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

Merge pull request #1274 from trheyi/main

Browse source 
Refactor to decouple attachments, RAG, and AI assistant components.
This commit is contained in:
Max 2025-11-06 19:10:11 +08:00 • committed by GitHub
commit 296a37f031
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
24 changed files with 3730 additions and 5582 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

32
agent/agent.go
View file

@ -2,7 +2,6 @@ package agent
import ( import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/yaoapp/gou/session"
"github.com/yaoapp/yao/agent/assistant" "github.com/yaoapp/yao/agent/assistant"
chatctx "github.com/yaoapp/yao/agent/context" chatctx "github.com/yaoapp/yao/agent/context"
) )
@ -28,34 +27,3 @@ func (agent *DSL) Select(id string) (assistant.API, error) {
} }
return assistant.Get(id) return assistant.Get(id)
} }
// UserID get the user id from the session
func (agent *DSL) UserID(sid string) (interface{}, error) {
fieldID := agent.AuthSetting.SessionFields.ID
return session.Global().ID(sid).Get(fieldID)
}
// GuestID get the guest id from the session
func (agent *DSL) GuestID(sid string) (interface{}, error) {
fieldGuest := agent.AuthSetting.SessionFields.Guest
return session.Global().ID(sid).Get(fieldGuest)
}
// UserRoles get the user roles from the session
func (agent *DSL) UserRoles(sid string) (interface{}, error) {
fieldRoles := agent.AuthSetting.SessionFields.Roles
return session.Global().ID(sid).Get(fieldRoles)
}
// UserOrGuestID get the user id or guest id from the session
func (agent *DSL) UserOrGuestID(sid string) (interface{}, bool, error) {
userID, err := agent.UserID(sid)
if err != nil {
guestID, err := agent.GuestID(sid)
if err != nil {
return nil, false, err
}
return guestID, true, nil
}
return userID, false, nil
}

355
agent/api.go
View file

