opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #812 from trheyi/main

Browse source 
Refactor Neo (60%)
This commit is contained in:
Max 2024-12-31 07:17:11 +08:00 • committed by GitHub
commit d55bfef374
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 2496 additions and 401 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

571
neo/api.go
View file

@ -12,6 +12,7 @@ import (
"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/yao/helper"
"github.com/yaoapp/yao/neo/conversation"
@ -36,32 +37,120 @@ func (neo *DSL) API(router *gin.Engine, path string) error {
router.OPTIONS(path+"/upload", neo.optionsHandler)
router.OPTIONS(path+"/download", neo.optionsHandler)
router.OPTIONS(path+"/mentions", neo.optionsHandler)
router.OPTIONS(path+"/generate", neo.optionsHandler)
router.OPTIONS(path+"/generate/title", neo.optionsHandler)
router.OPTIONS(path+"/generate/prompts", neo.optionsHandler)
router.OPTIONS(path+"/dangerous/clear_chats", neo.optionsHandler)
router.OPTIONS(path+"/assistants", neo.optionsHandler)
router.OPTIONS(path+"/assistants/:id", neo.optionsHandler)
// Register endpoints with middlewares
// Chat endpoint
// Example:
// curl -X GET 'http://localhost:5099/api/__yao/neo?content=Hello&chat_id=chat_123&context=previous_context&token=xxx'
// curl -X POST 'http://localhost:5099/api/__yao/neo' \
// -H 'Content-Type: application/json' \
// -d '{"content": "Hello", "chat_id": "chat_123", "context": "previous_context", "token": "xxx"}'
router.GET(path, append(middlewares, neo.handleChat)...)
router.POST(path, append(middlewares, neo.handleChat)...)
// Status check
// Status check endpoint
// Example:
// curl -X GET 'http://localhost:5099/api/__yao/neo/status?token=xxx'
router.GET(path+"/status", append(middlewares, neo.handleStatus)...)
// Chat api
// Assistant API endpoints
// List assistants example:
// curl -X GET 'http://localhost:5099/api/__yao/neo/assistants?page=1&pagesize=20&tags=tag1,tag2&token=xxx'
router.GET(path+"/assistants", append(middlewares, neo.handleAssistantList)...)
// Get assistant details example:
// curl -X GET 'http://localhost:5099/api/__yao/neo/assistants/assistant_123?token=xxx'
router.GET(path+"/assistants/:id", append(middlewares, neo.handleAssistantDetail)...)
// Create/Update assistant example:
// curl -X POST 'http://localhost:5099/api/__yao/neo/assistants' \
// -H 'Content-Type: application/json' \
// -d '{"name": "My Assistant", "type": "chat", "tags": ["tag1", "tag2"], "mentionable": true, "avatar": "path/to/avatar.png", "token": "xxx"}'
router.POST(path+"/assistants", append(middlewares, neo.handleAssistantSave)...)
// Delete assistant example:
// curl -X DELETE 'http://localhost:5099/api/__yao/neo/assistants/assistant_123?token=xxx'
router.DELETE(path+"/assistants/:id", append(middlewares, neo.handleAssistantDelete)...)
// Chat management endpoints
// List chats example:
// curl -X GET 'http://localhost:5099/api/__yao/neo/chats?page=1&pagesize=20&keywords=search+term&order=desc&token=xxx'
router.GET(path+"/chats", append(middlewares, neo.handleChatList)...)
// Get chat details example:
// curl -X GET 'http://localhost:5099/api/__yao/neo/chats/chat_123?token=xxx'
router.GET(path+"/chats/:id", append(middlewares, neo.handleChatDetail)...)
// Update chat example:
// curl -X POST 'http://localhost:5099/api/__yao/neo/chats/chat_123' \
// -H 'Content-Type: application/json' \
// -d '{"title": "New Title", "content": "Chat content for title generation", "token": "xxx"}'
router.POST(path+"/chats/:id", append(middlewares, neo.handleChatUpdate)...)
// Delete chat example:
// curl -X DELETE 'http://localhost:5099/api/__yao/neo/chats/chat_123?token=xxx'
router.DELETE(path+"/chats/:id", append(middlewares, neo.handleChatDelete)...)
// History api
// Chat history endpoint
// Example:
// curl -X GET 'http://localhost:5099/api/__yao/neo/history?chat_id=chat_123&token=xxx'
router.GET(path+"/history", append(middlewares, neo.handleChatHistory)...)
// File api
// File management endpoints
// Upload file example:
// curl -X POST 'http://localhost:5099/api/__yao/neo/upload?chat_id=chat_123&token=xxx' \
// -F 'file=@/path/to/file.txt'
router.POST(path+"/upload", append(middlewares, neo.handleUpload)...)
// Download file example:
// curl -X GET 'http://localhost:5099/api/__yao/neo/download?file_id=file_123&disposition=attachment&token=xxx' \
// -o downloaded_file.txt
router.GET(path+"/download", append(middlewares, neo.handleDownload)...)
// Mention api
// Mentions endpoint
// Example:
// curl -X GET 'http://localhost:5099/api/__yao/neo/mentions?keywords=assistant&token=xxx'
router.GET(path+"/mentions", append(middlewares, neo.handleMentions)...)
// Generation endpoints
// Generate custom content example:
// curl -X GET 'http://localhost:5099/api/__yao/neo/generate?content=Generate+something&type=custom&system_prompt=You+are+a+helpful+assistant&chat_id=chat_123&token=xxx'
// curl -X POST 'http://localhost:5099/api/__yao/neo/generate' \
// -H 'Content-Type: application/json' \
// -d '{"content": "Generate something", "type": "custom", "system_prompt": "You are a helpful assistant", "chat_id": "chat_123", "token": "xxx"}'
router.GET(path+"/generate", append(middlewares, neo.handleGenerateCustom)...)
router.POST(path+"/generate", append(middlewares, neo.handleGenerateCustom)...)
// Generate title example:
// curl -X GET 'http://localhost:5099/api/__yao/neo/generate/title?content=Chat+content&chat_id=chat_123&token=xxx'
// curl -X POST 'http://localhost:5099/api/__yao/neo/generate/title' \
// -H 'Content-Type: application/json' \
// -d '{"content": "Chat content", "chat_id": "chat_123", "token": "xxx"}'
router.GET(path+"/generate/title", append(middlewares, neo.handleGenerateTitle)...)
router.POST(path+"/generate/title", append(middlewares, neo.handleGenerateTitle)...)
// Generate prompts example:
// curl -X GET 'http://localhost:5099/api/__yao/neo/generate/prompts?content=Generate+prompts&chat_id=chat_123&token=xxx'
// curl -X POST 'http://localhost:5099/api/__yao/neo/generate/prompts' \
// -H 'Content-Type: application/json' \
// -d '{"content": "Generate prompts", "chat_id": "chat_123", "token": "xxx"}'
router.GET(path+"/generate/prompts", append(middlewares, neo.handleGeneratePrompts)...)
router.POST(path+"/generate/prompts", append(middlewares, neo.handleGeneratePrompts)...)
// Utility endpoints
// List connectors example:
// curl -X GET 'http://localhost:5099/api/__yao/neo/utility/connectors?token=xxx'
router.GET(path+"/utility/connectors", append(middlewares, neo.handleConnectors)...)
// Dangerous operations
// Dangerous operations
// Clear all chats example:
// curl -X DELETE 'http://localhost:5099/api/__yao/neo/dangerous/clear_chats?token=xxx'
router.DELETE(path+"/dangerous/clear_chats", append(middlewares, neo.handleChatsDeleteAll)...)
return nil
@ -383,49 +472,35 @@ func (neo *DSL) handleMentions(c *gin.Context) {
// Get keywords from query parameter
keywords := strings.ToLower(c.Query("keywords"))
mentions, err := neo.GetMentions(keywords)
mentionable := true
// Query mentionable assistants
filter := conversation.AssistantFilter{
Keywords: keywords,
Mentionable: &mentionable,
Page: 1,
PageSize: 20,
}
response, err := neo.Conversation.GetAssistants(filter)
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)
}
// Convert assistants to mentions
mentions := []Mention{}
for _, item := range response.Data {
mention := Mention{
ID: item["assistant_id"].(string),
Name: item["name"].(string),
Type: item["type"].(string),
Avatar: item["avatar"].(string),
}
testMentions = filtered
mentions = append(mentions, mention)
}
// Append test data to actual mentions
mentions = append(mentions, testMentions...)
c.JSON(200, map[string]interface{}{"data": mentions})
c.Done()
}
@ -462,7 +537,7 @@ func (neo *DSL) handleChatUpdate(c *gin.Context) {
ctx, cancel := NewContextWithCancel(sid, c.Query("chat_id"), "")
defer cancel()
title, err := neo.GenerateChatTitle(ctx, body.Content, c)
title, err := neo.GenerateChatTitle(ctx, body.Content, c, true)
if err != nil {
c.JSON(500, gin.H{"message": err.Error(), "code": 500})
c.Done()
@ -534,3 +609,417 @@ func (neo *DSL) handleChatsDeleteAll(c *gin.Context) {
c.JSON(200, gin.H{"message": "ok"})
c.Done()
}
// generateResponse is a helper struct to handle both SSE and HTTP responses
type generateResponse struct {
c *gin.Context
sid string
content string
result interface{}
err error
}
// validate checks common validation rules
func (r *generateResponse) validate() bool {
if r.sid == "" {
if strings.Contains(r.c.GetHeader("Accept"), "text/event-stream") {
r.c.Header("Content-Type", "text/event-stream;charset=utf-8")
r.c.Header("Cache-Control", "no-cache")
r.c.Header("Connection", "keep-alive")
msg := message.New().
Error("sid is required").
Done()
msg.Write(r.c.Writer)
} else {
r.c.JSON(400, gin.H{"message": "sid is required", "code": 400})
}
return false
}
if r.content == "" {
if strings.Contains(r.c.GetHeader("Accept"), "text/event-stream") {
r.c.Header("Content-Type", "text/event-stream;charset=utf-8")
r.c.Header("Cache-Control", "no-cache")
r.c.Header("Connection", "keep-alive")
msg := message.New().
Error("content is required").
Done()
msg.Write(r.c.Writer)
} else {
r.c.JSON(400, gin.H{"message": "content is required", "code": 400})
}
return false
}
return true
}
// send handles both SSE and HTTP responses
func (r *generateResponse) send(key string) {
if r.err != nil {
if strings.Contains(r.c.GetHeader("Accept"), "text/event-stream") {
r.c.Header("Content-Type", "text/event-stream;charset=utf-8")
r.c.Header("Cache-Control", "no-cache")
r.c.Header("Connection", "keep-alive")
msg := message.New().
Error(r.err.Error()).
Done()
msg.Write(r.c.Writer)
} else {
r.c.JSON(500, gin.H{"message": r.err.Error(), "code": 500})
}
return
}
if strings.Contains(r.c.GetHeader("Accept"), "text/event-stream") {
r.c.Header("Content-Type", "text/event-stream;charset=utf-8")
r.c.Header("Cache-Control", "no-cache")
r.c.Header("Connection", "keep-alive")
msg := message.New().
Map(gin.H{key: r.result}).
Done()
msg.Write(r.c.Writer)
} else {
r.c.JSON(200, gin.H{key: r.result})
}
}
// handleGenerateTitle handles generating a chat title
func (neo *DSL) handleGenerateTitle(c *gin.Context) {
var content string
if c.Request.Method == "GET" {
content = c.Query("content")
} else {
var body struct {
Content string `json:"content"`
}
if err := c.BindJSON(&body); err != nil {
// For SSE requests, send error message in SSE format
if strings.Contains(c.GetHeader("Accept"), "text/event-stream") {
c.Header("Content-Type", "text/event-stream;charset=utf-8")
c.Header("Cache-Control", "no-cache")
c.Header("Connection", "keep-alive")
msg := message.New().Error("invalid request body").Done()
msg.Write(c.Writer)
return
}
c.JSON(400, gin.H{"message": "invalid request body", "code": 400})
return
}
content = body.Content
}
resp := &generateResponse{
c: c,
sid: c.GetString("__sid"),
content: content,
}
// For SSE requests, set headers before validation
if strings.Contains(c.GetHeader("Accept"), "text/event-stream") {
c.Header("Content-Type", "text/event-stream;charset=utf-8")
c.Header("Cache-Control", "no-cache")
c.Header("Connection", "keep-alive")
}
if !resp.validate() {
return
}
ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "")
defer cancel()
// Use silent mode for regular HTTP requests, streaming for SSE
silent := !strings.Contains(c.GetHeader("Accept"), "text/event-stream")
resp.result, resp.err = neo.GenerateChatTitle(ctx, resp.content, c, silent)
resp.send("result")
}
// handleGeneratePrompts handles generating prompts
func (neo *DSL) handleGeneratePrompts(c *gin.Context) {
var content string
if c.Request.Method == "GET" {
content = c.Query("content")
} else {
var body struct {
Content string `json:"content"`
}
if err := c.BindJSON(&body); err != nil {
// For SSE requests, send error message in SSE format
if strings.Contains(c.GetHeader("Accept"), "text/event-stream") {
c.Header("Content-Type", "text/event-stream;charset=utf-8")
c.Header("Cache-Control", "no-cache")
c.Header("Connection", "keep-alive")
msg := message.New().Error("invalid request body").Done()
msg.Write(c.Writer)
return
}
c.JSON(400, gin.H{"message": "invalid request body", "code": 400})
return
}
content = body.Content
}
resp := &generateResponse{
c: c,
sid: c.GetString("__sid"),
content: content,
}
// For SSE requests, set headers before validation
if strings.Contains(c.GetHeader("Accept"), "text/event-stream") {
c.Header("Content-Type", "text/event-stream;charset=utf-8")
c.Header("Cache-Control", "no-cache")
c.Header("Connection", "keep-alive")
}
if !resp.validate() {
return
}
ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "")
defer cancel()
// Use silent mode for regular HTTP requests, streaming for SSE
silent := !strings.Contains(c.GetHeader("Accept"), "text/event-stream")
resp.result, resp.err = neo.GeneratePrompts(ctx, resp.content, c, silent)
resp.send("result")
}
// handleGenerateCustom handles generating custom content
func (neo *DSL) handleGenerateCustom(c *gin.Context) {
var content, genType, systemPrompt string
if c.Request.Method == "GET" {
content = c.Query("content")
genType = c.Query("type")
systemPrompt = c.Query("system_prompt")
} else {
var body struct {
Content string `json:"content"`
Type string `json:"type"`
SystemPrompt string `json:"system_prompt"`
}
if err := c.BindJSON(&body); err != nil {
c.JSON(400, gin.H{"message": "invalid request body", "code": 400})
return
}
content = body.Content
genType = body.Type
systemPrompt = body.SystemPrompt
}
resp := &generateResponse{
c: c,
sid: c.GetString("__sid"),
content: content,
}
if !resp.validate() {
return
}
// Additional validations for custom generation
if genType == "" {
c.JSON(400, gin.H{"message": "type is required", "code": 400})
return
}
if systemPrompt == "" {
c.JSON(400, gin.H{"message": "system_prompt is required", "code": 400})
return
}
ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "")
defer cancel()
// Use silent mode for regular HTTP requests, streaming for SSE
silent := !strings.Contains(c.GetHeader("Accept"), "text/event-stream")
resp.result, resp.err = neo.GenerateWithAI(ctx, resp.content, genType, systemPrompt, c, silent)
resp.send("result")
}
// handleAssistantList handles listing assistants
func (neo *DSL) handleAssistantList(c *gin.Context) {
// Parse filter parameters
filter := conversation.AssistantFilter{
Page: 1,
PageSize: 20,
}
// 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
}
}
// Parse tags
if tags := c.Query("tags"); tags != "" {
filter.Tags = strings.Split(tags, ",")
}
// Parse keywords
if keywords := c.Query("keywords"); keywords != "" {
filter.Keywords = keywords
}
// Parse connector
if connector := c.Query("connector"); connector != "" {
filter.Connector = connector
}
// Parse select fields
if selectFields := c.Query("select"); selectFields != "" {
filter.Select = strings.Split(selectFields, ",")
}
// Parse mentionable (support various boolean formats)
if mentionable := c.Query("mentionable"); mentionable != "" {
val := parseBoolValue(mentionable)
if val != nil {
filter.Mentionable = val
}
}
// Parse automated (support various boolean formats)
if automated := c.Query("automated"); automated != "" {
val := parseBoolValue(automated)
if val != nil {
filter.Automated = val
}
}
response, err := neo.Conversation.GetAssistants(filter)
if err != nil {
c.JSON(500, gin.H{"message": err.Error(), "code": 500})
c.Done()
return
}
c.JSON(200, response)
c.Done()
}
// parseBoolValue parses various string formats into a boolean pointer
// Supports: 1, 0, "1", "0", "true", "false", etc.
func parseBoolValue(value string) *bool {
value = strings.ToLower(strings.TrimSpace(value))
switch value {
case "1", "true", "yes", "on":
v := true
return &v
case "0", "false", "no", "off":
v := false
return &v
default:
return nil
}
}
// handleAssistantDetail handles getting a single assistant's details
func (neo *DSL) handleAssistantDetail(c *gin.Context) {
assistantID := c.Param("id")
if assistantID == "" {
c.JSON(400, gin.H{"message": "assistant id is required", "code": 400})
c.Done()
return
}
filter := conversation.AssistantFilter{
AssistantID: assistantID,
Page: 1,
PageSize: 1,
}
response, err := neo.Conversation.GetAssistants(filter)
if err != nil {
c.JSON(500, gin.H{"message": err.Error(), "code": 500})
c.Done()
return
}
if len(response.Data) == 0 {
c.JSON(404, gin.H{"message": "assistant not found", "code": 404})
c.Done()
return
}
c.JSON(200, map[string]interface{}{"data": response.Data[0]})
c.Done()
}
// handleAssistantSave handles creating or updating an assistant
func (neo *DSL) handleAssistantSave(c *gin.Context) {
var assistant map[string]interface{}
if err := c.BindJSON(&assistant); err != nil {
c.JSON(400, gin.H{"message": "invalid request body", "code": 400})
c.Done()
return
}
id, err := neo.Conversation.SaveAssistant(assistant)
if err != nil {
c.JSON(500, gin.H{"message": err.Error(), "code": 500})
c.Done()
return
}
// Update the assistant map with the returned ID if it's not already set
if _, ok := assistant["assistant_id"]; !ok {
assistant["assistant_id"] = id
}
c.JSON(200, gin.H{"message": "ok", "data": assistant})
c.Done()
}
// handleAssistantDelete handles deleting an assistant
func (neo *DSL) handleAssistantDelete(c *gin.Context) {
assistantID := c.Param("id")
if assistantID == "" {
c.JSON(400, gin.H{"message": "assistant id is required", "code": 400})
c.Done()
return
}
err := neo.Conversation.DeleteAssistant(assistantID)
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()
}
// handleConnectors handles listing connectors
func (neo *DSL) handleConnectors(c *gin.Context) {
options := []map[string]interface{}{}
// Filter and format connectors
for id, conn := range connector.Connectors {
if conn.Is(connector.OPENAI) || conn.Is(connector.MOAPI) {
setting := conn.Setting()
label := setting["label"]
if label == nil || label == "" {
label = setting["name"]
}
if label == nil || label == "" {
label = id
}
options = append(options, map[string]interface{}{
"label": label,
"value": id,
})
}
}
c.JSON(200, gin.H{"data": options})
c.Done()
}

