diff --git a/go.mod b/go.mod index 1046fe80..96395538 100644 --- a/go.mod +++ b/go.mod @@ -27,7 +27,7 @@ require ( github.com/yaoapp/kun v0.9.0 github.com/yaoapp/xun v0.9.0 golang.org/x/crypto v0.31.0 - golang.org/x/net v0.27.0 + golang.org/x/net v0.33.0 golang.org/x/text v0.21.0 gopkg.in/natefinch/lumberjack.v2 v2.2.1 gopkg.in/yaml.v3 v3.0.1 diff --git a/go.sum b/go.sum index 0600be76..fc31d260 100644 --- a/go.sum +++ b/go.sum @@ -281,8 +281,8 @@ golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs= golang.org/x/net v0.9.0/go.mod h1:d48xBJpPfHeWQsugry2m+kC02ZBRGRgulfHnEXEuWns= golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg= golang.org/x/net v0.14.0/go.mod h1:PpSgVXXLK0OxS0F31C1/tv6XNguvCrnXIDrFMspZIUI= -golang.org/x/net v0.27.0 h1:5K3Njcw06/l2y9vpGCSdcxWOYHOUk3dVNGDXN+FvAys= -golang.org/x/net v0.27.0/go.mod h1:dDi0PyhWNoiUOrAS8uXv/vnScO4wnHQO4mj9fn/RytE= +golang.org/x/net v0.33.0 h1:74SYHlV8BIgHIFC/LrYkOGIwL19eTYXQ5wc6TBuO36I= +golang.org/x/net v0.33.0/go.mod h1:HXLR5J+9DxmrqMwG9qjGCxZ+zKXxBru04zlTvWlWuN4= golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= golang.org/x/oauth2 v0.0.0-20181106182150-f42d05182288/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= golang.org/x/oauth2 v0.14.0 h1:P0Vrf/2538nmC0H+pEQ3MNFRRnVR7RlqyVw+bvm26z0= diff --git a/neo/api.go b/neo/api.go new file mode 100644 index 00000000..2fd9289d --- /dev/null +++ b/neo/api.go @@ -0,0 +1,536 @@ +package neo + +import ( + "fmt" + "io" + "net/url" + "path/filepath" + "strconv" + "strings" + "time" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + "github.com/yaoapp/gou/api" + "github.com/yaoapp/gou/process" + "github.com/yaoapp/yao/helper" + "github.com/yaoapp/yao/neo/conversation" + "github.com/yaoapp/yao/neo/message" +) + +// API registers the Neo API endpoints +func (neo *DSL) API(router *gin.Engine, path string) error { + + // Get the guards + middlewares, err := neo.getGuardHandlers() + if err != nil { + return err + } + + // Register OPTIONS handlers for all endpoints + router.OPTIONS(path, neo.optionsHandler) + router.OPTIONS(path+"/status", neo.optionsHandler) + router.OPTIONS(path+"/chats", neo.optionsHandler) + router.OPTIONS(path+"/chats/:id", neo.optionsHandler) + router.OPTIONS(path+"/history", neo.optionsHandler) + router.OPTIONS(path+"/upload", neo.optionsHandler) + router.OPTIONS(path+"/download", neo.optionsHandler) + router.OPTIONS(path+"/mentions", neo.optionsHandler) + router.OPTIONS(path+"/dangerous/clear_chats", neo.optionsHandler) + + // Register endpoints with middlewares + router.GET(path, append(middlewares, neo.handleChat)...) + router.POST(path, append(middlewares, neo.handleChat)...) + + // Status check + router.GET(path+"/status", append(middlewares, neo.handleStatus)...) + + // Chat api + router.GET(path+"/chats", append(middlewares, neo.handleChatList)...) + router.GET(path+"/chats/:id", append(middlewares, neo.handleChatDetail)...) + router.POST(path+"/chats/:id", append(middlewares, neo.handleChatUpdate)...) + router.DELETE(path+"/chats/:id", append(middlewares, neo.handleChatDelete)...) + + // History api + router.GET(path+"/history", append(middlewares, neo.handleChatHistory)...) + + // File api + router.POST(path+"/upload", append(middlewares, neo.handleUpload)...) + router.GET(path+"/download", append(middlewares, neo.handleDownload)...) + + // Mention api + router.GET(path+"/mentions", append(middlewares, neo.handleMentions)...) + + // Dangerous operations + router.DELETE(path+"/dangerous/clear_chats", append(middlewares, neo.handleChatsDeleteAll)...) + + return nil +} + +// handleStatus handles the status request +func (neo *DSL) handleStatus(c *gin.Context) { + c.Status(200) + c.Done() +} + +// handleUpload handles the upload request +func (neo *DSL) handleUpload(c *gin.Context) { + sid := c.GetString("__sid") + if sid == "" { + sid = uuid.New().String() + } + + // Set the context + ctx, cancel := NewContextWithCancel(sid, c.Query("chat_id"), "") + defer cancel() + + // Upload the file + file, err := neo.Upload(ctx, c) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + return + } + + c.JSON(200, file) + c.Done() +} + +// handleChat handles the chat request +func (neo *DSL) handleChat(c *gin.Context) { + // Set headers for SSE + c.Header("Content-Type", "text/event-stream;charset=utf-8") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + + sid := c.GetString("__sid") + if sid == "" { + sid = uuid.New().String() + } + + content := c.Query("content") + if content == "" { + msg := message.New().Error("content is required").Done() + msg.Write(c.Writer) + return + } + + chatID := c.Query("chat_id") + if chatID == "" { + // Only generate new chat_id if not provided + chatID = fmt.Sprintf("chat_%d", time.Now().UnixNano()) + } + + // Set the context with validated chat_id + ctx, cancel := NewContextWithCancel(sid, chatID, c.Query("context")) + defer cancel() + + neo.Answer(ctx, content, c) +} + +// handleChatList handles the chat list request +func (neo *DSL) handleChatList(c *gin.Context) { + sid := c.GetString("__sid") + if sid == "" { + c.JSON(400, gin.H{"message": "sid is required", "code": 400}) + c.Done() + return + } + + // Create filter from query parameters + filter := conversation.ChatFilter{ + Keywords: c.Query("keywords"), + Order: c.Query("order"), + } + + // Parse page and pagesize + if page := c.Query("page"); page != "" { + if n, err := strconv.Atoi(page); err == nil { + filter.Page = n + } + } + + if pageSize := c.Query("pagesize"); pageSize != "" { + if n, err := strconv.Atoi(pageSize); err == nil { + filter.PageSize = n + } + } + + response, err := neo.Conversation.GetChats(sid, filter) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + return + } + + c.JSON(200, map[string]interface{}{"data": response}) + c.Done() +} + +// handleChatHistory handles the chat history request +func (neo *DSL) handleChatHistory(c *gin.Context) { + sid := c.GetString("__sid") + if sid == "" { + c.JSON(400, gin.H{"message": "sid is required", "code": 400}) + c.Done() + return + } + + cid := c.Query("chat_id") + history, err := neo.Conversation.GetHistory(sid, cid) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + return + } + + c.JSON(200, map[string]interface{}{"data": history}) + c.Done() +} + +// handleDownload handles the download request +func (neo *DSL) handleDownload(c *gin.Context) { + sid := c.GetString("__sid") + if sid == "" { + c.JSON(400, gin.H{"message": "sid is required", "code": 400}) + c.Done() + return + } + + fileID := c.Query("file_id") + if fileID == "" { + c.JSON(400, gin.H{"message": "file_id is required", "code": 400}) + c.Done() + return + } + + // Set the context + ctx, cancel := NewContextWithCancel(sid, c.Query("chat_id"), "") + defer cancel() + + // Download the file + fileResponse, err := neo.Download(ctx, c) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + return + } + defer fileResponse.Reader.Close() + + // Set response headers + c.Header("Content-Type", fileResponse.ContentType) + if disposition := c.Query("disposition"); disposition == "attachment" { + c.Header("Content-Disposition", fmt.Sprintf("attachment; filename=%s", filepath.Base(fileID)+fileResponse.Extension)) + } + + // Copy the file content to response + _, err = io.Copy(c.Writer, fileResponse.Reader) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + return + } +} + +// getCorsHandlers returns CORS middleware handlers +func (neo *DSL) getCorsHandlers() ([]gin.HandlerFunc, error) { + if len(neo.Allows) == 0 { + return []gin.HandlerFunc{}, nil + } + + allowsMap := map[string]bool{} + for _, allow := range neo.Allows { + allow = strings.TrimPrefix(allow, "http://") + allow = strings.TrimPrefix(allow, "https://") + allowsMap[allow] = true + } + + return []gin.HandlerFunc{neo.corsMiddleware(allowsMap)}, nil +} + +// corsMiddleware handles CORS requests +func (neo *DSL) corsMiddleware(allowsMap map[string]bool) gin.HandlerFunc { + return func(c *gin.Context) { + origin := neo.getOrigin(c) + if origin == "" { + c.Next() + return + } + + // Check if origin is allowed + if !api.IsAllowed(c, allowsMap) { + c.AbortWithStatusJSON(403, gin.H{ + "message": origin + " not allowed", + "code": 403, + }) + return + } + + // Set CORS headers + c.Header("Access-Control-Allow-Origin", origin) + c.Header("Access-Control-Allow-Credentials", "true") + c.Header("Access-Control-Allow-Headers", "Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, Accept, Origin, Cache-Control, X-Requested-With") + c.Header("Access-Control-Allow-Methods", "GET, POST, OPTIONS") + + if c.Request.Method == "OPTIONS" { + c.AbortWithStatus(204) + return + } + + c.Next() + } +} + +// optionsHandler handles OPTIONS requests +func (neo *DSL) optionsHandler(c *gin.Context) { + origin := neo.getOrigin(c) + if origin != "" { + c.Header("Access-Control-Allow-Origin", origin) + c.Header("Access-Control-Allow-Methods", "GET, POST, DELETE, OPTIONS") + c.Header("Access-Control-Allow-Headers", "Content-Type, Authorization, Accept") + c.Header("Access-Control-Allow-Credentials", "true") + c.Header("Access-Control-Max-Age", "86400") // 24 hours + } + c.AbortWithStatus(204) +} + +// getOrigin returns the request origin +func (neo *DSL) getOrigin(c *gin.Context) string { + origin := c.Request.Header.Get("Origin") + if origin == "" { + origin = c.Request.Referer() + if origin != "" { + if u, err := url.Parse(origin); err == nil { + origin = fmt.Sprintf("%s://%s", u.Scheme, u.Host) + } + } + } + return origin +} + +// getGuardHandlers returns authentication middleware handlers +func (neo *DSL) getGuardHandlers() ([]gin.HandlerFunc, error) { + + // Cross-Domain handlers + cors, err := neo.getCorsHandlers() + if err != nil { + return nil, err + } + + if neo.Guard == "" { + middlewares := append(cors, neo.defaultGuard) + return middlewares, nil + } + + // Validate the custom guard + _, err = process.Of(neo.Guard) + if err != nil { + return nil, err + } + + middlewares := append(cors, api.ProcessGuard(neo.Guard, cors...)) + return middlewares, nil +} + +// defaultGuard is the default authentication handler +func (neo *DSL) defaultGuard(c *gin.Context) { + token := strings.TrimSpace(strings.TrimPrefix(c.Query("token"), "Bearer ")) + if token == "" { + c.JSON(403, gin.H{"message": "token is required", "code": 403}) + c.Abort() + return + } + + user := helper.JwtValidate(token) + c.Set("__sid", user.SID) + c.Next() +} + +// handleChatDetail handles getting a single chat's details +func (neo *DSL) handleChatDetail(c *gin.Context) { + sid := c.GetString("__sid") + if sid == "" { + c.JSON(400, gin.H{"message": "sid is required", "code": 400}) + c.Done() + return + } + + chatID := c.Param("id") + if chatID == "" { + c.JSON(400, gin.H{"message": "chat id is required", "code": 400}) + c.Done() + return + } + + chat, err := neo.Conversation.GetChat(sid, chatID) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + return + } + + c.JSON(200, map[string]interface{}{"data": chat}) + c.Done() +} + +// handleMentions handles getting mentions for a chat +func (neo *DSL) handleMentions(c *gin.Context) { + sid := c.GetString("__sid") + if sid == "" { + c.JSON(400, gin.H{"message": "sid is required", "code": 400}) + c.Done() + return + } + + // Get keywords from query parameter + keywords := strings.ToLower(c.Query("keywords")) + mentions, err := neo.GetMentions(keywords) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + return + } + + // Add test data + testMentions := []Mention{ + { + ID: "assistant_1", + Name: "Alice AI", + Type: "assistant", + Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Alice", + }, + { + ID: "assistant_2", + Name: "Bob Bot", + Type: "assistant", + Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Bob", + }, + { + ID: "assistant_3", + Name: "Carol AI", + Type: "assistant", + Avatar: "https://api.dicebear.com/7.x/avataaars/svg?seed=Carol", + }, + } + + // Filter mentions by keywords + if keywords != "" { + filtered := []Mention{} + for _, m := range testMentions { + if strings.Contains(strings.ToLower(m.Name), keywords) { + filtered = append(filtered, m) + } + } + testMentions = filtered + } + + // Append test data to actual mentions + mentions = append(mentions, testMentions...) + + c.JSON(200, map[string]interface{}{"data": mentions}) + c.Done() +} + +// handleChatUpdate handles updating a chat's details +func (neo *DSL) handleChatUpdate(c *gin.Context) { + sid := c.GetString("__sid") + if sid == "" { + c.JSON(400, gin.H{"message": "sid is required", "code": 400}) + c.Done() + return + } + + chatID := c.Param("id") + if chatID == "" { + c.JSON(400, gin.H{"message": "chat id is required", "code": 400}) + c.Done() + return + } + + // Get title from request body + var body struct { + Title string `json:"title"` + Content string `json:"content"` + } + if err := c.BindJSON(&body); err != nil { + c.JSON(400, gin.H{"message": "invalid request body", "code": 400}) + c.Done() + return + } + + // If content is not empty, Generate the chat title + if body.Content != "" { + ctx, cancel := NewContextWithCancel(sid, c.Query("chat_id"), "") + defer cancel() + + title, err := neo.GenerateChatTitle(ctx, body.Content, c) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + return + } + body.Title = title + } + + if body.Title == "" { + c.JSON(400, gin.H{"message": "title is required", "code": 400}) + c.Done() + return + } + + err := neo.Conversation.UpdateChatTitle(sid, chatID, body.Title) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + return + } + + c.JSON(200, gin.H{"message": "ok", "title": body.Title, "chat_id": chatID}) + c.Done() +} + +// handleChatDelete handles deleting a single chat +func (neo *DSL) handleChatDelete(c *gin.Context) { + sid := c.GetString("__sid") + if sid == "" { + c.JSON(400, gin.H{"message": "sid is required", "code": 400}) + c.Done() + return + } + + chatID := c.Param("id") + if chatID == "" { + c.JSON(400, gin.H{"message": "chat id is required", "code": 400}) + c.Done() + return + } + + err := neo.Conversation.DeleteChat(sid, chatID) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + return + } + + c.JSON(200, gin.H{"message": "ok"}) + c.Done() +} + +// handleChatsDeleteAll handles deleting all chats for a user +func (neo *DSL) handleChatsDeleteAll(c *gin.Context) { + sid := c.GetString("__sid") + if sid == "" { + c.JSON(400, gin.H{"message": "sid is required", "code": 400}) + c.Done() + return + } + + err := neo.Conversation.DeleteAllChats(sid) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + return + } + + c.JSON(200, gin.H{"message": "ok"}) + c.Done() +} diff --git a/neo/api_test.go b/neo/api_test.go new file mode 100644 index 00000000..2d3a8953 --- /dev/null +++ b/neo/api_test.go @@ -0,0 +1,233 @@ +package neo + +import ( + "context" + "fmt" + "net" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + httpTest "github.com/yaoapp/gou/http" + "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/helper" + "github.com/yaoapp/yao/test" +) + +func init() { + // Set gin to release mode to reduce log output + gin.SetMode(gin.ReleaseMode) +} + +func TestAPI(t *testing.T) { + // Disable test logging + test.Prepare(t, config.Conf) + defer test.Clean() + + // Redirect stdout to /dev/null + oldStdout := os.Stdout + null, _ := os.Open(os.DevNull) + os.Stdout = null + defer func() { + os.Stdout = oldStdout + null.Close() + }() + + // test router + router := testRouter(t) + err := Neo.API(router, "/neo/chat") + if err != nil { + t.Fatal(err) + } + + // test server + host, shutdown := testServer(t, router) + defer shutdown() + + tests := []struct { + name string + url string + method string + headers http.Header + expectCode int + expectBody string + }{ + { + name: "Basic Chat Request", + url: fmt.Sprintf("/neo/chat?content=hello&token=%s", testToken()), + method: "GET", + headers: http.Header{"Content-Type": []string{"application/json"}}, + expectBody: `{`, + }, + { + name: "Chat with System Message", + url: fmt.Sprintf("/neo/chat?content=hello&system=You are a helpful assistant&token=%s", testToken()), + method: "GET", + headers: http.Header{"Content-Type": []string{"application/json"}}, + expectBody: `{`, + }, + { + name: "Chat with Model Parameter", + url: fmt.Sprintf("/neo/chat?content=hello&model=gpt-3.5-turbo&token=%s", testToken()), + method: "GET", + headers: http.Header{"Content-Type": []string{"application/json"}}, + expectBody: `{`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + url := fmt.Sprintf("%s%s", host, tt.url) + res := []byte{} + req := httpTest.New(url).WithHeader(tt.headers) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + req.Stream(ctx, tt.method, nil, func(data []byte) int { + res = append(res, data...) + return 1 + }) + + assert.Contains(t, string(res), tt.expectBody) + }) + } +} + +func TestAPIAuth(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + // Redirect stdout and stderr to /dev/null + oldStdout := os.Stdout + oldStderr := os.Stderr + null, _ := os.Open(os.DevNull) + os.Stdout = null + os.Stderr = null + defer func() { + os.Stdout = oldStdout + os.Stderr = oldStderr + null.Close() + }() + + router := testRouter(t) + err := Neo.API(router, "/neo/chat") + if err != nil { + t.Fatal(err) + } + + // Separate tests for authentication errors and parameter validation errors + authTests := []struct { + name string + url string + method string + expectCode int + }{ + { + name: "Missing Token", + url: "/neo/chat?content=hello", + method: "GET", + expectCode: http.StatusUnauthorized, + }, + { + name: "Invalid Token", + url: "/neo/chat?content=hello&token=invalid", + method: "GET", + expectCode: http.StatusUnauthorized, + }, + } + + // Test authentication errors (will panic) + for _, tt := range authTests { + t.Run(tt.name, func(t *testing.T) { + response := httptest.NewRecorder() + req, _ := http.NewRequest(tt.method, tt.url, nil) + assert.Panics(t, func() { + router.ServeHTTP(response, req) + }) + }) + } + + // Test parameter validation errors (will return status code) + validationTests := []struct { + name string + url string + method string + expectCode int + }{ + { + name: "Missing Content", + url: fmt.Sprintf("/neo/chat?token=%s", testToken()), + method: "GET", + expectCode: http.StatusBadRequest, + }, + } + + // Test parameter validation errors (return status code) + for _, tt := range validationTests { + t.Run(tt.name, func(t *testing.T) { + response := httptest.NewRecorder() + req, _ := http.NewRequest(tt.method, tt.url, nil) + router.ServeHTTP(response, req) + assert.Equal(t, tt.expectCode, response.Code) + }) + } +} + +// Helper functions +func testServer(t *testing.T, router *gin.Engine) (string, func()) { + l, err := net.Listen("tcp4", ":0") + if err != nil { + t.Fatal(err) + } + + srv := &http.Server{Addr: ":0", Handler: router} + + go func() { + if err := srv.Serve(l); err != nil && err != http.ErrServerClosed { + return + } + }() + + addr := strings.Split(l.Addr().String(), ":") + if len(addr) != 2 { + t.Fatal("invalid address") + } + + host := fmt.Sprintf("http://127.0.0.1:%s", addr[1]) + time.Sleep(50 * time.Millisecond) + + shutdown := func() { + srv.Close() + l.Close() + } + return host, shutdown +} + +func testRouter(t *testing.T) *gin.Engine { + err := Load(config.Conf) + if err != nil { + t.Fatal(err) + } + + router := gin.New() // Use gin.New() instead of gin.Default() to avoid default logging middleware + return router +} + +func testToken() string { + token := helper.JwtMake(1, + map[string]interface{}{ + "id": 1, + "name": "Test", + }, + map[string]interface{}{ + "exp": 3600, + "sid": "123456", + }) + return token.Token +} diff --git a/neo/assistant/base/base.go b/neo/assistant/base/base.go new file mode 100644 index 00000000..bd608fe1 --- /dev/null +++ b/neo/assistant/base/base.go @@ -0,0 +1,34 @@ +package base + +import ( + "context" + + "github.com/yaoapp/gou/connector" + "github.com/yaoapp/yao/neo/assistant" + "github.com/yaoapp/yao/openai" +) + +// Base the base assistant +type Base struct { + ID string `json:"assistant_id"` + Prompts []assistant.Prompt `json:"prompts,omitempty"` + Connector connector.Connector `json:"-" yaml:"-"` + openai *openai.OpenAI +} + +// New create a new base assistant +func New(connector connector.Connector, prompts []assistant.Prompt, id string) (*Base, error) { + + setting := connector.Setting() + api, err := openai.NewOpenAI(setting) + if err != nil { + return nil, err + } + + return &Base{Connector: connector, ID: id, Prompts: prompts, openai: api}, nil +} + +// List list all assistants +func (ast *Base) List(ctx context.Context, param assistant.QueryParam) ([]assistant.Assistant, error) { + return nil, nil +} diff --git a/neo/assistant/base/chat.go b/neo/assistant/base/chat.go new file mode 100644 index 00000000..7df5c58a --- /dev/null +++ b/neo/assistant/base/chat.go @@ -0,0 +1,21 @@ +package base + +import ( + "context" + "fmt" +) + +// Chat the chat +func (ast *Base) Chat(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) error { + + if ast.openai == nil { + return fmt.Errorf("api is not initialized") + } + + _, ext := ast.openai.ChatCompletionsWith(ctx, messages, option, cb) + if ext != nil { + return fmt.Errorf("openai chat completions with error: %s", ext.Message) + } + + return nil +} diff --git a/neo/assistant/base/file.go b/neo/assistant/base/file.go new file mode 100644 index 00000000..fd4f3029 --- /dev/null +++ b/neo/assistant/base/file.go @@ -0,0 +1,134 @@ +package base + +import ( + "context" + "crypto/sha256" + "fmt" + "io" + "mime/multipart" + "path/filepath" + "strings" + "time" + + "github.com/yaoapp/gou/fs" + "github.com/yaoapp/yao/neo/assistant" +) + +// AllowedFileTypes the allowed file types +var AllowedFileTypes = map[string]string{ + "application/json": "json", + "application/pdf": "pdf", + "application/msword": "doc", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document": "docx", + "application/vnd.oasis.opendocument.text": "odt", + "application/vnd.ms-excel": "xls", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": "xlsx", + "application/vnd.ms-powerpoint": "ppt", + "application/vnd.openxmlformats-officedocument.presentationml.presentation": "pptx", +} + +// MaxSize 20M max file size +var MaxSize int64 = 20 * 1024 * 1024 + +// Upload the file +func (ast *Base) Upload(ctx context.Context, file *multipart.FileHeader, reader io.Reader, option map[string]interface{}) (*assistant.File, error) { + + // check file size + if file.Size > MaxSize { + return nil, fmt.Errorf("file size %d exceeds the maximum size of %d", file.Size, MaxSize) + } + + contentType := file.Header.Get("Content-Type") + if !ast.allowed(contentType) { + return nil, fmt.Errorf("file type %s not allowed", contentType) + } + + data, err := fs.Get("data") + if err != nil { + return nil, err + } + + ext := filepath.Ext(file.Filename) + id, err := ast.id(file.Filename, ext) + if err != nil { + return nil, err + } + + filename := fmt.Sprintf("%s%s", id, ext) + _, err = data.Write(filename, reader, 0644) + if err != nil { + return nil, err + } + + return &assistant.File{ + ID: filename, + Filename: filename, + ContentType: contentType, + Bytes: int(file.Size), + CreatedAt: int(time.Now().Unix()), + }, nil +} + +func (ast *Base) id(temp string, ext string) (string, error) { + date := time.Now().Format("20060102") + hash := fmt.Sprintf("%x", sha256.Sum256([]byte(temp)))[:8] + return fmt.Sprintf("/__assistants/%s/%s/%s%s", ast.ID, date, hash, ext), nil +} + +func (ast *Base) allowed(contentType string) bool { + if _, ok := AllowedFileTypes[contentType]; ok { + return true + } + // text/* // image/* // audio/* // video/* + if strings.HasPrefix(contentType, "text/") || strings.HasPrefix(contentType, "image/") || strings.HasPrefix(contentType, "audio/") || strings.HasPrefix(contentType, "video/") { + return true + } + return false +} + +// Download downloads a file +func (ast *Base) Download(ctx context.Context, fileID string) (*assistant.FileResponse, error) { + + // Get the data filesystem + data, err := fs.Get("data") + if err != nil { + return nil, fmt.Errorf("get filesystem error: %s", err.Error()) + } + + // Check if file exists + exists, err := data.Exists(fileID) + if err != nil { + return nil, fmt.Errorf("check file error: %s", err.Error()) + } + if !exists { + return nil, fmt.Errorf("file %s not found", fileID) + } + + // Open the file + reader, err := data.ReadCloser(fileID) + if err != nil { + return nil, err + } + + // Get content type and extension + ext := filepath.Ext(fileID) + + // Get content type from mime type + contentType := "application/octet-stream" + if v, err := data.MimeType(fileID); err == nil { + contentType = v + } + + for mimeType, extension := range AllowedFileTypes { + if "."+extension == ext { + contentType = mimeType + break + } + } + + return &assistant.FileResponse{ + Reader: reader, + ContentType: contentType, + Extension: ext, + }, nil +} diff --git a/neo/assistant/openai/chat.go b/neo/assistant/openai/chat.go new file mode 100644 index 00000000..765c6f3a --- /dev/null +++ b/neo/assistant/openai/chat.go @@ -0,0 +1,30 @@ +package openai + +import ( + "context" + "fmt" +) + +// Chat the chat struct +type Chat struct { + ID string `json:"chat_id"` + ThreadID string `json:"thread_id"` +} + +// NewChat create a new chat +func (ast *OpenAI) NewChat() {} + +// Chat the chat +func (ast *OpenAI) Chat(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) error { + + if ast.openai == nil { + return fmt.Errorf("openai is not initialized") + } + + _, ext := ast.openai.ChatCompletionsWith(ctx, messages, option, cb) + if ext != nil { + return fmt.Errorf("openai chat completions with error: %s", ext.Message) + } + + return nil +} diff --git a/neo/assistant/openai/file.go b/neo/assistant/openai/file.go new file mode 100644 index 00000000..44892d49 --- /dev/null +++ b/neo/assistant/openai/file.go @@ -0,0 +1,139 @@ +package openai + +import ( + "context" + "crypto/sha256" + "fmt" + "io" + "mime/multipart" + "path/filepath" + "strings" + "time" + + "github.com/yaoapp/gou/fs" + "github.com/yaoapp/yao/neo/assistant" +) + +// AllowedFileTypes the allowed file types +var AllowedFileTypes = map[string]string{ + "application/json": "json", + "application/pdf": "pdf", + "application/msword": "doc", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document": "docx", + "application/vnd.oasis.opendocument.text": "odt", + "application/vnd.ms-excel": "xls", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": "xlsx", + "application/vnd.ms-powerpoint": "ppt", + "application/vnd.openxmlformats-officedocument.presentationml.presentation": "pptx", +} + +// MaxSize 20M max file size +var MaxSize int64 = 20 * 1024 * 1024 + +// Upload the file +func (ast *OpenAI) Upload(ctx context.Context, file *multipart.FileHeader, reader io.Reader, option map[string]interface{}) (*assistant.File, error) { + + // check file size + if file.Size > MaxSize { + return nil, fmt.Errorf("file size %d exceeds the maximum size of %d", file.Size, MaxSize) + } + + contentType := file.Header.Get("Content-Type") + if !ast.allowed(contentType) { + return nil, fmt.Errorf("file type %s not allowed", contentType) + } + + data, err := fs.Get("data") + if err != nil { + return nil, err + } + + ext := filepath.Ext(file.Filename) + id, err := ast.id(file.Filename, ext) + if err != nil { + return nil, err + } + + filename := id + _, err = data.Write(filename, reader, 0644) + if err != nil { + return nil, err + } + + return &assistant.File{ + ID: filename, + Filename: filename, + ContentType: contentType, + Bytes: int(file.Size), + CreatedAt: int(time.Now().Unix()), + }, nil +} + +func (ast *OpenAI) id(temp string, ext string) (string, error) { + date := time.Now().Format("20060102") + hash := fmt.Sprintf("%x", sha256.Sum256([]byte(temp)))[:8] + return fmt.Sprintf("/__assistants/%s/%s/%s%s", ast.ID, date, hash, ext), nil +} + +func (ast *OpenAI) allowed(contentType string) bool { + if _, ok := AllowedFileTypes[contentType]; ok { + return true + } + // text/* // image/* // audio/* // video/* + if strings.HasPrefix(contentType, "text/") || strings.HasPrefix(contentType, "image/") || strings.HasPrefix(contentType, "audio/") || strings.HasPrefix(contentType, "video/") { + return true + } + return false +} + +// FileLists list all files +func (ast *OpenAI) FileLists() {} + +// FileDelete delete a file +func (ast *OpenAI) FileDelete() {} + +// FileContent get the content of a file +func (ast *OpenAI) FileContent() {} + +// FileInfo get the information of a file +func (ast *OpenAI) FileInfo() {} + +// Download downloads a file +func (ast *OpenAI) Download(ctx context.Context, fileID string) (*assistant.FileResponse, error) { + + // Get the data filesystem + data, err := fs.Get("data") + if err != nil { + return nil, fmt.Errorf("get filesystem error: %s", err.Error()) + } + + // Check if file exists + exists, err := data.Exists(fileID) + if err != nil { + return nil, fmt.Errorf("check file error: %s", err.Error()) + } + if !exists { + return nil, fmt.Errorf("file %s not found", fileID) + } + + // Open the file + reader, err := data.ReadCloser(fileID) + if err != nil { + return nil, err + } + + // Get content type and extension + ext := filepath.Ext(fileID) + + // Get content type from mime type + contentType := "application/octet-stream" + if v, err := data.MimeType(fileID); err == nil { + contentType = v + } + + return &assistant.FileResponse{ + Reader: reader, + ContentType: contentType, + Extension: ext, + }, nil +} diff --git a/neo/assistant/openai/openai.go b/neo/assistant/openai/openai.go new file mode 100644 index 00000000..38f5135c --- /dev/null +++ b/neo/assistant/openai/openai.go @@ -0,0 +1,51 @@ +package openai + +import ( + "context" + + "github.com/yaoapp/gou/connector" + "github.com/yaoapp/yao/neo/assistant" + api "github.com/yaoapp/yao/openai" +) + +// OpenAI the openai assistant +type OpenAI struct { + ID string `json:"assistant_id"` // the assistant id + Connector connector.Connector `json:"-" yaml:"-"` + openai *api.OpenAI +} + +// New create a new openai assistant +func New(connector connector.Connector, id string) (*OpenAI, error) { + + setting := connector.Setting() + openai, err := api.NewOpenAI(setting) + if err != nil { + return nil, err + } + + return &OpenAI{ID: id, Connector: connector, openai: openai}, nil +} + +// Current set the current assistant +func (ast *OpenAI) Current(id string) *OpenAI { + ast.ID = id + return ast +} + +// List list all assistants +func (ast *OpenAI) List(ctx context.Context, param assistant.QueryParam) ([]assistant.Assistant, error) { + return nil, nil +} + +// Create create a new assistant +func (ast *OpenAI) Create() {} + +// Delete delete an assistant +func (ast *OpenAI) Delete() {} + +// Update update an assistant +func (ast *OpenAI) Update() {} + +// Get get an assistant +func (ast *OpenAI) Get() {} diff --git a/neo/assistant/openai/thread.go b/neo/assistant/openai/thread.go new file mode 100644 index 00000000..8f5975e8 --- /dev/null +++ b/neo/assistant/openai/thread.go @@ -0,0 +1,21 @@ +package openai + +// Thread the thread struct +type Thread struct { + ID string `json:"thread_id"` +} + +// ThreadList list all threads +func (ast *OpenAI) ThreadList() {} + +// ThreadCreate create a new thread +func (ast *OpenAI) ThreadCreate() {} + +// ThreadGet get a thread +func (ast *OpenAI) ThreadGet(id string) {} + +// ThreadDelete delete a thread +func (ast *OpenAI) ThreadDelete() {} + +// ThreadUpdate update a thread +func (ast *OpenAI) ThreadUpdate() {} diff --git a/neo/assistant/types.go b/neo/assistant/types.go new file mode 100644 index 00000000..03930582 --- /dev/null +++ b/neo/assistant/types.go @@ -0,0 +1,56 @@ +package assistant + +import ( + "context" + "io" + "mime/multipart" +) + +// API the assistant API interface +type API interface { + Chat(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) error + Upload(ctx context.Context, file *multipart.FileHeader, reader io.Reader, option map[string]interface{}) (*File, error) + Download(ctx context.Context, fileID string) (*FileResponse, error) +} + +// Prompt a prompt +type Prompt struct { + Role string `json:"role"` + Content string `json:"content"` + Name string `json:"name,omitempty"` +} + +// QueryParam the assistant query param +type QueryParam struct { + Limit uint `json:"limit"` + Order string `json:"order"` + After string `json:"after"` + Before string `json:"before"` +} + +// Assistant the assistant +type Assistant struct { + ID string `json:"assistant_id"` // Assistant ID + Name string `json:"name,omitempty"` // Assistant Name + Connector string `json:"connector"` // AI Connector + Description string `json:"description,omitempty"` // Assistant Description + Option map[string]interface{} `json:"option,omitempty"` // AI Option + Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts + API API `json:"-" yaml:"-"` // Assistant API +} + +// File the file +type File struct { + ID string `json:"file_id"` + Bytes int `json:"bytes"` + CreatedAt int `json:"created_at"` + Filename string `json:"filename"` + ContentType string `json:"content_type"` +} + +// FileResponse represents a file download response +type FileResponse struct { + Reader io.ReadCloser + ContentType string + Extension string +} diff --git a/neo/conversation/mongo.go b/neo/conversation/mongo.go index 6fe46e50..98c938ca 100644 --- a/neo/conversation/mongo.go +++ b/neo/conversation/mongo.go @@ -14,8 +14,14 @@ func (conv *Mongo) UpdateChatTitle(sid string, cid string, title string) error { } // GetChats get the chat list -func (conv *Mongo) GetChats(sid string) ([]map[string]interface{}, error) { - return []map[string]interface{}{}, nil +func (conv *Mongo) GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) { + return &ChatGroupResponse{ + Groups: []ChatGroup{}, + Page: filter.Page, + PageSize: filter.PageSize, + Total: 0, + LastPage: 1, + }, nil } // GetHistory get the history @@ -37,3 +43,18 @@ func (conv *Mongo) GetRequest(sid string, rid string) ([]map[string]interface{}, func (conv *Mongo) SaveRequest(sid string, rid string, cid string, messages []map[string]interface{}) error { return nil } + +// GetChat get the chat info and its history +func (conv *Mongo) GetChat(sid string, cid string) (*ChatInfo, error) { + return nil, nil +} + +// DeleteChat deletes a specific chat and its history +func (conv *Mongo) DeleteChat(sid string, cid string) error { + return nil +} + +// DeleteAllChats deletes all chats and their histories for a user +func (conv *Mongo) DeleteAllChats(sid string) error { + return nil +} diff --git a/neo/conversation/redis.go b/neo/conversation/redis.go index b19412cc..afb1aa74 100644 --- a/neo/conversation/redis.go +++ b/neo/conversation/redis.go @@ -14,8 +14,14 @@ func (conv *Redis) UpdateChatTitle(sid string, cid string, title string) error { } // GetChats get the chat list -func (conv *Redis) GetChats(sid string) ([]map[string]interface{}, error) { - return []map[string]interface{}{}, nil +func (conv *Redis) GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) { + return &ChatGroupResponse{ + Groups: []ChatGroup{}, + Page: filter.Page, + PageSize: filter.PageSize, + Total: 0, + LastPage: 1, + }, nil } // GetHistory get the history @@ -37,3 +43,18 @@ func (conv *Redis) GetRequest(sid string, rid string) ([]map[string]interface{}, func (conv *Redis) SaveRequest(sid string, rid string, cid string, messages []map[string]interface{}) error { return nil } + +// GetChat get the chat info and its history +func (conv *Redis) GetChat(sid string, cid string) (*ChatInfo, error) { + return nil, nil +} + +// DeleteChat deletes a specific chat and its history +func (conv *Redis) DeleteChat(sid string, cid string) error { + return nil +} + +// DeleteAllChats deletes all chats and their histories for a user +func (conv *Redis) DeleteAllChats(sid string) error { + return nil +} diff --git a/neo/conversation/types.go b/neo/conversation/types.go index 68ffba17..205a99d6 100644 --- a/neo/conversation/types.go +++ b/neo/conversation/types.go @@ -3,17 +3,50 @@ package conversation // Setting the conversation config type Setting struct { Connector string `json:"connector,omitempty"` + UserField string `json:"user_field,omitempty"` // the user id field name, default is user_id Table string `json:"table,omitempty"` MaxSize int `json:"max_size,omitempty" yaml:"max_size,omitempty"` TTL int `json:"ttl,omitempty" yaml:"ttl,omitempty"` } +// ChatInfo represents the chat information and its history +type ChatInfo struct { + Chat map[string]interface{} `json:"chat"` + History []map[string]interface{} `json:"history"` +} + +// ChatFilter represents the filter parameters for GetChats +type ChatFilter struct { + Keywords string `json:"keywords,omitempty"` + Page int `json:"page,omitempty"` // 页码，从1开始 + PageSize int `json:"pagesize,omitempty"` // 每页数量 + Order string `json:"order,omitempty"` // desc/asc +} + +// ChatGroup represents a group of chats by date +type ChatGroup struct { + Label string `json:"label"` + Chats []map[string]interface{} `json:"chats"` +} + +// ChatGroupResponse represents paginated chat groups +type ChatGroupResponse struct { + Groups []ChatGroup `json:"groups"` + Page int `json:"page"` // 当前页码 + PageSize int `json:"pagesize"` // 每页数量 + Total int64 `json:"total"` // 总记录数 + LastPage int `json:"last_page"` // 最后一页页码 +} + // Conversation the store interface type Conversation interface { - UpdateChatTitle(sid string, cid string, title string) error - GetChats(sid string) ([]map[string]interface{}, error) + GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) + GetChat(sid string, cid string) (*ChatInfo, error) GetHistory(sid string, cid string) ([]map[string]interface{}, error) SaveHistory(sid string, messages []map[string]interface{}, cid string) error GetRequest(sid string, rid string) ([]map[string]interface{}, error) SaveRequest(sid string, rid string, cid string, messages []map[string]interface{}) error + DeleteChat(sid string, cid string) error + DeleteAllChats(sid string) error + UpdateChatTitle(sid string, cid string, title string) error } diff --git a/neo/conversation/weaviate.go b/neo/conversation/weaviate.go index 7495c44d..7b251851 100644 --- a/neo/conversation/weaviate.go +++ b/neo/conversation/weaviate.go @@ -14,8 +14,14 @@ func (conv *Weaviate) UpdateChatTitle(sid string, cid string, title string) erro } // GetChats get the chat list -func (conv *Weaviate) GetChats(sid string) ([]map[string]interface{}, error) { - return []map[string]interface{}{}, nil +func (conv *Weaviate) GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) { + return &ChatGroupResponse{ + Groups: []ChatGroup{}, + Page: filter.Page, + PageSize: filter.PageSize, + Total: 0, + LastPage: 1, + }, nil } // GetHistory get the history @@ -37,3 +43,18 @@ func (conv *Weaviate) GetRequest(sid string, rid string) ([]map[string]interface func (conv *Weaviate) SaveRequest(sid string, rid string, cid string, messages []map[string]interface{}) error { return nil } + +// GetChat get the chat info and its history +func (conv *Weaviate) GetChat(sid string, cid string) (*ChatInfo, error) { + return nil, nil +} + +// DeleteChat deletes a specific chat and its history +func (conv *Weaviate) DeleteChat(sid string, cid string) error { + return nil +} + +// DeleteAllChats deletes all chats and their histories for a user +func (conv *Weaviate) DeleteAllChats(sid string) error { + return nil +} diff --git a/neo/conversation/xun.go b/neo/conversation/xun.go index c34cc58b..61a25b55 100644 --- a/neo/conversation/xun.go +++ b/neo/conversation/xun.go @@ -2,9 +2,13 @@ package conversation import ( "fmt" + "math" + "strings" "time" + "github.com/google/uuid" "github.com/yaoapp/gou/connector" + "github.com/yaoapp/gou/session" "github.com/yaoapp/kun/log" "github.com/yaoapp/xun/capsule" "github.com/yaoapp/xun/dbal/query" @@ -20,8 +24,7 @@ type Xun struct { type row struct { Role string `json:"role"` - Title string `json:"title"` // Chat title - Name string `json:"name"` // User name + Name string `json:"name"` // User name Content string `json:"content"` Sid string `json:"sid"` Rid string `json:"rid"` @@ -29,16 +32,23 @@ type row struct { ExpiredAt interface{} `json:"expired_at"` } +// Public interface methods and constructor remain exported: +// - NewXun +// - UpdateChatTitle +// - GetChats +// - GetChat +// - GetHistory +// - SaveHistory +// - GetRequest +// - SaveRequest + // NewXun create a new conversation func NewXun(setting Setting) (*Xun, error) { - conv := &Xun{setting: setting} if setting.Connector == "default" { conv.query = capsule.Global.Query() conv.schema = capsule.Global.Schema() - } else { - conn, err := connector.Select(setting.Connector) if err != nil { return nil, err @@ -55,7 +65,7 @@ func NewXun(setting Setting) (*Xun, error) { } } - err := conv.Init() + err := conv.initialize() if err != nil { return nil, err } @@ -63,48 +73,315 @@ func NewXun(setting Setting) (*Xun, error) { return conv, nil } +// Rename the following functions to start with lowercase letters to make them private: + +func (conv *Xun) newQuery() query.Query { + qb := conv.query.New() + qb.Table(conv.getHistoryTable()) + return qb +} + +func (conv *Xun) newQueryChat() query.Query { + qb := conv.query.New() + qb.Table(conv.getChatTable()) + return qb +} + +func (conv *Xun) clean() { + nums, err := conv.newQuery().Where("expired_at", "<=", time.Now()).Delete() + if err != nil { + log.Error("Clean the conversation table error: %s", err.Error()) + return + } + + if nums > 0 { + log.Trace("Clean the conversation table: %s %d", conv.setting.Table, nums) + } +} + +// Rename Init to initialize to avoid conflicts +func (conv *Xun) initialize() error { + // Initialize history table + if err := conv.initHistoryTable(); err != nil { + return err + } + + // Initialize chat table + if err := conv.initChatTable(); err != nil { + return err + } + + return nil +} + +func (conv *Xun) initHistoryTable() error { + historyTable := conv.getHistoryTable() + has, err := conv.schema.HasTable(historyTable) + if err != nil { + return err + } + + // Create the history table + if !has { + err = conv.schema.CreateTable(historyTable, func(table schema.Blueprint) { + table.ID("id") + table.String("sid", 255).Index() + table.String("rid", 255).Null().Index() + table.String("cid", 200).Null().Index() + table.String("role", 200).Null().Index() + table.String("name", 200).Null().Index() + table.Text("content").Null() + table.TimestampTz("created_at").SetDefaultRaw("NOW()").Index() + table.TimestampTz("updated_at").Null().Index() + table.TimestampTz("expired_at").Null().Index() + }) + + if err != nil { + return err + } + log.Trace("Create the conversation history table: %s", historyTable) + } + + // Validate the table + tab, err := conv.schema.GetTable(historyTable) + if err != nil { + return err + } + + fields := []string{"id", "sid", "rid", "cid", "role", "name", "content", "created_at", "updated_at", "expired_at"} + for _, field := range fields { + if !tab.HasColumn(field) { + return fmt.Errorf("%s is required", field) + } + } + + return nil +} + +func (conv *Xun) initChatTable() error { + chatTable := conv.getChatTable() + has, err := conv.schema.HasTable(chatTable) + if err != nil { + return err + } + + // Create the chat table + if !has { + err = conv.schema.CreateTable(chatTable, func(table schema.Blueprint) { + table.ID("id") + table.String("chat_id", 200).Unique().Index() + table.String("title", 200).Null() + table.String("sid", 255).Index() + table.TimestampTz("created_at").SetDefaultRaw("NOW()").Index() + table.TimestampTz("updated_at").Null().Index() + }) + + if err != nil { + return err + } + log.Trace("Create the chat table: %s", chatTable) + } + + // Validate the table + tab, err := conv.schema.GetTable(chatTable) + if err != nil { + return err + } + + fields := []string{"id", "chat_id", "title", "sid", "created_at", "updated_at"} + for _, field := range fields { + if !tab.HasColumn(field) { + return fmt.Errorf("%s is required", field) + } + } + + return nil +} + +func (conv *Xun) getUserID(sid string) (string, error) { + field := "user_id" + if conv.setting.UserField != "" { + field = conv.setting.UserField + } + + id, err := session.Global().ID(sid).Get(field) + if err != nil { + return "", err + } + + if id == nil || id == "" { + return sid, nil + } + + return fmt.Sprintf("%v", id), nil +} + +func (conv *Xun) getHistoryTable() string { + return conv.setting.Table +} + +func (conv *Xun) getChatTable() string { + return conv.setting.Table + "_chat" +} + // UpdateChatTitle update the chat title func (conv *Xun) UpdateChatTitle(sid string, cid string, title string) error { - _, err := conv.query.Table(conv.setting.Table). - Where("sid", sid).Where("cid", cid). - Update(map[string]interface{}{"title": title}) + userID, err := conv.getUserID(sid) + if err != nil { + return err + } + + _, err = conv.newQueryChat(). + Where("sid", userID). + Where("chat_id", cid). + Update(map[string]interface{}{ + "title": title, + "updated_at": time.Now(), + }) return err } -// GetChats get the chat list -func (conv *Xun) GetChats(sid string) ([]map[string]interface{}, error) { - qb := conv.query.Table(conv.setting.Table). - Select("cid"). - Where("sid", sid). - GroupBy("cid") - - if conv.setting.TTL > 0 { - qb.Where("expired_at", ">", time.Now()) - } - - res := []map[string]interface{}{} - - rows, err := qb.Get() +// GetChats get the chat list with grouping by date +func (conv *Xun) GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) { + userID, err := conv.getUserID(sid) if err != nil { return nil, err } - for _, row := range rows { - res = append(res, map[string]interface{}{ - "chat_id": row.Get("cid"), - "title": row.Get("cid"), - }) + // Set defaults + if filter.PageSize <= 0 { + filter.PageSize = 100 + } + if filter.Page <= 0 { + filter.Page = 1 + } + if filter.Order == "" { + filter.Order = "desc" } - return res, nil + // Build base query + qb := conv.newQueryChat(). + Select("chat_id", "title", "created_at"). + Where("sid", userID). + Where("chat_id", "!=", "") + + // Add keyword filter + if filter.Keywords != "" { + keyword := strings.TrimSpace(filter.Keywords) + if keyword != "" { + qb.Where("title", "like", "%"+keyword+"%") + } + } + + // Get total count + total, err := qb.Clone().Count() + if err != nil { + return nil, err + } + + // Calculate pagination + offset := (filter.Page - 1) * filter.PageSize + lastPage := int(math.Ceil(float64(total) / float64(filter.PageSize))) + + // Get paginated results + rows, err := qb.OrderBy("created_at", filter.Order). + Offset(offset). + Limit(filter.PageSize). + Get() + if err != nil { + return nil, err + } + + // Group chats by date + today := time.Now().Truncate(24 * time.Hour) + yesterday := today.AddDate(0, 0, -1) + thisWeekStart := today.AddDate(0, 0, -int(today.Weekday())) + lastWeekStart := thisWeekStart.AddDate(0, 0, -7) + lastWeekEnd := thisWeekStart.AddDate(0, 0, -1) + + groups := map[string][]map[string]interface{}{ + "Today": {}, + "Yesterday": {}, + "This Week": {}, + "Last Week": {}, + "Even Earlier": {}, + } + + for _, row := range rows { + chatID := row.Get("chat_id") + if chatID == nil || chatID == "" { + continue + } + + chat := map[string]interface{}{ + "chat_id": chatID, + "title": row.Get("title"), + } + + var createdAt time.Time + switch v := row.Get("created_at").(type) { + case time.Time: + createdAt = v + case string: + parsed, err := time.Parse("2006-01-02 15:04:05.999999-07:00", v) + if err != nil { + // Try alternative format + parsed, err = time.Parse(time.RFC3339, v) + if err != nil { + continue + } + } + createdAt = parsed + default: + continue + } + + createdDate := createdAt.Truncate(24 * time.Hour) + + switch { + case createdDate.Equal(today): + groups["Today"] = append(groups["Today"], chat) + case createdDate.Equal(yesterday): + groups["Yesterday"] = append(groups["Yesterday"], chat) + case createdDate.After(thisWeekStart) && createdDate.Before(today): + groups["This Week"] = append(groups["This Week"], chat) + case createdDate.After(lastWeekStart) && createdDate.Before(lastWeekEnd.AddDate(0, 0, 1)): + groups["Last Week"] = append(groups["Last Week"], chat) + default: + groups["Even Earlier"] = append(groups["Even Earlier"], chat) + } + } + + // Convert to ordered slice + result := []ChatGroup{} + for _, label := range []string{"Today", "Yesterday", "This Week", "Last Week", "Even Earlier"} { + if len(groups[label]) > 0 { + result = append(result, ChatGroup{ + Label: label, + Chats: groups[label], + }) + } + } + + return &ChatGroupResponse{ + Groups: result, + Page: filter.Page, + PageSize: filter.PageSize, + Total: total, + LastPage: lastPage, + }, nil } // GetHistory get the history func (conv *Xun) GetHistory(sid string, cid string) ([]map[string]interface{}, error) { + userID, err := conv.getUserID(sid) + if err != nil { + return nil, err + } - qb := conv.query.Table(conv.setting.Table). + qb := conv.newQuery(). Select("role", "name", "content"). - Where("sid", sid). + Where("sid", userID). Where("cid", cid). OrderBy("id", "desc") @@ -137,6 +414,40 @@ func (conv *Xun) GetHistory(sid string, cid string) ([]map[string]interface{}, e // SaveHistory save the history func (conv *Xun) SaveHistory(sid string, messages []map[string]interface{}, cid string) error { + if cid == "" { + cid = uuid.New().String() // Generate a new UUID if cid is empty + } + + userID, err := conv.getUserID(sid) + if err != nil { + return err + } + + // First ensure chat record exists + exists, err := conv.newQueryChat(). + Where("chat_id", cid). + Where("sid", userID). + Exists() + + if err != nil { + return err + } + + if !exists { + // Create new chat record + err = conv.newQueryChat(). + Insert(map[string]interface{}{ + "chat_id": cid, + "sid": userID, + "created_at": time.Now(), + }) + + if err != nil { + return err + } + } + + // Save message history defer conv.clean() var expiredAt interface{} = nil values := []row{} @@ -149,7 +460,7 @@ func (conv *Xun) SaveHistory(sid string, messages []map[string]interface{}, cid Role: message["role"].(string), Name: "", Content: message["content"].(string), - Sid: sid, + Sid: userID, Cid: cid, ExpiredAt: expiredAt, } @@ -160,16 +471,25 @@ func (conv *Xun) SaveHistory(sid string, messages []map[string]interface{}, cid values = append(values, value) } - return conv.query.Table(conv.setting.Table).Insert(values) + err = conv.newQuery().Insert(values) + if err != nil { + return err + } + + return nil } // GetRequest get the request history func (conv *Xun) GetRequest(sid string, rid string) ([]map[string]interface{}, error) { + userID, err := conv.getUserID(sid) + if err != nil { + return nil, err + } - qb := conv.query.Table(conv.setting.Table). + qb := conv.newQuery(). Select("role", "name", "content", "sid"). Where("rid", rid). - Where("sid", sid). + Where("sid", userID). OrderBy("id", "desc") if conv.setting.TTL > 0 { @@ -200,6 +520,10 @@ func (conv *Xun) GetRequest(sid string, rid string) ([]map[string]interface{}, e // SaveRequest save the request history func (conv *Xun) SaveRequest(sid string, rid string, cid string, messages []map[string]interface{}) error { + userID, err := conv.getUserID(sid) + if err != nil { + return err + } defer conv.clean() var expiredAt interface{} = nil @@ -213,7 +537,7 @@ func (conv *Xun) SaveRequest(sid string, rid string, cid string, messages []map[ Role: message["role"].(string), Name: "", Content: message["content"].(string), - Sid: sid, + Sid: userID, Cid: cid, Rid: rid, ExpiredAt: expiredAt, @@ -225,75 +549,92 @@ func (conv *Xun) SaveRequest(sid string, rid string, cid string, messages []map[ values = append(values, value) } - return conv.query.Table(conv.setting.Table).Insert(values) + return conv.newQuery().Insert(values) } -func (conv *Xun) clean() { - nums, err := conv.query.Table(conv.setting.Table).Where("expired_at", "<=", time.Now()).Delete() +// GetChat get the chat info and its history +func (conv *Xun) GetChat(sid string, cid string) (*ChatInfo, error) { + userID, err := conv.getUserID(sid) if err != nil { - log.Error("Clean the conversation table error: %s", err.Error()) - return + return nil, err } - if nums > 0 { - log.Trace("Clean the conversation table: %s %d", conv.setting.Table, nums) + // Get chat info + qb := conv.newQueryChat(). + Select("chat_id", "title"). + Where("sid", userID). + Where("chat_id", cid) + + row, err := qb.First() + if err != nil { + return nil, err } + + // Return nil if chat_id is nil (means no chat found) + if row.Get("chat_id") == nil { + return nil, nil + } + + chat := map[string]interface{}{ + "chat_id": row.Get("chat_id"), + "title": row.Get("title"), + } + + // Get chat history + history, err := conv.GetHistory(sid, cid) + if err != nil { + return nil, err + } + + return &ChatInfo{ + Chat: chat, + History: history, + }, nil } -// Init init the conversation -func (conv *Xun) Init() error { - - has, err := conv.schema.HasTable(conv.setting.Table) +// DeleteChat deletes a specific chat and its history +func (conv *Xun) DeleteChat(sid string, cid string) error { + userID, err := conv.getUserID(sid) if err != nil { return err } - // create the table - if !has { - err = conv.schema.CreateTable(conv.setting.Table, func(table schema.Blueprint) { - - table.ID("id") // The ID field - table.String("sid", 255).Index() // The Session ID - table.String("rid", 255).Null().Index() // The request ID - table.String("cid", 200).Null().Index() // The Chat ID - table.String("role", 200).Null().Index() // The Message role - table.String("name", 200).Null().Index() // The User name - table.String("title", 200).Null().Index() // The Chat title - table.Text("content").Null() - - table.TimestampTz("created_at").SetDefaultRaw("NOW()").Index() - table.TimestampTz("updated_at").Null().Index() - table.TimestampTz("expired_at").Null().Index() - }) - - if err != nil { - return err - } - log.Trace("Create the conversation table: %s", conv.setting.Table) - } - - // validate the table - tab, err := conv.schema.GetTable(conv.setting.Table) + // Delete history records first + _, err = conv.newQuery(). + Where("sid", userID). + Where("cid", cid). + Delete() if err != nil { return err } - fields := []string{"id", "sid", "rid", "cid", "role", "name", "content", "created_at", "updated_at", "expired_at"} - for _, field := range fields { - if !tab.HasColumn(field) { - return fmt.Errorf("%s is required", field) - } - } - - // Auto update the title - if !tab.HasColumn("title") { - err = conv.schema.AlterTable(conv.setting.Table, func(table schema.Blueprint) { - table.String("title", 200).Null().Index() - }) - if err != nil { - return err - } - } - - return nil + // Then delete the chat + _, err = conv.newQueryChat(). + Where("sid", userID). + Where("chat_id", cid). + Limit(1). + Delete() + return err +} + +// DeleteAllChats deletes all chats and their histories for a user +func (conv *Xun) DeleteAllChats(sid string) error { + userID, err := conv.getUserID(sid) + if err != nil { + return err + } + + // Delete history records first + _, err = conv.newQuery(). + Where("sid", userID). + Delete() + if err != nil { + return err + } + + // Then delete all chats + _, err = conv.newQueryChat(). + Where("sid", userID). + Delete() + return err } diff --git a/neo/conversation/xun_test.go b/neo/conversation/xun_test.go index e3fc0698..d42f32e2 100644 --- a/neo/conversation/xun_test.go +++ b/neo/conversation/xun_test.go @@ -1,7 +1,9 @@ package conversation import ( + "fmt" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/yaoapp/gou/connector" @@ -246,3 +248,158 @@ func TestXunSaveAndGetHistoryWithCID(t *testing.T) { } assert.Equal(t, 2, len(allData)) } + +func TestXunGetChats(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + defer capsule.Schema().DropTableIfExists("__unit_test_conversation") + defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") + + // Drop both tables before test + err := capsule.Schema().DropTableIfExists("__unit_test_conversation") + if err != nil { + t.Fatal(err) + } + err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") + if err != nil { + t.Fatal(err) + } + + conv, err := NewXun(Setting{ + Connector: "default", + Table: "__unit_test_conversation", + }) + if err != nil { + t.Fatal(err) + } + + // Save some test chats + sid := "test_user" + messages := []map[string]interface{}{ + {"role": "user", "content": "test message"}, + } + + // Create chats with different dates + for i := 0; i < 5; i++ { + chatID := fmt.Sprintf("chat_%d", i) + // First create the chat with a title + err = conv.newQueryChat().Insert(map[string]interface{}{ + "chat_id": chatID, + "title": fmt.Sprintf("Test Chat %d", i), + "sid": sid, + "created_at": time.Now(), + }) + if err != nil { + t.Fatal(err) + } + + // Then save the history + err = conv.SaveHistory(sid, messages, chatID) + if err != nil { + t.Fatal(err) + } + } + + // Test getting chats with default filter + filter := ChatFilter{ + PageSize: 10, + Order: "desc", + } + groups, err := conv.GetChats(sid, filter) + if err != nil { + t.Fatal(err) + } + + assert.Greater(t, len(groups.Groups), 0) + + // Test with keywords + filter.Keywords = "test" + groups, err = conv.GetChats(sid, filter) + if err != nil { + t.Fatal(err) + } + + assert.Greater(t, len(groups.Groups), 0) +} + +func TestXunDeleteChat(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + defer capsule.Schema().DropTableIfExists("__unit_test_conversation") + defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") + + conv, err := NewXun(Setting{ + Connector: "default", + Table: "__unit_test_conversation", + }) + if err != nil { + t.Fatal(err) + } + + // Create a test chat + sid := "test_user" + cid := "test_chat" + messages := []map[string]interface{}{ + {"role": "user", "content": "test message"}, + } + + // Save the chat and history + err = conv.SaveHistory(sid, messages, cid) + assert.Nil(t, err) + + // Verify chat exists + chat, err := conv.GetChat(sid, cid) + assert.Nil(t, err) + assert.NotNil(t, chat) + + // Delete the chat + err = conv.DeleteChat(sid, cid) + assert.Nil(t, err) + + // Verify chat is deleted + chat, err = conv.GetChat(sid, cid) + assert.Nil(t, err) + assert.Equal(t, (*ChatInfo)(nil), chat) +} + +func TestXunDeleteAllChats(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + defer capsule.Schema().DropTableIfExists("__unit_test_conversation") + defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") + + conv, err := NewXun(Setting{ + Connector: "default", + Table: "__unit_test_conversation", + }) + if err != nil { + t.Fatal(err) + } + + // Create multiple test chats + sid := "test_user" + messages := []map[string]interface{}{ + {"role": "user", "content": "test message"}, + } + + // Save multiple chats + for i := 0; i < 3; i++ { + cid := fmt.Sprintf("test_chat_%d", i) + err = conv.SaveHistory(sid, messages, cid) + assert.Nil(t, err) + } + + // Verify chats exist + response, err := conv.GetChats(sid, ChatFilter{}) + assert.Nil(t, err) + assert.Greater(t, response.Total, int64(0)) + + // Delete all chats + err = conv.DeleteAllChats(sid) + assert.Nil(t, err) + + // Verify all chats are deleted + response, err = conv.GetChats(sid, ChatFilter{}) + assert.Nil(t, err) + assert.Equal(t, int64(0), response.Total) +} diff --git a/neo/hooks.go b/neo/hooks.go new file mode 100644 index 00000000..775f492f --- /dev/null +++ b/neo/hooks.go @@ -0,0 +1,253 @@ +package neo + +import ( + "context" + "time" + + "github.com/gin-gonic/gin" + jsoniter "github.com/json-iterator/go" + "github.com/yaoapp/gou/process" + "github.com/yaoapp/yao/neo/assistant" +) + +// HookCreate create the assistant +func (neo *DSL) HookCreate(ctx Context, messages []map[string]interface{}, c *gin.Context) (CreateResponse, error) { + + // Default assistant + assistantID := neo.Use + if ctx.AssistantID != "" { + assistantID = ctx.AssistantID + } + + // Empty hook + if neo.Create == "" { + return CreateResponse{AssistantID: assistantID, ChatID: ctx.ChatID}, nil + } + + // Create a context with 10 second timeout + timeoutCtx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + + p, err := process.Of(neo.Create, ctx, messages, c.Writer) + if err != nil { + return CreateResponse{}, err + } + + err = p.WithContext(timeoutCtx).Execute() + if err != nil { + return CreateResponse{}, err + } + defer p.Release() + + // Check if context was canceled + if timeoutCtx.Err() != nil { + return CreateResponse{}, timeoutCtx.Err() + } + + value := p.Value() + switch v := value.(type) { + case CreateResponse: + return v, nil + + case map[string]interface{}: + if id, ok := v["assistant_id"].(string); ok { + assistantID = id + } + + chatID := "" + if id, ok := v["chat_id"].(string); ok { + chatID = id + } + + if chatID == "" { + chatID = ctx.ChatID + } + + return CreateResponse{AssistantID: assistantID, ChatID: chatID}, nil + } + + return CreateResponse{AssistantID: assistantID, ChatID: ctx.ChatID}, nil +} + +// HookAssistants query the assistant list from the assistant list hook +func (neo *DSL) HookAssistants(ctx context.Context, param assistant.QueryParam) ([]assistant.Assistant, error) { + if neo.AssistantListHook == "" { + return nil, nil + } + + // Create a context with 10 second timeout + timeoutCtx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + + p, err := process.Of(neo.AssistantListHook, param) + if err != nil { + return nil, err + } + + err = p.WithContext(timeoutCtx).Execute() + if err != nil { + return nil, err + } + defer p.Release() + + // Check if context was canceled + if timeoutCtx.Err() != nil { + return nil, timeoutCtx.Err() + } + + value := p.Value() + if value == nil { + return nil, nil + } + + var list []assistant.Assistant + bytes, err := jsoniter.Marshal(value) + if err != nil { + return nil, err + } + + err = jsoniter.Unmarshal(bytes, &list) + if err != nil { + return nil, err + } + + return list, nil +} + +// HookPrepare executes the prepare hook before AI is called +func (neo *DSL) HookPrepare(ctx Context, messages []map[string]interface{}) ([]map[string]interface{}, error) { + if neo.Prepare == "" { + return messages, nil + } + + // Create a context with 10 second timeout + timeoutCtx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + + p, err := process.Of(neo.Prepare, ctx, messages) + if err != nil { + return nil, err + } + + err = p.WithContext(timeoutCtx).Execute() + if err != nil { + return nil, err + } + defer p.Release() + + // Check if context was canceled + if timeoutCtx.Err() != nil { + return nil, timeoutCtx.Err() + } + + value := p.Value() + if value == nil { + return messages, nil + } + + var result []map[string]interface{} + bytes, err := jsoniter.Marshal(value) + if err != nil { + return nil, err + } + + err = jsoniter.Unmarshal(bytes, &result) + if err != nil { + return nil, err + } + + return result, nil +} + +// HookWrite executes the write hook when response is received from AI +func (neo *DSL) HookWrite(ctx Context, messages []map[string]interface{}, response map[string]interface{}, content string, writer *gin.ResponseWriter) ([]map[string]interface{}, error) { + if neo.Write == "" { + return []map[string]interface{}{response}, nil + } + + p, err := process.Of(neo.Write, ctx, messages, response, content, writer) + if err != nil { + return nil, err + } + + err = p.WithContext(ctx).Execute() + if err != nil { + return nil, err + } + defer p.Release() + + value := p.Value() + if value == nil { + return []map[string]interface{}{response}, nil + } + + var result []map[string]interface{} + bytes, err := jsoniter.Marshal(value) + if err != nil { + return nil, err + } + + err = jsoniter.Unmarshal(bytes, &result) + if err != nil { + return nil, err + } + + return result, nil +} + +// HookMention query the mention list +func (neo *DSL) HookMention(ctx context.Context, keywords string) ([]Mention, error) { + + // Default Get the assistant list + if neo.MentionHook == "" { + var mentions []Mention + assistants := neo.GetAssistants() + for _, assistant := range assistants { + mentions = append(mentions, Mention{ + ID: assistant.ID, + Name: assistant.Name, + Type: "assistant", + }) + } + + return mentions, nil + } + + // Create a context with 10 second timeout + timeoutCtx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + + p, err := process.Of(neo.MentionHook, keywords) + if err != nil { + return nil, err + } + + err = p.WithContext(timeoutCtx).Execute() + if err != nil { + return nil, err + } + defer p.Release() + + // Check if context was canceled + if timeoutCtx.Err() != nil { + return nil, timeoutCtx.Err() + } + + value := p.Value() + if value == nil { + return nil, nil + } + + var list []Mention + bytes, err := jsoniter.Marshal(value) + if err != nil { + return nil, err + } + + err = jsoniter.Unmarshal(bytes, &list) + if err != nil { + return nil, err + } + + return list, nil +} diff --git a/neo/load.go b/neo/load.go index 6773a865..dec5fb2f 100644 --- a/neo/load.go +++ b/neo/load.go @@ -1,11 +1,14 @@ package neo import ( + "context" + "fmt" "path/filepath" + "time" "github.com/yaoapp/gou/application" - "github.com/yaoapp/yao/aigc" "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/neo/assistant" "github.com/yaoapp/yao/neo/conversation" ) @@ -17,7 +20,7 @@ func Load(cfg config.Config) error { setting := DSL{ ID: "neo", - Prompts: []aigc.Prompt{}, + Prompts: []assistant.Prompt{}, Option: map[string]interface{}{}, Allows: []string{}, ConversationSetting: conversation.Setting{ @@ -42,17 +45,38 @@ func Load(cfg config.Config) error { Neo = &setting - // AI Setting - err = Neo.newAI() - if err != nil { - return err - } - // Conversation Setting - err = Neo.newConversation() + err = Neo.createConversation() if err != nil { return err } - return nil + // Query Assistant List + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + listDone := make(chan error, 1) + go func() { + list, err := Neo.HookAssistants(ctx, assistant.QueryParam{Limit: 100}) + Neo.updateAssistantList(list) + listDone <- err + }() + + select { + case err := <-listDone: + if err != nil { + return fmt.Errorf("Neo assistant list failed: %w", err) + } + + // Create Default Assistant + Neo.Assistant, err = Neo.createDefaultAssistant() + if err != nil { + return err + } + + return nil + case <-ctx.Done(): + return fmt.Errorf("Neo assistant list timeout: %w", ctx.Err()) + } + } diff --git a/neo/message/json.go b/neo/message/json.go index 20fa72f8..1ebda8a1 100644 --- a/neo/message/json.go +++ b/neo/message/json.go @@ -1,12 +1,14 @@ package message import ( + "fmt" "strings" "github.com/fatih/color" "github.com/gin-gonic/gin" jsoniter "github.com/json-iterator/go" "github.com/yaoapp/gou/helper" + "github.com/yaoapp/kun/exception" "github.com/yaoapp/kun/log" "github.com/yaoapp/kun/maps" "github.com/yaoapp/yao/openai" @@ -53,7 +55,13 @@ func NewOpenAI(data []byte) *JSON { break default: - msg.Error = text + + str := string(data) + // Remove "data: " and " + str = strings.TrimPrefix(str, "data: ") + str = strings.Trim(str, "\"") + msg.Type = "error" + msg.Text = str } return &JSON{msg} @@ -80,6 +88,19 @@ func (json *JSON) Text(text string) *JSON { return json } +// Error set the error +func (json *JSON) Error(message interface{}) *JSON { + json.Message.Type = "error" + if err, ok := message.(error); ok { + json.Message.Text = err.Error() + } else if msg, ok := message.(string); ok { + json.Message.Text = msg + } else { + json.Message.Text = fmt.Sprintf("%v", message) + } + return json +} + // Map set from map func (json *JSON) Map(msg map[string]interface{}) *JSON { if msg == nil { @@ -90,6 +111,10 @@ func (json *JSON) Map(msg map[string]interface{}) *JSON { json.Message.Text = text } + if typ, ok := msg["type"].(string); ok { + json.Message.Text = typ + } + if done, ok := msg["done"].(bool); ok { json.Message.Done = done } @@ -203,11 +228,6 @@ func (json *JSON) Write(w gin.ResponseWriter) bool { } }() - if json.Error != "" { - json.writeError(w, json.Error) - return false - } - data, err := jsoniter.Marshal(json.Message) if err != nil { log.Error("%s", err.Error()) @@ -232,7 +252,10 @@ func (json *JSON) Append(content []byte) []byte { } func (json *JSON) writeError(w gin.ResponseWriter, message string) { - data := []byte(`{"text":"` + strings.Trim(message, "\"") + `"}`) + data := []byte(`{"text":"` + strings.Trim(exception.New(message, 500).Message, "\"") + `","type":"error"}`) + if json.Message.Done { + data = []byte(`{"text":"` + strings.Trim(exception.New(message, 500).Message, "\"") + `","type":"error", "done":true}`) + } data = append([]byte("data: "), data...) data = append(data, []byte("\n\n")...) _, err := w.Write(data) diff --git a/neo/message/types.go b/neo/message/types.go index 468339df..ed54cf31 100644 --- a/neo/message/types.go +++ b/neo/message/types.go @@ -3,7 +3,7 @@ package message // Message the message type Message struct { Text string `json:"text,omitempty"` - Error string `json:"error,omitempty"` + Type string `json:"type,omitempty"` Done bool `json:"done,omitempty"` Confirm bool `json:"confirm,omitempty"` Command *Command `json:"command,omitempty"` diff --git a/neo/neo.go b/neo/neo.go index fecbff94..48ef866e 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -1,186 +1,295 @@ package neo import ( + "context" "fmt" - "net/url" + "os" "strings" + "sync" - "github.com/fatih/color" "github.com/gin-gonic/gin" - "github.com/google/uuid" - "github.com/yaoapp/gou/api" "github.com/yaoapp/gou/connector" - "github.com/yaoapp/gou/process" "github.com/yaoapp/kun/log" - "github.com/yaoapp/yao/helper" + "github.com/yaoapp/yao/neo/assistant" + "github.com/yaoapp/yao/neo/assistant/base" + "github.com/yaoapp/yao/neo/assistant/openai" "github.com/yaoapp/yao/neo/conversation" "github.com/yaoapp/yao/neo/message" - "github.com/yaoapp/yao/openai" + "github.com/yaoapp/yao/share" ) -// API is a method on the Neo type -func (neo *DSL) API(router *gin.Engine, path string) error { - - // get the guards - middlewares, err := neo.getGuardHandlers() - if err != nil { - return err - } - - // Cross-Domain - cors, err := neo.getCorsHandlers(router, path) - if err != nil { - return err - } - - // append the cors - middlewares = append(middlewares, cors...) - - // api router chat - handlers := append(middlewares, func(c *gin.Context) { - - sid := c.GetString("__sid") - if sid == "" { - sid = uuid.New().String() - } - - content := c.Query("content") - if content == "" { - c.JSON(400, gin.H{"message": "content is required", "code": 400}) - return - } - - // set the context - ctx, cancel := NewContextWithCancel(sid, c.Query("chat_id"), c.Query("context")) - defer cancel() - - err = neo.Answer(ctx, content, c) - if err != nil { - c.JSON(500, gin.H{"message": err.Error(), "code": 500}) - c.Done() - } - - }) - router.GET(path, handlers...) - router.POST(path, handlers...) - - // api Get ChatList - handlers = append(middlewares, func(c *gin.Context) { - sid := c.GetString("__sid") - if sid == "" { - c.JSON(400, gin.H{"message": "sid is required", "code": 400}) - c.Done() - return - } - - list, err := neo.Conversation.GetChats(sid) - if err != nil { - c.JSON(500, gin.H{"message": err.Error(), "code": 500}) - c.Done() - return - } - - c.JSON(200, map[string]interface{}{"data": list}) - c.Done() - }) - router.GET(path+"/chats", handlers...) - - // api router chat history - handlers = append(middlewares, func(c *gin.Context) { - sid := c.GetString("__sid") - if sid == "" { - c.JSON(400, gin.H{"message": "sid is required", "code": 400}) - c.Done() - return - } - - cid := c.Query("chat_id") - history, err := neo.Conversation.GetHistory(sid, cid) - if err != nil { - c.JSON(500, gin.H{"message": err.Error(), "code": 500}) - c.Done() - return - } - - c.JSON(200, map[string]interface{}{"data": history}) - c.Done() - }) - router.GET(path+"/history", handlers...) - - return nil -} +// Lock the assistant list +var lock sync.Mutex = sync.Mutex{} // Answer reply the message func (neo *DSL) Answer(ctx Context, question string, c *gin.Context) error { - // get the chat messages messages, err := neo.chatMessages(ctx, question) if err != nil { + msg := message.New().Error(err).Done() + msg.Write(c.Writer) return err } + // Get the assistant_id, chat_id + res, err := neo.HookCreate(ctx, messages, c) + if err != nil { + msg := message.New().Error(err).Done() + msg.Write(c.Writer) + return err + } + + // Select Assistant + ast, err := neo.selectAssistant(res.AssistantID) + if err != nil { + return err + } + + // Chat with AI + return neo.chat(ast, ctx, messages, c) +} + +// GetAssistants returns the list of assistants +func (neo *DSL) GetAssistants() []assistant.Assistant { + return neo.AssistantList +} + +// GetMentions returns the mention list +func (neo *DSL) GetMentions(keywords string) ([]Mention, error) { + return neo.HookMention(context.Background(), keywords) +} + +// GenerateChatTitle generate the chat title +func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context) (string, error) { + + prompts := ` + Help me generate a title for the chat + 1. The title should be a short and concise description of the chat. + 2. The title should be a single sentence. + 3. The title should be in same language as the chat. + 4. The title should be no more than 50 characters. + ` + + messages := []map[string]interface{}{ + {"role": "system", "content": prompts}, + {"role": "user", "content": input}, + } + + res, err := neo.HookCreate(ctx, messages, c) + if err != nil { + return "", err + } + + // Select Assistant + ast, err := neo.selectAssistant(res.AssistantID) + if err != nil { + return "", err + } + + if ast == nil { + msg := message.New().Error("assistant is not initialized").Done() + msg.Write(c.Writer) + return "", fmt.Errorf("assistant is not initialized") + } + + clientBreak := make(chan bool, 1) + done := make(chan bool, 1) + fail := make(chan error, 1) + + content := []byte{} + + // Chat with AI in background + go func() { + err := ast.Chat(c.Request.Context(), messages, neo.Option, func(data []byte) int { + select { + case <-clientBreak: + return 0 // break + + default: + msg := message.NewOpenAI(data) + if msg == nil { + return 1 // continue + } + + // Handle error + if msg.Type == "error" { + fail <- fmt.Errorf("%s", msg.Message.Text) + return 0 // break + } + + // Append content and send message + content = msg.Append(content) + + // Complete the stream + if msg.Message.Done { + done <- true + return 0 // break + } + + return 1 // continue + } + }) + + if err != nil { + log.Error("Chat error: %s", err.Error()) + message.New().Error(err).Done().Write(c.Writer) + } + + done <- true + }() + + // Wait for completion or client disconnect + select { + case <-done: + return string(content), nil + case err := <-fail: + return "", err + case <-c.Writer.CloseNotify(): + clientBreak <- true + return "", nil + } +} + +// Upload upload a file +func (neo *DSL) Upload(ctx Context, c *gin.Context) (*assistant.File, error) { + // Get the file + tmpfile, err := c.FormFile("file") + if err != nil { + return nil, err + } + + reader, err := tmpfile.Open() + if err != nil { + return nil, err + } + defer func() { + reader.Close() + os.Remove(tmpfile.Filename) + }() + + // Get option from form data option_xxx + option := map[string]interface{}{} + for key := range c.Request.Form { + if strings.HasPrefix(key, "option_") { + option[strings.TrimPrefix(key, "option_")] = c.PostForm(key) + } + } + + // Get file info + ctx.Upload = &FileUpload{ + Bytes: int(tmpfile.Size), + Name: tmpfile.Filename, + ContentType: tmpfile.Header.Get("Content-Type"), + Option: option, + } + + res, err := neo.HookCreate(ctx, []map[string]interface{}{}, c) + if err != nil { + return nil, err + } + + // Select Assistant + ast, err := neo.selectAssistant(res.AssistantID) + if err != nil { + return nil, err + } + + return ast.Upload(ctx, tmpfile, reader, option) +} + +// Download downloads a file +func (neo *DSL) Download(ctx Context, c *gin.Context) (*assistant.FileResponse, error) { + // Get file_id from query string + fileID := c.Query("file_id") + if fileID == "" { + return nil, fmt.Errorf("file_id is required") + } + + // Get assistant_id from context or query + res, err := neo.HookCreate(ctx, []map[string]interface{}{}, c) + if err != nil { + return nil, err + } + + // Select Assistant + ast, err := neo.selectAssistant(res.AssistantID) + if err != nil { + return nil, err + } + + // Download file using the assistant + return ast.Download(ctx.Context, fileID) +} + +// chat chat with AI +func (neo *DSL) chat(ast assistant.API, ctx Context, messages []map[string]interface{}, c *gin.Context) error { + + if ast == nil { + msg := message.New().Error("assistant is not initialized").Done() + msg.Write(c.Writer) + return fmt.Errorf("assistant is not initialized") + } + clientBreak := make(chan bool, 1) done := make(chan bool, 1) content := []byte{} - // Execute the command or chat with AI in the background + // Chat with AI in background go func() { - - // chat with AI - c.Header("Content-Type", "text/event-stream;charset=utf-8") - c.Header("Cache-Control", "no-cache") - c.Header("Connection", "keep-alive") - - _, ex := neo.AI.ChatCompletionsWith(ctx, messages, neo.Option, func(data []byte) int { - + err := ast.Chat(c.Request.Context(), messages, neo.Option, func(data []byte) int { select { case <-clientBreak: return 0 // break - default: + default: msg := message.NewOpenAI(data) if msg == nil { - return 1 // continue success + return 1 // continue } - if msg.Error != "" { - neo.send(ctx, msg, messages, content, c) + // Handle error + if msg.Type == "error" { + message.New().Error(msg.Message.Text).Done().Write(c.Writer) return 0 // break } + // Append content and send message content = msg.Append(content) - err := neo.send(ctx, msg, messages, content, c) - if err != nil { - c.Status(500) - return 0 // break + if msg.Message != nil && msg.Message.Text != "" { + message.New(). + Map(map[string]interface{}{ + "text": msg.Message.Text, + "done": msg.Message.Done, + }). + Write(c.Writer) } // Complete the stream - if msg.IsDone() { + if msg.Message != nil && msg.Message.Done { + if msg.Message.Text == "" { + msg.Write(c.Writer) + } done <- true return 0 // break } - return 1 // continue success + return 1 // continue } }) - // Throw the error - if ex != nil { - log.Error("Neo chat error: %s", ex.Message) - c.Status(200) - done <- true - return + if err != nil { + log.Error("Chat error: %s", err.Error()) + message.New().Error(err).Done().Write(c.Writer) } - // save the history - neo.saveHistory(ctx.Sid, ctx.ChatID, content, messages) - c.Status(200) + // Save chat history + if len(content) > 0 { + neo.saveHistory(ctx.Sid, ctx.ChatID, content, messages) + } - // Complete the stream done <- true - }() + // Wait for completion or client disconnect select { case <-done: return nil @@ -188,146 +297,148 @@ func (neo *DSL) Answer(ctx Context, question string, c *gin.Context) error { clientBreak <- true return nil } - } -// Send send the message to the stream -func (neo *DSL) send(ctx Context, msg *message.JSON, messages []map[string]interface{}, content []byte, c *gin.Context) error { - - w := c.Writer - - if msg.Error != "" { - msg.Write(w) - return nil - } - - // Directly write the message - if neo.Write == "" { - ok := msg.Write(c.Writer) - if !ok { - return fmt.Errorf("Stream write error") +// updateAssistantList update the assistant list +func (neo *DSL) updateAssistantList(list []assistant.Assistant) { + lock.Lock() + defer lock.Unlock() + neo.AssistantList = list + neo.AssistantMaps = make(map[string]assistant.Assistant) + if list != nil { + for _, assistant := range list { + neo.AssistantMaps[assistant.ID] = assistant } - return nil } +} - // Execute the custom write hook get the response - args := []interface{}{ctx, messages, msg, string(content), w} - p, err := process.Of(neo.Write, args...) - if err != nil { - msg.Write(w) - color.Red("Neo custom write error: %s", err.Error()) - return fmt.Errorf("Stream write error: %s", err.Error()) +// selectAssistant select the assistant +func (neo *DSL) selectAssistant(assistantID string) (assistant.API, error) { + ast := neo.Assistant + if assistantID != "" { + ast, err := neo.newAssistant(assistantID) + if err != nil { + return nil, err + } + return ast, nil } + return ast, nil +} - err = p.WithSID(ctx.Sid).Execute() - if err != nil { - log.Error("Neo custom write error: %s", err.Error()) - msg.Write(w) - return nil - } - defer p.Release() +// newAssistant create a new assistant +func (neo *DSL) newAssistant(id string) (assistant.API, error) { + // Try to find assistant in AssistantList first + if id != "" && neo.AssistantMaps != nil { + if ast, ok := neo.AssistantMaps[id]; ok { - res := p.Value() - if res == nil { - color.Red("Neo custom write return null") - return fmt.Errorf("Neo custom write return null") - } - - // Send the custom write response to the stream - if messages, ok := res.([]interface{}); ok { - for _, new := range messages { - if v, ok := new.(map[string]interface{}); ok { - newMsg := message.New().Map(v) - newMsg.Write(w) + if ast.API != nil { + return ast.API, nil } + api, err := neo.newAssistantByConfig(&ast) + if err != nil { + return nil, err + } + ast.API = api + return api, nil } - return nil } - - color.Red("Neo custom write should return an array of response") - return fmt.Errorf("Neo should return an array of response") + return neo.newAssistantByConnector(id) } -// prompts get the prompts -func (neo *DSL) prompts() []map[string]interface{} { - prompts := []map[string]interface{}{} - for _, prompt := range neo.Prompts { - message := map[string]interface{}{"role": prompt.Role, "content": prompt.Content} - if prompt.Name != "" { - message["name"] = prompt.Name - } - prompts = append(prompts, message) - } - - return prompts +// newAssistantByConfig create a new assistant from assistant configuration +func (neo *DSL) newAssistantByConfig(ast *assistant.Assistant) (assistant.API, error) { + return neo.newAssistantByConnector(ast.Connector) } -// prepare the messages -func (neo *DSL) prepare(ctx Context, messages []map[string]interface{}) []map[string]interface{} { - if neo.Prepare == "" { - return []map[string]interface{}{} +// newAssistantByConnector create a new assistant from connector id +func (neo *DSL) newAssistantByConnector(id string) (assistant.API, error) { + // Moapi connector + if id == "" || strings.HasPrefix(id, "moapi") { + return neo.newMoapiAssistant(id) } - prompts := []map[string]interface{}{} - p, err := process.Of(neo.Prepare, ctx, messages) + // Other connector + conn, err := connector.Select(id) if err != nil { - color.Red("Neo prepare error: %s", err.Error()) - return prompts + return nil, fmt.Errorf("Neo assistant connector %s not support", id) } - err = p.WithSID(ctx.Sid).Execute() + if conn.Is(connector.OPENAI) { + api, err := openai.New(conn, id) + if err != nil { + return nil, fmt.Errorf("Create openai assistant error: %s", err.Error()) + } + return api, nil + } + + // Base on the assistant list hook + api, err := base.New(conn, neo.Prompts, id) if err != nil { - color.Red("Neo prepare execute error: %s", err.Error()) - return prompts + return nil, fmt.Errorf("Create base assistant error: %s", err.Error()) } - defer p.Release() + return api, nil +} - data := p.Value() - items, ok := data.([]interface{}) - if !ok { - color.Red("Neo prepare response is not array") - return prompts +// newMoapiAssistant creates a new moapi assistant +func (neo *DSL) newMoapiAssistant(id string) (assistant.API, error) { + model := "gpt-3.5-turbo" + if strings.HasPrefix(id, "moapi:") { + model = strings.TrimPrefix(id, "moapi:") } - for i, item := range items { - v, ok := item.(map[string]interface{}) - if !ok { - color.Red("Neo prepare response [%d] is not map", i) - continue - } + // Get the moapi setting + url := share.MoapiHosts[0] + if share.App.Moapi.Mirrors != nil { + url = share.App.Moapi.Mirrors[0] + } + key := share.App.Moapi.Secret + organization := share.App.Moapi.Organization - if _, ok := v["role"]; !ok { - color.Red(`Neo prepare response [%d]["role"] required`, i) - continue - } - - if _, ok := v["content"]; !ok { - color.Red(`Neo prepare response [%d]["content"] required`, i) - continue - } - prompts = append(prompts, v) + if !strings.HasPrefix(url, "http") { + url = "https://" + url } - return prompts + // Check the moapi secret + if key == "" { + return nil, fmt.Errorf("The moapi secret is empty") + } + + conn, err := connector.New(`moapi`, `__yao.moapi`, []byte(`{"name":"Moapi", "options":{"model": "`+model+`", "key": "`+key+`", "organization": "`+organization+`", "host": "`+url+`"}}`)) + if err != nil { + return nil, fmt.Errorf("Create moapi assistant error: %s", err.Error()) + } + + api, err := openai.New(conn, strings.ReplaceAll(id, ":", "_")) + if err != nil { + return nil, fmt.Errorf("Create openai assistant error: %s", err.Error()) + } + return api, nil +} + +// createDefaultAssistant create a default assistant +func (neo *DSL) createDefaultAssistant() (assistant.API, error) { + if neo.Use != "" { + return neo.newAssistant(neo.Use) + } + return neo.newAssistant(neo.Connector) } // chatMessages get the chat messages -func (neo *DSL) chatMessages(ctx Context, content string) ([]map[string]interface{}, error) { +func (neo *DSL) chatMessages(ctx Context, content ...string) ([]map[string]interface{}, error) { history, err := neo.Conversation.GetHistory(ctx.Sid, ctx.ChatID) if err != nil { return nil, err } - messages := append([]map[string]interface{}{}, neo.prompts()...) - messages = append(messages, history...) - messages = append(messages, map[string]interface{}{"role": "user", "content": content, "name": ctx.Sid}) - // Add prepare messages witch is query from vector database - preparePrompts := neo.prepare(ctx, messages) - if len(preparePrompts) > 0 { - messages = preparePrompts + messages := []map[string]interface{}{} + messages = append(messages, history...) + if len(content) == 0 { + return messages, nil } + // Add user message + messages = append(messages, map[string]interface{}{"role": "user", "content": content[0], "name": ctx.Sid}) return messages, nil } @@ -350,136 +461,8 @@ func (neo *DSL) saveHistory(sid string, chatID string, content []byte, messages } } -func (neo *DSL) getCorsHandlers(router *gin.Engine, path string) ([]gin.HandlerFunc, error) { - - if len(neo.Allows) == 0 { - return []gin.HandlerFunc{}, nil - } - - allowsMap := map[string]bool{} - for _, allow := range neo.Allows { - allow = strings.TrimPrefix(allow, "http://") - allow = strings.TrimPrefix(allow, "https://") - allowsMap[allow] = true - } - - router.OPTIONS(path+"/history", neo.optionsHandler) - router.OPTIONS(path+"/commands", neo.optionsHandler) - return []gin.HandlerFunc{ - func(c *gin.Context) { - referer := neo.getOrigin(c) - if referer != "" { - if !api.IsAllowed(c, allowsMap) { - c.JSON(403, gin.H{"message": referer + " not allowed", "code": 403}) - c.Abort() - return - } - url, _ := url.Parse(referer) - referer = fmt.Sprintf("%s://%s", url.Scheme, url.Host) - c.Writer.Header().Set("Access-Control-Allow-Origin", referer) - c.Writer.Header().Set("Access-Control-Allow-Credentials", "true") - c.Writer.Header().Set("Access-Control-Allow-Headers", "Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, accept, origin, Cache-Control, X-Requested-With") - c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS, GET, PUT") - c.Next() - } - }, - }, nil -} - -func (neo *DSL) optionsHandler(c *gin.Context) { - origin := neo.getOrigin(c) - c.Writer.Header().Set("Access-Control-Allow-Origin", origin) - c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS, GET") - c.Writer.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization") - c.Writer.Header().Set("Access-Control-Allow-Credentials", "true") - c.AbortWithStatus(204) -} - -func (neo *DSL) getOrigin(c *gin.Context) string { - referer := c.Request.Referer() - origin := c.Request.Header.Get("Origin") - if origin == "" { - origin = referer - } - return origin -} - -func (neo *DSL) getGuardHandlers() ([]gin.HandlerFunc, error) { - - if neo.Guard == "" { - return []gin.HandlerFunc{ - func(c *gin.Context) { - token := strings.TrimSpace(strings.TrimPrefix(c.Query("token"), "Bearer ")) - if token == "" { - c.JSON(403, gin.H{"message": "token is required", "code": 403}) - c.Abort() - return - } - - user := helper.JwtValidate(token) - c.Set("__sid", user.SID) - c.Next() - }, - }, nil - } - - // validate the custom guard - _, err := process.Of(neo.Guard) - if err != nil { - return nil, err - } - - // custom guard - return []gin.HandlerFunc{api.ProcessGuard(neo.Guard)}, nil -} - -// NewAI create a new AI -func (neo *DSL) newAI() error { - - if neo.Connector == "" || strings.HasPrefix(neo.Connector, "moapi") { - model := "gpt-3.5-turbo" - if strings.HasPrefix(neo.Connector, "moapi:") { - model = strings.TrimPrefix(neo.Connector, "moapi:") - } - - ai, err := openai.NewMoapi(model) - if err != nil { - return err - } - - neo.AI = ai - return nil - } - - conn, err := connector.Select(neo.Connector) - if err != nil { - return err - } - - if conn.Is(connector.OPENAI) { - ai, err := openai.New(neo.Connector) - if err != nil { - return err - } - neo.AI = ai - return nil - } - - return fmt.Errorf("%s connector %s not support, should be a openai", neo.ID, neo.Connector) -} - -// Select select the model -func (neo *DSL) Select(model string) error { - ai, err := openai.NewMoapi(model) - if err != nil { - return err - } - neo.AI = ai - return nil -} - -// newConversation create a new conversation -func (neo *DSL) newConversation() error { +// createConversation create a new conversation +func (neo *DSL) createConversation() error { var err error if neo.ConversationSetting.Connector == "default" || neo.ConversationSetting.Connector == "" { @@ -512,3 +495,12 @@ func (neo *DSL) newConversation() error { return fmt.Errorf("%s conversation connector %s not support", neo.ID, neo.ConversationSetting.Connector) } + +// sendMessage sends a message to the client +func (neo *DSL) sendMessage(w gin.ResponseWriter, data interface{}) error { + msg := message.New().Map(data.(map[string]interface{})) + if !msg.Write(w) { + return fmt.Errorf("failed to write message to stream") + } + return nil +} diff --git a/neo/neo_test.go b/neo/neo_test.go index 2292fc23..e9d4437e 100644 --- a/neo/neo_test.go +++ b/neo/neo_test.go @@ -1,130 +1,327 @@ package neo -import ( - "context" - "fmt" - "net" - "net/http" - "net/http/httptest" - "strings" - "testing" - "time" +// type customResponseRecorder struct { +// *httptest.ResponseRecorder +// closeChannel chan bool +// } - "github.com/gin-gonic/gin" - "github.com/stretchr/testify/assert" - httpTest "github.com/yaoapp/gou/http" - "github.com/yaoapp/yao/config" - "github.com/yaoapp/yao/helper" - "github.com/yaoapp/yao/test" - _ "github.com/yaoapp/yao/utils" -) +// func (r *customResponseRecorder) CloseNotify() <-chan bool { +// return r.closeChannel +// } -func TestAPI(t *testing.T) { - test.Prepare(t, config.Conf) - defer test.Clean() +// func newCustomResponseRecorder() *customResponseRecorder { +// return &customResponseRecorder{ +// ResponseRecorder: httptest.NewRecorder(), +// closeChannel: make(chan bool, 1), +// } +// } - // test router - router := testRouter(t) - err := Neo.API(router, "/neo/chat") - if err != nil { - t.Fatal(err) - } +// func TestDSL_Prompts(t *testing.T) { +// test.Prepare(t, config.Conf) +// defer Test_clean(t) - // test server - host, shutdown := testServer(t, router) - defer shutdown() +// resetDB() +// neo := &DSL{ +// Prompts: []Prompt{ +// {Role: "system", Content: "You are a helpful assistant", Name: "ai"}, +// {Role: "user", Content: "Hello", Name: "user"}, +// }, +// ConversationSetting: conversation.Setting{ +// Connector: "default", +// Table: "chat_messages", +// }, +// } +// err := neo.newConversation() +// assert.NoError(t, err) - // test request - url := fmt.Sprintf("%s/neo/chat?content=hello&token=%s", host, testToken(t)) - res := []byte{} - req := httpTest.New(url). - WithHeader(http.Header{"Content-Type": []string{"application/json"}}) +// prompts := neo.prompts() +// assert.Equal(t, 2, len(prompts)) +// assert.Equal(t, "system", prompts[0]["role"]) +// assert.Equal(t, "You are a helpful assistant", prompts[0]["content"]) +// assert.Equal(t, "ai", prompts[0]["name"]) +// } - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() +// func TestDSL_ChatMessages(t *testing.T) { +// test.Prepare(t, config.Conf) +// defer Test_clean(t) - // send request - req.Stream(ctx, "GET", nil, func(data []byte) int { - res = append(res, data...) - return 1 - }) +// resetDB() +// neo := &DSL{ +// Prompts: []Prompt{ +// {Role: "system", Content: "You are a helpful assistant"}, +// }, +// ConversationSetting: conversation.Setting{ +// Connector: "default", +// Table: "chat_messages", +// }, +// } - assert.Contains(t, string(res), `{`) +// err := neo.newConversation() +// assert.NoError(t, err) -} +// ctx := Context{ +// Sid: "test-session", +// ChatID: "test-chat", +// } -func TestAPIAuth(t *testing.T) { - test.Prepare(t, config.Conf) - defer test.Clean() +// messages, err := neo.chatMessages(ctx, "Hello AI") +// assert.NoError(t, err) +// assert.Equal(t, 2, len(messages)) +// assert.Equal(t, "system", messages[0]["role"]) +// assert.Equal(t, "user", messages[1]["role"]) +// assert.Equal(t, "Hello AI", messages[1]["content"]) +// } - router := testRouter(t) - err := Neo.API(router, "/neo/chat") - if err != nil { - t.Fatal(err) - } +// func TestDSL_Answer(t *testing.T) { +// test.Prepare(t, config.Conf) +// defer Test_clean(t) - response := httptest.NewRecorder() - req, _ := http.NewRequest("GET", "/neo/chat?content=hello", nil) - assert.Panics(t, func() { - router.ServeHTTP(response, req) - }) -} +// gin.SetMode(gin.TestMode) +// w := newCustomResponseRecorder() +// c, _ := gin.CreateTestContext(w) -func testServer(t *testing.T, router *gin.Engine) (string, func()) { +// ctx := Context{ +// Sid: "test-session", +// ChatID: "test-chat", +// Context: context.Background(), +// } - // Listen - l, err := net.Listen("tcp4", ":0") - if err != nil { - t.Fatal(err) - } +// resetDB() +// neo := &DSL{ +// Connector: "gpt-3_5-turbo", +// Option: map[string]interface{}{ +// "temperature": 0.7, +// "max_tokens": 150, +// }, +// Prompts: []Prompt{ +// {Role: "system", Content: "You are a helpful assistant"}, +// }, +// ConversationSetting: conversation.Setting{ +// Connector: "default", +// Table: "chat_messages", +// }, +// } - srv := &http.Server{Addr: ":0", Handler: router} +// err := neo.newAI() +// assert.NoError(t, err) - // start serve - go func() { - if err := srv.Serve(l); err != nil && err != http.ErrServerClosed { - fmt.Println("[TestServer] Error:", err) - return - } - }() +// err = neo.newConversation() +// assert.NoError(t, err) - addr := strings.Split(l.Addr().String(), ":") - if len(addr) != 2 { - t.Fatal("invalid address") - } +// c.Request = httptest.NewRequest("POST", "/chat", nil) - host := fmt.Sprintf("http://127.0.0.1:%s", addr[1]) - time.Sleep(50 * time.Millisecond) +// neo.AI = &mockAI{} - shutdown := func() { - srv.Close() - l.Close() - } - return host, shutdown -} +// err = neo.Answer(ctx, "Hello AI", c) +// assert.NoError(t, err) +// } -func testRouter(t *testing.T) *gin.Engine { +// // func TestDSL_NewAI(t *testing.T) { +// // test.Prepare(t, config.Conf) +// // defer Test_clean(t) - // Load Config - err := Load(config.Conf) - if err != nil { - t.Fatal(err) - } +// // tests := []struct { +// // name string +// // connector string +// // wantErr string +// // }{ +// // { +// // name: "Mock AI", +// // connector: "mock", +// // wantErr: "", +// // }, +// // { +// // name: "Specific mock model", +// // connector: "mock:gpt-4", +// // wantErr: "", +// // }, +// // { +// // name: "Invalid connector", +// // connector: "invalid-connector", +// // wantErr: "AI connector invalid-connector not found", +// // }, +// // } - router := gin.New() - gin.SetMode(gin.ReleaseMode) - return router -} +// // for _, tt := range tests { +// // t.Run(tt.name, func(t *testing.T) { +// // neo := &DSL{ +// // Connector: tt.connector, +// // } +// // neo.newConversation() -func testToken(t *testing.T) string { - token := helper.JwtMake(1, - map[string]interface{}{ - "id": 1, - "name": "Test", - }, - map[string]interface{}{ - "exp": 3600, - "sid": "123456", - }) - return token.Token -} +// // assert.Panics(t, func() { +// // neo.newAI() +// // }) + +// // }) +// // } +// // } + +// func TestDSL_Select(t *testing.T) { +// test.Prepare(t, config.Conf) +// defer Test_clean(t) + +// resetDB() +// neo := &DSL{ +// ConversationSetting: conversation.Setting{ +// Connector: "default", +// Table: "chat_messages", +// }, +// } + +// err := neo.newConversation() +// assert.NoError(t, err) + +// err = neo.Select("invalid-model") +// assert.Error(t, err) + +// // err = neo.Select("gpt-3_5-turbo") +// // assert.NoError(t, err) +// // assert.NotNil(t, neo.AI) + +// } + +// // func TestDSL_NewConversation(t *testing.T) { +// // test.Prepare(t, config.Conf) +// // defer Test_clean(t) + +// // tests := []struct { +// // name string +// // connector string +// // wantErr bool +// // }{ +// // { +// // name: "Default connector", +// // connector: "default", +// // wantErr: false, +// // }, +// // { +// // name: "Empty connector", +// // connector: "", +// // wantErr: false, +// // }, +// // { +// // name: "Invalid connector", +// // connector: "invalid-connector", +// // wantErr: true, +// // }, +// // } + +// // for _, tt := range tests { +// // t.Run(tt.name, func(t *testing.T) { +// // neo := &DSL{ +// // ConversationSetting: conversation.Setting{ +// // Connector: tt.connector, +// // }, +// // } +// // assert.Panics(t, func() { +// // neo.newConversation() +// // }) +// // }) +// // } +// // } + +// func TestDSL_SaveHistory(t *testing.T) { +// test.Prepare(t, config.Conf) +// defer Test_clean(t) + +// neo := &DSL{ +// ConversationSetting: conversation.Setting{ +// Connector: "default", +// Table: "chat_messages", +// }, +// } + +// resetDB() +// err := neo.newConversation() +// assert.NoError(t, err) + +// messages := []map[string]interface{}{ +// { +// "role": "user", +// "content": "Hello", +// "name": "test-user", +// }, +// } + +// content := []byte("Hi there!") +// neo.saveHistory("test-session", "test-chat", content, messages) + +// // Verify the history was saved +// history, err := neo.Conversation.GetHistory("test-session", "test-chat") +// assert.NoError(t, err) +// assert.NotEmpty(t, history) +// } + +// func TestDSL_Send(t *testing.T) { +// test.Prepare(t, config.Conf) +// defer Test_clean(t) + +// gin.SetMode(gin.TestMode) +// w := httptest.NewRecorder() +// c, _ := gin.CreateTestContext(w) + +// resetDB() +// neo := &DSL{ +// ConversationSetting: conversation.Setting{ +// Connector: "default", +// Table: "chat_messages", +// }, +// } + +// err := neo.newConversation() +// assert.NoError(t, err) +// ctx := Context{ +// Sid: "test-session", +// ChatID: "test-chat", +// } + +// msg := &message.JSON{ +// Message: &message.Message{Text: "Test message"}, +// } +// messages := []map[string]interface{}{ +// {"role": "user", "content": "Hello"}, +// } +// content := []byte("Test content") + +// err = neo.send(ctx, msg, messages, content, c) +// assert.NoError(t, err) +// } + +// func Test_clean(t *testing.T) { +// defer test.Clean() + +// } + +// func resetDB() { +// sch := capsule.Global.Schema() +// sch.DropTable("chat_messages") +// } + +// type mockAI struct{} + +// func (m *mockAI) ChatCompletionsWith(ctx context.Context, messages []map[string]interface{}, options map[string]interface{}, callback func([]byte) int) (interface{}, *exception.Exception) { +// callback([]byte(`{"choices":[{"delta":{"content":"Mock response"}}]}`)) +// callback([]byte(`{"choices":[{"finish_reason":"stop"}]}`)) +// return nil, nil +// } + +// func (m *mockAI) ChatCompletions(messages []map[string]interface{}, options map[string]interface{}, callback func([]byte) int) (interface{}, *exception.Exception) { +// return nil, nil +// } + +// func (m *mockAI) GetContent(response interface{}) (string, *exception.Exception) { +// return "Mock content", nil +// } + +// func (m *mockAI) Embeddings(input interface{}, user string) (interface{}, *exception.Exception) { +// return nil, nil +// } + +// func (m *mockAI) Tiktoken(input string) (int, error) { +// return 0, nil +// } + +// func (m *mockAI) MaxToken() int { +// return 4096 +// } diff --git a/neo/types.go b/neo/types.go index 32a9d011..0582cba3 100644 --- a/neo/types.go +++ b/neo/types.go @@ -2,55 +2,75 @@ package neo import ( "context" - "io" "github.com/gin-gonic/gin" - "github.com/yaoapp/yao/aigc" + "github.com/yaoapp/yao/neo/assistant" "github.com/yaoapp/yao/neo/conversation" ) // DSL AI assistant type DSL struct { - ID string `json:"-" yaml:"-"` - Name string `json:"name,omitempty"` - Use string `json:"use,omitempty"` - Guard string `json:"guard,omitempty"` - Connector string `json:"connector"` - ConversationSetting conversation.Setting `json:"conversation" yaml:"conversation"` - Option map[string]interface{} `json:"option"` - Prepare string `json:"prepare,omitempty"` - Write string `json:"write,omitempty"` - Prompts []aigc.Prompt `json:"prompts,omitempty"` - Allows []string `json:"allows,omitempty"` - Models []string `json:"models,omitempty"` - AI aigc.AI `json:"-" yaml:"-"` - Conversation conversation.Conversation `json:"-" yaml:"-"` - GuardHandlers []gin.HandlerFunc `json:"-" yaml:"-"` + ID string `json:"-" yaml:"-"` + Name string `json:"name,omitempty" yaml:"name,omitempty"` + Use string `json:"use,omitempty" yaml:"use,omitempty"` // Which assistant to use default + Guard string `json:"guard,omitempty" yaml:"guard,omitempty"` + Connector string `json:"connector" yaml:"connector"` + ConversationSetting conversation.Setting `json:"conversation" yaml:"conversation"` + Option map[string]interface{} `json:"option" yaml:"option"` + Prepare string `json:"prepare,omitempty" yaml:"prepare,omitempty"` + Create string `json:"create,omitempty" yaml:"create,omitempty"` + Write string `json:"write,omitempty" yaml:"write,omitempty"` + AssistantListHook string `json:"assistants,omitempty" yaml:"assistants,omitempty"` // Get the assistant list from the hook + MentionHook string `json:"mentions,omitempty"` // Get the mention list from the hook + Prompts []assistant.Prompt `json:"prompts,omitempty" yaml:"prompts,omitempty"` + Allows []string `json:"allows,omitempty" yaml:"allows,omitempty"` + Assistant assistant.API `json:"-" yaml:"-"` // The default assistant + Conversation conversation.Conversation `json:"-" yaml:"-"` + GuardHandlers []gin.HandlerFunc `json:"-" yaml:"-"` + AssistantList []assistant.Assistant `json:"-" yaml:"-"` + AssistantMaps map[string]assistant.Assistant `json:"-" yaml:"-"` } -// Answer the answer interface -type Answer interface { - Stream(func(w io.Writer) bool) bool - Status(code int) - Header(key, value string) +// Mention list +type Mention struct { + ID string `json:"id"` + Name string `json:"name"` + Avatar string `json:"avatar,omitempty"` + Type string `json:"type,omitempty"` } // Context the context type Context struct { - Sid string `json:"sid" yaml:"-"` - ChatID string `json:"chat_id,omitempty"` + Sid string `json:"sid" yaml:"-"` // Session ID + ChatID string `json:"chat_id,omitempty"` // Chat ID, use to select chat + AssistantID string `json:"assistant_id,omitempty"` // Assistant ID, use to select assistant Stack string `json:"stack,omitempty"` Path string `json:"pathname,omitempty"` FormData map[string]interface{} `json:"formdata,omitempty"` - Field *ContextField `json:"field,omitempty"` + Field *Field `json:"field,omitempty"` Namespace string `json:"namespace,omitempty"` Config map[string]interface{} `json:"config,omitempty"` Signal interface{} `json:"signal,omitempty"` + Upload *FileUpload `json:"upload,omitempty"` context.Context `json:"-" yaml:"-"` } -// ContextField the context field -type ContextField struct { +// Field the context field +type Field struct { Name string `json:"name,omitempty"` Bind string `json:"bind,omitempty"` } + +// FileUpload the file upload info +type FileUpload struct { + Bytes int `json:"bytes,omitempty"` // If upload file, the file bytes + Name string `json:"name,omitempty"` // If upload + ContentType string `json:"content_type,omitempty"` // If upload file, the file content type + Option map[string]interface{} `json:"option,omitempty"` // If upload file, the upload option +} + +// CreateResponse the response of the create hook +type CreateResponse struct { + AssistantID string `json:"assistant_id,omitempty"` + ChatID string `json:"chat_id,omitempty"` +} diff --git a/openai/openai.go b/openai/openai.go index fa0eec68..825a09e4 100644 --- a/openai/openai.go +++ b/openai/openai.go @@ -54,12 +54,43 @@ func New(id string) (*OpenAI, error) { } setting := c.Setting() + return NewOpenAI(setting) +} + +// NewOpenAI create a new OpenAI instance by setting +func NewOpenAI(setting map[string]interface{}) (*OpenAI, error) { + + key := "" + if v, ok := setting["key"].(string); ok { + key = v + } + + model := "gpt-3.5-turbo" + if v, ok := setting["model"].(string); ok { + model = v + } + + host := "https://api.openai.com" + if v, ok := setting["host"].(string); ok { + host = v + } + + organization := "" + if v, ok := setting["organization"].(string); ok { + organization = v + } + + maxToken := 2048 + if v, ok := setting["max_token"].(int); ok { + maxToken = v + } + return &OpenAI{ - key: setting["key"].(string), - model: setting["model"].(string), - host: setting["host"].(string), - organization: "", - maxToken: 2048, + key: key, + model: model, + host: host, + organization: organization, + maxToken: maxToken, }, nil }