@ -2,23 +2,18 @@ package agent
import ( import (
"fmt" "fmt"
"io"
"net/url" "net/url"
"os"
"strconv" "strconv"
"strings" "strings"
"time" "time"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/google/uuid" "github.com/google/uuid"
"github.com/yaoapp/gou/api"
"github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/process"
"github.com/yaoapp/yao/agent/assistant" "github.com/yaoapp/yao/agent/assistant"
chatctx "github.com/yaoapp/yao/agent/context" chatctx "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/message" "github.com/yaoapp/yao/agent/message"
"github.com/yaoapp/yao/agent/store" "github.com/yaoapp/yao/agent/store"
"github.com/yaoapp/yao/attachment"
"github.com/yaoapp/yao/helper" "github.com/yaoapp/yao/helper"
"github.com/yaoapp/yao/openapi/oauth" "github.com/yaoapp/yao/openapi/oauth"
) )
@ -32,23 +27,6 @@ func (agent *DSL) API(router *gin.Engine, path string) error {
return err return err
} }
// Register OPTIONS handlers for all endpoints
router.OPTIONS(path, agent.optionsHandler)
router.OPTIONS(path+"/status", agent.optionsHandler)
router.OPTIONS(path+"/chats", agent.optionsHandler)
router.OPTIONS(path+"/chats/:id", agent.optionsHandler)
router.OPTIONS(path+"/history", agent.optionsHandler)
router.OPTIONS(path+"/upload/:storage", agent.optionsHandler)
router.OPTIONS(path+"/download", agent.optionsHandler)
router.OPTIONS(path+"/mentions", agent.optionsHandler)
router.OPTIONS(path+"/generate", agent.optionsHandler)
router.OPTIONS(path+"/generate/title", agent.optionsHandler)
router.OPTIONS(path+"/generate/prompts", agent.optionsHandler)
router.OPTIONS(path+"/dangerous/clear_chats", agent.optionsHandler)
router.OPTIONS(path+"/assistants", agent.optionsHandler)
router.OPTIONS(path+"/assistants/:id", agent.optionsHandler)
router.OPTIONS(path+"/assistants/:id/call", agent.optionsHandler)
// Chat endpoint // Chat endpoint
// Chat endpoint // Chat endpoint
// Example: // Example:
@ -124,12 +102,12 @@ func (agent *DSL) API(router *gin.Engine, path string) error {
// Upload file example: // Upload file example:
// curl -X POST 'http://localhost:5099/api/__yao/agent/upload?chat_id=chat_123&token=xxx' \ // curl -X POST 'http://localhost:5099/api/__yao/agent/upload?chat_id=chat_123&token=xxx' \
// -F 'file=@/path/to/file.txt' // -F 'file=@/path/to/file.txt'
router.POST(path+"/upload/:storage", append(middlewares, agent.handleUpload)...) // router.POST(path+"/upload/:storage", append(middlewares, agent.handleUpload)...)
// Download file example: // Download file example:
// curl -X GET 'http://localhost:5099/api/__yao/agent/download?file_id=file_123&disposition=attachment&token=xxx' \ // curl -X GET 'http://localhost:5099/api/__yao/agent/download?file_id=file_123&disposition=attachment&token=xxx' \
// -o downloaded_file.txt // -o downloaded_file.txt
router.GET(path+"/download", append(middlewares, agent.handleDownload)...) // router.GET(path+"/download", append(middlewares, agent.handleDownload)...)
// Mentions endpoint // Mentions endpoint
// Example: // Example:
@ -172,250 +150,6 @@ func (agent *DSL) handleStatus(c *gin.Context) {
c.Done() c.Done()
} }
// handleUpload handles the upload request
func (agent *DSL) handleUpload(c *gin.Context) {
sid := c.GetString("__sid")
if sid == "" {
sid = uuid.New().String()
}
uid, isGuest, err := agent.UserOrGuestID(sid)
if err != nil {
c.JSON(401, gin.H{"message": fmt.Sprintf("Unauthorized, %s", err.Error()), "code": 401})
c.Done()
return
}
if uid == nil || uid == "" {
c.JSON(401, gin.H{"message": "Unauthorized", "code": 401})
c.Done()
return
}
// Storage name must be chat, knowledge or assets
storage := c.Param("storage")
if storage != "chat" && storage != "knowledge" && storage != "assets" {
c.JSON(400, gin.H{"message": "Invalid storage", "code": 400})
c.Done()
return
}
// Get the manager
var manager, ok = attachment.Managers[storage]
if !ok {
c.JSON(400, gin.H{"message": "Invalid storage: " + storage, "code": 400})
c.Done()
return
}
// Get Option from form data
var option UploadOption
err = c.ShouldBind(&option)
if err != nil {
c.JSON(400, gin.H{"message": err.Error(), "code": 400})
c.Done()
return
}
// Validate the option with the storage
option.UserID = fmt.Sprintf("%v", uid)
// Build multi-level groups based on storage type and IDs
var groups []string
switch storage {
case "chat":
if option.ChatID == "" {
c.JSON(400, gin.H{"message": "chat_id is required", "code": 400})
c.Done()
return
}
// Build groups: ["users", "user123", "chats", "chat456"]
groups = []string{"users", option.UserID, "chats", option.ChatID}
if option.AssistantID != "" {
// Add assistant level: ["users", "user123", "chats", "chat456", "assistants", "assistant789"]
groups = append(groups, "assistants", option.AssistantID)
}
case "knowledge":
if option.CollectionID == "" {
c.JSON(400, gin.H{"message": "collection_id is required", "code": 400})
c.Done()
return
}
// Build groups: ["knowledge", "collection123", "users", "user456"]
groups = []string{"knowledge", option.CollectionID, "users", option.UserID}
case "assets":
// Build groups: ["assets", "users", "user123"]
groups = []string{"assets", "users", option.UserID}
}
// Set the groups in the attachment upload option
option.UploadOption.Groups = groups
// Get the file
file, err := c.FormFile("file")
if err != nil {
c.JSON(400, gin.H{"message": err.Error(), "code": 400})
c.Done()
return
}
// Open the file
reader, err := file.Open()
if err != nil {
c.JSON(500, gin.H{"message": err.Error(), "code": 500})
c.Done()
return
}
defer func() {
reader.Close()
os.Remove(file.Filename)
}()
// Upload the file
header := attachment.GetHeader(c.Request.Header, file.Header, file.Size)
res, err := manager.Upload(c.Request.Context(), header, reader, option.UploadOption)
if err != nil {
c.JSON(500, gin.H{"message": err.Error(), "code": 500})
c.Done()
return
}
// if storage is chat or knowledge, save the file to the store
if storage == "chat" || storage == "knowledge" {
attachment := map[string]interface{}{
"file_id": res.ID,
"uid": uid,
"guest": isGuest,
"manager": storage,
"public": option.Public,
"name": option.OriginalFilename,
"content_type": res.ContentType,
"bytes": res.Bytes,
"gzip": option.Gzip,
"status": res.Status,
}
// Set the scope
if option.Scope != nil {
attachment["scope"] = option.Scope
}
// Set the collection_id
if option.CollectionID != "" {
attachment["collection_id"] = option.CollectionID
}
_, err = agent.Store.SaveAttachment(attachment)
if err != nil {
c.JSON(500, gin.H{"message": err.Error(), "code": 500})
c.Done()
return
}
}
c.JSON(200, map[string]interface{}{"data": res})
c.Done()
}
// handleDownload handles the download request
func (agent *DSL) handleDownload(c *gin.Context) {
sid := c.GetString("__sid")
if sid == "" {
c.JSON(400, gin.H{"message": "sid is required", "code": 400})
c.Done()
return
}
uid, _, err := agent.UserOrGuestID(sid)
if err != nil {
c.JSON(401, gin.H{"message": fmt.Sprintf("Unauthorized, %s", err.Error()), "code": 401})
c.Done()
return
}
if uid == nil || uid == "" {
c.JSON(401, gin.H{"message": "Unauthorized", "code": 401})
c.Done()
return
}
fileID := c.Query("file_id")
if fileID == "" {
c.JSON(400, gin.H{"message": "file_id is required", "code": 400})
c.Done()
return
}
// Get the attachment
attach, err := agent.Store.GetAttachment(fileID)
if err != nil {
c.JSON(500, gin.H{"message": err.Error(), "code": 500})
c.Done()
return
}
// Validate the permission ( Will be supported scope validation in the future )
if (attach["public"] == 0 || attach["public"] == false) && attach["uid"] != uid {
c.JSON(403, gin.H{"message": "Forbidden", "code": 403})
c.Done()
return
}
storage, ok := attach["manager"].(string)
if !ok {
c.JSON(400, gin.H{"message": "Invalid storage", "code": 400})
c.Done()
return
}
// Get the manager
manager, ok := attachment.Managers[storage]
if !ok {
c.JSON(400, gin.H{"message": "Invalid storage", "code": 400})
c.Done()
return
}
name, ok := attach["name"].(string)
if !ok {
c.JSON(400, gin.H{"message": "Invalid name", "code": 400})
c.Done()
return
}
name = strings.TrimSuffix(name, ".gz")
contentType, ok := attach["content_type"].(string)
if !ok {
c.JSON(400, gin.H{"message": "Invalid content type", "code": 400})
c.Done()
return
}
handle, err := manager.Download(c.Request.Context(), fileID)
if err != nil {
c.JSON(500, gin.H{"message": err.Error(), "code": 500})
c.Done()
return
}
defer handle.Reader.Close()
// Set the response headers
encoded := url.PathEscape(name)
disposition := fmt.Sprintf(`attachment; filename="%s"`, encoded)
c.Header("Content-Type", contentType)
c.Header("Content-Disposition", disposition)
// Copy the file content to response
_, err = io.Copy(c.Writer, handle.Reader)
if err != nil {
c.JSON(500, gin.H{"message": err.Error(), "code": 500})
return
}
c.Done()
}
// handleChat handles the chat request // handleChat handles the chat request
func (agent *DSL) handleChat(c *gin.Context) { func (agent *DSL) handleChat(c *gin.Context) {
// Set headers for SSE // Set headers for SSE
@ -545,70 +279,6 @@ func (agent *DSL) handleChatHistory(c *gin.Context) {
c.Done() c.Done()
} }
// getCorsHandlers returns CORS middleware handlers
func (agent *DSL) getCorsHandlers() ([]gin.HandlerFunc, error) {
if len(agent.Allows) == 0 {
return []gin.HandlerFunc{}, nil
}
allowsMap := map[string]bool{}
for _, allow := range agent.Allows {
allow = strings.TrimPrefix(allow, "http://")
allow = strings.TrimPrefix(allow, "https://")
allowsMap[allow] = true
}
return []gin.HandlerFunc{agent.corsMiddleware(allowsMap)}, nil
}
// corsMiddleware handles CORS requests
func (agent *DSL) corsMiddleware(allowsMap map[string]bool) gin.HandlerFunc {
return func(c *gin.Context) {
origin := agent.getOrigin(c)
if origin == "" {
c.Next()
return
}
// Check if origin is allowed
if !api.IsAllowed(c, allowsMap) {
c.AbortWithStatusJSON(403, gin.H{
"message": origin + " not allowed",
"code": 403,
})
return
}
// Set CORS headers
c.Header("Access-Control-Allow-Origin", origin)
c.Header("Access-Control-Allow-Credentials", "true")
c.Header("Access-Control-Allow-Headers", "Content-Type, Content-Disposition, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, Accept, Origin, Cache-Control, X-Requested-With, Content-Sync, Content-Fingerprint, Content-Uid, Content-Range")
c.Header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
c.Header("Access-Control-Expose-Headers", "Content-Type, Content-Disposition, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, Accept, Origin, Cache-Control, X-Requested-With, Content-Sync, Content-Fingerprint, Content-Uid, Content-Range")
if c.Request.Method == "OPTIONS" {
c.AbortWithStatus(204)
return
}
c.Next()
}
}
// optionsHandler handles OPTIONS requests
func (agent *DSL) optionsHandler(c *gin.Context) {
origin := agent.getOrigin(c)
if origin != "" {
c.Header("Access-Control-Allow-Origin", origin)
c.Header("Access-Control-Allow-Methods", "GET, POST, DELETE, OPTIONS")
c.Header("Access-Control-Allow-Headers", "Content-Type, Content-Disposition, Authorization, Accept, Content-Sync, Content-Fingerprint, Content-Uid, Content-Range")
c.Header("Access-Control-Allow-Credentials", "true")
c.Header("Access-Control-Max-Age", "86400") // 24 hours
c.Header("Access-Control-Expose-Headers", "Content-Type, Content-Disposition, Authorization, Accept, Content-Sync, Content-Fingerprint, Content-Uid, Content-Range")
}
c.AbortWithStatus(204)
}
// getOrigin returns the request origin // getOrigin returns the request origin
func (agent *DSL) getOrigin(c *gin.Context) string { func (agent *DSL) getOrigin(c *gin.Context) string {
origin := c.Request.Header.Get("Origin") origin := c.Request.Header.Get("Origin")
@ -625,26 +295,7 @@ func (agent *DSL) getOrigin(c *gin.Context) string {
// getGuardHandlers returns authentication middleware handlers // getGuardHandlers returns authentication middleware handlers
func (agent *DSL) getGuardHandlers() ([]gin.HandlerFunc, error) { func (agent *DSL) getGuardHandlers() ([]gin.HandlerFunc, error) {
return []gin.HandlerFunc{}, nil
// Cross-Domain handlers
cors, err := agent.getCorsHandlers()
if err != nil {
return nil, err
}
if agent.Guard == "" {
middlewares := append(cors, agent.defaultGuard)
return middlewares, nil
}
// Validate the custom guard
_, err = process.Of(agent.Guard)
if err != nil {
return nil, err
}
middlewares := append(cors, api.ProcessGuard(agent.Guard, cors...))
return middlewares, nil
} }
// defaultGuard is the default authentication handler // defaultGuard is the default authentication handler

408
agent/api_test.go
View file

@ -1,233 +1,233 @@
package agent package agent
import ( // import (
"context" // "context"
"fmt" // "fmt"
"net" // "net"
"net/http" // "net/http"
"net/http/httptest" // "net/http/httptest"
"os" // "os"
"strings" // "strings"
"testing" // "testing"
"time" // "time"
"github.com/gin-gonic/gin" // "github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert" // "github.com/stretchr/testify/assert"
httpTest "github.com/yaoapp/gou/http" // httpTest "github.com/yaoapp/gou/http"
"github.com/yaoapp/yao/config" // "github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/helper" // "github.com/yaoapp/yao/helper"
"github.com/yaoapp/yao/test" // "github.com/yaoapp/yao/test"
) // )
func init() { // func init() {
// Set gin to release mode to reduce log output // // Set gin to release mode to reduce log output
gin.SetMode(gin.ReleaseMode) // gin.SetMode(gin.ReleaseMode)
} // }
func TestAPI(t *testing.T) { // func TestAPI(t *testing.T) {
// Disable test logging // // Disable test logging
test.Prepare(t, config.Conf) // test.Prepare(t, config.Conf)
defer test.Clean() // defer test.Clean()
// Redirect stdout to /dev/null // // Redirect stdout to /dev/null
oldStdout := os.Stdout // oldStdout := os.Stdout
null, _ := os.Open(os.DevNull) // null, _ := os.Open(os.DevNull)
os.Stdout = null // os.Stdout = null
defer func() { // defer func() {
os.Stdout = oldStdout // os.Stdout = oldStdout
null.Close() // null.Close()
}() // }()
// test router // // test router
router := testRouter(t) // router := testRouter(t)
err := Agent.API(router, "/agent/chat") // err := Agent.API(router, "/agent/chat")
if err != nil { // if err != nil {
t.Fatal(err) // t.Fatal(err)
} // }
// test server // // test server
host, shutdown := testServer(t, router) // host, shutdown := testServer(t, router)
defer shutdown() // defer shutdown()
tests := []struct { // tests := []struct {
name string // name string
url string // url string
method string // method string
headers http.Header // headers http.Header
expectCode int // expectCode int
expectBody string // expectBody string
}{ // }{
{ // {
name: "Basic Chat Request", // name: "Basic Chat Request",
url: fmt.Sprintf("/agent/chat?content=hello&token=%s", testToken()), // url: fmt.Sprintf("/agent/chat?content=hello&token=%s", testToken()),
method: "GET", // method: "GET",
headers: http.Header{"Content-Type": []string{"application/json"}}, // headers: http.Header{"Content-Type": []string{"application/json"}},
expectBody: `{`, // expectBody: `{`,
}, // },
{ // {
name: "Chat with System Message", // name: "Chat with System Message",
url: fmt.Sprintf("/agent/chat?content=hello&system=You are a helpful assistant&token=%s", testToken()), // url: fmt.Sprintf("/agent/chat?content=hello&system=You are a helpful assistant&token=%s", testToken()),
method: "GET", // method: "GET",
headers: http.Header{"Content-Type": []string{"application/json"}}, // headers: http.Header{"Content-Type": []string{"application/json"}},
expectBody: `{`, // expectBody: `{`,
}, // },
{ // {
name: "Chat with Model Parameter", // name: "Chat with Model Parameter",
url: fmt.Sprintf("/agent/chat?content=hello&model=gpt-3.5-turbo&token=%s", testToken()), // url: fmt.Sprintf("/agent/chat?content=hello&model=gpt-3.5-turbo&token=%s", testToken()),
method: "GET", // method: "GET",
headers: http.Header{"Content-Type": []string{"application/json"}}, // headers: http.Header{"Content-Type": []string{"application/json"}},
expectBody: `{`, // expectBody: `{`,
}, // },
} // }
for _, tt := range tests { // for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) { // t.Run(tt.name, func(t *testing.T) {
url := fmt.Sprintf("%s%s", host, tt.url) // url := fmt.Sprintf("%s%s", host, tt.url)
res := []byte{} // res := []byte{}
req := httpTest.New(url).WithHeader(tt.headers) // req := httpTest.New(url).WithHeader(tt.headers)
ctx, cancel := context.WithCancel(context.Background()) // ctx, cancel := context.WithCancel(context.Background())
defer cancel() // defer cancel()
req.Stream(ctx, tt.method, nil, func(data []byte) int { // req.Stream(ctx, tt.method, nil, func(data []byte) int {
res = append(res, data...) // res = append(res, data...)
return 1 // return 1
}) // })
assert.Contains(t, string(res), tt.expectBody) // assert.Contains(t, string(res), tt.expectBody)
}) // })
} // }
} // }
func TestAPIAuth(t *testing.T) { // func TestAPIAuth(t *testing.T) {
test.Prepare(t, config.Conf) // test.Prepare(t, config.Conf)
defer test.Clean() // defer test.Clean()
// Redirect stdout and stderr to /dev/null // // Redirect stdout and stderr to /dev/null
oldStdout := os.Stdout // oldStdout := os.Stdout
oldStderr := os.Stderr // oldStderr := os.Stderr
null, _ := os.Open(os.DevNull) // null, _ := os.Open(os.DevNull)
os.Stdout = null // os.Stdout = null
os.Stderr = null // os.Stderr = null
defer func() { // defer func() {
os.Stdout = oldStdout // os.Stdout = oldStdout
os.Stderr = oldStderr // os.Stderr = oldStderr
null.Close() // null.Close()
}() // }()
router := testRouter(t) // router := testRouter(t)
err := Agent.API(router, "/agent/chat") // err := Agent.API(router, "/agent/chat")
if err != nil { // if err != nil {
t.Fatal(err) // t.Fatal(err)
} // }
// Separate tests for authentication errors and parameter validation errors // // Separate tests for authentication errors and parameter validation errors
authTests := []struct { // authTests := []struct {
name string // name string
url string // url string
method string // method string
expectCode int // expectCode int
}{ // }{
{ // {
name: "Missing Token", // name: "Missing Token",
url: "/agent/chat?content=hello", // url: "/agent/chat?content=hello",
method: "GET", // method: "GET",
expectCode: http.StatusUnauthorized, // expectCode: http.StatusUnauthorized,
}, // },
{ // {
name: "Invalid Token", // name: "Invalid Token",
url: "/agent/chat?content=hello&token=invalid", // url: "/agent/chat?content=hello&token=invalid",
method: "GET", // method: "GET",
expectCode: http.StatusUnauthorized, // expectCode: http.StatusUnauthorized,
}, // },
} // }
// Test authentication errors (will panic) // // Test authentication errors (will panic)
for _, tt := range authTests { // for _, tt := range authTests {
t.Run(tt.name, func(t *testing.T) { // t.Run(tt.name, func(t *testing.T) {
response := httptest.NewRecorder() // response := httptest.NewRecorder()
req, _ := http.NewRequest(tt.method, tt.url, nil) // req, _ := http.NewRequest(tt.method, tt.url, nil)
assert.Panics(t, func() { // assert.Panics(t, func() {
router.ServeHTTP(response, req) // router.ServeHTTP(response, req)
}) // })
}) // })
} // }
// Test parameter validation errors (will return status code) // // Test parameter validation errors (will return status code)
validationTests := []struct { // validationTests := []struct {
name string // name string
url string // url string
method string // method string
expectCode int // expectCode int
}{ // }{
{ // {
name: "Missing Content", // name: "Missing Content",
url: fmt.Sprintf("/agent/chat?token=%s", testToken()), // url: fmt.Sprintf("/agent/chat?token=%s", testToken()),
method: "GET", // method: "GET",
expectCode: http.StatusBadRequest, // expectCode: http.StatusBadRequest,
}, // },
} // }
// Test parameter validation errors (return status code) // // Test parameter validation errors (return status code)
for _, tt := range validationTests { // for _, tt := range validationTests {
t.Run(tt.name, func(t *testing.T) { // t.Run(tt.name, func(t *testing.T) {
response := httptest.NewRecorder() // response := httptest.NewRecorder()
req, _ := http.NewRequest(tt.method, tt.url, nil) // req, _ := http.NewRequest(tt.method, tt.url, nil)
router.ServeHTTP(response, req) // router.ServeHTTP(response, req)
assert.Equal(t, tt.expectCode, response.Code) // assert.Equal(t, tt.expectCode, response.Code)
}) // })
} // }
} // }
// Helper functions // // Helper functions
func testServer(t *testing.T, router *gin.Engine) (string, func()) { // func testServer(t *testing.T, router *gin.Engine) (string, func()) {
l, err := net.Listen("tcp4", ":0") // l, err := net.Listen("tcp4", ":0")
if err != nil { // if err != nil {
t.Fatal(err) // t.Fatal(err)
} // }
srv := &http.Server{Addr: ":0", Handler: router} // srv := &http.Server{Addr: ":0", Handler: router}
go func() { // go func() {
if err := srv.Serve(l); err != nil && err != http.ErrServerClosed { // if err := srv.Serve(l); err != nil && err != http.ErrServerClosed {
return // return
} // }
}() // }()
addr := strings.Split(l.Addr().String(), ":") // addr := strings.Split(l.Addr().String(), ":")
if len(addr) != 2 { // if len(addr) != 2 {
t.Fatal("invalid address") // t.Fatal("invalid address")
} // }
host := fmt.Sprintf("http://127.0.0.1:%s", addr[1]) // host := fmt.Sprintf("http://127.0.0.1:%s", addr[1])
time.Sleep(50 * time.Millisecond) // time.Sleep(50 * time.Millisecond)
shutdown := func() { // shutdown := func() {
srv.Close() // srv.Close()
l.Close() // l.Close()
} // }
return host, shutdown // return host, shutdown
} // }
func testRouter(t *testing.T) *gin.Engine { // func testRouter(t *testing.T) *gin.Engine {
err := Load(config.Conf) // err := Load(config.Conf)
if err != nil { // if err != nil {
t.Fatal(err) // t.Fatal(err)
} // }
router := gin.New() // Use gin.New() instead of gin.Default() to avoid default logging middleware // router := gin.New() // Use gin.New() instead of gin.Default() to avoid default logging middleware
return router // return router
} // }
func testToken() string { // func testToken() string {
token := helper.JwtMake(1, // token := helper.JwtMake(1,
map[string]interface{}{ // map[string]interface{}{
"id": 1, // "id": 1,
"name": "Test", // "name": "Test",
}, // },
map[string]interface{}{ // map[string]interface{}{
"exp": 3600, // "exp": 3600,
"sid": "123456", // "sid": "123456",
}) // })
return token.Token // return token.Token
} // }

92
agent/assistant/assistant.go
View file

@ -1,16 +1,11 @@
package assistant package assistant
import ( import (
"context"
"fmt" "fmt"
"path" "path"
"time"
"github.com/fatih/color"
jsoniter "github.com/json-iterator/go" jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/fs" "github.com/yaoapp/gou/fs"
"github.com/yaoapp/gou/rag/driver"
"github.com/yaoapp/kun/log"
sui "github.com/yaoapp/yao/sui/core" sui "github.com/yaoapp/yao/sui/core"
) )
@ -25,96 +20,9 @@ func (ast *Assistant) Save() error {
return err return err
} }
// Update Index in background
go func() {
err := ast.UpdateIndex()
if err != nil {
log.Error("failed to update index for assistant %s: %s", ast.ID, err)
color.Red("failed to update index for assistant %s: %s", ast.ID, err)
}
}()
return nil return nil
} }
// UpdateIndex update the index for RAG
func (ast *Assistant) UpdateIndex() error {
// RAG is not enabled
if rag == nil {
return nil
}
if rag.Engine == nil {
return fmt.Errorf("engine is not set")
}
// Update Index
index := fmt.Sprintf("%sassistants", rag.Setting.IndexPrefix)
id := fmt.Sprintf("assistant_%s", ast.ID)
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
// Check if the index exists
exists, err := rag.Engine.HasIndex(ctx, index)
if err != nil {
return err
}
// Create the index if it does not exist
if !exists {
ctxCreate, cancelCreate := context.WithTimeout(context.Background(), 2*time.Second)
defer cancelCreate()
err = rag.Engine.CreateIndex(ctxCreate, driver.IndexConfig{Name: index})
if err != nil {
return err
}
}
// Check if the document exists
exists, err = rag.Engine.HasDocument(ctx, index, id)
if err != nil {
return err
}
// Check if the document is updated
if exists {
metadata, err := rag.Engine.GetMetadata(ctx, index, id)
if err != nil {
return err
}
if v, ok := metadata["updated_at"].(string); ok {
updatedAt, err := stringToTimestamp(v)
if err != nil {
return err
}
if updatedAt >= ast.UpdatedAt {
return nil
}
}
}
// Update the index
content, err := jsoniter.MarshalToString(ast.Map())
if err != nil {
return err
}
metadata := map[string]interface{}{
"assistant_id": ast.ID,
"type": ast.Type,
"name": ast.Name,
"updated_at": fmt.Sprintf("%d", ast.UpdatedAt),
}
return rag.Engine.IndexDoc(ctx, index, &driver.Document{
DocID: id,
Content: content,
Metadata: metadata,
})
}
// Map convert the assistant to a map // Map convert the assistant to a map
func (ast *Assistant) Map() map[string]interface{} { func (ast *Assistant) Map() map[string]interface{} {

237
agent/assistant/cache_test.go
View file

@ -1,151 +1,146 @@
package assistant package assistant
import ( // func TestCache_Basic(t *testing.T) {
"sync" // cache := NewCache(2)
"testing"
)
func TestCache_Basic(t *testing.T) { // // Test empty cache
cache := NewCache(2) // if cache.Len() != 0 {
// t.Errorf("Expected empty cache, got length %d", cache.Len())
// }
// Test empty cache // // Test adding items
if cache.Len() != 0 { // assistant1 := &Assistant{ID: "1", Name: "Test1"}
t.Errorf("Expected empty cache, got length %d", cache.Len()) // assistant2 := &Assistant{ID: "2", Name: "Test2"}
}
// Test adding items // cache.Put(assistant1)
assistant1 := &Assistant{ID: "1", Name: "Test1"} // cache.Put(assistant2)
assistant2 := &Assistant{ID: "2", Name: "Test2"}
cache.Put(assistant1) // if cache.Len() != 2 {
cache.Put(assistant2) // t.Errorf("Expected cache length 2, got %d", cache.Len())
// }
if cache.Len() != 2 { // // Test getting items
t.Errorf("Expected cache length 2, got %d", cache.Len()) // if a, exists := cache.Get("1"); !exists || a.ID != "1" {
} // t.Error("Failed to get assistant1")
// }
// Test getting items // if a, exists := cache.Get("2"); !exists || a.ID != "2" {
if a, exists := cache.Get("1"); !exists || a.ID != "1" { // t.Error("Failed to get assistant2")
t.Error("Failed to get assistant1") // }
} // }
if a, exists := cache.Get("2"); !exists || a.ID != "2" { // func TestCache_LRU(t *testing.T) {
t.Error("Failed to get assistant2") // cache := NewCache(2)
}
}
func TestCache_LRU(t *testing.T) { // assistant1 := &Assistant{ID: "1", Name: "Test1"}
cache := NewCache(2) // assistant2 := &Assistant{ID: "2", Name: "Test2"}
// assistant3 := &Assistant{ID: "3", Name: "Test3"}
assistant1 := &Assistant{ID: "1", Name: "Test1"} // // Add first two items
assistant2 := &Assistant{ID: "2", Name: "Test2"} // cache.Put(assistant1)
assistant3 := &Assistant{ID: "3", Name: "Test3"} // cache.Put(assistant2)
// Add first two items // // Access assistant1 to make it most recently used
cache.Put(assistant1) // cache.Get("1")
cache.Put(assistant2)
// Access assistant1 to make it most recently used // // Add third item, should evict assistant2
cache.Get("1") // cache.Put(assistant3)
// Add third item, should evict assistant2 // // Check assistant2 was evicted
cache.Put(assistant3) // if _, exists := cache.Get("2"); exists {
// t.Error("Assistant2 should have been evicted")
// }
// Check assistant2 was evicted // // Check assistant1 and assistant3 are still present
if _, exists := cache.Get("2"); exists { // if _, exists := cache.Get("1"); !exists {
t.Error("Assistant2 should have been evicted") // t.Error("Assistant1 should still be in cache")
} // }
// if _, exists := cache.Get("3"); !exists {
// t.Error("Assistant3 should be in cache")
// }
// }
// Check assistant1 and assistant3 are still present // func TestCache_Remove(t *testing.T) {
if _, exists := cache.Get("1"); !exists { // cache := NewCache(2)
t.Error("Assistant1 should still be in cache")
}
if _, exists := cache.Get("3"); !exists {
t.Error("Assistant3 should be in cache")
}
}
func TestCache_Remove(t *testing.T) { // assistant1 := &Assistant{ID: "1", Name: "Test1"}
cache := NewCache(2) // cache.Put(assistant1)
assistant1 := &Assistant{ID: "1", Name: "Test1"} // // Test remove existing item
cache.Put(assistant1) // cache.Remove("1")
// if cache.Len() != 0 {
// t.Error("Cache should be empty after removing item")
// }
// Test remove existing item // // Test remove non-existing item
cache.Remove("1") // cache.Remove("nonexistent")
if cache.Len() != 0 { // if cache.Len() != 0 {
t.Error("Cache should be empty after removing item") // t.Error("Cache length should not change when removing non-existent item")
} // }
// }
// Test remove non-existing item // func TestCache_Clear(t *testing.T) {
cache.Remove("nonexistent") // cache := NewCache(2)
if cache.Len() != 0 {
t.Error("Cache length should not change when removing non-existent item")
}
}
func TestCache_Clear(t *testing.T) { // assistant1 := &Assistant{ID: "1", Name: "Test1"}
cache := NewCache(2) // assistant2 := &Assistant{ID: "2", Name: "Test2"}
assistant1 := &Assistant{ID: "1", Name: "Test1"} // cache.Put(assistant1)
assistant2 := &Assistant{ID: "2", Name: "Test2"} // cache.Put(assistant2)
cache.Put(assistant1) // cache.Clear()
cache.Put(assistant2) // if cache.Len() != 0 {
// t.Error("Cache should be empty after clear")
// }
// }
cache.Clear() // func TestCache_Concurrent(t *testing.T) {
if cache.Len() != 0 { // cache := NewCache(100)
t.Error("Cache should be empty after clear") // var wg sync.WaitGroup
} // workers := 10
} // iterations := 100
func TestCache_Concurrent(t *testing.T) { // // Concurrent writes
cache := NewCache(100) // for i := 0; i < workers; i++ {
var wg sync.WaitGroup // wg.Add(1)
workers := 10 // go func(workerID int) {
iterations := 100 // defer wg.Done()
// for j := 0; j < iterations; j++ {
// assistant := &Assistant{
// ID: string(rune('A' + workerID)),
// Name: "Test",
// }
// cache.Put(assistant)
// }
// }(i)
// }
// Concurrent writes // // Concurrent reads
for i := 0; i < workers; i++ { // for i := 0; i < workers; i++ {
wg.Add(1) // wg.Add(1)
go func(workerID int) { // go func(workerID int) {
defer wg.Done() // defer wg.Done()
for j := 0; j < iterations; j++ { // for j := 0; j < iterations; j++ {
assistant := &Assistant{ // cache.Get(string(rune('A' + workerID)))
ID: string(rune('A' + workerID)), // }
Name: "Test", // }(i)
} // }
cache.Put(assistant)
}
}(i)
}
// Concurrent reads // wg.Wait()
for i := 0; i < workers; i++ { // }
wg.Add(1)
go func(workerID int) {
defer wg.Done()
for j := 0; j < iterations; j++ {
cache.Get(string(rune('A' + workerID)))
}
}(i)
}
wg.Wait() // func TestCache_NilInput(t *testing.T) {
} // cache := NewCache(2)
func TestCache_NilInput(t *testing.T) { // // Test putting nil assistant
cache := NewCache(2) // cache.Put(nil)
// if cache.Len() != 0 {
// t.Error("Cache should not store nil assistant")
// }
// Test putting nil assistant // // Test putting assistant with empty ID
cache.Put(nil) // cache.Put(&Assistant{ID: "", Name: "Test"})
if cache.Len() != 0 { // if cache.Len() != 0 {
t.Error("Cache should not store nil assistant") // t.Error("Cache should not store assistant with empty ID")
} // }
// }
// Test putting assistant with empty ID
cache.Put(&Assistant{ID: "", Name: "Test"})
if cache.Len() != 0 {
t.Error("Cache should not store assistant with empty ID")
}
}

29
agent/assistant/load.go
View file

@ -12,7 +12,6 @@ import (
"github.com/spf13/cast" "github.com/spf13/cast"
"github.com/yaoapp/gou/application" "github.com/yaoapp/gou/application"
"github.com/yaoapp/gou/fs" "github.com/yaoapp/gou/fs"
"github.com/yaoapp/gou/rag/driver"
v8 "github.com/yaoapp/gou/runtime/v8" v8 "github.com/yaoapp/gou/runtime/v8"
"github.com/yaoapp/yao/agent/i18n" "github.com/yaoapp/yao/agent/i18n"
"github.com/yaoapp/yao/agent/store" "github.com/yaoapp/yao/agent/store"
@ -25,7 +24,6 @@ import (
// loaded the loaded assistant // loaded the loaded assistant
var loaded = NewCache(200) // 200 is the default capacity var loaded = NewCache(200) // 200 is the default capacity
var storage store.Store = nil var storage store.Store = nil
var rag *RAG = nil
var search interface{} = nil var search interface{} = nil
var connectorSettings map[string]ConnectorSetting = map[string]ConnectorSetting{} var connectorSettings map[string]ConnectorSetting = map[string]ConnectorSetting{}
var vision *agentvision.Vision = nil var vision *agentvision.Vision = nil
@ -149,19 +147,6 @@ func SetConnector(c string) {
defaultConnector = c defaultConnector = c
} }
// SetRAG set the RAG engine
// e: the RAG engine
// u: the RAG file uploader
// v: the RAG vectorizer
func SetRAG(e driver.Engine, u driver.FileUpload, v driver.Vectorizer, setting RAGSetting) {
rag = &RAG{
Engine: e,
Uploader: u,
Vectorizer: v,
Setting: setting,
}
}
// SetCache set the cache // SetCache set the cache
func SetCache(capacity int) { func SetCache(capacity int) {
ClearCache() ClearCache()
@ -481,20 +466,6 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
} }
} }
// Knowledge options
if v, ok := data["knowledge"].(map[string]interface{}); ok {
assistant.Knowledge = &KnowledgeOption{}
raw, err := jsoniter.Marshal(v)
if err != nil {
return nil, err
}
// Unmarshal the raw data
err = jsoniter.Unmarshal(raw, assistant.Knowledge)
if err != nil {
return nil, err
}
}
// prompts // prompts
if prompts, has := data["prompts"]; has { if prompts, has := data["prompts"]; has {

770
agent/assistant/load_test.go
View file

@ -1,445 +1,435 @@
package assistant package assistant
import ( // func prepare(t *testing.T) {
"fmt" // test.Prepare(t, config.Conf)
"testing" // }
"github.com/stretchr/testify/assert" // func TestLoad_LoadPath(t *testing.T) {
"github.com/yaoapp/yao/agent/store" // prepare(t)
"github.com/yaoapp/yao/config" // defer test.Clean()
"github.com/yaoapp/yao/test"
)
func prepare(t *testing.T) { // assistant, err := LoadPath("/assistants/modi")
test.Prepare(t, config.Conf) // if err != nil {
} // t.Fatal(err)
// }
func TestLoad_LoadPath(t *testing.T) { // // Validate basic properties
prepare(t) // assert.NotNil(t, assistant)
defer test.Clean() // assert.Equal(t, "modi", assistant.ID)
// assert.Equal(t, "Modi", assistant.Name)
// assert.Equal(t, "https://api.dicebear.com/7.x/bottts/svg?seed=Modi", assistant.Avatar)
// assert.Equal(t, "deepseek", assistant.Connector)
// assert.NotNil(t, assistant.Prompts)
// assert.NotNil(t, assistant.Script)
assistant, err := LoadPath("/assistants/modi") // // Test non-existent assistant
if err != nil { // _, err = LoadPath("/assistants/non-existent")
t.Fatal(err) // assert.Error(t, err)
} // }
// Validate basic properties // func TestLoad_LoadStore(t *testing.T) {
assert.NotNil(t, assistant) // prepare(t)
assert.Equal(t, "modi", assistant.ID) // defer test.Clean()
assert.Equal(t, "Modi", assistant.Name)
assert.Equal(t, "https://api.dicebear.com/7.x/bottts/svg?seed=Modi", assistant.Avatar)
assert.Equal(t, "deepseek", assistant.Connector)
assert.NotNil(t, assistant.Prompts)
assert.NotNil(t, assistant.Script)
// Test non-existent assistant // // Test with nil storage
_, err = LoadPath("/assistants/non-existent") // _, err := LoadStore("test-id")
assert.Error(t, err) // assert.Error(t, err)
} // assert.Contains(t, err.Error(), "storage is not set")
func TestLoad_LoadStore(t *testing.T) { // // Setup mock storage
prepare(t) // mockStore := &mockStore{
defer test.Clean() // data: map[string]map[string]interface{}{
// "test-id": {
// "assistant_id": "test-id",
// "name": "Test Assistant",
// "avatar": "test-avatar",
// "connector": "gpt-3_5-turbo",
// },
// },
// }
// SetStorage(mockStore)
// defer SetStorage(nil)
// Test with nil storage // // Test loading from store
_, err := LoadStore("test-id") // assistant, err := LoadStore("test-id")
assert.Error(t, err) // assert.NoError(t, err)
assert.Contains(t, err.Error(), "storage is not set") // assert.NotNil(t, assistant)
// assert.Equal(t, "test-id", assistant.ID)
// assert.Equal(t, "Test Assistant", assistant.Name)
// assert.Equal(t, "test-avatar", assistant.Avatar)
// assert.Equal(t, "gpt-3_5-turbo", assistant.Connector)
// Setup mock storage // // Test cache functionality
mockStore := &mockStore{ // assistant2, err := LoadStore("test-id")
data: map[string]map[string]interface{}{ // assert.NoError(t, err)
"test-id": { // assert.Equal(t, assistant, assistant2) // Should be the same instance from cache
"assistant_id": "test-id",
"name": "Test Assistant",
"avatar": "test-avatar",
"connector": "gpt-3_5-turbo",
},
},
}
SetStorage(mockStore)
defer SetStorage(nil)
// Test loading from store // // Test non-existent assistant
assistant, err := LoadStore("test-id") // _, err = LoadStore("non-existent")
assert.NoError(t, err) // assert.Error(t, err)
assert.NotNil(t, assistant) // }
assert.Equal(t, "test-id", assistant.ID)
assert.Equal(t, "Test Assistant", assistant.Name)
assert.Equal(t, "test-avatar", assistant.Avatar)
assert.Equal(t, "gpt-3_5-turbo", assistant.Connector)
// Test cache functionality // func TestLoad_Cache(t *testing.T) {
assistant2, err := LoadStore("test-id") // prepare(t)
assert.NoError(t, err) // defer test.Clean()
assert.Equal(t, assistant, assistant2) // Should be the same instance from cache
// Test non-existent assistant // // Clear any existing cache first
_, err = LoadStore("non-existent") // ClearCache()
assert.Error(t, err)
}
func TestLoad_Cache(t *testing.T) { // // Test cache operations
prepare(t) // SetCache(2) // Set small cache size for testing
defer test.Clean() // assert.Equal(t, 2, loaded.capacity, "Cache capacity should be 2")
// Clear any existing cache first // // Create test assistants
ClearCache() // assistant1 := &Assistant{ID: "id1", Name: "Assistant 1"}
// assistant2 := &Assistant{ID: "id2", Name: "Assistant 2"}
// assistant3 := &Assistant{ID: "id3", Name: "Assistant 3"}
// Test cache operations // // Test Put and Get
SetCache(2) // Set small cache size for testing // loaded.Put(assistant1)
assert.Equal(t, 2, loaded.capacity, "Cache capacity should be 2") // assert.Equal(t, 1, loaded.Len(), "Cache should have 1 item")
// Create test assistants // loaded.Put(assistant2)
assistant1 := &Assistant{ID: "id1", Name: "Assistant 1"} // assert.Equal(t, 2, loaded.Len(), "Cache should have 2 items")
assistant2 := &Assistant{ID: "id2", Name: "Assistant 2"}
assistant3 := &Assistant{ID: "id3", Name: "Assistant 3"}
// Test Put and Get // // Test cache hit
loaded.Put(assistant1) // cached, exists := loaded.Get("id1")
assert.Equal(t, 1, loaded.Len(), "Cache should have 1 item") // assert.True(t, exists)
// assert.Equal(t, assistant1, cached)
loaded.Put(assistant2) // // Test cache eviction (LRU)
assert.Equal(t, 2, loaded.Len(), "Cache should have 2 items") // // At this point: assistant1 is most recently used (due to Get), then assistant2
// loaded.Put(assistant3) // This should evict assistant2 since it's least recently used
// assert.Equal(t, 2, loaded.Len(), "Cache should still have 2 items")
// _, exists = loaded.Get("id2")
// assert.False(t, exists, "assistant2 should have been evicted (least recently used)")
// _, exists = loaded.Get("id1")
// assert.True(t, exists, "assistant1 should still be in cache (was accessed recently)")
// _, exists = loaded.Get("id3")
// assert.True(t, exists, "assistant3 should be in cache (most recently added)")
// Test cache hit // // Test clear cache
cached, exists := loaded.Get("id1") // ClearCache()
assert.True(t, exists) // assert.Nil(t, loaded)
assert.Equal(t, assistant1, cached)
// Test cache eviction (LRU) // // Test setting new cache capacity
// At this point: assistant1 is most recently used (due to Get), then assistant2 // SetCache(100)
loaded.Put(assistant3) // This should evict assistant2 since it's least recently used // assert.NotNil(t, loaded)
assert.Equal(t, 2, loaded.Len(), "Cache should still have 2 items") // }
_, exists = loaded.Get("id2")
assert.False(t, exists, "assistant2 should have been evicted (least recently used)")
_, exists = loaded.Get("id1")
assert.True(t, exists, "assistant1 should still be in cache (was accessed recently)")
_, exists = loaded.Get("id3")
assert.True(t, exists, "assistant3 should be in cache (most recently added)")
// Test clear cache // func TestLoad_Validate(t *testing.T) {
ClearCache() // tests := []struct {
assert.Nil(t, loaded) // name string
// ast *Assistant
// wantErr bool
// }{
// {
// name: "valid assistant",
// ast: &Assistant{
// ID: "test-id",
// Name: "Test Assistant",
// Connector: "test-connector",
// },
// wantErr: false,
// },
// {
// name: "missing id",
// ast: &Assistant{
// Name: "Test Assistant",
// Connector: "test-connector",
// },
// wantErr: true,
// },
// {
// name: "missing name",
// ast: &Assistant{
// ID: "test-id",
// Connector: "test-connector",
// },
// wantErr: true,
// },
// {
// name: "missing connector",
// ast: &Assistant{
// ID: "test-id",
// Name: "Test Assistant",
// },
// wantErr: true,
// },
// }
// Test setting new cache capacity // for _, tt := range tests {
SetCache(100) // t.Run(tt.name, func(t *testing.T) {
assert.NotNil(t, loaded) // err := tt.ast.Validate()
} // if (err != nil) != tt.wantErr {
// t.Errorf("Assistant.Validate() error = %v, wantErr %v", err, tt.wantErr)
// }
// })
// }
// }
func TestLoad_Validate(t *testing.T) { // func TestLoad_Clone(t *testing.T) {
tests := []struct { // // Create a test assistant with all fields populated
name string // original := &Assistant{
ast *Assistant // ID: "test-id",
wantErr bool // Type: "test-type",
}{ // Name: "Test Assistant",
{ // Avatar: "test-avatar",
name: "valid assistant", // Connector: "test-connector",
ast: &Assistant{ // Path: "test-path",
ID: "test-id", // BuiltIn: true,
Name: "Test Assistant", // Sort: 1,
Connector: "test-connector", // Description: "test description",
}, // Tags: []string{"tag1", "tag2"},
wantErr: false, // Readonly: true,
}, // Mentionable: true,
{ // Automated: true,
name: "missing id", // Options: map[string]interface{}{"key": "value"},
ast: &Assistant{ // Prompts: []Prompt{{Role: "system", Content: "test"}},
Name: "Test Assistant", // Workflow: map[string]interface{}{"step": "test"},
Connector: "test-connector", // }
},
wantErr: true,
},
{
name: "missing name",
ast: &Assistant{
ID: "test-id",
Connector: "test-connector",
},
wantErr: true,
},
{
name: "missing connector",
ast: &Assistant{
ID: "test-id",
Name: "Test Assistant",
},
wantErr: true,
},
}
for _, tt := range tests { // // Clone the assistant
t.Run(tt.name, func(t *testing.T) { // clone := original.Clone()
err := tt.ast.Validate()
if (err != nil) != tt.wantErr {
t.Errorf("Assistant.Validate() error = %v, wantErr %v", err, tt.wantErr)
}
})
}
}
func TestLoad_Clone(t *testing.T) { // // Verify all fields are correctly cloned
// Create a test assistant with all fields populated // assert.Equal(t, original.ID, clone.ID)
original := &Assistant{ // assert.Equal(t, original.Type, clone.Type)
ID: "test-id", // assert.Equal(t, original.Name, clone.Name)
Type: "test-type", // assert.Equal(t, original.Avatar, clone.Avatar)
Name: "Test Assistant", // assert.Equal(t, original.Connector, clone.Connector)
Avatar: "test-avatar", // assert.Equal(t, original.Path, clone.Path)
Connector: "test-connector", // assert.Equal(t, original.BuiltIn, clone.BuiltIn)
Path: "test-path", // assert.Equal(t, original.Sort, clone.Sort)
BuiltIn: true, // assert.Equal(t, original.Description, clone.Description)
Sort: 1, // assert.Equal(t, original.Tags, clone.Tags)
Description: "test description", // assert.Equal(t, original.Readonly, clone.Readonly)
Tags: []string{"tag1", "tag2"}, // assert.Equal(t, original.Mentionable, clone.Mentionable)
Readonly: true, // assert.Equal(t, original.Automated, clone.Automated)
Mentionable: true, // assert.Equal(t, original.Options, clone.Options)
Automated: true, // assert.Equal(t, original.Prompts, clone.Prompts)
Options: map[string]interface{}{"key": "value"}, // assert.Equal(t, original.Workflow, clone.Workflow)
Prompts: []Prompt{{Role: "system", Content: "test"}},
Workflow: map[string]interface{}{"step": "test"},
}
// Clone the assistant // // Verify deep copy by modifying original
clone := original.Clone() // original.Tags[0] = "modified"
// original.Options["key"] = "modified"
// original.Workflow["step"] = "modified"
// assert.NotEqual(t, original.Tags[0], clone.Tags[0])
// assert.NotEqual(t, original.Options["key"], clone.Options["key"])
// assert.NotEqual(t, original.Workflow["step"], clone.Workflow["step"])
// Verify all fields are correctly cloned // // Test nil case
assert.Equal(t, original.ID, clone.ID) // var nilAssistant *Assistant
assert.Equal(t, original.Type, clone.Type) // assert.Nil(t, nilAssistant.Clone())
assert.Equal(t, original.Name, clone.Name) // }
assert.Equal(t, original.Avatar, clone.Avatar)
assert.Equal(t, original.Connector, clone.Connector)
assert.Equal(t, original.Path, clone.Path)
assert.Equal(t, original.BuiltIn, clone.BuiltIn)
assert.Equal(t, original.Sort, clone.Sort)
assert.Equal(t, original.Description, clone.Description)
assert.Equal(t, original.Tags, clone.Tags)
assert.Equal(t, original.Readonly, clone.Readonly)
assert.Equal(t, original.Mentionable, clone.Mentionable)
assert.Equal(t, original.Automated, clone.Automated)
assert.Equal(t, original.Options, clone.Options)
assert.Equal(t, original.Prompts, clone.Prompts)
assert.Equal(t, original.Workflow, clone.Workflow)
// Verify deep copy by modifying original // func TestLoad_Update(t *testing.T) {
original.Tags[0] = "modified" // // Create a test assistant
original.Options["key"] = "modified" // ast := &Assistant{
original.Workflow["step"] = "modified" // ID: "test-id",
assert.NotEqual(t, original.Tags[0], clone.Tags[0]) // Name: "Original Name",
assert.NotEqual(t, original.Options["key"], clone.Options["key"]) // Connector: "original-connector",
assert.NotEqual(t, original.Workflow["step"], clone.Workflow["step"]) // }
// Test nil case // // Test updating various fields
var nilAssistant *Assistant // updates := map[string]interface{}{
assert.Nil(t, nilAssistant.Clone()) // "name": "Updated Name",
} // "avatar": "updated-avatar",
// "description": "Updated description",
// "connector": "updated-connector",
// "type": "updated-type",
// "sort": 2,
// "mentionable": true,
// "automated": true,
// "tags": []string{"new-tag"},
// "options": map[string]interface{}{"new": "value"},
// }
func TestLoad_Update(t *testing.T) { // err := ast.Update(updates)
// Create a test assistant // assert.NoError(t, err)
ast := &Assistant{
ID: "test-id",
Name: "Original Name",
Connector: "original-connector",
}
// Test updating various fields // // Verify updates
updates := map[string]interface{}{ // assert.Equal(t, "Updated Name", ast.Name)
"name": "Updated Name", // assert.Equal(t, "updated-avatar", ast.Avatar)
"avatar": "updated-avatar", // assert.Equal(t, "Updated description", ast.Description)
"description": "Updated description", // assert.Equal(t, "updated-connector", ast.Connector)
"connector": "updated-connector", // assert.Equal(t, "updated-type", ast.Type)
"type": "updated-type", // assert.Equal(t, 2, ast.Sort)
"sort": 2, // assert.True(t, ast.Mentionable)
"mentionable": true, // assert.True(t, ast.Automated)
"automated": true, // assert.Equal(t, []string{"new-tag"}, ast.Tags)
"tags": []string{"new-tag"}, // assert.Equal(t, map[string]interface{}{"new": "value"}, ast.Options)
"options": map[string]interface{}{"new": "value"},
}
err := ast.Update(updates) // // Test nil assistant
assert.NoError(t, err) // var nilAssistant *Assistant
// err = nilAssistant.Update(updates)
// assert.Error(t, err)
// Verify updates // // Test invalid update that would make the assistant invalid
assert.Equal(t, "Updated Name", ast.Name) // invalidUpdates := map[string]interface{}{
assert.Equal(t, "updated-avatar", ast.Avatar) // "name": "",
assert.Equal(t, "Updated description", ast.Description) // }
assert.Equal(t, "updated-connector", ast.Connector) // err = ast.Update(invalidUpdates)
assert.Equal(t, "updated-type", ast.Type) // assert.Error(t, err)
assert.Equal(t, 2, ast.Sort) // }
assert.True(t, ast.Mentionable)
assert.True(t, ast.Automated)
assert.Equal(t, []string{"new-tag"}, ast.Tags)
assert.Equal(t, map[string]interface{}{"new": "value"}, ast.Options)
// Test nil assistant // func TestLoadBuiltIn(t *testing.T) {
var nilAssistant *Assistant // prepare(t)
err = nilAssistant.Update(updates) // defer test.Clean()
assert.Error(t, err)
// Test invalid update that would make the assistant invalid // // Clear any existing cache and storage
invalidUpdates := map[string]interface{}{ // ClearCache()
"name": "", // SetStorage(nil)
}
err = ast.Update(invalidUpdates)
assert.Error(t, err)
}
func TestLoadBuiltIn(t *testing.T) { // // Create a mock store to verify built-in assistants are saved
prepare(t) // mockStore := &mockStore{
defer test.Clean() // data: make(map[string]map[string]interface{}),
// }
// SetStorage(mockStore)
// SetCache(100)
// Clear any existing cache and storage // // Test loading built-in assistants
ClearCache() // err := LoadBuiltIn()
SetStorage(nil) // assert.NoError(t, err)
// Create a mock store to verify built-in assistants are saved // // Verify Modi assistant was loaded
mockStore := &mockStore{ // assistant, exists := loaded.Get("modi")
data: make(map[string]map[string]interface{}), // assert.True(t, exists, "Modi assistant should be loaded in cache")
} // if exists {
SetStorage(mockStore) // assert.Equal(t, "modi", assistant.ID)
SetCache(100) // assert.Equal(t, "Modi", assistant.Name)
// assert.Equal(t, "deepseek", assistant.Connector)
// assert.True(t, assistant.BuiltIn)
// assert.True(t, assistant.Readonly)
// assert.NotNil(t, assistant.Prompts)
// assert.NotNil(t, assistant.Script)
// }
// Test loading built-in assistants // }
err := LoadBuiltIn()
assert.NoError(t, err)
// Verify Modi assistant was loaded // // mockStore implements store.Store interface for testing
assistant, exists := loaded.Get("modi") // type mockStore struct {
assert.True(t, exists, "Modi assistant should be loaded in cache") // data map[string]map[string]interface{}
if exists { // }
assert.Equal(t, "modi", assistant.ID)
assert.Equal(t, "Modi", assistant.Name)
assert.Equal(t, "deepseek", assistant.Connector)
assert.True(t, assistant.BuiltIn)
assert.True(t, assistant.Readonly)
assert.NotNil(t, assistant.Prompts)
assert.NotNil(t, assistant.Script)
}
} // func (m *mockStore) GetAssistant(id string, locale ...string) (map[string]interface{}, error) {
// if data, ok := m.data[id]; ok {
// return data, nil
// }
// return nil, fmt.Errorf("assistant not found: %s", id)
// }
// mockStore implements store.Store interface for testing // // Add other required interface methods with empty implementations
type mockStore struct { // func (m *mockStore) GetThread(id string) (map[string]interface{}, error) { return nil, nil }
data map[string]map[string]interface{} // func (m *mockStore) GetMessage(id string) (map[string]interface{}, error) { return nil, nil }
} // func (m *mockStore) GetFile(id string) (map[string]interface{}, error) { return nil, nil }
// func (m *mockStore) CreateAssistant(data map[string]interface{}) (map[string]interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) CreateThread(data map[string]interface{}) (map[string]interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) CreateMessage(data map[string]interface{}) (map[string]interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) CreateFile(data map[string]interface{}) (map[string]interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) UpdateAssistant(id string, data map[string]interface{}) error { return nil }
// func (m *mockStore) UpdateThread(id string, data map[string]interface{}) error { return nil }
// func (m *mockStore) UpdateMessage(id string, data map[string]interface{}) error { return nil }
// func (m *mockStore) UpdateFile(id string, data map[string]interface{}) error { return nil }
// func (m *mockStore) DeleteAssistant(id string) error { return nil }
// func (m *mockStore) DeleteThread(id string) error { return nil }
// func (m *mockStore) DeleteMessage(id string) error { return nil }
// func (m *mockStore) DeleteFile(id string) error { return nil }
// func (m *mockStore) ListAssistants(query map[string]interface{}) ([]map[string]interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) ListThreads(query map[string]interface{}) ([]map[string]interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) ListMessages(query map[string]interface{}) ([]map[string]interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) ListFiles(query map[string]interface{}) ([]map[string]interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) DeleteAllChats(id string) error { return nil }
// func (m *mockStore) DeleteChat(id string, chatID string) error { return nil }
// func (m *mockStore) GetAssistants(filter store.AssistantFilter, locale ...string) (*store.AssistantResponse, error) {
// return nil, nil
// }
// func (m *mockStore) GetChat(id string, chatID string, locale ...string) (*store.ChatInfo, error) {
// return nil, nil
// }
// func (m *mockStore) GetChatWithFilter(id string, chatID string, filter store.ChatFilter, locale ...string) (*store.ChatInfo, error) {
// return nil, nil
// }
// func (m *mockStore) GetChats(id string, filter store.ChatFilter, locale ...string) (*store.ChatGroupResponse, error) {
// return nil, nil
// }
// func (m *mockStore) GetHistory(id string, chatID string, locale ...string) ([]map[string]interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) GetHistoryWithFilter(id string, chatID string, filter store.ChatFilter, locale ...string) ([]map[string]interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) SaveAssistant(assistant map[string]interface{}) (interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error {
// return nil
// }
// func (m *mockStore) UpdateChatTitle(sid string, cid string, title string) error { return nil }
// func (m *mockStore) DeleteAssistants(filter store.AssistantFilter) (int64, error) { return 0, nil }
// func (m *mockStore) GetAssistantTags(locale ...string) ([]store.Tag, error) {
// return []store.Tag{}, nil
// }
func (m *mockStore) GetAssistant(id string, locale ...string) (map[string]interface{}, error) { // // Attachment related methods
if data, ok := m.data[id]; ok { // func (m *mockStore) SaveAttachment(attachment map[string]interface{}) (interface{}, error) {
return data, nil // return attachment["file_id"], nil
} // }
return nil, fmt.Errorf("assistant not found: %s", id)
}
// Add other required interface methods with empty implementations // func (m *mockStore) DeleteAttachment(fileID string) error {
func (m *mockStore) GetThread(id string) (map[string]interface{}, error) { return nil, nil } // return nil
func (m *mockStore) GetMessage(id string) (map[string]interface{}, error) { return nil, nil } // }
func (m *mockStore) GetFile(id string) (map[string]interface{}, error) { return nil, nil }
func (m *mockStore) CreateAssistant(data map[string]interface{}) (map[string]interface{}, error) {
return nil, nil
}
func (m *mockStore) CreateThread(data map[string]interface{}) (map[string]interface{}, error) {
return nil, nil
}
func (m *mockStore) CreateMessage(data map[string]interface{}) (map[string]interface{}, error) {
return nil, nil
}
func (m *mockStore) CreateFile(data map[string]interface{}) (map[string]interface{}, error) {
return nil, nil
}
func (m *mockStore) UpdateAssistant(id string, data map[string]interface{}) error { return nil }
func (m *mockStore) UpdateThread(id string, data map[string]interface{}) error { return nil }
func (m *mockStore) UpdateMessage(id string, data map[string]interface{}) error { return nil }
func (m *mockStore) UpdateFile(id string, data map[string]interface{}) error { return nil }
func (m *mockStore) DeleteAssistant(id string) error { return nil }
func (m *mockStore) DeleteThread(id string) error { return nil }
func (m *mockStore) DeleteMessage(id string) error { return nil }
func (m *mockStore) DeleteFile(id string) error { return nil }
func (m *mockStore) ListAssistants(query map[string]interface{}) ([]map[string]interface{}, error) {
return nil, nil
}
func (m *mockStore) ListThreads(query map[string]interface{}) ([]map[string]interface{}, error) {
return nil, nil
}
func (m *mockStore) ListMessages(query map[string]interface{}) ([]map[string]interface{}, error) {
return nil, nil
}
func (m *mockStore) ListFiles(query map[string]interface{}) ([]map[string]interface{}, error) {
return nil, nil
}
func (m *mockStore) DeleteAllChats(id string) error { return nil }
func (m *mockStore) DeleteChat(id string, chatID string) error { return nil }
func (m *mockStore) GetAssistants(filter store.AssistantFilter, locale ...string) (*store.AssistantResponse, error) {
return nil, nil
}
func (m *mockStore) GetChat(id string, chatID string, locale ...string) (*store.ChatInfo, error) {
return nil, nil
}
func (m *mockStore) GetChatWithFilter(id string, chatID string, filter store.ChatFilter, locale ...string) (*store.ChatInfo, error) {
return nil, nil
}
func (m *mockStore) GetChats(id string, filter store.ChatFilter, locale ...string) (*store.ChatGroupResponse, error) {
return nil, nil
}
func (m *mockStore) GetHistory(id string, chatID string, locale ...string) ([]map[string]interface{}, error) {
return nil, nil
}
func (m *mockStore) GetHistoryWithFilter(id string, chatID string, filter store.ChatFilter, locale ...string) ([]map[string]interface{}, error) {
return nil, nil
}
func (m *mockStore) SaveAssistant(assistant map[string]interface{}) (interface{}, error) {
return nil, nil
}
func (m *mockStore) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error {
return nil
}
func (m *mockStore) UpdateChatTitle(sid string, cid string, title string) error { return nil }
func (m *mockStore) DeleteAssistants(filter store.AssistantFilter) (int64, error) { return 0, nil }
func (m *mockStore) GetAssistantTags(locale ...string) ([]store.Tag, error) {
return []store.Tag{}, nil
}
// Attachment related methods // func (m *mockStore) GetAttachments(filter store.AttachmentFilter, locale ...string) (*store.AttachmentResponse, error) {
func (m *mockStore) SaveAttachment(attachment map[string]interface{}) (interface{}, error) { // return &store.AttachmentResponse{}, nil
return attachment["file_id"], nil // }
}
func (m *mockStore) DeleteAttachment(fileID string) error { // func (m *mockStore) GetAttachment(fileID string, locale ...string) (map[string]interface{}, error) {
return nil // return nil, nil
} // }
func (m *mockStore) GetAttachments(filter store.AttachmentFilter, locale ...string) (*store.AttachmentResponse, error) { // func (m *mockStore) DeleteAttachments(filter store.AttachmentFilter) (int64, error) {
return &store.AttachmentResponse{}, nil // return 0, nil
} // }
func (m *mockStore) GetAttachment(fileID string, locale ...string) (map[string]interface{}, error) { // // Knowledge related methods
return nil, nil // func (m *mockStore) SaveKnowledge(knowledge map[string]interface{}) (interface{}, error) {
} // return knowledge["collection_id"], nil
// }
func (m *mockStore) DeleteAttachments(filter store.AttachmentFilter) (int64, error) { // func (m *mockStore) DeleteKnowledge(collectionID string) error {
return 0, nil // return nil
} // }
// Knowledge related methods // func (m *mockStore) GetKnowledges(filter store.KnowledgeFilter, locale ...string) (*store.KnowledgeResponse, error) {
func (m *mockStore) SaveKnowledge(knowledge map[string]interface{}) (interface{}, error) { // return &store.KnowledgeResponse{}, nil
return knowledge["collection_id"], nil // }
}
func (m *mockStore) DeleteKnowledge(collectionID string) error { // func (m *mockStore) GetKnowledge(collectionID string, locale ...string) (map[string]interface{}, error) {
return nil // return nil, nil
} // }
func (m *mockStore) GetKnowledges(filter store.KnowledgeFilter, locale ...string) (*store.KnowledgeResponse, error) { // func (m *mockStore) DeleteKnowledges(filter store.KnowledgeFilter) (int64, error) {
return &store.KnowledgeResponse{}, nil // return 0, nil
} // }
func (m *mockStore) GetKnowledge(collectionID string, locale ...string) (map[string]interface{}, error) { // // Close closes the store and releases any resources
return nil, nil // func (m *mockStore) Close() error {
} // return nil
// }
func (m *mockStore) DeleteKnowledges(filter store.KnowledgeFilter) (int64, error) {
return 0, nil
}
// Close closes the store and releases any resources
func (m *mockStore) Close() error {
return nil
}

24
agent/assistant/types.go
View file

@ -5,7 +5,6 @@ import (
"io" "io"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/yaoapp/gou/rag/driver"
v8 "github.com/yaoapp/gou/runtime/v8" v8 "github.com/yaoapp/gou/runtime/v8"
chatctx "github.com/yaoapp/yao/agent/context" chatctx "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n" "github.com/yaoapp/yao/agent/i18n"
@ -76,34 +75,12 @@ type NextAction struct {
Payload map[string]interface{} `json:"payload,omitempty"` Payload map[string]interface{} `json:"payload,omitempty"`
} }
// RAG the RAG interface
type RAG struct {
Engine driver.Engine
Uploader driver.FileUpload
Vectorizer driver.Vectorizer
Setting RAGSetting
}
// SearchOption the search option // SearchOption the search option
type SearchOption struct { type SearchOption struct {
WebSearch *bool `json:"web_search,omitempty" yaml:"web_search,omitempty"` // Whether to search the web WebSearch *bool `json:"web_search,omitempty" yaml:"web_search,omitempty"` // Whether to search the web
Knowledge *bool `json:"knowledge,omitempty" yaml:"knowledge,omitempty"` // Whether to search the knowledge Knowledge *bool `json:"knowledge,omitempty" yaml:"knowledge,omitempty"` // Whether to search the knowledge
} }
// KnowledgeOption the knowledge option
type KnowledgeOption struct {
Collections []string `json:"collections,omitempty" yaml:"collections,omitempty"` // The Global Collections
ChunkingMethod string `json:"chunking_method,omitempty" yaml:"chunking_method,omitempty"`
ChunkSize int `json:"chunk_size,omitempty" yaml:"chunk_size,omitempty"`
ChunkOverlap int `json:"chunk_overlap,omitempty" yaml:"chunk_overlap,omitempty"`
SearchMethod string `json:"search_method,omitempty" yaml:"search_method,omitempty"`
}
// RAGSetting the RAG setting
type RAGSetting struct {
IndexPrefix string `json:"index_prefix" yaml:"index_prefix"`
}
// Prompt a prompt // Prompt a prompt
type Prompt struct { type Prompt struct {
Role string `json:"role"` Role string `json:"role"`
@ -141,7 +118,6 @@ type Assistant struct {
Placeholder *Placeholder `json:"placeholder,omitempty"` // Assistant Placeholder Placeholder *Placeholder `json:"placeholder,omitempty"` // Assistant Placeholder
Locales i18n.Map `json:"locales,omitempty"` // Assistant Locales Locales i18n.Map `json:"locales,omitempty"` // Assistant Locales
Search *SearchOption `json:"search,omitempty" yaml:"search,omitempty"` // Whether this assistant supports search Search *SearchOption `json:"search,omitempty" yaml:"search,omitempty"` // Whether this assistant supports search
Knowledge *KnowledgeOption `json:"knowledge,omitempty" yaml:"knowledge,omitempty"` // Whether this assistant supports knowledge
CreatedAt int64 `json:"created_at"` // Creation timestamp CreatedAt int64 `json:"created_at"` // Creation timestamp
UpdatedAt int64 `json:"updated_at"` // Last update timestamp UpdatedAt int64 `json:"updated_at"` // Last update timestamp
Script *v8.Script `json:"-" yaml:"-"` // Assistant Script Script *v8.Script `json:"-" yaml:"-"` // Assistant Script

167
agent/load.go
View file

@ -6,11 +6,9 @@ import (
"github.com/yaoapp/gou/application" "github.com/yaoapp/gou/application"
"github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/model"
"github.com/yaoapp/yao/agent/assistant" "github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/agent/i18n" "github.com/yaoapp/yao/agent/i18n"
"github.com/yaoapp/yao/agent/store" "github.com/yaoapp/yao/agent/store"
"github.com/yaoapp/yao/attachment"
"github.com/yaoapp/yao/config" "github.com/yaoapp/yao/config"
) )
@ -22,10 +20,9 @@ func Load(cfg config.Config) error {
setting := DSL{ setting := DSL{
ID: "agent", ID: "agent",
Allows: []string{},
StoreSetting: store.Setting{ StoreSetting: store.Setting{
Prefix: "yao_agent_", MaxSize: 20,
Connector: "default", TTL: 90 * 24 * 60 * 60, // 90 days in seconds
}, },
} }
@ -78,18 +75,6 @@ func Load(cfg config.Config) error {
return err return err
} }
// Initialize Auth
err = initAuth()
if err != nil {
return err
}
// Initialize Upload
err = initUpload()
if err != nil {
return err
}
// Initialize Assistant // Initialize Assistant
err = initAssistant() err = initAssistant()
if err != nil { if err != nil {
@ -99,154 +84,6 @@ func Load(cfg config.Config) error {
return nil return nil
} }
// initAuth initialize the auth
func initAuth() error {
if Agent.AuthSetting == nil {
Agent.AuthSetting = &Auth{
Models: &AuthModels{User: "admin.user", Guest: "guest"},
Fields: &AuthFields{ID: "id", Roles: "roles", Permission: "permission"},
SessionFields: &AuthSessionFields{ID: "user_id", Roles: "user_roles", Guest: "guest_id"},
}
}
if Agent.AuthSetting.Models == nil {
Agent.AuthSetting.Models = &AuthModels{User: "admin.user", Guest: "guest"}
}
if Agent.AuthSetting.Fields == nil {
Agent.AuthSetting.Fields = &AuthFields{ID: "id", Roles: "roles", Permission: "permission"}
}
if Agent.AuthSetting.SessionFields == nil {
Agent.AuthSetting.SessionFields = &AuthSessionFields{ID: "user_id", Roles: "user_roles", Guest: "guest_id"}
}
if Agent.AuthSetting.Models.User == "" {
Agent.AuthSetting.Models.User = "admin.user"
}
if Agent.AuthSetting.Models.Guest == "" {
Agent.AuthSetting.Models.Guest = "guest"
}
if Agent.AuthSetting.Fields.Roles == "" {
Agent.AuthSetting.Fields.Roles = "roles"
}
if Agent.AuthSetting.Fields.Permission == "" {
Agent.AuthSetting.Fields.Permission = "permission"
}
if Agent.AuthSetting.Fields.ID == "" {
Agent.AuthSetting.Fields.ID = "id"
}
if Agent.AuthSetting.Fields.ID == "" {
Agent.AuthSetting.Fields.ID = "id"
}
if Agent.AuthSetting.SessionFields.ID == "" {
Agent.AuthSetting.SessionFields.ID = "user_id"
}
if Agent.AuthSetting.SessionFields.Roles == "" {
Agent.AuthSetting.SessionFields.Roles = "user_roles"
}
if Agent.AuthSetting.SessionFields.Guest == "" {
Agent.AuthSetting.SessionFields.Guest = "guest_id"
}
// Validate User Model and Fields
if !model.Exists(Agent.AuthSetting.Models.User) {
return fmt.Errorf("model %s not found", Agent.AuthSetting.Models.User)
}
user := model.Select(Agent.AuthSetting.Models.User)
shouldHave := []string{Agent.AuthSetting.Fields.ID, Agent.AuthSetting.Fields.Roles, Agent.AuthSetting.Fields.Permission}
for _, name := range shouldHave {
if _, has := user.Columns[name]; !has {
return fmt.Errorf("model %s should have column %s", Agent.AuthSetting.Models.User, name)
}
}
return nil
}
// initUpload initialize the upload
func initUpload() error {
if Agent.UploadSetting == nil {
_, err := attachment.RegisterDefault("chat")
if err != nil {
return err
}
_, err = attachment.RegisterDefault("knowledge")
if err != nil {
return err
}
return nil
}
// If the chat upload setting is not set, use the default chat upload setting.
if Agent.UploadSetting.Chat == nil {
_, err := attachment.RegisterDefault("chat")
if err != nil {
return err
}
}
// Use the chat upload setting for knowledge upload, if the knowledge upload setting is not set.
if Agent.UploadSetting.Knowledge == nil {
if Agent.UploadSetting.Chat == nil {
_, err := attachment.RegisterDefault("knowledge")
if err != nil {
return err
}
} else {
_, err := attachment.Register("knowledge", Agent.UploadSetting.Chat.Driver, *Agent.UploadSetting.Chat)
if err != nil {
return err
}
}
}
// Use custom chat upload setting
if Agent.UploadSetting.Chat != nil {
Agent.UploadSetting.Chat.ReplaceEnv(config.Conf.DataRoot)
_, err := attachment.Register("chat", Agent.UploadSetting.Chat.Driver, *Agent.UploadSetting.Chat) // Register the chat upload manager
if err != nil {
return err
}
}
// Use custom knowledge upload setting
if Agent.UploadSetting.Knowledge != nil {
Agent.UploadSetting.Knowledge.ReplaceEnv(config.Conf.DataRoot)
_, err := attachment.Register("knowledge", Agent.UploadSetting.Knowledge.Driver, *Agent.UploadSetting.Knowledge)
if err != nil {
return err
}
}
// Use the chat upload setting for asset upload, if the asset upload setting is not set. (public assets)
if Agent.UploadSetting.Assets == nil {
_, err := attachment.RegisterDefault("assets")
if err != nil {
return err
}
}
// Use custom asset upload setting
if Agent.UploadSetting.Assets != nil {
Agent.UploadSetting.Assets.ReplaceEnv(config.Conf.DataRoot)
_, err := attachment.Register("assets", Agent.UploadSetting.Assets.Driver, *Agent.UploadSetting.Assets)
if err != nil {
return err
}
}
return nil
}
// initGlobalI18n initialize the global i18n // initGlobalI18n initialize the global i18n
func initGlobalI18n() error { func initGlobalI18n() error {
locales, err := i18n.GetLocales("agent") locales, err := i18n.GetLocales("agent")

32
agent/load_test.go
View file

@ -1,24 +1,16 @@
package agent package agent
import ( // func TestLoad(t *testing.T) {
"testing" // test.Prepare(t, config.Conf)
// defer test.Clean()
"github.com/stretchr/testify/assert" // err := Load(config.Conf)
"github.com/yaoapp/yao/config" // if err != nil {
"github.com/yaoapp/yao/test" // t.Fatal(err)
) // }
// check(t)
// }
func TestLoad(t *testing.T) { // func check(t *testing.T) {
test.Prepare(t, config.Conf) // assert.NotNil(t, Agent)
defer test.Clean() // }
err := Load(config.Conf)
if err != nil {
t.Fatal(err)
}
check(t)
}
func check(t *testing.T) {
assert.NotNil(t, Agent)
}

124
agent/process.go
View file

@ -1,15 +1,12 @@
package agent package agent
import ( import (
"context"
"encoding/json"
"fmt" "fmt"
"strconv" "strconv"
"strings" "strings"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/yaoapp/gou/process" "github.com/yaoapp/gou/process"
"github.com/yaoapp/gou/rag/driver"
"github.com/yaoapp/kun/exception" "github.com/yaoapp/kun/exception"
"github.com/yaoapp/yao/agent/message" "github.com/yaoapp/yao/agent/message"
"github.com/yaoapp/yao/agent/store" "github.com/yaoapp/yao/agent/store"
@ -163,131 +160,10 @@ func processAssistantMatch(process *process.Process) interface{} {
} }
} }
// Force Using sotre
forceStore := false
if store, has := params["store"]; has {
switch v := store.(type) {
case bool:
forceStore = v
case int:
forceStore = v == 1
case string:
forceStore = v == "true" || v == "1"
}
}
// Rag Support match using RAG
if Agent.RAG != nil && !forceStore {
return assistantMatchRAG(content, params)
}
// Match using Store // Match using Store
return assistantMatchStore(content, params) return assistantMatchStore(content, params)
} }
func assistantMatchRAG(content interface{}, params map[string]interface{}) interface{} {
if Agent == nil {
exception.New("Agent is not initialized", 500).Throw()
}
// Convert content to JSON string
var contentStr string
switch v := content.(type) {
case string:
contentStr = v
case []byte:
contentStr = string(v)
default:
bytes, err := json.Marshal(v)
if err != nil {
exception.New("Failed to convert content to JSON: %s", 500, err.Error()).Throw()
}
contentStr = string(bytes)
}
// Get limit from params
limit := 20 // default limit
if v, has := params["limit"]; has {
switch lv := v.(type) {
case int:
limit = lv
case string:
limitInt, err := strconv.Atoi(lv)
if err == nil {
limit = limitInt
}
}
}
// Get min_score from params
minScore := 0.0 // default min_score
if v, has := params["min_score"]; has {
switch lv := v.(type) {
case float64:
minScore = lv
case float32:
minScore = float64(lv)
case int:
minScore = float64(lv)
case string:
if score, err := strconv.ParseFloat(lv, 64); err == nil {
minScore = score
}
}
}
ctx := context.Background()
// Get vectors using vectorizer
vectors, err := Agent.RAG.Vectorizer().Vectorize(ctx, contentStr)
if err != nil {
exception.New("Failed to encode content: %s", 500, err.Error()).Throw()
}
// Search using RAG engine
opts := driver.VectorSearchOptions{
TopK: limit,
MinScore: minScore,
QueryText: contentStr,
}
index := fmt.Sprintf("%sassistants", Agent.RAG.Setting().IndexPrefix)
results, err := Agent.RAG.Engine().Search(ctx, index, vectors, opts)
if err != nil {
exception.New("Failed to search with RAG: %s", 500, err.Error()).Throw()
}
// Convert results to assistant data array
ids := []string{}
// Collect IDs from search results
for _, result := range results {
if result.Metadata != nil {
if id, ok := result.Metadata["assistant_id"].(string); ok {
ids = append(ids, id)
}
}
}
// If no IDs found, return empty array
if len(ids) == 0 {
return []map[string]interface{}{}
}
// Fetch complete assistant data from store using AssistantIDs
filter := store.AssistantFilter{
AssistantIDs: ids,
Page: 1,
PageSize: len(ids),
}
res, err := Agent.Store.GetAssistants(filter)
if err != nil {
exception.New("get assistants error: %s", 500, err).Throw()
}
return res.Data
}
// parseAssistantFilter parse common filter parameters // parseAssistantFilter parse common filter parameters
func parseAssistantFilter(params map[string]interface{}) store.AssistantFilter { func parseAssistantFilter(params map[string]interface{}) store.AssistantFilter {
filter := store.AssistantFilter{} filter := store.AssistantFilter{}

1022
agent/process_test.go
View file

File diff suppressed because it is too large Load diff

120
agent/rag/rag.go
View file

@ -1,120 +0,0 @@
package rag
import (
"fmt"
"os"
"strings"
"github.com/yaoapp/gou/rag"
"github.com/yaoapp/gou/rag/driver"
)
// RAG the RAG instance
type RAG struct {
setting Setting
engine driver.Engine
vectorizer driver.Vectorizer
fileUpload driver.FileUpload
}
// parseEnvValue parse environment variable if the value starts with $ENV.
func parseEnvValue(value string) string {
if strings.HasPrefix(value, "$ENV.") {
envKey := strings.TrimPrefix(value, "$ENV.")
if envVal := os.Getenv(envKey); envVal != "" {
return envVal
}
}
return value
}
// convertOptions convert interface{} options map to string map and parse environment variables
func convertOptions(options map[string]interface{}) map[string]string {
converted := make(map[string]string)
for k, v := range options {
if str, ok := v.(string); ok {
converted[k] = parseEnvValue(str)
}
}
return converted
}
// New create a new RAG instance
func New(setting Setting) (*RAG, error) {
if setting.Engine.Driver == "" {
return nil, fmt.Errorf("engine driver is required")
}
if setting.Vectorizer.Driver == "" {
return nil, fmt.Errorf("vectorizer driver is required")
}
// Set default values
if setting.Upload.ChunkSize == 0 {
setting.Upload.ChunkSize = 1024
}
if setting.Upload.ChunkOverlap == 0 {
setting.Upload.ChunkOverlap = 256
}
if setting.IndexPrefix == "" {
setting.IndexPrefix = "yao_agent_"
}
// Convert options map for vectorizer and handle environment variables
vectorizerOpts := convertOptions(setting.Vectorizer.Options)
// Create vectorizer
vectorizer, err := rag.NewVectorizer(setting.Vectorizer.Driver, driver.VectorizeConfig{
Model: vectorizerOpts["model"],
Options: vectorizerOpts,
})
if err != nil {
return nil, fmt.Errorf("create vectorizer: %v", err)
}
// Convert options map for engine and handle environment variables
engineOpts := convertOptions(setting.Engine.Options)
// Create engine
engine, err := rag.NewEngine(setting.Engine.Driver, driver.IndexConfig{
Options: engineOpts,
}, vectorizer)
if err != nil {
return nil, fmt.Errorf("create engine: %v", err)
}
// Create file upload
fileUpload, err := rag.NewFileUpload(setting.Engine.Driver, engine, vectorizer)
if err != nil {
return nil, fmt.Errorf("create file upload: %v", err)
}
return &RAG{
setting: setting,
engine: engine,
vectorizer: vectorizer,
fileUpload: fileUpload,
}, nil
}
// Setting get the RAG settings
func (rag *RAG) Setting() Setting {
return rag.setting
}
// Engine get the vector database engine
func (rag *RAG) Engine() driver.Engine {
return rag.engine
}
// Vectorizer get the text vectorizer
func (rag *RAG) Vectorizer() driver.Vectorizer {
return rag.vectorizer
}
// FileUpload get the file upload handler
func (rag *RAG) FileUpload() driver.FileUpload {
return rag.fileUpload
}

29
agent/rag/types.go
View file

@ -1,29 +0,0 @@
package rag
// Setting RAG settings
type Setting struct {
Engine Engine `json:"engine" yaml:"engine"`
Vectorizer Vectorizer `json:"vectorizer" yaml:"vectorizer"`
Upload Upload `json:"upload" yaml:"upload"`
IndexPrefix string `json:"index_prefix" yaml:"index_prefix"`
}
// Engine the vector database engine settings
type Engine struct {
Driver string `json:"driver" yaml:"driver"`
Options map[string]interface{} `json:"options" yaml:"options"`
}
// Vectorizer the text vectorizer settings
type Vectorizer struct {
Driver string `json:"driver" yaml:"driver"`
Options map[string]interface{} `json:"options" yaml:"options"`
}
// Upload the file upload settings
type Upload struct {
Async bool `json:"async" yaml:"async"`
AllowedTypes []string `json:"allowed_types" yaml:"allowed_types"`
ChunkSize int `json:"chunk_size" yaml:"chunk_size"`
ChunkOverlap int `json:"chunk_overlap" yaml:"chunk_overlap"`
}

59
agent/store/types.go
View file

@ -3,11 +3,10 @@ package store
// Setting represents the conversation configuration structure // Setting represents the conversation configuration structure
// Used to configure basic conversation parameters including connector, user field, table name, etc. // Used to configure basic conversation parameters including connector, user field, table name, etc.
type Setting struct { type Setting struct {
Connector string `json:"connector,omitempty"` // Name of the connector used to specify data storage method Connector string `json:"connector,omitempty" yaml:"connector,omitempty"` // Connector name, default is "default"
UserField string `json:"user_field,omitempty"` // User ID field name, defaults to "user_id" MaxSize int `json:"max_size,omitempty" yaml:"max_size,omitempty"` // Maximum storage size limit, default is 20
Prefix string `json:"prefix,omitempty"` // Database table name prefix TTL int `json:"ttl,omitempty" yaml:"ttl,omitempty"` // Time To Live in seconds, default is 90 * 24 * 60 * 60 (90 days)
MaxSize int `json:"max_size,omitempty" yaml:"max_size,omitempty"` // Maximum storage size limit Options map[string]interface{} `json:"optional,omitempty" yaml:"optional,omitempty"` // The options for the store
TTL int `json:"ttl,omitempty" yaml:"ttl,omitempty"` // Time To Live in seconds
} }
// ChatInfo represents the chat information structure // ChatInfo represents the chat information structure
@ -169,56 +168,6 @@ type Store interface {
// Returns: Number of deleted records and potential error // Returns: Number of deleted records and potential error
DeleteAssistants(filter AssistantFilter) (int64, error) DeleteAssistants(filter AssistantFilter) (int64, error)
// SaveAttachment saves attachment information
// attachment: Attachment information
// Returns: Attachment ID and potential error
SaveAttachment(attachment map[string]interface{}) (interface{}, error)
// DeleteAttachment deletes an attachment
// fileID: Attachment file ID
// Returns: Potential error
DeleteAttachment(fileID string) error
// GetAttachments retrieves a list of attachments
// filter: Filter conditions
// Returns: Paginated attachment list and potential error
GetAttachments(filter AttachmentFilter, locale ...string) (*AttachmentResponse, error)
// GetAttachment retrieves a single attachment by file ID
// fileID: Attachment file ID
// Returns: Attachment information and potential error
GetAttachment(fileID string, locale ...string) (map[string]interface{}, error)
// DeleteAttachments deletes attachments based on filter conditions
// filter: Filter conditions
// Returns: Number of deleted records and potential error
DeleteAttachments(filter AttachmentFilter) (int64, error)
// SaveKnowledge saves knowledge collection information
// knowledge: Knowledge collection information
// Returns: Collection ID and potential error
SaveKnowledge(knowledge map[string]interface{}) (interface{}, error)
// DeleteKnowledge deletes a knowledge collection
// collectionID: Knowledge collection ID
// Returns: Potential error
DeleteKnowledge(collectionID string) error
// GetKnowledges retrieves a list of knowledge collections
// filter: Filter conditions
// Returns: Paginated knowledge collection list and potential error
GetKnowledges(filter KnowledgeFilter, locale ...string) (*KnowledgeResponse, error)
// GetKnowledge retrieves a single knowledge collection by ID
// collectionID: Knowledge collection ID
// Returns: Knowledge collection information and potential error
GetKnowledge(collectionID string, locale ...string) (map[string]interface{}, error)
// DeleteKnowledges deletes knowledge collections based on filter conditions
// filter: Filter conditions
// Returns: Number of deleted records and potential error
DeleteKnowledges(filter KnowledgeFilter) (int64, error)
// Close closes the store and releases any resources // Close closes the store and releases any resources
// Returns: Potential error // Returns: Potential error
Close() error Close() error

771
agent/store/xun.go
View file

@ -9,7 +9,6 @@ import (
"github.com/google/uuid" "github.com/google/uuid"
jsoniter "github.com/json-iterator/go" jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/session"
"github.com/yaoapp/kun/log" "github.com/yaoapp/kun/log"
"github.com/yaoapp/xun/capsule" "github.com/yaoapp/xun/capsule"
"github.com/yaoapp/xun/dbal/query" "github.com/yaoapp/xun/dbal/query"
@ -112,7 +111,7 @@ func (conv *Xun) clean() {
} }
if nums > 0 { if nums > 0 {
log.Trace("Clean the conversation table: %s %d", conv.setting.Prefix, nums) log.Trace("Clean the conversation table: %d", nums)
} }
} }
@ -136,7 +135,7 @@ func (conv *Xun) startAutoClean() {
} }
}() }()
log.Trace("Started automatic cleanup for: %s", conv.setting.Prefix) log.Trace("Started automatic cleanup")
} }
// stopAutoClean stops the automatic cleanup routine // stopAutoClean stops the automatic cleanup routine
@ -151,7 +150,7 @@ func (conv *Xun) stopAutoClean() {
conv.cleanStop = nil conv.cleanStop = nil
} }
log.Trace("Stopped automatic cleanup for: %s", conv.setting.Prefix) log.Trace("Stopped automatic cleanup")
} }
// Close stops the automatic cleanup and closes resources // Close stops the automatic cleanup and closes resources
@ -162,31 +161,22 @@ func (conv *Xun) Close() error {
// Rename Init to initialize to avoid conflicts // Rename Init to initialize to avoid conflicts
func (conv *Xun) initialize() error { func (conv *Xun) initialize() error {
// Initialize history table
if err := conv.initHistoryTable(); err != nil {
return err
}
// Initialize chat table // Initialize chat table
if err := conv.initChatTable(); err != nil { if err := conv.initChatTable(); err != nil {
return err return err
} }
// Initialize history table
if err := conv.initHistoryTable(); err != nil {
return err
}
// Initialize assistant table // Initialize assistant table
if err := conv.initAssistantTable(); err != nil { if err := conv.initAssistantTable(); err != nil {
return err return err
} }
// Initialize attachment table
if err := conv.initAttachmentTable(); err != nil {
return err
}
// Initialize knowledge table
if err := conv.initKnowledgeTable(); err != nil {
return err
}
// Start automatic cleanup if TTL is enabled // Start automatic cleanup if TTL is enabled
if conv.setting.TTL > 0 { if conv.setting.TTL > 0 {
conv.startAutoClean() conv.startAutoClean()
@ -345,153 +335,21 @@ func (conv *Xun) initAssistantTable() error {
return nil return nil
} }
func (conv *Xun) initAttachmentTable() error {
attachmentTable := conv.getAttachmentTable()
has, err := conv.schema.HasTable(attachmentTable)
if err != nil {
return err
}
// Create the attachment table
if !has {
err = conv.schema.CreateTable(attachmentTable, func(table schema.Blueprint) {
table.ID("id")
table.String("file_id", 255).Unique().Index()
table.String("uid", 255).Index()
table.Boolean("guest").SetDefault(false).Index()
table.String("manager", 200).Index()
table.String("content_type", 200).Index()
table.String("name", 500).Index()
table.Boolean("public").SetDefault(false).Index()
table.JSON("scope").Null()
table.Boolean("gzip").SetDefault(false).Index()
table.BigInteger("bytes").Index()
table.String("collection_id", 200).Null().Index()
table.Enum("status", []string{"uploading", "uploaded", "indexing", "indexed", "upload_failed", "index_failed"}).SetDefault("uploading").Index() // Status field enum
table.String("progress", 200).Null() // Progress information
table.String("error", 600).Null() // Error information
table.TimestampTz("created_at").SetDefaultRaw("CURRENT_TIMESTAMP").Index()
table.TimestampTz("updated_at").Null().Index()
})
if err != nil {
return err
}
log.Trace("Create the attachment table: %s", attachmentTable)
}
// Validate the table
tab, err := conv.schema.GetTable(attachmentTable)
if err != nil {
return err
}
fields := []string{"id", "file_id", "uid", "guest", "manager", "content_type", "name", "public", "scope", "gzip", "bytes", "collection_id", "status", "progress", "error", "created_at", "updated_at"}
for _, field := range fields {
if !tab.HasColumn(field) {
return fmt.Errorf("%s is required", field)
}
}
return nil
}
func (conv *Xun) initKnowledgeTable() error {
knowledgeTable := conv.getKnowledgeTable()
has, err := conv.schema.HasTable(knowledgeTable)
if err != nil {
return err
}
// Create the knowledge table
if !has {
err = conv.schema.CreateTable(knowledgeTable, func(table schema.Blueprint) {
table.ID("id")
table.String("collection_id", 200).Unique().Index()
table.String("name", 200).Index()
table.String("description", 600).Null().Index() // knowledge description
table.String("uid", 255).Index()
table.Boolean("public").SetDefault(false).Index()
table.JSON("scope").Null()
table.Boolean("readonly").SetDefault(false).Index()
table.JSON("option").Null()
table.Boolean("system").SetDefault(false).Index()
table.Integer("sort").SetDefault(9999).Index() // knowledge sort order
table.String("cover", 500).Null()
table.TimestampTz("created_at").SetDefaultRaw("CURRENT_TIMESTAMP").Index()
table.TimestampTz("updated_at").Null().Index()
})
if err != nil {
return err
}
log.Trace("Create the knowledge table: %s", knowledgeTable)
}
// Validate the table
tab, err := conv.schema.GetTable(knowledgeTable)
if err != nil {
return err
}
fields := []string{"id", "collection_id", "name", "description", "uid", "public", "scope", "readonly", "option", "system", "sort", "cover", "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) { func (conv *Xun) getUserID(sid string) (string, error) {
field := "user_id" // TODO: get the user id from the authentication system
if conv.setting.UserField != "" { return "guest", nil
field = conv.setting.UserField
}
id, err := session.Global().ID(sid).Get(field)
if err != nil {
return "", err
}
if id == nil || id == "" {
return sid, nil
}
return fmt.Sprintf("%v", id), nil
} }
func (conv *Xun) getHistoryTable() string { func (conv *Xun) getHistoryTable() string {
return conv.setting.Prefix + "history" return "__yao.agent.history"
} }
func (conv *Xun) getChatTable() string { func (conv *Xun) getChatTable() string {
return conv.setting.Prefix + "chat" return "__yao.agent.chat"
} }
func (conv *Xun) getAssistantTable() string { func (conv *Xun) getAssistantTable() string {
return conv.setting.Prefix + "assistant" return "__yao.agent.assistant"
}
func (conv *Xun) getAttachmentTable() string {
return conv.setting.Prefix + "attachment"
}
func (conv *Xun) getKnowledgeTable() string {
return conv.setting.Prefix + "knowledge"
}
func (conv *Xun) newQueryAttachment() query.Query {
qb := conv.query.New()
qb.Table(conv.getAttachmentTable())
return qb
}
func (conv *Xun) newQueryKnowledge() query.Query {
qb := conv.query.New()
qb.Table(conv.getKnowledgeTable())
return qb
} }
// UpdateChatTitle update the chat title // UpdateChatTitle update the chat title
@ -1648,606 +1506,3 @@ func (conv *Xun) GenerateAssistantID() (string, error) {
return "", fmt.Errorf("failed to generate unique ID after %d attempts", maxAttempts) return "", fmt.Errorf("failed to generate unique ID after %d attempts", maxAttempts)
} }
// SaveAttachment saves attachment information
func (conv *Xun) SaveAttachment(attachment map[string]interface{}) (interface{}, error) {
// Validate required fields
requiredFields := []string{"file_id", "uid", "manager", "content_type", "name"}
for _, field := range requiredFields {
if _, ok := attachment[field]; !ok {
return nil, fmt.Errorf("field %s is required", field)
}
if attachment[field] == nil || attachment[field] == "" {
return nil, fmt.Errorf("field %s cannot be empty", field)
}
}
// Create a copy of the attachment map to avoid modifying the original
attachmentCopy := make(map[string]interface{})
for k, v := range attachment {
attachmentCopy[k] = v
}
// Process JSON fields
jsonFields := []string{"scope"}
for _, field := range jsonFields {
if val, ok := attachmentCopy[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 {
attachmentCopy[field] = parsed
}
}
}
}
// Check if attachment exists
exists, err := conv.query.New().
Table(conv.getAttachmentTable()).
Where("file_id", attachmentCopy["file_id"]).
Exists()
if err != nil {
return nil, err
}
// Convert JSON fields to strings for storage
for _, field := range jsonFields {
if val, ok := attachmentCopy[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)
}
attachmentCopy[field] = jsonStr
}
}
// Update or insert
if exists {
attachmentCopy["updated_at"] = time.Now()
_, err := conv.query.New().
Table(conv.getAttachmentTable()).
Where("file_id", attachmentCopy["file_id"]).
Update(attachmentCopy)
if err != nil {
return nil, err
}
return attachmentCopy["file_id"], nil
}
attachmentCopy["created_at"] = time.Now()
err = conv.query.New().
Table(conv.getAttachmentTable()).
Insert(attachmentCopy)
if err != nil {
return nil, err
}
return attachmentCopy["file_id"], nil
}
// DeleteAttachment deletes an attachment by file_id
func (conv *Xun) DeleteAttachment(fileID string) error {
// Check if attachment exists
exists, err := conv.query.New().
Table(conv.getAttachmentTable()).
Where("file_id", fileID).
Exists()
if err != nil {
return err
}
if !exists {
return fmt.Errorf("attachment %s not found", fileID)
}
_, err = conv.query.New().
Table(conv.getAttachmentTable()).
Where("file_id", fileID).
Delete()
return err
}
// GetAttachments retrieves attachments with pagination and filtering
func (conv *Xun) GetAttachments(filter AttachmentFilter, locale ...string) (*AttachmentResponse, error) {
qb := conv.query.New().
Table(conv.getAttachmentTable())
// Apply UID filter if provided
if filter.UID != "" {
qb.Where("uid", filter.UID)
}
// Apply guest filter if provided
if filter.Guest != nil {
qb.Where("guest", *filter.Guest)
}
// Apply manager filter if provided
if filter.Manager != "" {
qb.Where("manager", filter.Manager)
}
// Apply content_type filter if provided
if filter.ContentType != "" {
qb.Where("content_type", filter.ContentType)
}
// Apply name filter if provided
if filter.Name != "" {
qb.Where("name", "like", fmt.Sprintf("%%%s%%", filter.Name))
}
// Apply public filter if provided
if filter.Public != nil {
qb.Where("public", *filter.Public)
}
// Apply gzip filter if provided
if filter.Gzip != nil {
qb.Where("gzip", *filter.Gzip)
}
// Apply collection_id filter if provided
if filter.CollectionID != "" {
qb.Where("collection_id", filter.CollectionID)
}
// Apply status filter if provided
if filter.Status != "" {
qb.Where("status", filter.Status)
}
// Apply keyword filter if provided
if filter.Keywords != "" {
qb.Where("name", "like", fmt.Sprintf("%%%s%%", filter.Keywords))
}
// 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{"scope"}
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 &AttachmentResponse{
Data: data,
Page: filter.Page,
PageSize: filter.PageSize,
PageCnt: totalPages,
Next: nextPage,
Prev: prevPage,
Total: total,
}, nil
}
// GetAttachment retrieves a single attachment by file_id
func (conv *Xun) GetAttachment(fileID string, locale ...string) (map[string]interface{}, error) {
row, err := conv.query.New().
Table(conv.getAttachmentTable()).
Where("file_id", fileID).
First()
if err != nil {
return nil, err
}
if row == nil {
return nil, fmt.Errorf("attachment %s not found", fileID)
}
data := row.ToMap()
if data == nil || len(data) == 0 {
return nil, fmt.Errorf("the attachment %s is empty", fileID)
}
// Parse JSON fields
jsonFields := []string{"scope"}
conv.parseJSONFields(data, jsonFields)
return data, nil
}
// DeleteAttachments deletes attachments based on filter conditions
func (conv *Xun) DeleteAttachments(filter AttachmentFilter) (int64, error) {
qb := conv.query.New().
Table(conv.getAttachmentTable())
// Apply UID filter if provided
if filter.UID != "" {
qb.Where("uid", filter.UID)
}
// Apply guest filter if provided
if filter.Guest != nil {
qb.Where("guest", *filter.Guest)
}
// Apply manager filter if provided
if filter.Manager != "" {
qb.Where("manager", filter.Manager)
}
// Apply content_type filter if provided
if filter.ContentType != "" {
qb.Where("content_type", filter.ContentType)
}
// Apply name filter if provided
if filter.Name != "" {
qb.Where("name", "like", fmt.Sprintf("%%%s%%", filter.Name))
}
// Apply public filter if provided
if filter.Public != nil {
qb.Where("public", *filter.Public)
}
// Apply gzip filter if provided
if filter.Gzip != nil {
qb.Where("gzip", *filter.Gzip)
}
// Apply collection_id filter if provided
if filter.CollectionID != "" {
qb.Where("collection_id", filter.CollectionID)
}
// Apply status filter if provided
if filter.Status != "" {
qb.Where("status", filter.Status)
}
// Apply keyword filter if provided
if filter.Keywords != "" {
qb.Where("name", "like", fmt.Sprintf("%%%s%%", filter.Keywords))
}
// Execute delete and return number of deleted records
return qb.Delete()
}
// SaveKnowledge saves knowledge collection information
func (conv *Xun) SaveKnowledge(knowledge map[string]interface{}) (interface{}, error) {
// Validate required fields
requiredFields := []string{"collection_id", "name", "uid"}
for _, field := range requiredFields {
if _, ok := knowledge[field]; !ok {
return nil, fmt.Errorf("field %s is required", field)
}
if knowledge[field] == nil || knowledge[field] == "" {
return nil, fmt.Errorf("field %s cannot be empty", field)
}
}
// Create a copy of the knowledge map to avoid modifying the original
knowledgeCopy := make(map[string]interface{})
for k, v := range knowledge {
knowledgeCopy[k] = v
}
// Process JSON fields
jsonFields := []string{"scope", "option"}
for _, field := range jsonFields {
if val, ok := knowledgeCopy[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 {
knowledgeCopy[field] = parsed
}
}
}
}
// Check if knowledge exists
exists, err := conv.query.New().
Table(conv.getKnowledgeTable()).
Where("collection_id", knowledgeCopy["collection_id"]).
Exists()
if err != nil {
return nil, err
}
// Convert JSON fields to strings for storage
for _, field := range jsonFields {
if val, ok := knowledgeCopy[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)
}
knowledgeCopy[field] = jsonStr
}
}
// Update or insert
if exists {
knowledgeCopy["updated_at"] = time.Now()
_, err := conv.query.New().
Table(conv.getKnowledgeTable()).
Where("collection_id", knowledgeCopy["collection_id"]).
Update(knowledgeCopy)
if err != nil {
return nil, err
}
return knowledgeCopy["collection_id"], nil
}
knowledgeCopy["created_at"] = time.Now()
err = conv.query.New().
Table(conv.getKnowledgeTable()).
Insert(knowledgeCopy)
if err != nil {
return nil, err
}
return knowledgeCopy["collection_id"], nil
}
// DeleteKnowledge deletes a knowledge collection by collection_id
func (conv *Xun) DeleteKnowledge(collectionID string) error {
// Check if knowledge exists
exists, err := conv.query.New().
Table(conv.getKnowledgeTable()).
Where("collection_id", collectionID).
Exists()
if err != nil {
return err
}
if !exists {
return fmt.Errorf("knowledge collection %s not found", collectionID)
}
_, err = conv.query.New().
Table(conv.getKnowledgeTable()).
Where("collection_id", collectionID).
Delete()
return err
}
// GetKnowledges retrieves knowledge collections with pagination and filtering
func (conv *Xun) GetKnowledges(filter KnowledgeFilter, locale ...string) (*KnowledgeResponse, error) {
qb := conv.query.New().
Table(conv.getKnowledgeTable())
// Apply UID filter if provided
if filter.UID != "" {
qb.Where("uid", filter.UID)
}
// Apply name filter if provided
if filter.Name != "" {
qb.Where("name", "like", fmt.Sprintf("%%%s%%", filter.Name))
}
// 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 public filter if provided
if filter.Public != nil {
qb.Where("public", *filter.Public)
}
// Apply readonly filter if provided
if filter.Readonly != nil {
qb.Where("readonly", *filter.Readonly)
}
// Apply system filter if provided
if filter.System != nil {
qb.Where("system", *filter.System)
}
// 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("sort", "asc").
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{"scope", "option"}
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 &KnowledgeResponse{
Data: data,
Page: filter.Page,
PageSize: filter.PageSize,
PageCnt: totalPages,
Next: nextPage,
Prev: prevPage,
Total: total,
}, nil
}
// GetKnowledge retrieves a single knowledge collection by collection_id
func (conv *Xun) GetKnowledge(collectionID string, locale ...string) (map[string]interface{}, error) {
row, err := conv.query.New().
Table(conv.getKnowledgeTable()).
Where("collection_id", collectionID).
First()
if err != nil {
return nil, err
}
if row == nil {
return nil, fmt.Errorf("knowledge collection %s not found", collectionID)
}
data := row.ToMap()
if data == nil || len(data) == 0 {
return nil, fmt.Errorf("the knowledge collection %s is empty", collectionID)
}
// Parse JSON fields
jsonFields := []string{"scope", "option"}
conv.parseJSONFields(data, jsonFields)
return data, nil
}
// DeleteKnowledges deletes knowledge collections based on filter conditions
func (conv *Xun) DeleteKnowledges(filter KnowledgeFilter) (int64, error) {
qb := conv.query.New().
Table(conv.getKnowledgeTable())
// Apply UID filter if provided
if filter.UID != "" {
qb.Where("uid", filter.UID)
}
// Apply name filter if provided
if filter.Name != "" {
qb.Where("name", "like", fmt.Sprintf("%%%s%%", filter.Name))
}
// 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 public filter if provided
if filter.Public != nil {
qb.Where("public", *filter.Public)
}
// Apply readonly filter if provided
if filter.Readonly != nil {
qb.Where("readonly", *filter.Readonly)
}
// Apply system filter if provided
if filter.System != nil {
qb.Where("system", *filter.System)
}
// Execute delete and return number of deleted records
return qb.Delete()
}