4
neo/assistant/base/chat.go → neo/assistant/local/chat.go
View file

@ -1,4 +1,4 @@
package base
package local
import (
"context"
@ -6,7 +6,7 @@ import (
)
// Chat the chat
func (ast *Base) Chat(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) error {
func (ast *Local) 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")

10
neo/assistant/base/file.go → neo/assistant/local/file.go
View file

@ -1,4 +1,4 @@
package base
package local
import (
"context"
@ -31,7 +31,7 @@ var AllowedFileTypes = map[string]string{
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) {
func (ast *Local) Upload(ctx context.Context, file *multipart.FileHeader, reader io.Reader, option map[string]interface{}) (*assistant.File, error) {
// check file size
if file.Size > MaxSize {
@ -69,13 +69,13 @@ func (ast *Base) Upload(ctx context.Context, file *multipart.FileHeader, reader
}, nil
}
func (ast *Base) id(temp string, ext string) (string, error) {
func (ast *Local) 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 {
func (ast *Local) allowed(contentType string) bool {
if _, ok := AllowedFileTypes[contentType]; ok {
return true
}
@ -87,7 +87,7 @@ func (ast *Base) allowed(contentType string) bool {
}
// Download downloads a file
func (ast *Base) Download(ctx context.Context, fileID string) (*assistant.FileResponse, error) {
func (ast *Local) Download(ctx context.Context, fileID string) (*assistant.FileResponse, error) {
// Get the data filesystem
data, err := fs.Get("data")

14
neo/assistant/base/base.go → neo/assistant/local/local.go
View file

@ -1,4 +1,4 @@
package base
package local
import (
"context"
@ -8,16 +8,16 @@ import (
"github.com/yaoapp/yao/openai"
)
// Base the base assistant
type Base struct {
// Local the local assistant
type Local 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) {
// New create a new local assistant
func New(connector connector.Connector, prompts []assistant.Prompt, id string) (*Local, error) {
setting := connector.Setting()
api, err := openai.NewOpenAI(setting)
@ -25,10 +25,10 @@ func New(connector connector.Connector, prompts []assistant.Prompt, id string) (
return nil, err
}
return &Base{Connector: connector, ID: id, Prompts: prompts, openai: api}, nil
return &Local{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) {
func (ast *Local) List(ctx context.Context, param assistant.QueryParam) ([]assistant.Assistant, error) {
return nil, nil
}

17
neo/assistant/types.go
View file

@ -30,13 +30,16 @@ type QueryParam struct {
// 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
ID string `json:"assistant_id"` // Assistant ID
Type string `json:"type,omitempty"` // Assistant Type, default is assistant
Name string `json:"name,omitempty"` // Assistant Name
Avatar string `json:"avatar,omitempty"` // Assistant Avatar
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
Flows []map[string]interface{} `json:"flows,omitempty"` // Assistant Flows
API API `json:"-" yaml:"-"` // Assistant API
}
// File the file

67
neo/conversation/mongo.go
View file

@ -1,60 +1,59 @@
package conversation
// Mongo conversation
// Mongo represents a MongoDB-based conversation storage
type Mongo struct{}
// NewMongo create a new conversation
// NewMongo creates a new MongoDB conversation storage
func NewMongo() *Mongo {
return &Mongo{}
}
// UpdateChatTitle update the chat title
func (conv *Mongo) UpdateChatTitle(sid string, cid string, title string) error {
return nil
// GetChats retrieves a list of chats
func (m *Mongo) GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) {
return &ChatGroupResponse{}, nil
}
// GetChats get the chat list
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
// GetChat retrieves a single chat's information
func (m *Mongo) GetChat(sid string, cid string) (*ChatInfo, error) {
return &ChatInfo{}, nil
}
// GetHistory get the history
func (conv *Mongo) GetHistory(sid string, cid string) ([]map[string]interface{}, error) {
// GetHistory retrieves chat history
func (m *Mongo) GetHistory(sid string, cid string) ([]map[string]interface{}, error) {
return []map[string]interface{}{}, nil
}
// SaveHistory save the history
func (conv *Mongo) SaveHistory(sid string, messages []map[string]interface{}, cid string) error {
// SaveHistory saves chat history
func (m *Mongo) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error {
return nil
}
// GetRequest get the request
func (conv *Mongo) GetRequest(sid string, rid string) ([]map[string]interface{}, error) {
return nil, nil
}
// SaveRequest save the request
func (conv *Mongo) SaveRequest(sid string, rid string, cid string, messages []map[string]interface{}) error {
// DeleteChat deletes a single chat
func (m *Mongo) DeleteChat(sid string, cid string) 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 {
// DeleteAllChats deletes all chats
func (m *Mongo) DeleteAllChats(sid string) error {
return nil
}
// DeleteAllChats deletes all chats and their histories for a user
func (conv *Mongo) DeleteAllChats(sid string) error {
// UpdateChatTitle updates chat title
func (m *Mongo) UpdateChatTitle(sid string, cid string, title string) error {
return nil
}
// SaveAssistant saves assistant information
func (m *Mongo) SaveAssistant(assistant map[string]interface{}) (interface{}, error) {
return assistant["assistant_id"], nil
}
// DeleteAssistant deletes an assistant
func (m *Mongo) DeleteAssistant(assistantID string) error {
return nil
}
// GetAssistants retrieves a list of assistants
func (m *Mongo) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) {
return &AssistantResponse{}, nil
}

67
neo/conversation/redis.go
View file

@ -1,60 +1,59 @@
package conversation
// Redis conversation
// Redis represents a Redis-based conversation storage
type Redis struct{}
// NewRedis create a new conversation
// NewRedis creates a new Redis conversation storage
func NewRedis() *Redis {
return &Redis{}
}
// UpdateChatTitle update the chat title
func (conv *Redis) UpdateChatTitle(sid string, cid string, title string) error {
return nil
// GetChats retrieves a list of chats
func (r *Redis) GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) {
return &ChatGroupResponse{}, nil
}
// GetChats get the chat list
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
// GetChat retrieves a single chat's information
func (r *Redis) GetChat(sid string, cid string) (*ChatInfo, error) {
return &ChatInfo{}, nil
}
// GetHistory get the history
func (conv *Redis) GetHistory(sid string, cid string) ([]map[string]interface{}, error) {
// GetHistory retrieves chat history
func (r *Redis) GetHistory(sid string, cid string) ([]map[string]interface{}, error) {
return []map[string]interface{}{}, nil
}
// SaveHistory save the history
func (conv *Redis) SaveHistory(sid string, messages []map[string]interface{}, cid string) error {
// SaveHistory saves chat history
func (r *Redis) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error {
return nil
}
// GetRequest get the request
func (conv *Redis) GetRequest(sid string, rid string) ([]map[string]interface{}, error) {
return nil, nil
}
// SaveRequest save the request
func (conv *Redis) SaveRequest(sid string, rid string, cid string, messages []map[string]interface{}) error {
// DeleteChat deletes a single chat
func (r *Redis) DeleteChat(sid string, cid string) 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 {
// DeleteAllChats deletes all chats
func (r *Redis) DeleteAllChats(sid string) error {
return nil
}
// DeleteAllChats deletes all chats and their histories for a user
func (conv *Redis) DeleteAllChats(sid string) error {
// UpdateChatTitle updates chat title
func (r *Redis) UpdateChatTitle(sid string, cid string, title string) error {
return nil
}
// SaveAssistant saves assistant information
func (r *Redis) SaveAssistant(assistant map[string]interface{}) (interface{}, error) {
return assistant["assistant_id"], nil
}
// DeleteAssistant deletes an assistant
func (r *Redis) DeleteAssistant(assistantID string) error {
return nil
}
// GetAssistants retrieves a list of assistants
func (r *Redis) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) {
return &AssistantResponse{}, nil
}

149
neo/conversation/types.go
View file

@ -1,52 +1,133 @@
package conversation
// Setting the conversation config
// Setting represents the conversation configuration structure
// Used to configure basic conversation parameters including connector, user field, table name, etc.
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"`
Connector string `json:"connector,omitempty"` // Name of the connector used to specify data storage method
UserField string `json:"user_field,omitempty"` // User ID field name, defaults to "user_id"
Table string `json:"table,omitempty"` // Database table name
MaxSize int `json:"max_size,omitempty" yaml:"max_size,omitempty"` // Maximum storage size limit
TTL int `json:"ttl,omitempty" yaml:"ttl,omitempty"` // Time To Live in seconds
}
// ChatInfo represents the chat information and its history
// ChatInfo represents the chat information structure
// Contains basic information and history for a single chat
type ChatInfo struct {
Chat map[string]interface{} `json:"chat"`
History []map[string]interface{} `json:"history"`
Chat map[string]interface{} `json:"chat"` // Basic chat information
History []map[string]interface{} `json:"history"` // Chat history records
}
// ChatFilter represents the filter parameters for GetChats
// ChatFilter represents the chat filter structure
// Used for filtering and pagination when retrieving chat lists
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
Keywords string `json:"keywords,omitempty"` // Keyword search
Page int `json:"page,omitempty"` // Page number, starting from 1
PageSize int `json:"pagesize,omitempty"` // Number of items per page
Order string `json:"order,omitempty"` // Sort order: desc/asc
}
// ChatGroup represents a group of chats by date
// ChatGroup represents the chat group structure
// Groups chats by date
type ChatGroup struct {
Label string `json:"label"`
Chats []map[string]interface{} `json:"chats"`
Label string `json:"label"` // Group label (typically a date)
Chats []map[string]interface{} `json:"chats"` // List of chats in this group
}
// ChatGroupResponse represents paginated chat groups
// ChatGroupResponse represents the paginated chat group response
// Contains paginated chat group information
type ChatGroupResponse struct {
Groups []ChatGroup `json:"groups"`
Page int `json:"page"` // 当前页码
PageSize int `json:"pagesize"` // 每页数量
Total int64 `json:"total"` // 总记录数
LastPage int `json:"last_page"` // 最后一页页码
Groups []ChatGroup `json:"groups"` // List of chat groups
Page int `json:"page"` // Current page number
PageSize int `json:"pagesize"` // Items per page
Total int64 `json:"total"` // Total number of records
LastPage int `json:"last_page"` // Last page number
}
// Conversation the store interface
type Conversation interface {
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
// AssistantFilter represents the assistant filter structure
// Used for filtering and pagination when retrieving assistant lists
type AssistantFilter struct {
Tags []string `json:"tags,omitempty"` // Filter by tags
Keywords string `json:"keywords,omitempty"` // Search in name and description
Connector string `json:"connector,omitempty"` // Filter by connector
AssistantID string `json:"assistant_id,omitempty"` // Filter by assistant ID
Mentionable *bool `json:"mentionable,omitempty"` // Filter by mentionable status
Automated *bool `json:"automated,omitempty"` // Filter by automation status
Page int `json:"page,omitempty"` // Page number, starting from 1
PageSize int `json:"pagesize,omitempty"` // Items per page
Select []string `json:"select,omitempty"` // Fields to return, returns all fields if empty
}
// AssistantResponse represents the assistant response structure
// Used for returning paginated assistant lists
type AssistantResponse struct {
Data []map[string]interface{} `json:"data"` // The paginated data
Page int `json:"page"` // Current page number
PageSize int `json:"pagesize"` // Number of items per page
PageCnt int `json:"pagecnt"` // Total number of pages
Next int `json:"next"` // Next page number
Prev int `json:"prev"` // Previous page number
Total int64 `json:"total"` // Total number of items
}
// Conversation defines the conversation storage interface
// Provides basic operations required for conversation management
type Conversation interface {
// GetChats retrieves a list of chats
// sid: Session ID
// filter: Filter conditions
// Returns: Grouped chat list and potential error
GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error)
// GetChat retrieves a single chat's information
// sid: Session ID
// cid: Chat ID
// Returns: Chat information and potential error
GetChat(sid string, cid string) (*ChatInfo, error)
// GetHistory retrieves chat history
// sid: Session ID
// cid: Chat ID
// Returns: History record list and potential error
GetHistory(sid string, cid string) ([]map[string]interface{}, error)
// SaveHistory saves chat history
// sid: Session ID
// messages: Message list
// cid: Chat ID
// context: Context information
// Returns: Potential error
SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error
// DeleteChat deletes a single chat
// sid: Session ID
// cid: Chat ID
// Returns: Potential error
DeleteChat(sid string, cid string) error
// DeleteAllChats deletes all chats
// sid: Session ID
// Returns: Potential error
DeleteAllChats(sid string) error
// UpdateChatTitle updates chat title
// sid: Session ID
// cid: Chat ID
// title: New title
// Returns: Potential error
UpdateChatTitle(sid string, cid string, title string) error
// SaveAssistant saves assistant information
// assistant: Assistant information
// Returns: Potential error
SaveAssistant(assistant map[string]interface{}) (interface{}, error)
// DeleteAssistant deletes an assistant
// assistantID: Assistant ID
// Returns: Potential error
DeleteAssistant(assistantID string) error
// GetAssistants retrieves a list of assistants
// filter: Filter conditions
// Returns: Paginated assistant list and potential error
GetAssistants(filter AssistantFilter) (*AssistantResponse, error)
}

67
neo/conversation/weaviate.go
View file

@ -1,60 +1,59 @@
package conversation
// Weaviate Database conversation
// Weaviate represents a Weaviate-based conversation storage
type Weaviate struct{}
// NewWeaviate create a new conversation
// NewWeaviate creates a new Weaviate conversation storage
func NewWeaviate() *Weaviate {
return &Weaviate{}
}
// UpdateChatTitle update the chat title
func (conv *Weaviate) UpdateChatTitle(sid string, cid string, title string) error {
return nil
// GetChats retrieves a list of chats
func (w *Weaviate) GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) {
return &ChatGroupResponse{}, nil
}
// GetChats get the chat list
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
// GetChat retrieves a single chat's information
func (w *Weaviate) GetChat(sid string, cid string) (*ChatInfo, error) {
return &ChatInfo{}, nil
}
// GetHistory get the history
func (conv *Weaviate) GetHistory(sid string, cid string) ([]map[string]interface{}, error) {
// GetHistory retrieves chat history
func (w *Weaviate) GetHistory(sid string, cid string) ([]map[string]interface{}, error) {
return []map[string]interface{}{}, nil
}
// SaveHistory save the history
func (conv *Weaviate) SaveHistory(sid string, messages []map[string]interface{}, cid string) error {
// SaveHistory saves chat history
func (w *Weaviate) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error {
return nil
}
// GetRequest get the request
func (conv *Weaviate) GetRequest(sid string, rid string) ([]map[string]interface{}, error) {
return nil, nil
}
// SaveRequest save the request
func (conv *Weaviate) SaveRequest(sid string, rid string, cid string, messages []map[string]interface{}) error {
// DeleteChat deletes a single chat
func (w *Weaviate) DeleteChat(sid string, cid string) 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 {
// DeleteAllChats deletes all chats
func (w *Weaviate) DeleteAllChats(sid string) error {
return nil
}
// DeleteAllChats deletes all chats and their histories for a user
func (conv *Weaviate) DeleteAllChats(sid string) error {
// UpdateChatTitle updates chat title
func (w *Weaviate) UpdateChatTitle(sid string, cid string, title string) error {
return nil
}
// SaveAssistant saves assistant information
func (w *Weaviate) SaveAssistant(assistant map[string]interface{}) (interface{}, error) {
return assistant["assistant_id"], nil
}
// DeleteAssistant deletes an assistant
func (w *Weaviate) DeleteAssistant(assistantID string) error {
return nil
}
// GetAssistants retrieves a list of assistants
func (w *Weaviate) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) {
return &AssistantResponse{}, nil
}

511
neo/conversation/xun.go
View file

@ -7,6 +7,7 @@ import (
"time"
"github.com/google/uuid"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/session"
"github.com/yaoapp/kun/log"
@ -15,32 +16,34 @@ import (
"github.com/yaoapp/xun/dbal/schema"
)
// Xun Database conversation
// Package conversation provides functionality for managing chat conversations and assistants.
// Xun implements the Conversation interface using a database backend.
// It provides functionality for:
// - Managing chat conversations and their message histories
// - Organizing chats with pagination and date-based grouping
// - Handling chat metadata like titles and creation dates
// - Managing AI assistants with their configurations and metadata
// - Supporting data expiration through TTL settings
type Xun struct {
query query.Query
schema schema.Schema
setting Setting
}
type row struct {
Role string `json:"role"`
Name string `json:"name"` // User name
Content string `json:"content"`
Sid string `json:"sid"`
Rid string `json:"rid"`
Cid string `json:"cid"` // Chat ID from chat history
ExpiredAt interface{} `json:"expired_at"`
}
// Public interface methods and constructor remain exported:
// - NewXun
// - UpdateChatTitle
// - GetChats
// - GetChat
// - GetHistory
// - SaveHistory
// - GetRequest
// - SaveRequest
// Public interface methods:
//
// NewXun creates a new conversation instance with the given settings
// UpdateChatTitle updates the title of a specific chat
// GetChats retrieves a paginated list of chats grouped by date
// GetChat retrieves a specific chat and its message history
// GetHistory retrieves the message history for a specific chat
// SaveHistory saves new messages to a chat's history
// DeleteChat deletes a specific chat and its history
// DeleteAllChats deletes all chats and their histories for a user
// SaveAssistant creates or updates an assistant
// DeleteAssistant deletes an assistant by assistant_id
// GetAssistants retrieves a paginated list of assistants with filtering
// NewXun create a new conversation
func NewXun(setting Setting) (*Xun, error) {
@ -111,6 +114,11 @@ func (conv *Xun) initialize() error {
return err
}
// Initialize assistant table
if err := conv.initAssistantTable(); err != nil {
return err
}
return nil
}
@ -126,11 +134,12 @@ func (conv *Xun) initHistoryTable() error {
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("uid", 255).Null().Index()
table.String("role", 200).Null().Index()
table.String("name", 200).Null().Index()
table.Text("content").Null()
table.JSON("context").Null()
table.TimestampTz("created_at").SetDefaultRaw("NOW()").Index()
table.TimestampTz("updated_at").Null().Index()
table.TimestampTz("expired_at").Null().Index()
@ -148,7 +157,7 @@ func (conv *Xun) initHistoryTable() error {
return err
}
fields := []string{"id", "sid", "rid", "cid", "role", "name", "content", "created_at", "updated_at", "expired_at"}
fields := []string{"id", "sid", "cid", "uid", "role", "name", "content", "context", "created_at", "updated_at", "expired_at"}
for _, field := range fields {
if !tab.HasColumn(field) {
return fmt.Errorf("%s is required", field)
@ -198,6 +207,59 @@ func (conv *Xun) initChatTable() error {
return nil
}
func (conv *Xun) initAssistantTable() error {
assistantTable := conv.getAssistantTable()
has, err := conv.schema.HasTable(assistantTable)
if err != nil {
return err
}
// Create the assistant table
if !has {
err = conv.schema.CreateTable(assistantTable, func(table schema.Blueprint) {
table.ID("id")
table.String("assistant_id", 200).Unique().Index()
table.String("type", 200).SetDefault("assistant").Index() // default is assistant
table.String("name", 200).Null() // assistant name
table.String("avatar", 200).Null() // assistant avatar
table.String("connector", 200).NotNull() // assistant connector
table.Text("description").Null() // assistant description
table.JSON("options").Null() // assistant options
table.JSON("prompts").Null() // assistant prompts
table.JSON("flows").Null() // assistant flows
table.JSON("files").Null() // assistant files
table.JSON("functions").Null() // assistant functions
table.JSON("tags").Null() // assistant tags
table.Boolean("readonly").SetDefault(false).Index() // assistant readonly
table.JSON("permissions").Null() // assistant permissions
table.Boolean("automated").SetDefault(true).Index() // assistant autoable
table.Boolean("mentionable").SetDefault(true).Index() // Whether this assistant can appear in @ mention list
table.TimestampTz("created_at").SetDefaultRaw("NOW()").Index()
table.TimestampTz("updated_at").Null().Index()
})
if err != nil {
return err
}
log.Trace("Create the assistant table: %s", assistantTable)
}
// Validate the table
tab, err := conv.schema.GetTable(assistantTable)
if err != nil {
return err
}
fields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "options", "prompts", "flows", "files", "functions", "tags", "mentionable", "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 != "" {
@ -217,13 +279,17 @@ func (conv *Xun) getUserID(sid string) (string, error) {
}
func (conv *Xun) getHistoryTable() string {
return conv.setting.Table
return conv.setting.Table + "_history"
}
func (conv *Xun) getChatTable() string {
return conv.setting.Table + "_chat"
}
func (conv *Xun) getAssistantTable() string {
return conv.setting.Table + "_assistant"
}
// UpdateChatTitle update the chat title
func (conv *Xun) UpdateChatTitle(sid string, cid string, title string) error {
userID, err := conv.getUserID(sid)
@ -380,7 +446,7 @@ func (conv *Xun) GetHistory(sid string, cid string) ([]map[string]interface{}, e
}
qb := conv.newQuery().
Select("role", "name", "content").
Select("role", "name", "content", "context", "uid", "created_at", "updated_at").
Where("sid", userID).
Where("cid", cid).
OrderBy("id", "desc")
@ -401,18 +467,23 @@ func (conv *Xun) GetHistory(sid string, cid string) ([]map[string]interface{}, e
res := []map[string]interface{}{}
for _, row := range rows {
res = append([]map[string]interface{}{{
"role": row.Get("role"),
"name": row.Get("name"),
"content": row.Get("content"),
}}, res...)
message := map[string]interface{}{
"role": row.Get("role"),
"name": row.Get("name"),
"content": row.Get("content"),
"context": row.Get("context"),
"uid": row.Get("uid"),
"created_at": row.Get("created_at"),
"updated_at": row.Get("updated_at"),
}
res = append([]map[string]interface{}{message}, res...)
}
return res, nil
}
// SaveHistory save the history
func (conv *Xun) SaveHistory(sid string, messages []map[string]interface{}, cid string) error {
func (conv *Xun) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error {
if cid == "" {
cid = uuid.New().String() // Generate a new UUID if cid is empty
@ -450,24 +521,49 @@ func (conv *Xun) SaveHistory(sid string, messages []map[string]interface{}, cid
// Save message history
defer conv.clean()
var expiredAt interface{} = nil
values := []row{}
values := []map[string]interface{}{}
if conv.setting.TTL > 0 {
expiredAt = time.Now().Add(time.Duration(conv.setting.TTL) * time.Second)
}
now := time.Now()
for _, message := range messages {
value := row{
Role: message["role"].(string),
Name: "",
Content: message["content"].(string),
Sid: userID,
Cid: cid,
ExpiredAt: expiredAt,
// Type assertion safety checks
role, ok := message["role"].(string)
if !ok {
return fmt.Errorf("invalid role type in message: %v", message["role"])
}
if message["name"] != nil {
value.Name = message["name"].(string)
content, ok := message["content"].(string)
if !ok {
return fmt.Errorf("invalid content type in message: %v", message["content"])
}
var contextRaw interface{} = nil
if context != nil {
contextRaw, err = jsoniter.MarshalToString(context)
if err != nil {
return err
}
}
value := map[string]interface{}{
"role": role,
"name": "",
"content": content,
"sid": userID,
"cid": cid,
"uid": userID,
"context": contextRaw,
"created_at": now,
"updated_at": nil,
"expired_at": expiredAt,
}
if name, ok := message["name"].(string); ok {
value["name"] = name
}
values = append(values, value)
}
@ -479,79 +575,6 @@ func (conv *Xun) SaveHistory(sid string, messages []map[string]interface{}, cid
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.newQuery().
Select("role", "name", "content", "sid").
Where("rid", rid).
Where("sid", userID).
OrderBy("id", "desc")
if conv.setting.TTL > 0 {
qb.Where("expired_at", ">", time.Now())
}
limit := 20
if conv.setting.MaxSize > 0 {
limit = conv.setting.MaxSize
}
rows, err := qb.Limit(limit).Get()
if err != nil {
return nil, err
}
res := []map[string]interface{}{}
for _, row := range rows {
res = append([]map[string]interface{}{{
"role": row.Get("role"),
"name": row.Get("name"),
"content": row.Get("content"),
}}, res...)
}
return res, nil
}
// 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
values := []row{}
if conv.setting.TTL > 0 {
expiredAt = time.Now().Add(time.Duration(conv.setting.TTL) * time.Second)
}
for _, message := range messages {
value := row{
Role: message["role"].(string),
Name: "",
Content: message["content"].(string),
Sid: userID,
Cid: cid,
Rid: rid,
ExpiredAt: expiredAt,
}
if message["name"] != nil {
value.Name = message["name"].(string)
}
values = append(values, value)
}
return conv.newQuery().Insert(values)
}
// GetChat get the chat info and its history
func (conv *Xun) GetChat(sid string, cid string) (*ChatInfo, error) {
userID, err := conv.getUserID(sid)
@ -638,3 +661,265 @@ func (conv *Xun) DeleteAllChats(sid string) error {
Delete()
return err
}
// processJSONField processes a field that should be stored as JSON string
func (conv *Xun) processJSONField(field interface{}) (interface{}, error) {
if field == nil {
return nil, nil
}
switch v := field.(type) {
case string:
return v, nil
default:
jsonStr, err := jsoniter.MarshalToString(v)
if err != nil {
return nil, fmt.Errorf("failed to marshal %v to JSON: %v", field, err)
}
return jsonStr, nil
}
}
// parseJSONFields parses JSON string fields into their corresponding Go types
func (conv *Xun) parseJSONFields(data map[string]interface{}, fields []string) {
for _, field := range fields {
if val := data[field]; val != nil {
if strVal, ok := val.(string); ok && strVal != "" {
var parsed interface{}
if err := jsoniter.UnmarshalFromString(strVal, &parsed); err == nil {
data[field] = parsed
}
}
}
}
}
// SaveAssistant saves assistant information
func (conv *Xun) SaveAssistant(assistant map[string]interface{}) (interface{}, error) {
// Validate required fields
requiredFields := []string{"name", "type", "connector"}
for _, field := range requiredFields {
if _, ok := assistant[field]; !ok {
return nil, fmt.Errorf("field %s is required", field)
}
if assistant[field] == nil || assistant[field] == "" {
return nil, fmt.Errorf("field %s cannot be empty", field)
}
}
// Create a copy of the assistant map to avoid modifying the original
assistantCopy := make(map[string]interface{})
for k, v := range assistant {
assistantCopy[k] = v
}
// Process JSON fields
jsonFields := []string{"tags", "options", "prompts", "flows", "files", "functions", "permissions"}
for _, field := range jsonFields {
if val, ok := assistantCopy[field]; ok && val != nil {
// If it's a string, try to parse it first
if strVal, ok := val.(string); ok && strVal != "" {
var parsed interface{}
if err := jsoniter.UnmarshalFromString(strVal, &parsed); err == nil {
assistantCopy[field] = parsed
}
}
}
}
// Generate assistant_id if not provided
if _, ok := assistantCopy["assistant_id"]; !ok {
assistantCopy["assistant_id"] = uuid.New().String()
}
// Check if assistant exists
exists, err := conv.query.New().
Table(conv.getAssistantTable()).
Where("assistant_id", assistantCopy["assistant_id"]).
Exists()
if err != nil {
return nil, err
}
// Convert JSON fields to strings for storage
for _, field := range jsonFields {
if val, ok := assistantCopy[field]; ok && val != nil {
jsonStr, err := jsoniter.MarshalToString(val)
if err != nil {
return nil, fmt.Errorf("failed to marshal %s to JSON: %v", field, err)
}
assistantCopy[field] = jsonStr
}
}
// Update or insert
if exists {
_, err := conv.query.New().
Table(conv.getAssistantTable()).
Where("assistant_id", assistantCopy["assistant_id"]).
Update(assistantCopy)
if err != nil {
return nil, err
}
return assistantCopy["assistant_id"], nil
}
err = conv.query.New().
Table(conv.getAssistantTable()).
Insert(assistantCopy)
if err != nil {
return nil, err
}
return assistantCopy["assistant_id"], nil
}
// DeleteAssistant deletes an assistant by assistant_id
func (conv *Xun) DeleteAssistant(assistantID string) error {
// Check if assistant exists
exists, err := conv.query.New().
Table(conv.getAssistantTable()).
Where("assistant_id", assistantID).
Exists()
if err != nil {
return err
}
if !exists {
return fmt.Errorf("assistant %s not found", assistantID)
}
_, err = conv.query.New().
Table(conv.getAssistantTable()).
Where("assistant_id", assistantID).
Delete()
return err
}
// GetAssistants retrieves assistants with pagination and filtering
func (conv *Xun) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) {
qb := conv.query.New().
Table(conv.getAssistantTable())
// Apply tag filter if provided
if filter.Tags != nil && len(filter.Tags) > 0 {
qb.Where(func(qb query.Query) {
for i, tag := range filter.Tags {
// For each tag, we need to match it as part of a JSON array
// This will match both single tag arrays ["tag1"] and multi-tag arrays ["tag1","tag2"]
pattern := fmt.Sprintf("%%\"%s\"%%", tag)
if i == 0 {
qb.Where("tags", "like", pattern)
} else {
qb.OrWhere("tags", "like", pattern)
}
}
})
}
// Apply keyword filter if provided
if filter.Keywords != "" {
qb.Where(func(qb query.Query) {
qb.Where("name", "like", fmt.Sprintf("%%%s%%", filter.Keywords)).
OrWhere("description", "like", fmt.Sprintf("%%%s%%", filter.Keywords))
})
}
// Apply connector filter if provided
if filter.Connector != "" {
qb.Where("connector", filter.Connector)
}
// Apply assistant_id filter if provided
if filter.AssistantID != "" {
qb.Where("assistant_id", filter.AssistantID)
}
// Apply mentionable filter if provided
if filter.Mentionable != nil {
qb.Where("mentionable", *filter.Mentionable)
}
// Apply automated filter if provided
if filter.Automated != nil {
qb.Where("automated", *filter.Automated)
}
// Set defaults for pagination
if filter.PageSize <= 0 {
filter.PageSize = 20
}
if filter.Page <= 0 {
filter.Page = 1
}
// Get total count
total, err := qb.Clone().Count()
if err != nil {
return nil, err
}
// Calculate pagination
offset := (filter.Page - 1) * filter.PageSize
totalPages := int(math.Ceil(float64(total) / float64(filter.PageSize)))
nextPage := filter.Page + 1
if nextPage > totalPages {
nextPage = 0
}
prevPage := filter.Page - 1
if prevPage < 1 {
prevPage = 0
}
// Apply select fields if provided
if filter.Select != nil && len(filter.Select) > 0 {
selectFields := make([]interface{}, len(filter.Select))
for i, field := range filter.Select {
selectFields[i] = field
}
qb.Select(selectFields...)
}
// Get paginated results
rows, err := qb.OrderBy("created_at", "desc").
Offset(offset).
Limit(filter.PageSize).
Get()
if err != nil {
return nil, err
}
// Convert rows to map slice and parse JSON fields
data := make([]map[string]interface{}, len(rows))
jsonFields := []string{"tags", "options", "prompts", "flows", "files", "functions", "permissions"}
for i, row := range rows {
data[i] = row
// Only parse JSON fields if they are selected or no select filter is provided
if filter.Select == nil || len(filter.Select) == 0 {
conv.parseJSONFields(data[i], jsonFields)
} else {
// Parse only selected JSON fields
selectedJSONFields := []string{}
for _, field := range jsonFields {
for _, selected := range filter.Select {
if selected == field {
selectedJSONFields = append(selectedJSONFields, field)
break
}
}
}
if len(selectedJSONFields) > 0 {
conv.parseJSONFields(data[i], selectedJSONFields)
}
}
}
return &AssistantResponse{
Data: data,
Page: filter.Page,
PageSize: filter.PageSize,
PageCnt: totalPages,
Next: nextPage,
Prev: prevPage,
Total: total,
}, nil
}

621
neo/conversation/xun_test.go
View file

@ -5,6 +5,7 @@ import (
"testing"
"time"
jsoniter "github.com/json-iterator/go"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/xun/capsule"
@ -15,13 +16,28 @@ import (
func TestNewXunDefault(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_history")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
err := capsule.Schema().DropTableIfExists("__unit_test_conversation")
err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
if err != nil {
t.Fatal(err)
}
err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
if err != nil {
t.Fatal(err)
}
err = capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
if err != nil {
t.Fatal(err)
}
// Add a small delay to ensure table is created
time.Sleep(100 * time.Millisecond)
conv, err := NewXun(Setting{
Connector: "default",
Table: "__unit_test_conversation",
@ -32,35 +48,59 @@ func TestNewXunDefault(t *testing.T) {
return
}
has, err := capsule.Schema().HasTable("__unit_test_conversation")
// Check history table
has, err := capsule.Schema().HasTable("__unit_test_conversation_history")
if err != nil {
t.Fatal(err)
}
assert.Equal(t, true, has)
// validate the table
tab, err := conv.schema.GetTable(conv.setting.Table)
// Check chat table
has, err = capsule.Schema().HasTable("__unit_test_conversation_chat")
if err != nil {
t.Fatal(err)
}
assert.Equal(t, true, has)
// Check assistant table
has, err = capsule.Schema().HasTable("__unit_test_conversation_assistant")
if err != nil {
t.Fatal(err)
}
assert.Equal(t, true, has)
// validate the history table
tab, err := conv.schema.GetTable(conv.getHistoryTable())
if err != nil {
t.Fatal(err)
}
fields := []string{"id", "sid", "cid", "rid", "role", "name", "content", "created_at", "updated_at", "expired_at"}
fields := []string{"id", "sid", "cid", "uid", "role", "name", "content", "context", "created_at", "updated_at", "expired_at"}
for _, field := range fields {
assert.Equal(t, true, tab.HasColumn(field))
}
conv, err = NewXun(Setting{
Connector: "default",
Table: "__unit_test_conversation",
})
has, err = capsule.Schema().HasTable("__unit_test_conversation")
// validate the chat table
tab, err = conv.schema.GetTable(conv.getChatTable())
if err != nil {
t.Fatal(err)
}
assert.Equal(t, true, has)
chatFields := []string{"id", "chat_id", "title", "sid", "created_at", "updated_at"}
for _, field := range chatFields {
assert.Equal(t, true, tab.HasColumn(field))
}
// validate the assistant table
tab, err = conv.schema.GetTable(conv.getAssistantTable())
if err != nil {
t.Fatal(err)
}
assistantFields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "options", "prompts", "flows", "files", "functions", "tags", "readonly", "permissions", "automated", "mentionable", "created_at", "updated_at"}
for _, field := range assistantFields {
assert.Equal(t, true, tab.HasColumn(field))
}
}
func TestNewXunConnector(t *testing.T) {
@ -77,9 +117,17 @@ func TestNewXunConnector(t *testing.T) {
t.Fatal(err)
}
defer sch.DropTableIfExists("__unit_test_conversation")
defer sch.DropTableIfExists("__unit_test_conversation_history")
defer sch.DropTableIfExists("__unit_test_conversation_chat")
defer sch.DropTableIfExists("__unit_test_conversation_assistant")
sch.DropTableIfExists("__unit_test_conversation_history")
sch.DropTableIfExists("__unit_test_conversation_chat")
sch.DropTableIfExists("__unit_test_conversation_assistant")
// Add a small delay to ensure table is created
time.Sleep(100 * time.Millisecond)
sch.DropTableIfExists("__unit_test_conversation")
conv, err := NewXun(Setting{
Connector: "mysql",
Table: "__unit_test_conversation",
@ -90,44 +138,73 @@ func TestNewXunConnector(t *testing.T) {
return
}
has, err := sch.HasTable("__unit_test_conversation")
// Check history table
has, err := sch.HasTable("__unit_test_conversation_history")
if err != nil {
t.Fatal(err)
}
assert.Equal(t, true, has)
// validate the table
tab, err := conv.schema.GetTable(conv.setting.Table)
// Check chat table
has, err = sch.HasTable("__unit_test_conversation_chat")
if err != nil {
t.Fatal(err)
}
assert.Equal(t, true, has)
// Check assistant table
has, err = sch.HasTable("__unit_test_conversation_assistant")
if err != nil {
t.Fatal(err)
}
assert.Equal(t, true, has)
// validate the history table
tab, err := conv.schema.GetTable(conv.getHistoryTable())
if err != nil {
t.Fatal(err)
}
fields := []string{"id", "sid", "cid", "rid", "role", "name", "content", "created_at", "updated_at", "expired_at"}
fields := []string{"id", "sid", "cid", "uid", "role", "name", "content", "context", "created_at", "updated_at", "expired_at"}
for _, field := range fields {
assert.Equal(t, true, tab.HasColumn(field))
}
conv, err = NewXun(Setting{
Connector: "default",
Table: "__unit_test_conversation",
})
has, err = sch.HasTable("__unit_test_conversation")
// validate the chat table
tab, err = conv.schema.GetTable(conv.getChatTable())
if err != nil {
t.Fatal(err)
}
assert.Equal(t, true, has)
chatFields := []string{"id", "chat_id", "title", "sid", "created_at", "updated_at"}
for _, field := range chatFields {
assert.Equal(t, true, tab.HasColumn(field))
}
// validate the assistant table
tab, err = conv.schema.GetTable(conv.getAssistantTable())
if err != nil {
t.Fatal(err)
}
assistantFields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "options", "prompts", "flows", "files", "functions", "tags", "readonly", "permissions", "automated", "mentionable", "created_at", "updated_at"}
for _, field := range assistantFields {
assert.Equal(t, true, tab.HasColumn(field))
}
}
func TestXunSaveAndGetHistory(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_history")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
err := capsule.Schema().DropTableIfExists("__unit_test_conversation")
err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
if err != nil {
t.Fatal(err)
}
err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
if err != nil {
t.Fatal(err)
}
@ -143,7 +220,7 @@ func TestXunSaveAndGetHistory(t *testing.T) {
err = conv.SaveHistory("123456", []map[string]interface{}{
{"role": "user", "name": "user1", "content": "hello"},
{"role": "assistant", "name": "user1", "content": "Hello there, how"},
}, cid)
}, cid, nil)
assert.Nil(t, err)
// get the history
@ -154,44 +231,18 @@ func TestXunSaveAndGetHistory(t *testing.T) {
assert.Equal(t, 2, len(data))
}
func TestXunSaveAndGetRequest(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
defer capsule.Schema().DropTableIfExists("__unit_test_conversation")
err := capsule.Schema().DropTableIfExists("__unit_test_conversation")
if err != nil {
t.Fatal(err)
}
conv, err := NewXun(Setting{
Connector: "default",
Table: "__unit_test_conversation",
TTL: 3600,
})
// save the history
err = conv.SaveRequest("123456", "912836", "test.command", []map[string]interface{}{
{"role": "user", "name": "user1", "content": "hello"},
{"role": "assistant", "name": "user1", "content": "Hello there, how"},
})
assert.Nil(t, err)
// get the history
data, err := conv.GetRequest("123456", "912836")
if err != nil {
t.Fatal(err)
}
assert.Equal(t, 2, len(data))
}
func TestXunSaveAndGetHistoryWithCID(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_history")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
err := capsule.Schema().DropTableIfExists("__unit_test_conversation")
err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
if err != nil {
t.Fatal(err)
}
err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
if err != nil {
t.Fatal(err)
}
@ -209,7 +260,7 @@ func TestXunSaveAndGetHistoryWithCID(t *testing.T) {
{"role": "user", "name": "user1", "content": "hello"},
{"role": "assistant", "name": "assistant1", "content": "Hi! How can I help you?"},
}
err = conv.SaveHistory(sid, messages, cid)
err = conv.SaveHistory(sid, messages, cid, nil)
assert.Nil(t, err)
// get the history for specific cid
@ -224,7 +275,7 @@ func TestXunSaveAndGetHistoryWithCID(t *testing.T) {
moreMessages := []map[string]interface{}{
{"role": "user", "name": "user1", "content": "another message"},
}
err = conv.SaveHistory(sid, moreMessages, anotherCID)
err = conv.SaveHistory(sid, moreMessages, anotherCID, nil)
assert.Nil(t, err)
// get history for the first cid - should still be 2 messages
@ -252,11 +303,11 @@ func TestXunSaveAndGetHistoryWithCID(t *testing.T) {
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_history")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
// Drop both tables before test
err := capsule.Schema().DropTableIfExists("__unit_test_conversation")
err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
if err != nil {
t.Fatal(err)
}
@ -294,7 +345,7 @@ func TestXunGetChats(t *testing.T) {
}
// Then save the history
err = conv.SaveHistory(sid, messages, chatID)
err = conv.SaveHistory(sid, messages, chatID, nil)
if err != nil {
t.Fatal(err)
}
@ -325,7 +376,7 @@ func TestXunGetChats(t *testing.T) {
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_history")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
conv, err := NewXun(Setting{
@ -344,7 +395,7 @@ func TestXunDeleteChat(t *testing.T) {
}
// Save the chat and history
err = conv.SaveHistory(sid, messages, cid)
err = conv.SaveHistory(sid, messages, cid, nil)
assert.Nil(t, err)
// Verify chat exists
@ -365,7 +416,7 @@ func TestXunDeleteChat(t *testing.T) {
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_history")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
conv, err := NewXun(Setting{
@ -385,7 +436,7 @@ func TestXunDeleteAllChats(t *testing.T) {
// Save multiple chats
for i := 0; i < 3; i++ {
cid := fmt.Sprintf("test_chat_%d", i)
err = conv.SaveHistory(sid, messages, cid)
err = conv.SaveHistory(sid, messages, cid, nil)
assert.Nil(t, err)
}
@ -403,3 +454,423 @@ func TestXunDeleteAllChats(t *testing.T) {
assert.Nil(t, err)
assert.Equal(t, int64(0), response.Total)
}
func TestXunAssistantCRUD(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
// Drop assistant table before test
err := capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
if err != nil {
t.Fatal(err)
}
// Add a small delay to ensure table is created
time.Sleep(100 * time.Millisecond)
conv, err := NewXun(Setting{
Connector: "default",
Table: "__unit_test_conversation",
})
if err != nil {
t.Fatal(err)
}
// Test creating a new assistant with different JSON field formats
// Test case 1: JSON fields as strings
tagsJSON := `["tag1", "tag2", "tag3"]`
optionsJSON := `{"model": "gpt-4"}`
assistant := map[string]interface{}{
"name": "Test Assistant",
"type": "assistant",
"avatar": "https://example.com/avatar.png",
"connector": "openai",
"description": "Test Description",
"tags": tagsJSON,
"options": optionsJSON,
"mentionable": true,
"automated": true,
}
// Test SaveAssistant (Create) with string JSON
v, err := conv.SaveAssistant(assistant)
assert.Nil(t, err)
assistantID := v.(string)
assert.NotEmpty(t, assistantID)
// Test case 2: JSON fields as native types
assistant2 := map[string]interface{}{
"name": "Test Assistant 2",
"type": "assistant",
"avatar": "https://example.com/avatar2.png",
"connector": "openai",
"description": "Test Description 2",
"tags": []string{"tag1", "tag2", "tag3"},
"options": map[string]interface{}{"model": "gpt-4"},
"prompts": []string{"prompt1", "prompt2"},
"flows": []string{"flow1", "flow2"},
"files": []string{"file1", "file2"},
"functions": []map[string]interface{}{{"name": "func1"}, {"name": "func2"}},
"permissions": map[string]interface{}{"read": true, "write": true},
"mentionable": true,
"automated": true,
}
// Test SaveAssistant (Create) with native types
v, err = conv.SaveAssistant(assistant2)
assert.Nil(t, err)
assistant2ID := v.(string)
assert.NotEmpty(t, assistant2ID)
// Test case 3: Test with nil JSON fields
assistant3 := map[string]interface{}{
"name": "Test Assistant 3",
"type": "assistant",
"connector": "openai",
"description": "Test Description 3",
"tags": nil,
"options": nil,
"prompts": nil,
"flows": nil,
"files": nil,
"functions": nil,
"permissions": nil,
"mentionable": true,
"automated": true,
}
// Test SaveAssistant (Create) with nil fields
v, err = conv.SaveAssistant(assistant3)
assert.Nil(t, err)
assistant3ID := v.(string)
assert.NotEmpty(t, assistant3ID)
// Test GetAssistants to verify JSON fields are properly stored
resp, err := conv.GetAssistants(AssistantFilter{})
assert.Nil(t, err)
assert.Equal(t, 3, len(resp.Data))
// Verify first assistant (string JSON)
found := false
for _, item := range resp.Data {
if item["assistant_id"].(string) == assistantID {
found = true
// Now we expect parsed JSON values instead of JSON strings
assert.Equal(t, []interface{}{"tag1", "tag2", "tag3"}, item["tags"])
assert.Equal(t, map[string]interface{}{"model": "gpt-4"}, item["options"])
break
}
}
assert.True(t, found)
// Verify second assistant (native types converted to JSON)
found = false
for _, item := range resp.Data {
if item["assistant_id"].(string) == assistant2ID {
found = true
// Now we expect parsed JSON values directly
assert.Equal(t, []interface{}{"tag1", "tag2", "tag3"}, item["tags"])
assert.Equal(t, map[string]interface{}{"model": "gpt-4"}, item["options"])
// Verify other JSON fields
assert.Equal(t, []interface{}{"prompt1", "prompt2"}, item["prompts"])
assert.Equal(t, []interface{}{"flow1", "flow2"}, item["flows"])
assert.Equal(t, []interface{}{"file1", "file2"}, item["files"])
assert.Equal(t,
[]interface{}{
map[string]interface{}{"name": "func1"},
map[string]interface{}{"name": "func2"},
},
item["functions"])
assert.Equal(t,
map[string]interface{}{
"read": true,
"write": true,
},
item["permissions"])
break
}
}
assert.True(t, found)
// Verify third assistant (nil fields)
found = false
for _, item := range resp.Data {
if item["assistant_id"].(string) == assistant3ID {
found = true
assert.Nil(t, item["tags"])
assert.Nil(t, item["options"])
assert.Nil(t, item["prompts"])
assert.Nil(t, item["flows"])
assert.Nil(t, item["files"])
assert.Nil(t, item["functions"])
assert.Nil(t, item["permissions"])
break
}
}
assert.True(t, found)
// Test updating with mixed JSON formats
assistant2["assistant_id"] = assistant2ID
_, err = conv.SaveAssistant(assistant2)
assert.Nil(t, err)
// Verify update
resp, err = conv.GetAssistants(AssistantFilter{})
assert.Nil(t, err)
for _, item := range resp.Data {
if item["assistant_id"].(string) == assistant2ID {
// Now we expect parsed JSON values
assert.Equal(t, []interface{}{"tag1", "tag2", "tag3"}, item["tags"])
assert.Equal(t, map[string]interface{}{"model": "gpt-4"}, item["options"])
break
}
}
// Test DeleteAssistant
err = conv.DeleteAssistant(assistantID)
assert.Nil(t, err)
err = conv.DeleteAssistant(assistant2ID)
assert.Nil(t, err)
err = conv.DeleteAssistant(assistant3ID)
assert.Nil(t, err)
resp, err = conv.GetAssistants(AssistantFilter{})
assert.Nil(t, err)
assert.Equal(t, 0, len(resp.Data))
}
func TestXunAssistantPagination(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
// Drop assistant table before test
err := capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
if err != nil {
t.Fatal(err)
}
// Add a small delay to ensure table is created
time.Sleep(100 * time.Millisecond)
conv, err := NewXun(Setting{
Connector: "default",
Table: "__unit_test_conversation",
})
if err != nil {
t.Fatal(err)
}
// Create multiple assistants for pagination testing
mentionable := true
automated := true
for i := 0; i < 25; i++ {
tagsJSON, err := jsoniter.MarshalToString([]string{fmt.Sprintf("tag%d", i%5)})
if err != nil {
t.Fatal(err)
}
// Alternate mentionable and automated flags
if i%2 == 0 {
mentionable = !mentionable
}
if i%3 == 0 {
automated = !automated
}
assistant := map[string]interface{}{
"name": fmt.Sprintf("Assistant %d", i),
"type": "assistant",
"connector": fmt.Sprintf("connector%d", i%3),
"description": fmt.Sprintf("Description %d", i),
"tags": tagsJSON,
"mentionable": mentionable,
"automated": automated,
}
_, err = conv.SaveAssistant(assistant)
assert.Nil(t, err)
}
// Test first page
resp, err := conv.GetAssistants(AssistantFilter{
Page: 1,
PageSize: 10,
})
assert.Nil(t, err)
assert.Equal(t, 10, len(resp.Data))
assert.Equal(t, int64(25), resp.Total)
assert.Equal(t, 3, resp.PageCnt)
assert.Equal(t, 2, resp.Next)
assert.Equal(t, 0, resp.Prev)
// Test second page
resp, err = conv.GetAssistants(AssistantFilter{
Page: 2,
PageSize: 10,
})
assert.Nil(t, err)
assert.Equal(t, 10, len(resp.Data))
assert.Equal(t, 3, resp.Next)
assert.Equal(t, 1, resp.Prev)
// Test last page
resp, err = conv.GetAssistants(AssistantFilter{
Page: 3,
PageSize: 10,
})
assert.Nil(t, err)
assert.Equal(t, 5, len(resp.Data))
assert.Equal(t, 0, resp.Next)
assert.Equal(t, 2, resp.Prev)
// Test filtering with tags
resp, err = conv.GetAssistants(AssistantFilter{
Tags: []string{"tag0"},
Page: 1,
PageSize: 10,
})
assert.Nil(t, err)
assert.Equal(t, 5, len(resp.Data))
// Test filtering with keywords
resp, err = conv.GetAssistants(AssistantFilter{
Keywords: "Assistant 1",
Page: 1,
PageSize: 10,
})
assert.Nil(t, err)
assert.Greater(t, len(resp.Data), 0)
// Test filtering with connector
resp, err = conv.GetAssistants(AssistantFilter{
Connector: "connector0",
Page: 1,
PageSize: 10,
})
assert.Nil(t, err)
assert.Greater(t, len(resp.Data), 0)
// Test filtering with mentionable
mentionableTrue := true
resp, err = conv.GetAssistants(AssistantFilter{
Mentionable: &mentionableTrue,
Page: 1,
PageSize: 10,
})
assert.Nil(t, err)
assert.Greater(t, len(resp.Data), 0)
// Test filtering with automated
automatedTrue := true
resp, err = conv.GetAssistants(AssistantFilter{
Automated: &automatedTrue,
Page: 1,
PageSize: 10,
})
assert.Nil(t, err)
assert.Greater(t, len(resp.Data), 0)
// Test filtering by assistant_id
// First get an assistant_id from previous results
firstAssistantID := resp.Data[0]["assistant_id"].(string)
// Test exact match with assistant_id
resp, err = conv.GetAssistants(AssistantFilter{
AssistantID: firstAssistantID,
Page: 1,
PageSize: 10,
})
assert.Nil(t, err)
assert.Equal(t, 1, len(resp.Data))
assert.Equal(t, firstAssistantID, resp.Data[0]["assistant_id"])
// Test assistant_id with other filters
resp, err = conv.GetAssistants(AssistantFilter{
AssistantID: firstAssistantID,
Select: []string{"name", "assistant_id", "description"},
Page: 1,
PageSize: 10,
})
assert.Nil(t, err)
assert.Equal(t, 1, len(resp.Data))
assert.Equal(t, firstAssistantID, resp.Data[0]["assistant_id"])
// Verify only selected fields are returned
assert.Contains(t, resp.Data[0], "name")
assert.Contains(t, resp.Data[0], "assistant_id")
assert.Contains(t, resp.Data[0], "description")
assert.NotContains(t, resp.Data[0], "tags")
assert.NotContains(t, resp.Data[0], "options")
// Test non-existent assistant_id
resp, err = conv.GetAssistants(AssistantFilter{
AssistantID: "non-existent-id",
Page: 1,
PageSize: 10,
})
assert.Nil(t, err)
assert.Equal(t, 0, len(resp.Data))
// Test combined filters
resp, err = conv.GetAssistants(AssistantFilter{
Tags: []string{"tag0"},
Keywords: "Assistant",
Connector: "connector0",
Mentionable: &mentionableTrue,
Automated: &automatedTrue,
Page: 1,
PageSize: 10,
})
assert.Nil(t, err)
// Test filtering with select fields
resp, err = conv.GetAssistants(AssistantFilter{
Select: []string{"name", "description", "tags"},
Page: 1,
PageSize: 10,
})
assert.Nil(t, err)
assert.Equal(t, 10, len(resp.Data))
// Verify only selected fields are returned
for _, item := range resp.Data {
// These fields should exist
assert.Contains(t, item, "name")
assert.Contains(t, item, "description")
assert.Contains(t, item, "tags")
// These fields should not exist
assert.NotContains(t, item, "options")
assert.NotContains(t, item, "prompts")
assert.NotContains(t, item, "flows")
assert.NotContains(t, item, "files")
assert.NotContains(t, item, "functions")
assert.NotContains(t, item, "permissions")
}
// Test filtering with select fields and other filters combined
resp, err = conv.GetAssistants(AssistantFilter{
Tags: []string{"tag0"},
Keywords: "Assistant",
Select: []string{"name", "tags"},
Page: 1,
PageSize: 10,
})
assert.Nil(t, err)
// Verify only selected fields are returned
for _, item := range resp.Data {
// These fields should exist
assert.Contains(t, item, "name")
assert.Contains(t, item, "tags")
// These fields should not exist
assert.NotContains(t, item, "description")
assert.NotContains(t, item, "options")
assert.NotContains(t, item, "prompts")
assert.NotContains(t, item, "flows")
assert.NotContains(t, item, "files")
assert.NotContains(t, item, "functions")
assert.NotContains(t, item, "permissions")
}
}

74
neo/neo.go
View file

@ -11,7 +11,7 @@ import (
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/yao/neo/assistant"
"github.com/yaoapp/yao/neo/assistant/base"
"github.com/yaoapp/yao/neo/assistant/local"
"github.com/yaoapp/yao/neo/assistant/openai"
"github.com/yaoapp/yao/neo/conversation"
"github.com/yaoapp/yao/neo/message"
@ -58,20 +58,50 @@ 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) {
// GeneratePrompts generate prompts for the AI assistant
func (neo *DSL) GeneratePrompts(ctx Context, input string, c *gin.Context, silent ...bool) (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.
Optimize the prompts for the AI assistant
1. Optimize prompts based on the user's input
2. The prompts should be clear and specific
3. The prompts should be in the same language as the input
4. Keep the prompts concise but comprehensive
5. DO NOT ASK USER FOR MORE INFORMATION, JUST GENERATE PROMPTS
6. DO NOT ANSWER THE QUESTION, JUST GENERATE PROMPTS
`
isSilent := false
if len(silent) > 0 {
isSilent = silent[0]
}
return neo.GenerateWithAI(ctx, input, "prompts", prompts, c, isSilent)
}
// GenerateChatTitle generate the chat title
func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context, silent ...bool) (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.
`
isSilent := false
if len(silent) > 0 {
isSilent = silent[0]
}
return neo.GenerateWithAI(ctx, input, "title", prompts, c, isSilent)
}
// GenerateWithAI generate content with AI, type can be "title", "prompts", etc.
func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, systemPrompt string, c *gin.Context, silent bool) (string, error) {
messages := []map[string]interface{}{
{"role": "system", "content": prompts},
{"role": "user", "content": input},
{"role": "system", "content": systemPrompt},
{
"role": "user",
"content": input,
"type": messageType,
"name": ctx.Sid,
},
}
res, err := neo.HookCreate(ctx, messages, c)
@ -119,8 +149,21 @@ func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context) (st
// Append content and send message
content = msg.Append(content)
// Only send real-time messages if not in silent mode
if !silent && 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.Message.Done {
if !silent && msg.Message.Text == "" {
msg.Write(c.Writer)
}
done <- true
return 0 // break
}
@ -131,7 +174,9 @@ func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context) (st
if err != nil {
log.Error("Chat error: %s", err.Error())
message.New().Error(err).Done().Write(c.Writer)
if !silent {
message.New().Error(err).Done().Write(c.Writer)
}
}
done <- true
@ -372,9 +417,9 @@ func (neo *DSL) newAssistantByConnector(id string) (assistant.API, error) {
}
// Base on the assistant list hook
api, err := base.New(conn, neo.Prompts, id)
api, err := local.New(conn, neo.Prompts, id)
if err != nil {
return nil, fmt.Errorf("Create base assistant error: %s", err.Error())
return nil, fmt.Errorf("Create local assistant error: %s", err.Error())
}
return api, nil
}
@ -453,6 +498,7 @@ func (neo *DSL) saveHistory(sid string, chatID string, content []byte, messages
{"role": "assistant", "content": string(content), "name": sid},
},
chatID,
nil,
)
if err != nil {

167
neo/process.go
View file

@ -1,15 +1,32 @@
package neo
import (
"fmt"
"strconv"
"github.com/gin-gonic/gin"
"github.com/yaoapp/gou/process"
"github.com/yaoapp/kun/exception"
"github.com/yaoapp/yao/neo/conversation"
"github.com/yaoapp/yao/neo/message"
)
// GetNeo returns the Neo instance
func GetNeo() *DSL {
if Neo == nil {
exception.New("Neo is not initialized", 500).Throw()
}
return Neo
}
func init() {
process.RegisterGroup("neo", map[string]process.Handler{
"write": ProcessWrite,
"write": ProcessWrite,
"assistant.create": processAssistantCreate,
"assistant.save": processAssistantSave,
"assistant.delete": processAssistantDelete,
"assistant.search": processAssistantSearch,
"assistant.find": processAssistantFind,
})
}
@ -39,3 +56,151 @@ func ProcessWrite(process *process.Process) interface{} {
return nil
}
// processAssistantCreate process the assistant create request
func processAssistantCreate(process *process.Process) interface{} {
process.ValidateArgNums(1)
data := process.ArgsMap(0)
neo := GetNeo()
if neo.Conversation == nil {
exception.New("Neo conversation is not initialized", 500).Throw()
}
id, err := neo.Conversation.SaveAssistant(data)
if err != nil {
exception.New("Failed to create assistant: %s", 500, err.Error()).Throw()
}
return id
}
// processAssistantSave process the assistant save request
func processAssistantSave(process *process.Process) interface{} {
process.ValidateArgNums(1)
data := process.ArgsMap(0)
neo := GetNeo()
if neo.Conversation == nil {
exception.New("Neo conversation is not initialized", 500).Throw()
}
id, err := neo.Conversation.SaveAssistant(data)
if err != nil {
exception.New("Failed to save assistant: %s", 500, err.Error()).Throw()
}
return id
}
// processAssistantDelete process the assistant delete request
func processAssistantDelete(process *process.Process) interface{} {
process.ValidateArgNums(1)
assistantID := process.ArgsString(0)
neo := GetNeo()
if neo.Conversation == nil {
exception.New("Neo conversation is not initialized", 500).Throw()
}
err := neo.Conversation.DeleteAssistant(assistantID)
if err != nil {
exception.New("Failed to delete assistant: %s", 500, err.Error()).Throw()
}
return gin.H{"message": "ok"}
}
// processAssistantSearch process the assistant search request
func processAssistantSearch(process *process.Process) interface{} {
params := process.ArgsMap(0)
filter := conversation.AssistantFilter{}
// Parse page and pagesize
if page, ok := params["page"]; ok {
pageStr := fmt.Sprintf("%v", page)
if pageInt, err := strconv.Atoi(pageStr); err == nil {
filter.Page = pageInt
}
}
if pagesize, ok := params["pagesize"]; ok {
pagesizeStr := fmt.Sprintf("%v", pagesize)
if pagesizeInt, err := strconv.Atoi(pagesizeStr); err == nil {
filter.PageSize = pagesizeInt
}
}
// Parse tags
if tags, ok := params["tags"]; ok {
switch v := tags.(type) {
case []interface{}:
filter.Tags = make([]string, len(v))
for i, tag := range v {
filter.Tags[i] = fmt.Sprintf("%v", tag)
}
case []string:
filter.Tags = v
}
}
// Parse keywords
if keywords, ok := params["keywords"].(string); ok {
filter.Keywords = keywords
}
// Parse connector
if connector, ok := params["connector"].(string); ok {
filter.Connector = connector
}
// Parse mentionable
if mentionable, ok := params["mentionable"].(bool); ok {
filter.Mentionable = &mentionable
}
// Parse automated
if automated, ok := params["automated"].(bool); ok {
filter.Automated = &automated
}
// Get assistants
neo := GetNeo()
if neo.Conversation == nil {
exception.New("Neo conversation is not initialized", 500).Throw()
}
res, err := neo.Conversation.GetAssistants(filter)
if err != nil {
exception.New("get assistants error: %s", 500, err).Throw()
}
return res
}
// processAssistantFind process the assistant find request
func processAssistantFind(process *process.Process) interface{} {
process.ValidateArgNums(1)
assistantID := process.ArgsString(0)
neo := GetNeo()
if neo.Conversation == nil {
exception.New("Neo conversation is not initialized", 500).Throw()
}
filter := conversation.AssistantFilter{
AssistantID: assistantID,
Page: 1,
PageSize: 1,
}
res, err := neo.Conversation.GetAssistants(filter)
if err != nil {
exception.New("Failed to find assistant: %s", 500, err.Error()).Throw()
}
if len(res.Data) == 0 {
exception.New("Assistant not found: %s", 404, assistantID).Throw()
}
return res.Data[0]
}

513
neo/process_test.go Normal file
View file

@ -0,0 +1,513 @@
package neo
import (
"fmt"
"testing"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/gou/process"
"github.com/yaoapp/kun/any"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/test"
)
func prepare(t *testing.T) {
test.Prepare(t, config.Conf)
err := Load(config.Conf)
if err != nil {
t.Fatal(err)
}
// Clean up the test data before each test
p, err := process.Of("neo.assistant.search", map[string]interface{}{
"page": 1,
"pagesize": 1000, // Use a large page size to get all records
})
if err != nil {
t.Fatal(err)
}
output, err := p.Exec()
if err != nil {
t.Fatal(err)
}
res := any.Of(output).Map()
items := res.Get("data")
if items != nil {
for _, item := range items.([]map[string]interface{}) {
assistantID := item["assistant_id"].(string)
p, err = process.Of("neo.assistant.delete", assistantID)
if err != nil {
t.Fatal(err)
}
_, err = p.Exec()
if err != nil {
t.Fatal(err)
}
}
}
// Verify cleanup
p, err = process.Of("neo.assistant.search")
if err != nil {
t.Fatal(err)
}
output, err = p.Exec()
if err != nil {
t.Fatal(err)
}
res = any.Of(output).Map()
total := res.Get("total")
if total != nil && any.Of(total).CInt() > 0 {
t.Fatalf("Failed to clean up test data, %d records remaining", any.Of(total).CInt())
}
check(t)
}
func TestProcessAssistantCRUD(t *testing.T) {
prepare(t)
defer test.Clean()
// Create an assistant with string JSON fields
tagsJSON := `["tag1", "tag2", "tag3"]`
optionsJSON := `{"model": "gpt-4"}`
assistant := map[string]interface{}{
"name": "Test Assistant",
"type": "assistant",
"avatar": "https://example.com/avatar.png",
"connector": "openai",
"description": "Test Description",
"tags": tagsJSON,
"options": optionsJSON,
"mentionable": true,
"automated": true,
}
// Test processAssistantCreate with string JSON
p, err := process.Of("neo.assistant.create", assistant)
if err != nil {
t.Fatal(err)
}
output, err := p.Exec()
if err != nil {
t.Fatal(err)
}
assistantID := output
assert.NotNil(t, assistantID)
// Test processAssistantFind
p, err = process.Of("neo.assistant.find", assistantID)
if err != nil {
t.Fatal(err)
}
output, err = p.Exec()
if err != nil {
t.Fatal(err)
}
foundAssistant := output.(map[string]interface{})
assert.Equal(t, assistantID, foundAssistant["assistant_id"])
assert.Equal(t, "Test Assistant", foundAssistant["name"])
assert.Equal(t, []interface{}{"tag1", "tag2", "tag3"}, foundAssistant["tags"])
assert.Equal(t, map[string]interface{}{"model": "gpt-4"}, foundAssistant["options"])
// Test processAssistantFind with non-existent ID
p, err = process.Of("neo.assistant.find", "non-existent-id")
if err != nil {
t.Fatal(err)
}
_, err = p.Exec()
assert.NotNil(t, err)
assert.Contains(t, err.Error(), "Assistant not found")
// Test with native type JSON fields
assistant2 := map[string]interface{}{
"name": "Test Assistant 2",
"type": "assistant",
"avatar": "https://example.com/avatar2.png",
"connector": "openai",
"description": "Test Description 2",
"tags": []string{"tag1", "tag2", "tag3"},
"options": map[string]interface{}{"model": "gpt-4"},
"prompts": []string{"prompt1", "prompt2"},
"flows": []string{"flow1", "flow2"},
"files": []string{"file1", "file2"},
"functions": []map[string]interface{}{{"name": "func1"}, {"name": "func2"}},
"permissions": map[string]interface{}{"read": true, "write": true},
"mentionable": true,
"automated": true,
}
// Test processAssistantCreate with native types
p, err = process.Of("neo.assistant.create", assistant2)
if err != nil {
t.Fatal(err)
}
output, err = p.Exec()
if err != nil {
t.Fatal(err)
}
assistant2ID := output
assert.NotNil(t, assistant2ID)
// Test with nil JSON fields
assistant3 := map[string]interface{}{
"name": "Test Assistant 3",
"type": "assistant",
"connector": "openai",
"description": "Test Description 3",
"tags": nil,
"options": nil,
"prompts": nil,
"flows": nil,
"files": nil,
"functions": nil,
"permissions": nil,
"mentionable": true,
"automated": true,
}
// Test processAssistantCreate with nil fields
p, err = process.Of("neo.assistant.create", assistant3)
if err != nil {
t.Fatal(err)
}
output, err = p.Exec()
if err != nil {
t.Fatal(err)
}
assistant3ID := output
assert.NotNil(t, assistant3ID)
// Test processAssistantSearch to verify all assistants
p, err = process.Of("neo.assistant.search")
if err != nil {
t.Fatal(err)
}
output, err = p.Exec()
if err != nil {
t.Fatal(err)
}
searchRes := any.Of(output).Map()
total := searchRes.Get("total")
if total == nil {
total = int64(0)
}
assert.Equal(t, int64(3), total)
items := searchRes.Get("data")
if items == nil {
items = []map[string]interface{}{}
}
assert.Equal(t, 3, len(items.([]map[string]interface{})))
// Verify each assistant's JSON fields
for _, item := range items.([]map[string]interface{}) {
switch item["assistant_id"].(string) {
case assistantID:
assert.Equal(t, []interface{}{"tag1", "tag2", "tag3"}, item["tags"])
assert.Equal(t, map[string]interface{}{"model": "gpt-4"}, item["options"])
case assistant2ID:
assert.Equal(t, []interface{}{"tag1", "tag2", "tag3"}, item["tags"])
assert.Equal(t, map[string]interface{}{"model": "gpt-4"}, item["options"])
assert.Equal(t, []interface{}{"prompt1", "prompt2"}, item["prompts"])
assert.Equal(t, []interface{}{"flow1", "flow2"}, item["flows"])
assert.Equal(t, []interface{}{"file1", "file2"}, item["files"])
assert.Equal(t,
[]interface{}{
map[string]interface{}{"name": "func1"},
map[string]interface{}{"name": "func2"},
},
item["functions"])
assert.Equal(t,
map[string]interface{}{
"read": true,
"write": true,
},
item["permissions"])
case assistant3ID:
assert.Nil(t, item["tags"])
assert.Nil(t, item["options"])
assert.Nil(t, item["prompts"])
assert.Nil(t, item["flows"])
assert.Nil(t, item["files"])
assert.Nil(t, item["functions"])
assert.Nil(t, item["permissions"])
}
}
// Test updating with mixed JSON formats
assistant2["assistant_id"] = assistant2ID
assistant2["tags"] = `["tag4", "tag5"]`
assistant2["options"] = map[string]interface{}{"model": "gpt-3.5"}
p, err = process.Of("neo.assistant.save", assistant2)
if err != nil {
t.Fatal(err)
}
output, err = p.Exec()
if err != nil {
t.Fatal(err)
}
savedID := output
assert.NotNil(t, savedID)
// Double check with a new search
p, err = process.Of("neo.assistant.search")
if err != nil {
t.Fatal(err)
}
output, err = p.Exec()
if err != nil {
t.Fatal(err)
}
searchRes = any.Of(output).Map()
items = searchRes.Get("data")
found := false
for _, item := range items.([]map[string]interface{}) {
if item["assistant_id"].(string) == assistant2ID {
found = true
assert.Equal(t, []interface{}{"tag4", "tag5"}, item["tags"])
assert.Equal(t, map[string]interface{}{"model": "gpt-3.5"}, item["options"])
break
}
}
assert.True(t, found)
// Test processAssistantDelete
p, err = process.Of("neo.assistant.delete", assistantID)
if err != nil {
t.Fatal(err)
}
output, err = p.Exec()
if err != nil {
t.Fatal(err)
}
deleteRes := any.Of(output).Map()
assert.Equal(t, "ok", deleteRes.Get("message"))
// Delete remaining assistants
p, err = process.Of("neo.assistant.delete", assistant2ID)
if err != nil {
t.Fatal(err)
}
_, err = p.Exec()
assert.Nil(t, err)
p, err = process.Of("neo.assistant.delete", assistant3ID)
if err != nil {
t.Fatal(err)
}
_, err = p.Exec()
assert.Nil(t, err)
// Verify all assistants are deleted
p, err = process.Of("neo.assistant.search")
if err != nil {
t.Fatal(err)
}
output, err = p.Exec()
if err != nil {
t.Fatal(err)
}
searchRes = any.Of(output).Map()
total = searchRes.Get("total")
if total == nil {
total = int64(0)
}
assert.Equal(t, int64(0), total)
}
func TestProcessAssistantSearchPagination(t *testing.T) {
prepare(t)
defer test.Clean()
// Create multiple assistants for pagination testing
for i := 0; i < 25; i++ {
assistant := map[string]interface{}{
"name": fmt.Sprintf("Assistant %d", i),
"type": "assistant",
"connector": fmt.Sprintf("connector%d", i%3),
"description": fmt.Sprintf("Description %d", i),
"tags": []string{fmt.Sprintf("tag%d", i%5)},
"mentionable": i%2 == 0,
"automated": i%3 == 0,
}
p, err := process.Of("neo.assistant.create", assistant)
if err != nil {
t.Fatal(err)
}
_, err = p.Exec()
if err != nil {
t.Fatal(err)
}
}
// Test first page
p, err := process.Of("neo.assistant.search", map[string]interface{}{
"page": 1,
"pagesize": 10,
})
if err != nil {
t.Fatal(err)
}
output, err := p.Exec()
if err != nil {
t.Fatal(err)
}
res := any.Of(output).Map()
total := res.Get("total")
if total == nil {
total = int64(0)
}
assert.Equal(t, int64(25), total)
items := res.Get("data")
if items == nil {
items = []map[string]interface{}{}
}
assert.Equal(t, 10, len(items.([]map[string]interface{})))
pageCnt := res.Get("pagecnt")
if pageCnt == nil {
pageCnt = 1
}
assert.Equal(t, 3, pageCnt)
// Test second page
p, err = process.Of("neo.assistant.search", map[string]interface{}{
"page": 2,
"pagesize": 10,
})
if err != nil {
t.Fatal(err)
}
output, err = p.Exec()
if err != nil {
t.Fatal(err)
}
res = any.Of(output).Map()
items = res.Get("data")
if items == nil {
items = []map[string]interface{}{}
}
assert.Equal(t, 10, len(items.([]map[string]interface{})))
// Test last page
p, err = process.Of("neo.assistant.search", map[string]interface{}{
"page": 3,
"pagesize": 10,
})
if err != nil {
t.Fatal(err)
}
output, err = p.Exec()
if err != nil {
t.Fatal(err)
}
res = any.Of(output).Map()
items = res.Get("data")
if items == nil {
items = []map[string]interface{}{}
}
assert.Equal(t, 5, len(items.([]map[string]interface{})))
// Test filtering with tags
p, err = process.Of("neo.assistant.search", map[string]interface{}{
"tags": []string{"tag0"},
"page": 1,
"pagesize": 10,
})
if err != nil {
t.Fatal(err)
}
output, err = p.Exec()
if err != nil {
t.Fatal(err)
}
res = any.Of(output).Map()
items = res.Get("data")
if items == nil {
items = []map[string]interface{}{}
}
assert.Equal(t, 5, len(items.([]map[string]interface{})))
}
func TestProcessAssistantValidation(t *testing.T) {
prepare(t)
defer test.Clean()
// Test missing required fields
p, err := process.Of("neo.assistant.create", map[string]interface{}{})
if err != nil {
t.Fatal(err)
}
_, err = p.Exec()
assert.NotNil(t, err)
// Test invalid assistant ID for delete
p, err = process.Of("neo.assistant.delete", "non-existent-id")
if err != nil {
t.Fatal(err)
}
_, err = p.Exec()
assert.NotNil(t, err)
// Test invalid assistant ID for find
p, err = process.Of("neo.assistant.find", "non-existent-id")
if err != nil {
t.Fatal(err)
}
_, err = p.Exec()
assert.NotNil(t, err)
assert.Contains(t, err.Error(), "Assistant not found")
// Test invalid page number
p, err = process.Of("neo.assistant.search", map[string]interface{}{
"page": -1,
"pagesize": 10,
})
if err != nil {
t.Fatal(err)
}
output, err := p.Exec()
assert.Nil(t, err)
res := any.Of(output).Map()
total := res.Get("total")
if total == nil {
total = int64(0)
}
assert.Equal(t, int64(0), total)
}

45
test/utils.go
View file

@ -4,6 +4,7 @@ import (
"fmt"
"os"
"path/filepath"
"regexp"
"strings"
"testing"
"time"
@ -82,6 +83,50 @@ func Prepare(t *testing.T, cfg config.Config, rootEnv ...string) {
// cfg.DataRoot = filepath.Join(root, "data")
// }
var appData []byte
var appFile string
// Read app setting
if has, _ := application.App.Exists("app.yao"); has {
appFile = "app.yao"
appData, err = application.App.Read("app.yao")
if err != nil {
t.Fatal(err)
}
} else if has, _ := application.App.Exists("app.jsonc"); has {
appFile = "app.jsonc"
appData, err = application.App.Read("app.jsonc")
if err != nil {
t.Fatal(err)
}
} else if has, _ := application.App.Exists("app.json"); has {
appFile = "app.json"
appData, err = application.App.Read("app.json")
if err != nil {
t.Fatal(err)
}
} else {
t.Fatal(fmt.Errorf("app.yao or app.jsonc or app.json does not exists"))
}
// Replace $ENV with os.Getenv
var envRe = regexp.MustCompile(`\$ENV\.([0-9a-zA-Z_-]+)`)
appData = envRe.ReplaceAllFunc(appData, func(s []byte) []byte {
key := string(s[5:])
val := os.Getenv(key)
if val == "" {
return s
}
return []byte(val)
})
share.App = share.AppInfo{}
err = application.Parse(appFile, appData, &share.App)
if err != nil {
t.Fatal(err)
}
utils.Init()
dbconnect(t, cfg)
load(t, cfg)