From 0a9c0167ddbeab13232d648d550bedbb2562e3f3 Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Sun, 12 Nov 2023 04:11:56 +0800
Subject: [PATCH] [add] chat.completions ( dev )

---
 moapi/api.go     |   9 ++++
 moapi/process.go | 104 +++++++++++++++++++++++++++++++++++++++++++++++
 openai/openai.go |  10 +++++
 3 files changed, 123 insertions(+)

diff --git a/moapi/api.go b/moapi/api.go
index a89bc7aa..92b86647 100644
--- a/moapi/api.go
+++ b/moapi/api.go
@@ -16,6 +16,15 @@ var dsl = []byte(`
 			"process": "moapi.images.Generations",
 			"in": ["$payload.model", "$payload.prompt", ":payload"],
 			"out": { "status": 200, "type": "application/json" }
+		},
+		
+		{
+			"path": "/chat/completions",
+			"guard": "query-jwt",
+			"method": "GET",
+			"process": "moapi.chat.Completions",
+			"processHandler": true,
+			"out": { "status": 200, "type": "text/event-stream" }
 		}
 	]
 }
diff --git a/moapi/process.go b/moapi/process.go
index a5c70f62..b5271434 100644
--- a/moapi/process.go
+++ b/moapi/process.go
@@ -1,6 +1,13 @@
 package moapi
 
 import (
+	"context"
+	"io"
+	"net/http"
+	"strings"
+
+	"github.com/gin-gonic/gin"
+	jsoniter "github.com/json-iterator/go"
 	"github.com/yaoapp/gou/process"
 	"github.com/yaoapp/kun/exception"
 	"github.com/yaoapp/kun/utils"
@@ -10,6 +17,7 @@ import (
 func init() {
 	process.RegisterGroup("moapi", map[string]process.Handler{
 		"images.generations": ImagesGenerations,
+		"chat.completions":   ChatCompletions,
 	})
 }
 
@@ -46,3 +54,99 @@ func ImagesGenerations(process *process.Process) interface{} {
 
 	return res
 }
+
+// ChatCompletions chat completions
+func ChatCompletions(process *process.Process) interface{} {
+
+	return func(c *gin.Context) {
+
+		option := map[string]interface{}{}
+		query := c.Query("payload")
+		err := jsoniter.UnmarshalFromString(query, &option)
+		if err != nil {
+			exception.New("ChatCompletions error: %s", 400, err).Throw()
+		}
+
+		// option := payload
+		// model := "gpt-3.5-turbo"
+		// messages := []map[string]interface{}{
+		// 	{
+		// 		"role":    "system",
+		// 		"content": "You are a helpful assistant.",
+		// 	},
+		// 	{
+		// 		"role":    "user",
+		// 		"content": "Hello!",
+		// 	},
+		// // }
+
+		// option["messages"] = messages
+		// option["model"] = model
+
+		delete(option, "context")
+		model, ok := option["model"].(string)
+		if !ok || model == "" {
+			exception.New("ChatCompletions error: model is required", 400).Throw()
+		}
+
+		ai, err := openai.NewMoapi(model)
+		if err != nil {
+			exception.New("ChatCompletions error: %s", 400, err).Throw()
+		}
+
+		if v, ok := option["stream"].(bool); ok && v {
+
+			chanStream := make(chan []byte, 1)
+			chanError := make(chan error, 1)
+
+			defer func() {
+				close(chanStream)
+				close(chanError)
+			}()
+
+			ctx, cancel := context.WithCancel(c.Request.Context())
+			defer cancel()
+
+			go ai.Stream(ctx, "/v1/chat/completions", option, func(data []byte) int {
+
+				if (string(data)) == "\n" || string(data) == "" {
+					return 1 // HandlerReturnOk
+				}
+
+				chanStream <- data
+				if strings.HasSuffix(string(data), "[DONE]") {
+					return 0 // HandlerReturnBreak0
+				}
+				return 1 // HandlerReturnOk
+			})
+
+			c.Header("Content-Type", "text/event-stream")
+			c.Stream(func(w io.Writer) bool {
+				select {
+				case err := <-chanError:
+					if err != nil {
+						c.JSON(http.StatusInternalServerError, err.Error())
+					}
+					return false
+
+				case msg := <-chanStream:
+
+					if string(msg) == "\n" {
+						return true
+					}
+
+					message := strings.TrimLeft(string(msg), "data: ")
+					c.SSEvent("message", message)
+					return true
+
+				case <-ctx.Done():
+					return false
+				}
+			})
+
+			return
+		}
+
+		return
+	}
+}
diff --git a/openai/openai.go b/openai/openai.go
index 560ad9ac..9ac7f94c 100644
--- a/openai/openai.go
+++ b/openai/openai.go
@@ -298,6 +298,16 @@ func (openai OpenAI) GetContent(response interface{}) (string, *exception.Except
 	return "", exception.New("response format error, %#v", 500, response)
 }
 
+// Post post request
+func (openai OpenAI) Post(path string, payload map[string]interface{}) (interface{}, *exception.Exception) {
+	return openai.post(path, payload)
+}
+
+// Stream post request
+func (openai OpenAI) Stream(ctx context.Context, path string, payload map[string]interface{}, cb func(data []byte) int) *exception.Exception {
+	return openai.stream(ctx, path, payload, cb)
+}
+
 // post post request
 func (openai OpenAI) post(path string, payload map[string]interface{}) (interface{}, *exception.Exception) {