3736
agent/store/xun_test.go
View file

File diff suppressed because it is too large Load diff

59
agent/types.go
View file

@ -3,10 +3,8 @@ package agent
import ( import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/yaoapp/yao/agent/assistant" "github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/agent/rag"
"github.com/yaoapp/yao/agent/store" "github.com/yaoapp/yao/agent/store"
"github.com/yaoapp/yao/agent/vision" "github.com/yaoapp/yao/agent/vision"
"github.com/yaoapp/yao/attachment"
) )
// DSL AI assistant // DSL AI assistant
@ -16,9 +14,9 @@ type DSL struct {
// =============================== // ===============================
Use *Use `json:"use,omitempty" yaml:"use,omitempty"` // Which assistant to use default, title, prompt Use *Use `json:"use,omitempty" yaml:"use,omitempty"` // Which assistant to use default, title, prompt
StoreSetting store.Setting `json:"store" yaml:"store"` // The store setting of the assistant StoreSetting store.Setting `json:"store" yaml:"store"` // The store setting of the assistant
AuthSetting *Auth `json:"auth,omitempty" yaml:"auth,omitempty"` // Authenticate Settings // AuthSetting *Auth `json:"auth,omitempty" yaml:"auth,omitempty"` // Authenticate Settings
UploadSetting *Upload `json:"upload,omitempty" yaml:"upload,omitempty"` // Upload Settings // UploadSetting *Upload `json:"upload,omitempty" yaml:"upload,omitempty"` // Upload Settings
KnowledgeSetting *Knowledge `json:"knowledge,omitempty" yaml:"knowledge,omitempty"` // Knowledge base Settings // KnowledgeSetting *Knowledge `json:"knowledge,omitempty" yaml:"knowledge,omitempty"` // Knowledge base Settings
// Global External Settings - connectors, tools, etc. // Global External Settings - connectors, tools, etc.
// =============================== // ===============================
@ -26,15 +24,14 @@ type DSL struct {
// Agent API Settings // Agent API Settings
// ===============================s // ===============================s
Guard string `json:"guard,omitempty" yaml:"guard,omitempty"` // The guard of the assistant // Guard string `json:"guard,omitempty" yaml:"guard,omitempty"` // The guard of the assistant
Allows []string `json:"allows,omitempty" yaml:"allows,omitempty"` // The allowed domains of the assistant // Allows []string `json:"allows,omitempty" yaml:"allows,omitempty"` // The allowed domains of the assistant
// Internal // Internal
// =============================== // ===============================
ID string `json:"-" yaml:"-"` // The id of the instance ID string `json:"-" yaml:"-"` // The id of the instance
Assistant assistant.API `json:"-" yaml:"-"` // The default assistant Assistant assistant.API `json:"-" yaml:"-"` // The default assistant
Store store.Store `json:"-" yaml:"-"` // The store of the assistant Store store.Store `json:"-" yaml:"-"` // The store of the assistant
RAG *rag.RAG `json:"-" yaml:"-"`
Vision *vision.Vision `json:"-" yaml:"-"` Vision *vision.Vision `json:"-" yaml:"-"`
GuardHandlers []gin.HandlerFunc `json:"-" yaml:"-"` GuardHandlers []gin.HandlerFunc `json:"-" yaml:"-"`
} }
@ -78,52 +75,6 @@ type AuthFields struct {
Permission string `json:"permission,omitempty" yaml:"permission,omitempty"` // the field name of the user permission, default is permission Permission string `json:"permission,omitempty" yaml:"permission,omitempty"` // the field name of the user permission, default is permission
} }
// Upload the upload setting
// ===============================
type Upload struct {
Chat *attachment.ManagerOption `json:"chat,omitempty" yaml:"chat,omitempty"` // Chat conversation upload setting, if not set use the local and root path is `/attachments`.
Assets *attachment.ManagerOption `json:"assets,omitempty" yaml:"assets,omitempty"` // Asset upload setting, if not set use the chat upload setting.
Knowledge *attachment.ManagerOption `json:"knowledge,omitempty" yaml:"knowledge,omitempty"` // Knowledge base upload setting, if not set use the chat upload setting.
}
// UploadOption the upload option
type UploadOption struct {
attachment.UploadOption
Public bool `json:"public,omitempty" yaml:"public,omitempty, form:public"` // The public of the file, default is false
Scope interface{} `json:"scope,omitempty" yaml:"scope,omitempty, form:scope"` // The scope of the file, default is private
CollectionID string `json:"collection_id,omitempty" yaml:"collection_id,omitempty, form:collection_id"` // The collection id of the file, default is empty
Knowledge bool `json:"knowledge,omitempty" form:"knowledge"` // Push to knowledge base, Optional, default is false
ChatID string `json:"chat_id,omitempty" form:"chat_id"` // Chat ID, Optional
AssistantID string `json:"assistant_id,omitempty" form:"assistant_id"` // Assistant ID, Optional
UserID string `json:"user_id,omitempty"` // User ID, Optional (used to build Groups)
}
// Knowledge base Settings
// ===============================
type Knowledge struct {
Vector KnowledgeVector `json:"vector" yaml:"vector"` // The vector database driver
Graph KnowledgeGraph `json:"graph" yaml:"graph"` // The graph database driver
Vectorizer KnowledgeVectorizer `json:"vectorizer" yaml:"vectorizer"` // The vectorizer driver
}
// KnowledgeVectorizer the knowledge vectorizer
type KnowledgeVectorizer struct {
Driver string `json:"driver" yaml:"driver"`
Options map[string]interface{} `json:"options" yaml:"options"`
}
// KnowledgeVector the knowledge vector
type KnowledgeVector struct {
Driver string `json:"driver" yaml:"driver"`
Options map[string]interface{} `json:"options" yaml:"options"`
}
// KnowledgeGraph the knowledge graph
type KnowledgeGraph struct {
Driver string `json:"driver" yaml:"driver"`
Options map[string]interface{} `json:"options" yaml:"options"`
}
// Mention Structure // Mention Structure
// =============================== // ===============================
type Mention struct { type Mention struct {

1000
agent/vision/vision_test.go
View file

File diff suppressed because it is too large Load diff

288
data/bindata.go
View file

File diff suppressed because it is too large Load diff

22
widgets/app/app.go
View file

@ -564,28 +564,6 @@ func processXgen(process *process.Process) interface{} {
// Available connectors // Available connectors
agentConfig["connectors"] = connector.AIConnectors agentConfig["connectors"] = connector.AIConnectors
// Available storages
agentConfig["storages"] = map[string]interface{}{
"chat": map[string]interface{}{
"max_size": agent.Agent.UploadSetting.Chat.MaxSize,
"chunk_size": agent.Agent.UploadSetting.Chat.ChunkSize,
"allowed_types": agent.Agent.UploadSetting.Chat.AllowedTypes,
"gzip": agent.Agent.UploadSetting.Chat.Gzip,
},
"assets": map[string]interface{}{
"max_size": agent.Agent.UploadSetting.Assets.MaxSize,
"chunk_size": agent.Agent.UploadSetting.Assets.ChunkSize,
"allowed_types": agent.Agent.UploadSetting.Assets.AllowedTypes,
"gzip": agent.Agent.UploadSetting.Assets.Gzip,
},
"knowledge": map[string]interface{}{
"max_size": agent.Agent.UploadSetting.Knowledge.MaxSize,
"chunk_size": agent.Agent.UploadSetting.Knowledge.ChunkSize,
"allowed_types": agent.Agent.UploadSetting.Knowledge.AllowedTypes,
"gzip": agent.Agent.UploadSetting.Knowledge.Gzip,
},
}
} }
// OpenAPI Settings // OpenAPI Settings

2
yao/models/agent/assistant.mod.yao
View file

@ -203,5 +203,5 @@
"comment": "Index for assistant sorting and automation" "comment": "Index for assistant sorting and automation"
} }
], ],
"option": { "timestamps": true, "soft_deletes": false } "option": { "timestamps": true, "soft_deletes": false, "permission": true }
} }

2
yao/models/agent/chat.mod.yao
View file

@ -90,5 +90,5 @@
"comment": "Index for silent mode filtering" "comment": "Index for silent mode filtering"
} }
], ],
"option": { "timestamps": true, "soft_deletes": false } "option": { "timestamps": true, "soft_deletes": false, "permission": true }
} }

2
yao/models/agent/history.mod.yao
View file

@ -170,5 +170,5 @@
"comment": "Index for expiration and cleanup" "comment": "Index for expiration and cleanup"
} }
], ],
"option": { "timestamps": true, "soft_deletes": false } "option": { "timestamps": true, "soft_deletes": false, "permission": true }
} }