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

@@ -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()
+}

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
+}

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 {
 

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
 }