From 1c078a5d2e3cd09b0e42d7ee15b31b7da36233d1 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 13 Dec 2024 11:17:00 +0800 Subject: [PATCH] Refactor Neo DSL tests and API structure by removing unused handlers, enhancing conversation management, and implementing new test cases for prompts, chat messages, and history saving. Introduce custom response recorder for improved testing and mock AI for simulating responses. --- neo/api.go | 198 ++++++++++++++++++++++ neo/api_test.go | 233 +++++++++++++++++++++++++ neo/neo.go | 179 -------------------- neo/neo_test.go | 439 +++++++++++++++++++++++++++++++++++------------- 4 files changed, 757 insertions(+), 292 deletions(-) create mode 100644 neo/api.go create mode 100644 neo/api_test.go 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 }