diff --git a/neo/api.go b/neo/api.go new file mode 100644 index 00000000..c9cd19b1 --- /dev/null +++ b/neo/api.go @@ -0,0 +1,198 @@ +package neo + +import ( + "fmt" + "net/url" + "strings" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + "github.com/yaoapp/gou/api" + "github.com/yaoapp/gou/process" + "github.com/yaoapp/yao/helper" +) + +// 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 + } + + // Cross-Domain handlers + cors, err := neo.getCorsHandlers(router, path) + if err != nil { + return err + } + + // Append cors handlers + middlewares = append(middlewares, cors...) + + // Register chat endpoint + router.GET(path, append(middlewares, neo.handleChat)...) + router.POST(path, append(middlewares, neo.handleChat)...) + + // Register chat list endpoint + router.GET(path+"/chats", append(middlewares, neo.handleChatList)...) + + // Register chat history endpoint + router.GET(path+"/history", append(middlewares, neo.handleChatHistory)...) + + return nil +} + +// handleChat handles the chat request +func (neo *DSL) handleChat(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() + } +} + +// 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 + } + + 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() +} + +// 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() +} + +// getCorsHandlers returns CORS middleware handlers +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{neo.corsMiddleware(allowsMap)}, nil +} + +// corsMiddleware handles CORS requests +func (neo *DSL) corsMiddleware(allowsMap map[string]bool) gin.HandlerFunc { + return 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() + } + } +} + +// optionsHandler handles OPTIONS requests +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) +} + +// getOrigin returns the request origin +func (neo *DSL) getOrigin(c *gin.Context) string { + referer := c.Request.Referer() + origin := c.Request.Header.Get("Origin") + if origin == "" { + origin = referer + } + return origin +} + +// getGuardHandlers returns authentication middleware handlers +func (neo *DSL) getGuardHandlers() ([]gin.HandlerFunc, error) { + if neo.Guard == "" { + return []gin.HandlerFunc{neo.defaultGuard}, nil + } + + // Validate the custom guard + _, err := process.Of(neo.Guard) + if err != nil { + return nil, err + } + + // Return custom guard + return []gin.HandlerFunc{api.ProcessGuard(neo.Guard)}, 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() +} 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/neo.go b/neo/neo.go index fecbff94..13007ba1 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -2,114 +2,18 @@ package neo import ( "fmt" - "net/url" "strings" "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/conversation" "github.com/yaoapp/yao/neo/message" "github.com/yaoapp/yao/openai" ) -// 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 -} - // Answer reply the message func (neo *DSL) Answer(ctx Context, question string, c *gin.Context) error { // get the chat messages @@ -350,89 +254,6 @@ 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 { diff --git a/neo/neo_test.go b/neo/neo_test.go index 2292fc23..91bacd33 100644 --- a/neo/neo_test.go +++ b/neo/neo_test.go @@ -2,129 +2,342 @@ package neo import ( "context" - "fmt" - "net" - "net/http" "net/http/httptest" - "strings" "testing" - "time" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" - httpTest "github.com/yaoapp/gou/http" + "github.com/yaoapp/kun/exception" + "github.com/yaoapp/xun/capsule" + "github.com/yaoapp/yao/aigc" "github.com/yaoapp/yao/config" - "github.com/yaoapp/yao/helper" + "github.com/yaoapp/yao/neo/conversation" + "github.com/yaoapp/yao/neo/message" "github.com/yaoapp/yao/test" - _ "github.com/yaoapp/yao/utils" ) -func TestAPI(t *testing.T) { +type customResponseRecorder struct { + *httptest.ResponseRecorder + closeChannel chan bool +} + +func (r *customResponseRecorder) CloseNotify() <-chan bool { + return r.closeChannel +} + +func newCustomResponseRecorder() *customResponseRecorder { + return &customResponseRecorder{ + ResponseRecorder: httptest.NewRecorder(), + closeChannel: make(chan bool, 1), + } +} + +func TestDSL_Prompts(t *testing.T) { test.Prepare(t, config.Conf) - defer test.Clean() + defer Test_clean(t) - // 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() - - // 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"}}) - - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - // send request - req.Stream(ctx, "GET", nil, func(data []byte) int { - res = append(res, data...) - return 1 - }) - - assert.Contains(t, string(res), `{`) - -} - -func TestAPIAuth(t *testing.T) { - test.Prepare(t, config.Conf) - defer test.Clean() - - router := testRouter(t) - err := Neo.API(router, "/neo/chat") - if err != nil { - t.Fatal(err) - } - - response := httptest.NewRecorder() - req, _ := http.NewRequest("GET", "/neo/chat?content=hello", nil) - assert.Panics(t, func() { - router.ServeHTTP(response, req) - }) -} - -func testServer(t *testing.T, router *gin.Engine) (string, func()) { - - // Listen - l, err := net.Listen("tcp4", ":0") - if err != nil { - t.Fatal(err) - } - - srv := &http.Server{Addr: ":0", Handler: router} - - // start serve - go func() { - if err := srv.Serve(l); err != nil && err != http.ErrServerClosed { - fmt.Println("[TestServer] Error:", err) - 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 { - - // Load Config - err := Load(config.Conf) - if err != nil { - t.Fatal(err) - } - - router := gin.New() - gin.SetMode(gin.ReleaseMode) - return router -} - -func testToken(t *testing.T) string { - token := helper.JwtMake(1, - map[string]interface{}{ - "id": 1, - "name": "Test", + resetDB() + neo := &DSL{ + Prompts: []aigc.Prompt{ + {Role: "system", Content: "You are a helpful assistant", Name: "ai"}, + {Role: "user", Content: "Hello", Name: "user"}, }, - map[string]interface{}{ - "exp": 3600, - "sid": "123456", - }) - return token.Token + ConversationSetting: conversation.Setting{ + Connector: "default", + Table: "chat_messages", + }, + } + err := neo.newConversation() + assert.NoError(t, err) + + 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"]) +} + +func TestDSL_ChatMessages(t *testing.T) { + test.Prepare(t, config.Conf) + defer Test_clean(t) + + resetDB() + neo := &DSL{ + Prompts: []aigc.Prompt{ + {Role: "system", Content: "You are a helpful assistant"}, + }, + ConversationSetting: conversation.Setting{ + Connector: "default", + Table: "chat_messages", + }, + } + + err := neo.newConversation() + assert.NoError(t, err) + + ctx := Context{ + Sid: "test-session", + ChatID: "test-chat", + } + + 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"]) +} + +func TestDSL_Answer(t *testing.T) { + test.Prepare(t, config.Conf) + defer Test_clean(t) + + gin.SetMode(gin.TestMode) + w := newCustomResponseRecorder() + c, _ := gin.CreateTestContext(w) + + ctx := Context{ + Sid: "test-session", + ChatID: "test-chat", + Context: context.Background(), + } + + resetDB() + neo := &DSL{ + Connector: "gpt-3_5-turbo", + Option: map[string]interface{}{ + "temperature": 0.7, + "max_tokens": 150, + }, + Prompts: []aigc.Prompt{ + {Role: "system", Content: "You are a helpful assistant"}, + }, + ConversationSetting: conversation.Setting{ + Connector: "default", + Table: "chat_messages", + }, + } + + err := neo.newAI() + assert.NoError(t, err) + + err = neo.newConversation() + assert.NoError(t, err) + + c.Request = httptest.NewRequest("POST", "/chat", nil) + + neo.AI = &mockAI{} + + err = neo.Answer(ctx, "Hello AI", c) + assert.NoError(t, err) +} + +// func TestDSL_NewAI(t *testing.T) { +// test.Prepare(t, config.Conf) +// defer Test_clean(t) + +// 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", +// }, +// } + +// for _, tt := range tests { +// t.Run(tt.name, func(t *testing.T) { +// neo := &DSL{ +// Connector: tt.connector, +// } +// neo.newConversation() + +// 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 }