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

Refactor agent and assistant structure for improved clarity and performance

Browse source 
- Removed deprecated API and vision components, streamlining the agent's architecture.
- Updated the Load function to utilize a new agentDSL variable, enhancing the management of assistant capabilities.
- Enhanced context handling by introducing a message metadata store for thread-safe operations, improving message tracking and management.
- Refactored context methods to eliminate deprecated fields, ensuring cleaner and more maintainable code.
- Improved documentation and comments throughout the codebase to clarify changes and enhance developer understanding.
This commit is contained in:
Max 2025-11-30 11:47:25 +08:00
parent ab240443f8
commit 0fe0843b43
39 changed files with 184 additions and 8192 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

38
agent/api/agent.go
View file

@ -1,38 +0,0 @@
package api
import (
"github.com/gin-gonic/gin"
"github.com/yaoapp/yao/agent/assistant"
chatctx "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/types"
)
// Agent the agent AI assistant
var Agent *API
// API the agent API
type API struct {
*types.DSL
}
// Answer reply the message
func (agent *API) Answer(ctx chatctx.Context, question string, c *gin.Context) error {
var err error
var ast assistant.API = Agent.Assistant
if ctx.AssistantID != "" {
ast, err = agent.Select(ctx.AssistantID)
if err != nil {
return err
}
}
_, err = ast.Execute(c, ctx, question, nil)
return err
}
// Select select an assistant
func (agent *API) Select(id string) (assistant.API, error) {
if id == "" {
return Agent.Assistant, nil
}
return assistant.Get(id)
}

1061
agent/api/api.go
View file

File diff suppressed because it is too large Load diff

233
agent/api/api_test.go
View file

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

1286
agent/assistant/api.go
View file

File diff suppressed because it is too large Load diff

56
agent/assistant/assistant.go
View file

@ -11,6 +11,62 @@ import (
sui "github.com/yaoapp/yao/sui/core" sui "github.com/yaoapp/yao/sui/core"
) )
// Get get the assistant by id
func Get(id string) (*Assistant, error) {
return LoadStore(id)
}
// GetByConnector get the assistant by connector
func GetByConnector(connector string, name string) (*Assistant, error) {
id := "connector:" + connector
assistant, exists := loaded.Get(id)
if exists {
return assistant, nil
}
data := map[string]interface{}{
"assistant_id": id,
"connector": connector,
"description": "Default assistant for " + connector,
"name": name,
"type": "assistant",
}
assistant, err := loadMap(data)
if err != nil {
return nil, err
}
loaded.Put(assistant)
return assistant, nil
}
// GetPlaceholder returns the placeholder of the assistant
func (ast *Assistant) GetPlaceholder(locale string) *store.Placeholder {
prompts := []string{}
if ast.Placeholder.Prompts != nil {
prompts = i18n.Translate(ast.ID, locale, ast.Placeholder.Prompts).([]string)
}
title := i18n.Translate(ast.ID, locale, ast.Placeholder.Title).(string)
description := i18n.Translate(ast.ID, locale, ast.Placeholder.Description).(string)
return &store.Placeholder{
Title: title,
Description: description,
Prompts: prompts,
}
}
// GetName returns the name of the assistant
func (ast *Assistant) GetName(locale string) string {
return i18n.Translate(ast.ID, locale, ast.Name).(string)
}
// GetDescription returns the description of the assistant
func (ast *Assistant) GetDescription(locale string) string {
return i18n.Translate(ast.ID, locale, ast.Description).(string)
}
// Save save the assistant // Save save the assistant
func (ast *Assistant) Save() error { func (ast *Assistant) Save() error {
if storage == nil { if storage == nil {

361
agent/assistant/call.go
View file

@ -1,361 +0,0 @@
package assistant
import (
"context"
"fmt"
"github.com/fatih/color"
"github.com/google/uuid"
"github.com/yaoapp/gou/runtime/v8/bridge"
"github.com/yaoapp/kun/log"
chatctx "github.com/yaoapp/yao/agent/context"
chatMessage "github.com/yaoapp/yao/agent/message"
"rogchap.com/v8go"
)
// objectCall is the object for the call function
type objectCall struct{}
// OptionsCall is the options for the call function
type OptionsCall struct {
Retry OptionsCallRetry `json:"retry,omitempty"` // Retry options
Options map[string]interface{} `json:"options,omitempty"` // LLM API options
Silent bool `json:"silent,omitempty"` // Silent mode, default is true
}
// OptionsCallRetry is the retry options for the call function
type OptionsCallRetry struct {
Times int `json:"times,omitempty"` // Retry times, default is 3
Delay int `json:"delay,omitempty"` // Retry delay, default is 200
DelayMax int `json:"delay_max,omitempty"` // Retry delay max, default is 5000
Prompt string `json:"prompt,omitempty"` // Retry prompt, default is "Please fix the error. \n {{ error }}"
}
// allowedEvents is the allowed events for the call function
var allowedEvents = map[string]bool{
"done": true,
"retry": true,
"message": true,
}
var callProps = []string{
"assistant_id",
"input",
"options",
"retry_times",
}
// jsNewPlan create a plan object and return it
func jsCall(info *v8go.FunctionCallbackInfo) *v8go.Value {
args := info.Args()
if len(args) < 2 {
return bridge.JsException(info.Context(), "Run requires at least two arguments")
}
options := v8go.Undefined(info.Context().Isolate())
if len(args) > 2 {
options = args[2]
}
// Export the object
obj := &objectCall{}
objectTmpl := obj.ExportObject(info)
this, err := objectTmpl.NewInstance(info.Context())
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
// Copy global properties
global := info.This()
for _, prop := range objectProperties {
if !global.Has(prop) {
continue
}
value, err := global.Get(prop)
if err != nil {
return bridge.JsException(info.Context(), fmt.Sprintf("Failed to get property %s: %s", prop, err.Error()))
}
this.Set(prop, value)
}
this.Set("assistant_id", args[0])
this.Set("input", args[1])
this.Set("options", options)
this.Set("retry_times", int32(1))
return this.Value
}
// ExportObject Export as a FS Object
func (obj *objectCall) ExportObject(info *v8go.FunctionCallbackInfo) *v8go.ObjectTemplate {
tmpl := v8go.NewObjectTemplate(info.Context().Isolate())
tmpl.Set("On", v8go.NewFunctionTemplate(info.Context().Isolate(), obj.on)) // On the call
tmpl.Set("Run", v8go.NewFunctionTemplate(info.Context().Isolate(), obj.run)) // Run the call
return tmpl
}
// on bind the callback to the call object
func (obj *objectCall) on(info *v8go.FunctionCallbackInfo) *v8go.Value {
args := info.Args()
if len(args) < 2 {
return bridge.JsException(info.Context(), "On requires at least one argument")
}
if !args[0].IsString() {
return bridge.JsException(info.Context(), "The first argument should be a string")
}
name := args[0].String()
if !allowedEvents[name] {
return bridge.JsException(info.Context(), fmt.Sprintf("Invalid event %s", name))
}
cb := args[1]
if !cb.IsFunction() {
return bridge.JsException(info.Context(), fmt.Sprintf("The second argument should be a function for event %s", name))
}
this := info.This()
this.Set(fmt.Sprintf("on_%s", name), cb)
return this.Value
}
// run run the call
func (obj *objectCall) run(info *v8go.FunctionCallbackInfo) *v8go.Value {
this := info.This()
args := info.Args()
global, err := getGlobal(info.Context(), this)
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
goArgs := []interface{}{}
jsArgs := []v8go.Valuer{}
if len(args) > 0 {
for _, arg := range args {
v, err := bridge.GoValue(arg, info.Context())
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
goArgs = append(goArgs, v)
jsArgs = append(jsArgs, arg)
}
}
// Get the assistant id
jsAssistantID, err := this.Get("assistant_id")
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
assistantID := jsAssistantID.String()
// Get the input
jsInput, err := this.Get("input")
if err != nil {
return bridge.JsException(info.Context(), fmt.Sprintf("Failed to get the input: %s", err.Error()))
}
input, err := bridge.GoValue(jsInput, info.Context())
if err != nil {
return bridge.JsException(info.Context(), fmt.Sprintf("Failed to unmarshal the input: %s", err.Error()))
}
// Get the retry input
if this.Has("retry_input") {
jsRetryInput, err := this.Get("retry_input")
if err != nil {
return bridge.JsException(info.Context(), fmt.Sprintf("Failed to get the retry input: %s", err.Error()))
}
input, err = bridge.GoValue(jsRetryInput, info.Context())
if err != nil {
return bridge.JsException(info.Context(), fmt.Sprintf("Failed to unmarshal the retry input: %s", err.Error()))
}
}
// Options
options := OptionsCall{
// Retry: OptionsCallRetry{
// Times: 3,
// Delay: 200,
// DelayMax: 1000,
// Prompt: "{{ input }}\n**Answer is not correct, please try again.**\nError:\n{{ error }} \nAssistant's last answer:\n{{ output }}",
// },
Silent: true,
Options: map[string]interface{}{}, // LLM API options
}
// Get the options
if this.Has("options") {
jsOptions, err := this.Get("options")
if err != nil {
return bridge.JsException(info.Context(), fmt.Sprintf("Failed to get the options: %s", err.Error()))
}
// Check if the options is undefined
if !jsOptions.IsUndefined() {
err = bridge.Unmarshal(jsOptions, &options)
if err != nil {
return bridge.JsException(info.Context(), fmt.Sprintf("Failed to unmarshal the options: %s", err.Error()))
}
}
}
// Get the assistant
newAst, err := Get(assistantID)
if err != nil {
return bridge.JsException(info.Context(), fmt.Sprintf("Failed to get the assistant: %s", err.Error()))
}
// Get the message event ( it will be used for the message event )
eventMessage := ""
goCallProps := map[string]interface{}{}
if this.Has("on_message") {
jsEventMessage, err := this.Get("on_message")
if err != nil {
return bridge.JsException(info.Context(), fmt.Sprintf("Failed to get the message: %s", err.Error()))
}
eventMessage = jsEventMessage.String()
for _, prop := range callProps {
if this.Has(prop) {
value, err := this.Get(prop)
if err != nil {
return bridge.JsException(info.Context(), fmt.Sprintf("Failed to get the %s property: %s", prop, err.Error()))
}
goValue, err := bridge.GoValue(value, info.Context())
if err != nil {
return bridge.JsException(info.Context(), fmt.Sprintf("Failed to get the %s property: %s", prop, err.Error()))
}
goCallProps[prop] = goValue
}
}
}
// Update the chat context
var chatCtx chatctx.Context = global.ChatContext
chatCtx.AssistantID = assistantID
chatCtx.ChatID = fmt.Sprintf("call_%s", uuid.New().String()) // New chat id
chatCtx.Silent = options.Silent
chatCtx.Referer = chatctx.RefererScript // Set the referer to hookscript
chatCtx.Args = goArgs // Arguments for call
// Define the callback function
var cb func(msg *chatMessage.Message) = nil
var output = []chatMessage.Message{}
cb = func(msg *chatMessage.Message) {
output = append(output, *msg)
if eventMessage != "" {
err := obj.triggerAnonymous(chatCtx, global, goCallProps, eventMessage, goArgs, msg)
if err != nil {
color.Red("Failed to trigger the message event: %s", err.Error())
log.Error("Failed to trigger the message event: %s", err.Error())
return
}
}
}
// Execute the assistant
result, err := newAst.Execute(global.GinContext, chatCtx, input, options.Options, cb) // Execute the assistant
if err != nil {
// result, err = obj.retry(jsArgs, err, input, output, info, options)
// if err != nil {
// return bridge.JsException(info.Context(), err.Error())
// }
return bridge.JsException(info.Context(), err.Error())
}
// Copy props
for name, value := range goCallProps {
info.Context().Global().Set(name, value)
}
// Trigger the done event
doneResult, err := obj.trigger(info, "done", jsArgs...)
if err != nil {
// result, err = obj.retry(jsArgs, err, input, output, info, options)
// if err != nil {
// return bridge.JsException(info.Context(), err.Error())
// }
return bridge.JsException(info.Context(), err.Error())
}
// Return the done result
if doneResult != nil && !doneResult.IsUndefined() {
return doneResult
}
// Return Value
switch v := result.(type) {
case *v8go.Value:
return v
case error:
return bridge.JsException(info.Context(), v.Error())
}
// Return Value
jsResult, err := bridge.JsValue(info.Context(), result)
if err != nil {
return bridge.JsException(info.Context(), fmt.Sprintf("Failed to get the result: %s", err.Error()))
}
return jsResult
}
func (obj *objectCall) triggerAnonymous(chatCtx chatctx.Context, global *GlobalVariables, goCallProps map[string]interface{}, source string, bindArgs []interface{}, fnArgs ...interface{}) error {
ctx, err := global.Assistant.Script.NewContext(global.ChatContext.Sid, nil)
if err != nil {
return err
}
defer ctx.Close()
// Update Context
global.Assistant.InitObject(ctx, global.GinContext, chatCtx, global.Contents)
// Copy props
for k, v := range goCallProps {
ctx.WithGlobal(k, v)
}
// Add the args
ctx.WithGlobal("args", bindArgs)
_, err = ctx.CallAnonymousWith(context.Background(), source, fnArgs...)
if err != nil {
return err
}
return nil
}
// trigger trigger the callback
func (obj *objectCall) trigger(info *v8go.FunctionCallbackInfo, name string, fnArgs ...v8go.Valuer) (*v8go.Value, error) {
// Try to get the callback
this := info.This()
if this.Has(fmt.Sprintf("on_%s", name)) {
event, err := this.Get(fmt.Sprintf("on_%s", name))
if err != nil {
return nil, err
}
if event.IsFunction() {
cb, err := event.AsFunction()
if err != nil {
return nil, err
}
result, err := cb.Call(this, fnArgs...)
if err != nil {
return nil, err
}
return result, nil
}
}
return nil, nil
}

7
agent/assistant/hook/script.go
View file

@ -12,7 +12,12 @@ func (s *Script) Execute(ctx *context.Context, method string, args ...interface{
return nil, nil return nil, nil
} }
scriptCtx, err := s.NewContext(ctx.Sid, nil) var sid = ""
if ctx.Authorized != nil {
sid = ctx.Authorized.SessionID
}
scriptCtx, err := s.NewContext(sid, nil)
if err != nil { if err != nil {
return nil, err return nil, err
} }

316
agent/assistant/hooks.go
View file

@ -1,316 +0,0 @@
package assistant
import (
"context"
"fmt"
"os"
"strings"
"time"
"github.com/gin-gonic/gin"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/kun/log"
chatctx "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/message"
chatMessage "github.com/yaoapp/yao/agent/message"
)
// HookCreate create a new assistant
func (ast *Assistant) HookCreate(c *gin.Context, context chatctx.Context, input []chatMessage.Message, options map[string]interface{}, contents *chatMessage.Contents) (*ResHookInit, error) {
// Create timeout context
ctx := ast.createBackgroundContext()
v, err := ast.call(ctx, "Create", c, contents, context, input, options)
if err != nil {
if err.Error() == HookErrorMethodNotFound {
return nil, nil
}
return nil, err
}
response := &ResHookInit{Result: nil}
switch v := v.(type) {
case map[string]interface{}:
if res, ok := v["assistant_id"].(string); ok {
response.AssistantID = res
}
if res, ok := v["chat_id"].(string); ok {
response.ChatID = res
}
// input
if input, has := v["input"]; has {
raw, _ := jsoniter.MarshalToString(input)
vv := []message.Message{}
err := jsoniter.UnmarshalFromString(raw, &vv)
if err != nil {
return nil, err
}
response.Input = vv
}
// result
if result, has := v["result"]; has {
response.Result = result
}
if res, ok := v["next"].(map[string]interface{}); ok {
response.Next = &NextAction{}
if name, ok := res["action"].(string); ok {
response.Next.Action = name
}
if payload, ok := res["payload"].(map[string]interface{}); ok {
response.Next.Payload = payload
}
}
case string:
response.AssistantID = v
response.ChatID = context.ChatID
case nil:
response.AssistantID = ast.ID
response.ChatID = context.ChatID
}
return response, nil
}
// HookRetry Handle retry of assistant response
func (ast *Assistant) HookRetry(c *gin.Context, context chatctx.Context, input []message.Message, contents *chatMessage.Contents, errmsg string) (interface{}, error) {
ctx := ast.createBackgroundContext()
output := []message.Data{}
if len(input) < 1 {
return "", fmt.Errorf("no input")
}
var lastInput message.Message = input[len(input)-1]
for _, data := range contents.Data {
if data.Type == "think" {
continue
}
output = append(output, data)
}
v, err := ast.call(ctx, "Retry", c, contents, context, lastInput.String(), output, errmsg)
if err != nil {
if err.Error() == HookErrorMethodNotFound {
return nil, nil
}
return nil, err
}
switch v := v.(type) {
case string, bool:
return v, nil
case map[string]interface{}:
// Has Action
if _, has := v["action"]; has {
var next NextAction
raw, _ := jsoniter.MarshalToString(v)
err := jsoniter.UnmarshalFromString(raw, &next)
if err != nil {
return nil, err
}
return &next, nil
}
// Ignore the error, and return the specific result
return v, nil
}
return nil, nil
}
// HookDone Handle completion of assistant response
func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []message.Message, contents *chatMessage.Contents) (*ResHookDone, error) {
// Create timeout context
ctx := ast.createBackgroundContext()
// format the output
// 1. Remove thinking message
// 2. Parse the tool call message content
output := []message.Data{}
if contents != nil && contents.Data != nil {
for _, data := range contents.Data {
if data.Type == "think" {
continue
}
// parse the tool call message content
if data.Type == "tool" && data.Props != nil {
props := map[string]interface{}{}
if text, ok := data.Props["text"].(string); ok {
// Extract the content between <tool> and </tool> tags more reliably
startTag := "<tool>"
endTag := "</tool>"
startIndex := strings.Index(text, startTag)
if startIndex != -1 {
// Find the content after <tool>
content := text[startIndex+len(startTag):]
endIndex := strings.LastIndex(content, endTag)
if endIndex != -1 {
// Extract the content between tags
text = content[:endIndex]
text = strings.TrimSpace(text)
if os.Getenv("YAO_AGENT_PRINT_TOOL_CALL") == "true" {
log.Trace("[TOOL CALL] %s", text)
}
}
}
// Parse the text into props
err := ParseJSON(text, &props)
if err != nil {
props["error"] = fmt.Sprintf("Can not parse the tool call: %s\n--original--\n%s", err.Error(), text)
}
}
output = append(output, message.Data{Type: "tool", Props: props})
continue
}
output = append(output, data)
}
}
v, err := ast.call(ctx, "Done", c, contents, context, input, output)
if err != nil {
if err.Error() == HookErrorMethodNotFound {
return nil, nil
}
return nil, err
}
response := &ResHookDone{Input: input, Output: contents.Data}
switch v := v.(type) {
case map[string]interface{}:
if res, ok := v["output"].(string); ok {
vv := []message.Data{}
err := jsoniter.UnmarshalFromString(res, &vv)
if err != nil {
return nil, err
}
response.Output = vv
}
if res, ok := v["output"].([]interface{}); ok {
vv := []message.Data{}
raw, _ := jsoniter.MarshalToString(res)
err := jsoniter.UnmarshalFromString(raw, &vv)
if err != nil {
return nil, err
}
response.Output = vv
}
// has result
if res, has := v["result"]; has {
response.Result = res
}
if res, ok := v["next"].(map[string]interface{}); ok {
response.Next = &NextAction{}
if name, ok := res["action"].(string); ok {
response.Next.Action = name
}
if payload, ok := res["payload"].(map[string]interface{}); ok {
response.Next.Payload = payload
}
}
case string:
vv := []message.Data{}
err := jsoniter.UnmarshalFromString(v, &vv)
if err != nil {
return nil, err
}
response.Output = vv
}
return response, nil
}
// HookFail Handle failure of assistant response
func (ast *Assistant) HookFail(c *gin.Context, context chatctx.Context, input []message.Message, err error, contents *chatMessage.Contents) (*ResHookFail, error) {
// Create timeout context
ctx, cancel := ast.createTimeoutContext(5 * time.Second)
defer cancel()
v, callErr := ast.call(ctx, "Fail", c, contents, context, input, err.Error())
if callErr != nil {
if callErr.Error() == HookErrorMethodNotFound {
return nil, nil
}
return nil, callErr
}
response := &ResHookFail{
Input: input,
Output: contents.Text(),
Error: err.Error(),
}
switch v := v.(type) {
case map[string]interface{}:
if res, ok := v["output"].(string); ok {
response.Output = res
}
if res, ok := v["error"].(string); ok {
response.Error = res
}
if res, ok := v["next"].(map[string]interface{}); ok {
response.Next = &NextAction{}
if name, ok := res["action"].(string); ok {
response.Next.Action = name
}
if payload, ok := res["payload"].(map[string]interface{}); ok {
response.Next.Payload = payload
}
}
case string:
response.Output = v
}
return response, nil
}
// createTimeoutContext creates a timeout context with 5 seconds timeout
func (ast *Assistant) createTimeoutContext(time time.Duration) (context.Context, context.CancelFunc) {
ctx, cancel := context.WithTimeout(context.Background(), time)
return ctx, cancel
}
// createBackgroundContext creates a background context
func (ast *Assistant) createBackgroundContext() context.Context {
return context.Background()
}
// Call the script method
func (ast *Assistant) call(ctx context.Context, method string, c *gin.Context, contents *chatMessage.Contents, context chatctx.Context, args ...any) (interface{}, error) {
if ast.Script == nil {
return nil, nil
}
scriptCtx, err := ast.Script.NewContext(context.Sid, nil)
if err != nil {
return nil, err
}
defer scriptCtx.Close()
// Initialize the object, add the global variables, methods to the script context
ast.InitObject(scriptCtx, c, context, contents)
// Check if the method exists
if !scriptCtx.Global().Has(method) {
return nil, fmt.Errorf(HookErrorMethodNotFound)
}
// Call the method directly in the current thread
if scriptCtx != nil {
return scriptCtx.CallWith(ctx, method, args...)
}
return nil, nil
}

8
agent/assistant/load.go
View file

@ -17,7 +17,6 @@ import (
"github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n" "github.com/yaoapp/yao/agent/i18n"
store "github.com/yaoapp/yao/agent/store/types" store "github.com/yaoapp/yao/agent/store/types"
agentvision "github.com/yaoapp/yao/agent/vision"
"github.com/yaoapp/yao/openai" "github.com/yaoapp/yao/openai"
"github.com/yaoapp/yao/share" "github.com/yaoapp/yao/share"
"gopkg.in/yaml.v3" "gopkg.in/yaml.v3"
@ -28,7 +27,6 @@ var loaded = NewCache(200) // 200 is the default capacity
var storage store.Store = nil var storage store.Store = nil
var search interface{} = nil var search interface{} = nil
var modelCapabilities map[string]ModelCapabilities = map[string]ModelCapabilities{} var modelCapabilities map[string]ModelCapabilities = map[string]ModelCapabilities{}
var vision *agentvision.Vision = nil
var defaultConnector string = "" // default connector var defaultConnector string = "" // default connector
var globalUses *context.Uses = nil // global uses configuration from agent.yml var globalUses *context.Uses = nil // global uses configuration from agent.yml
@ -132,11 +130,6 @@ func SetStorage(s store.Store) {
storage = s storage = s
} }
// SetVision set the vision
func SetVision(v *agentvision.Vision) {
vision = v
}
// SetModelCapabilities set the model capabilities configuration // SetModelCapabilities set the model capabilities configuration
func SetModelCapabilities(capabilities map[string]ModelCapabilities) { func SetModelCapabilities(capabilities map[string]ModelCapabilities) {
modelCapabilities = capabilities modelCapabilities = capabilities
@ -710,7 +703,6 @@ func (ast *Assistant) initialize() error {
return err return err
} }
defer scriptCtx.Close() defer scriptCtx.Close()
ast.initHook = scriptCtx.Global().Has("init")
} }
return nil return nil

1
agent/assistant/next.go
View file

@ -54,7 +54,6 @@ func (ast *Assistant) handleDelegation(
delegatedCtx := &agentContext.Context{ delegatedCtx := &agentContext.Context{
Context: ctx.Context, Context: ctx.Context,
Locale: ctx.Locale, Locale: ctx.Locale,
Sid: ctx.Sid,
Stack: ctx.Stack, // Maintain the call stack Stack: ctx.Stack, // Maintain the call stack
Authorized: ctx.Authorized, Authorized: ctx.Authorized,
Metadata: ctx.Metadata, Metadata: ctx.Metadata,

376
agent/assistant/object.go
View file

@ -1,376 +0,0 @@
package assistant
import (
"fmt"
"strings"
"github.com/gin-gonic/gin"
v8 "github.com/yaoapp/gou/runtime/v8"
"github.com/yaoapp/gou/runtime/v8/bridge"
chatctx "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/message"
chatMessage "github.com/yaoapp/yao/agent/message"
sui "github.com/yaoapp/yao/sui/core"
"rogchap.com/v8go"
)
// objectProperties is the properties of the assistant object
var objectProperties = []string{
"__yao_agent_global",
"assistant",
"context",
"Plan",
"Send",
"Call",
"Assets",
"Set",
"Get",
"Del",
"Clear",
}
// GlobalVariables is the global variables for the assistant
type GlobalVariables struct {
Assistant *Assistant
Contents *chatMessage.Contents
GinContext *gin.Context
ChatContext chatctx.Context
}
// JsValue return the javascript value of the global variables
func (global *GlobalVariables) JsValue(ctx *v8go.Context) (*v8go.Value, error) {
return v8go.NewExternal(ctx.Isolate(), global)
}
// InitObject add the global variables and methods to the script context
func (ast *Assistant) InitObject(v8ctx *v8.Context, c *gin.Context, context chatctx.Context, contents *chatMessage.Contents) {
// Add global variables to the script context
global := &GlobalVariables{
Assistant: ast,
Contents: contents,
GinContext: c,
ChatContext: context,
}
// Add global variables to the script context
v8ctx.WithGlobal("__yao_agent_global", global)
// Add assistant to the script context
v8ctx.WithGlobal("assistant", ast.Map())
v8ctx.WithGlobal("context", context.Map())
// Add methods to the script contexts
v8ctx.WithFunction("Send", jsSend)
v8ctx.WithFunction("Assets", jsAssets)
v8ctx.WithFunction("MakeCall", jsCall) // Create a new call object
v8ctx.WithFunction("MakePlan", jsPlan) // Create a new plan object
// Shared space methods
v8ctx.WithFunction("Set", jsSet)
v8ctx.WithFunction("Get", jsGet)
v8ctx.WithFunction("Del", jsDel)
v8ctx.WithFunction("Clear", jsClear)
// Template methods
v8ctx.WithFunction("Replace", jsReplace)
}
// jsSet function, set a value to the shared space
func jsSet(info *v8go.FunctionCallbackInfo) *v8go.Value {
global, err := global(info)
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
if global.ChatContext.Space == nil {
return bridge.JsException(info.Context(), "Shared space is not set")
}
args := info.Args()
if len(args) < 2 {
return bridge.JsException(info.Context(), "Set requires at least two arguments")
}
if !args[0].IsString() {
return bridge.JsException(info.Context(), "Set requires a valid key")
}
// Validate the key
key := args[0].String()
if key == "" {
return bridge.JsException(info.Context(), "Set requires a valid key")
}
// Validate the value
value, err := bridge.GoValue(args[1], info.Context())
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
// Set the value
err = global.ChatContext.Space.Set(key, value)
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
return nil
}
// jsGet function, get a value from the shared space
func jsGet(info *v8go.FunctionCallbackInfo) *v8go.Value {
global, err := global(info)
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
if global.ChatContext.Space == nil {
return bridge.JsException(info.Context(), "Shared space is not set")
}
args := info.Args()
if len(args) < 1 {
return bridge.JsException(info.Context(), "Get requires at least one argument")
}
if !args[0].IsString() {
return bridge.JsException(info.Context(), "Get requires a valid key")
}
// Get the key
key := args[0].String()
if key == "" {
return bridge.JsException(info.Context(), "Get requires a valid key")
}
// Get the value
value, err := global.ChatContext.Space.Get(key)
if err != nil {
// If the key is not found, return null
if strings.Contains(err.Error(), "not found") {
return v8go.Null(info.Context().Isolate())
}
return bridge.JsException(info.Context(), err.Error())
}
jsValue, err := bridge.JsValue(info.Context(), value)
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
return jsValue
}
// jsDel function, delete a value from the shared space
func jsDel(info *v8go.FunctionCallbackInfo) *v8go.Value {
global, err := global(info)
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
if global.ChatContext.Space == nil {
return bridge.JsException(info.Context(), "Shared space is not set")
}
args := info.Args()
if len(args) < 1 {
return bridge.JsException(info.Context(), "Get requires at least one argument")
}
if !args[0].IsString() {
return bridge.JsException(info.Context(), "Get requires a valid key")
}
// Get the key
key := args[0].String()
if key == "" {
return bridge.JsException(info.Context(), "Get requires a valid key")
}
err = global.ChatContext.Space.Delete(key)
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
return nil
}
func jsClear(info *v8go.FunctionCallbackInfo) *v8go.Value {
global, err := global(info)
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
if global.ChatContext.Space == nil {
return bridge.JsException(info.Context(), "Shared space is not set")
}
err = global.ChatContext.Space.Clear()
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
return nil
}
// jsAssets function, get the assets content
func jsAssets(info *v8go.FunctionCallbackInfo) *v8go.Value {
global, err := global(info)
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
// Get the message
args := info.Args()
if len(args) < 1 {
return bridge.JsException(info.Context(), "Assets requires at least one argument")
}
// Get the name
name := args[0].String()
data := map[string]interface{}{}
if len(args) > 1 {
raw, err := bridge.GoValue(args[1], info.Context())
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
v, ok := raw.(map[string]interface{})
if !ok {
return bridge.JsException(info.Context(), "Assets requires a map")
}
data = v
}
content, err := global.Assistant.Assets(name, data)
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
jsContent, err := bridge.JsValue(info.Context(), content)
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
return jsContent
}
// jsSend function, send a message to the http stream connection
func jsSend(info *v8go.FunctionCallbackInfo) *v8go.Value {
// Get the message
args := info.Args()
if len(args) < 1 {
return bridge.JsException(info.Context(), "SendMessage requires at least one argument")
}
input, err := bridge.GoValue(args[0], info.Context())
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
global, err := global(info)
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
// Save history by default
saveHistory := true
if len(args) > 1 && args[1].IsBoolean() {
saveHistory = args[1].Boolean()
}
switch v := input.(type) {
case string:
// Check if the message is json
msg, err := message.NewString(v)
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
// Set the role to assistant
if msg.Role == "" {
msg.Role = "assistant"
}
// Append the message to the contents
if saveHistory {
msg.AppendTo(global.Contents)
}
msg.Write(global.GinContext.Writer)
return nil
case map[string]interface{}:
msg := message.New().Map(v)
if msg.Role == "" {
msg.Role = "assistant"
}
// Append the message to the contents
if saveHistory {
msg.AppendTo(global.Contents)
}
msg.Write(global.GinContext.Writer)
return nil
default:
return bridge.JsException(info.Context(), "Send requires a string or a map")
}
}
func jsReplace(info *v8go.FunctionCallbackInfo) *v8go.Value {
args := info.Args()
if len(args) < 2 {
return bridge.JsException(info.Context(), "Replace requires at least two arguments")
}
if !args[0].IsString() {
return bridge.JsException(info.Context(), "the first argument must be a string")
}
tmpl := args[0].String()
raw, err := bridge.GoValue(args[1], info.Context())
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
data, ok := raw.(map[string]interface{})
if !ok {
return bridge.JsException(info.Context(), "the second argument must be a map")
}
replaced, _ := sui.Data(data).Replace(tmpl)
jsReplaced, err := bridge.JsValue(info.Context(), replaced)
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
return jsReplaced
}
// global get the global variables
func global(info *v8go.FunctionCallbackInfo) (global *GlobalVariables, err error) {
return getGlobal(info.Context(), info.This())
}
func getGlobal(ctx *v8go.Context, obj *v8go.Object) (global *GlobalVariables, err error) {
jsGlobal, err := obj.Get("__yao_agent_global")
if err != nil {
return nil, err
}
// Convert to go interface
goGlobal, err := bridge.GoValue(jsGlobal, ctx)
if err != nil {
return nil, err
}
global, ok := goGlobal.(*GlobalVariables)
if !ok {
return nil, fmt.Errorf("global is not a valid GlobalVariables. %#v", goGlobal)
}
return global, nil
}

140
agent/assistant/plan.go
View file

@ -1,140 +0,0 @@
package assistant
import (
"context"
"fmt"
"github.com/fatih/color"
"github.com/yaoapp/gou/runtime/v8/bridge"
v8plan "github.com/yaoapp/gou/runtime/v8/objects/plan"
"rogchap.com/v8go"
)
// TaskFn is the task function
func TaskFn(plan_id string, task_id string, source bool, method string, args ...interface{}) (interface{}, error) {
if !source {
return v8plan.DefaultTaskFn(plan_id, task_id, source, method, args...)
}
// Data
plan, err := v8plan.GetPlan(plan_id)
if err != nil {
return nil, err
}
global, ok := plan.Data().(*GlobalVariables)
if !ok {
return nil, fmt.Errorf("plan data is not a GlobalVariables")
}
if global.Assistant == nil {
return nil, fmt.Errorf("assistant is not set")
}
if global.Assistant.Script == nil {
return nil, fmt.Errorf("script is not set")
}
scriptCtx, err := global.Assistant.Script.NewContext(global.ChatContext.Sid, nil)
if err != nil {
return nil, err
}
defer scriptCtx.Close()
// Initialize the object
global.Assistant.InitObject(scriptCtx, global.GinContext, global.ChatContext, global.Contents)
fnargs := []interface{}{plan_id, task_id}
fnargs = append(fnargs, args...)
// Execute the anonymous function
return scriptCtx.CallAnonymousWith(context.Background(), method, fnargs...)
}
// SubscribeFn is the default subscribe function
func SubscribeFn(plan_id string, key string, value interface{}, source bool, method string, args ...interface{}) {
if !source {
v8plan.DefaultSubscribeFn(plan_id, key, value, source, method, args...)
return
}
// Data
plan, err := v8plan.GetPlan(plan_id)
if err != nil {
color.Red("Subscribe Failed to get the plan: %s", err.Error())
return
}
global, ok := plan.Data().(*GlobalVariables)
if !ok {
color.Red("Subscribe Failed: plan data is not a GlobalVariables")
return
}
if global.Assistant == nil {
color.Red("Subscribe Failed: assistant is not set")
return
}
if global.Assistant.Script == nil {
color.Red("Subscribe Failed: script is not set")
return
}
scriptCtx, err := global.Assistant.Script.NewContext(global.ChatContext.Sid, nil)
if err != nil {
color.Red("Subscribe Failed: Failed to create the script context: %s", err.Error())
return
}
defer scriptCtx.Close()
fnargs := []interface{}{plan_id, key, value}
fnargs = append(fnargs, args...)
// Initialize the object
global.Assistant.InitObject(scriptCtx, global.GinContext, global.ChatContext, global.Contents)
_, err = scriptCtx.CallAnonymousWith(context.Background(), method, fnargs...)
if err != nil {
return
}
}
// jsNewPlan create a plan object and return it
func jsPlan(info *v8go.FunctionCallbackInfo) *v8go.Value {
global, err := global(info)
if err != nil {
return bridge.JsException(info.Context(), err.Error())
}
obj := newPlanObject()
args := info.Args()
if len(args) < 1 {
return bridge.JsException(info.Context(), "the first parameter should be a string")
}
if !args[0].IsString() {
return bridge.JsException(info.Context(), "the first parameter should be a string")
}
id := args[0].String()
objectTmpl := obj.ExportObject(info.Context().Isolate())
plan, err := objectTmpl.NewInstance(info.Context())
if err != nil {
return bridge.JsException(info.Context(), fmt.Sprintf("failed to create plan object %s", err.Error()))
}
return obj.NewInstance(id, plan, global)
}
func newPlanObject() *v8plan.Object {
obj := v8plan.New(v8plan.Options{
TaskFn: TaskFn,
SubscribeFn: SubscribeFn,
})
return obj
}

156
agent/assistant/tool.go
View file

@ -1,156 +0,0 @@
package assistant
import (
"fmt"
jsoniter "github.com/json-iterator/go"
store "github.com/yaoapp/yao/agent/store/types"
)
// Tool represents a tool
type Tool struct {
Type string `json:"type"`
Function struct {
Name string `json:"name"`
Description string `json:"description"`
Parameters Parameter `json:"parameters"`
Strict bool `json:"strict,omitempty"`
} `json:"function"`
}
// SchemaProperty represents a JSON Schema property
type SchemaProperty struct {
Type string `json:"type,omitempty"`
Description string `json:"description,omitempty"`
Items *Parameter `json:"items,omitempty"`
OneOf []SchemaProperty `json:"oneOf,omitempty"`
Enum []interface{} `json:"enum,omitempty"`
}
// Parameter represents the parameters field in function calling format
type Parameter struct {
Type string `json:"type,omitempty"`
Properties map[string]SchemaProperty `json:"properties,omitempty"`
Description string `json:"description,omitempty"`
Required []string `json:"required,omitempty"`
AdditionalProperties bool `json:"additionalProperties,omitempty"`
Strict bool `json:"strict,omitempty"`
OneOf []SchemaProperty `json:"oneOf,omitempty"`
Enum []interface{} `json:"enum,omitempty"`
}
// Example returns a formatted example of how to use this tool
func (tool Tool) Example() string {
return fmt.Sprintf("<tool>\n{\"function\":\"%s\",\"arguments\":%s}\n</tool>",
tool.Function.Name,
jsoniter.Wrap(tool.ExampleArguments()).ToString())
}
// ExampleArguments generates example arguments for the tool based on parameter types
func (tool Tool) ExampleArguments() map[string]interface{} {
args := map[string]interface{}{}
// Handle the root parameter object
if tool.Function.Parameters.Type == "object" && tool.Function.Parameters.Properties != nil {
for name, prop := range tool.Function.Parameters.Properties {
args[name] = generateExampleValue(name, prop)
}
}
return args
}
// generateExampleValue creates an example value for a parameter
func generateExampleValue(name string, prop SchemaProperty) interface{} {
if len(prop.OneOf) > 0 {
// Return the first non-null type example value from oneOf
for _, subProp := range prop.OneOf {
if subProp.Type != "null" {
return generateExampleValue(name, subProp)
}
}
return nil
}
// If enum is defined, return the first enum value
if len(prop.Enum) > 0 {
return prop.Enum[0]
}
switch prop.Type {
case "string":
return fmt.Sprintf("<%s:string>", name)
case "number":
return fmt.Sprintf("<%s:number>", name)
case "integer":
return fmt.Sprintf("<%s:integer>", name)
case "boolean":
return fmt.Sprintf("<%s:boolean>", name)
case "object":
return fmt.Sprintf("<%s:object>", name)
case "array":
return fmt.Sprintf("<%s:array>", name)
case "null":
return nil
default:
return fmt.Sprintf("<%s>", name)
}
}
// ToRuntimeTool converts store.Tool to assistant.Tool (OpenAI format)
func ToRuntimeTool(storeTool store.Tool) (Tool, error) {
var tool Tool
// Marshal and unmarshal to convert between formats
raw, err := jsoniter.Marshal(storeTool)
if err != nil {
return tool, fmt.Errorf("failed to marshal store tool: %w", err)
}
// Try to unmarshal as OpenAI format first
err = jsoniter.Unmarshal(raw, &tool)
if err == nil && tool.Function.Name != "" {
return tool, nil
}
// If it's a simple format, convert it
tool.Type = "function"
if storeTool.Type != "" {
tool.Type = storeTool.Type
}
tool.Function.Name = storeTool.Name
tool.Function.Description = storeTool.Description
// Convert parameters
if storeTool.Parameters != nil {
raw, err := jsoniter.Marshal(storeTool.Parameters)
if err != nil {
return tool, fmt.Errorf("failed to marshal parameters: %w", err)
}
var params Parameter
err = jsoniter.Unmarshal(raw, &params)
if err != nil {
return tool, fmt.Errorf("failed to unmarshal parameters: %w", err)
}
tool.Function.Parameters = params
}
return tool, nil
}
// ToRuntimeTools converts []store.Tool to []assistant.Tool
func ToRuntimeTools(storeTools []store.Tool) ([]Tool, error) {
if storeTools == nil {
return nil, nil
}
tools := make([]Tool, 0, len(storeTools))
for _, storeTool := range storeTools {
tool, err := ToRuntimeTool(storeTool)
if err != nil {
return nil, err
}
tools = append(tools, tool)
}
return tools, nil
}

91
agent/assistant/types.go
View file

@ -1,14 +1,10 @@
package assistant package assistant
import ( import (
"context"
"io"
"github.com/gin-gonic/gin"
jsoniter "github.com/json-iterator/go" jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/yao/agent/assistant/hook" "github.com/yaoapp/yao/agent/assistant/hook"
chatctx "github.com/yaoapp/yao/agent/context" chatctx "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/message"
outputMessage "github.com/yaoapp/yao/agent/output/message" outputMessage "github.com/yaoapp/yao/agent/output/message"
store "github.com/yaoapp/yao/agent/store/types" store "github.com/yaoapp/yao/agent/store/types"
api "github.com/yaoapp/yao/openai" api "github.com/yaoapp/yao/openai"
@ -21,56 +17,7 @@ const (
// API the assistant API interface // API the assistant API interface
type API interface { type API interface {
Chat(ctx context.Context, messages []message.Message, option map[string]interface{}, cb func(data []byte) int) error
GetPlaceholder(locale string) *store.Placeholder GetPlaceholder(locale string) *store.Placeholder
Execute(c *gin.Context, ctx chatctx.Context, input interface{}, options map[string]interface{}, callback ...interface{}) (interface{}, error)
Call(c *gin.Context, payload APIPayload) (interface{}, error)
}
// APIPayload the API payload
type APIPayload struct {
Sid string `json:"sid"`
Name string `json:"name"`
Args []interface{} `json:"args,omitempty"`
}
// ResHookInit the response of the init hook
type ResHookInit struct {
AssistantID string `json:"assistant_id,omitempty"`
ChatID string `json:"chat_id,omitempty"`
Next *NextAction `json:"next,omitempty"`
Input []message.Message `json:"input,omitempty"`
Options map[string]interface{} `json:"options,omitempty"`
Result any `json:"result,omitempty"`
}
// ResHookStream the response of the stream hook
type ResHookStream struct {
Silent bool `json:"silent,omitempty"` // Whether to suppress the output
Next *NextAction `json:"next,omitempty"` // The next action
Output []message.Data `json:"output,omitempty"` // The output
}
// ResHookDone the response of the done hook
type ResHookDone struct {
Next *NextAction `json:"next,omitempty"`
Input []message.Message `json:"input,omitempty"`
Output []message.Data `json:"output,omitempty"`
Result any `json:"result,omitempty"`
}
// ResHookFail the response of the fail hook
type ResHookFail struct {
Next *NextAction `json:"next,omitempty"`
Input []message.Message `json:"input,omitempty"`
Output string `json:"output,omitempty"`
Error string `json:"error,omitempty"`
}
// NextAction the next action
type NextAction struct {
Action string `json:"action"`
Payload map[string]interface{} `json:"payload,omitempty"`
} }
// SearchOption the search option // SearchOption the search option
@ -79,21 +26,6 @@ type SearchOption struct {
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
} }
// Prompt a prompt
type Prompt struct {
Role string `json:"role"`
Content string `json:"content"`
Name string `json:"name,omitempty"`
}
// QueryParam the assistant query param
type QueryParam struct {
Limit uint `json:"limit"`
Order string `json:"order"`
After string `json:"after"`
Before string `json:"before"`
}
// Assistant the assistant // Assistant the assistant
type Assistant struct { type Assistant struct {
store.AssistantModel store.AssistantModel
@ -106,8 +38,6 @@ type Assistant struct {
search bool // Whether this assistant supports search search bool // Whether this assistant supports search
vision bool // Whether this assistant supports vision vision bool // Whether this assistant supports vision
// toolCalls bool // Whether this assistant supports tool_calls // toolCalls bool // Whether this assistant supports tool_calls
initHook bool // Whether this assistant has an init hook
runtimeTools []Tool // Converted tools for business logic (OpenAI format)
} }
// ModelCapabilities defines the capabilities of a language model // ModelCapabilities defines the capabilities of a language model
@ -148,25 +78,6 @@ var VisionCapableModels = map[string]bool{
"gpt-4o-mini": true, // Custom OpenAI compatible model - mini version "gpt-4o-mini": true, // Custom OpenAI compatible model - mini version
} }
// File the file
type File struct {
ID string `json:"file_id"`
Bytes int `json:"bytes"`
CreatedAt int `json:"created_at"`
Filename string `json:"filename"`
ContentType string `json:"content_type"`
Description string `json:"description,omitempty"` // Vision analysis result or other description
URL string `json:"url,omitempty"` // Vision URL for vision-capable models
DocIDs []string `json:"doc_ids,omitempty"` // RAG document IDs
}
// FileResponse represents a file download response
type FileResponse struct {
Reader io.ReadCloser
ContentType string
Extension string
}
// MCPTool represents a simplified MCP tool for building LLM requests // MCPTool represents a simplified MCP tool for building LLM requests
// This is an internal representation used when collecting tools from MCP servers // This is an internal representation used when collecting tools from MCP servers
// and preparing them for the LLM's tool calling interface // and preparing them for the LLM's tool calling interface

48
agent/context/JSAPI.md
View file

@ -123,20 +123,22 @@ const image_id = ctx.Send({
```javascript ```javascript
// Scenario 1: Simple messages without block grouping (most common) // Scenario 1: Simple messages without block grouping (most common)
function Next(ctx, response) { function Next(ctx, payload) {
const { completion } = payload;
// Each message is independent // Each message is independent
const loading_id = ctx.Send({ const loading_id = ctx.Send({
type: "loading", type: "loading",
props: { message: "Thinking..." } props: { message: "Thinking..." }
}); });
// Call LLM... // Process completion...
const result = Process("llms.chat", {...}); const result = completion.content;
// Replace loading with result // Replace loading with result
ctx.Replace(loading_id, { ctx.Replace(loading_id, {
type: "text", type: "text",
props: { content: result.content } props: { content: result }
}); });
} }
@ -154,14 +156,14 @@ function Create(ctx, messages) {
} }
// Scenario 3: LLM response + follow-up card in same block // Scenario 3: LLM response + follow-up card in same block
function Next(ctx, response) { function Next(ctx, payload) {
const { completion } = payload;
const block_id = ctx.BlockID(); const block_id = ctx.BlockID();
// LLM response // LLM response
const result = Process("llms.chat", {...});
ctx.Send({ ctx.Send({
type: "text", type: "text",
props: { content: result.content }, props: { content: completion.content },
block_id: block_id block_id: block_id
}); });
@ -948,9 +950,18 @@ Here's a comprehensive example using various Context API features:
```javascript ```javascript
/** /**
* Next Hook - Process LLM response and enhance with tools * Next Hook - Process LLM response and enhance with tools
* @param {Context} ctx - Agent context
* @param {Object} payload - Hook payload
* @param {Array} payload.messages - Messages sent to the assistant
* @param {Object} payload.completion - Completion response from LLM
* @param {Array} payload.tools - Tool call results
* @param {string} payload.error - Error message if failed
*/ */
function Next(ctx, messages, completion, tools) { function Next(ctx, payload) {
try { try {
// Destructure payload
const { messages, completion, tools, error } = payload;
// Create trace node for custom processing // Create trace node for custom processing
const process_node = ctx.Trace.Add( const process_node = ctx.Trace.Add(
{ completion, tools }, { completion, tools },
@ -1001,14 +1012,11 @@ function Next(ctx, messages, completion, tools) {
// Return enhanced response // Return enhanced response
return { return {
data: enhanced_response, data: enhanced_response,
done: true, metadata: { processed: true },
}; };
} catch (error) { } catch (error) {
ctx.Trace.Error("Processing failed", { error: error.message }); ctx.Trace.Error("Processing failed", { error: error.message });
throw error; throw error;
} finally {
// Optional: Manual cleanup
ctx.Release();
} }
} }
``` ```
@ -1043,13 +1051,17 @@ For TypeScript projects, the Context types are automatically inferred. You can a
```typescript ```typescript
import { Context, Message, TraceNodeOption } from "@yaoapps/types"; import { Context, Message, TraceNodeOption } from "@yaoapps/types";
function Next( interface NextPayload {
ctx: Context, messages: Message[];
messages: Message[], completion: any;
completion: any, tools: any[];
tools: any[] error?: string;
): any { }
function Next(ctx: Context, payload: NextPayload): any {
// Your code with full type checking // Your code with full type checking
const { messages, completion, tools, error } = payload;
// ...
} }
``` ```

20
agent/context/context.go
View file

@ -35,6 +35,7 @@ func New(parent context.Context, authorized *types.AuthorizedInfo, chatID, paylo
Space: plan.NewMemorySharedSpace(), Space: plan.NewMemorySharedSpace(),
ChatID: chatID, ChatID: chatID,
IDGenerator: message.NewIDGenerator(), // Initialize ID generator for this context IDGenerator: message.NewIDGenerator(), // Initialize ID generator for this context
messageMetadata: newMessageMetadataStore(), // Initialize message metadata store
} }
if payload == "" { if payload == "" {
@ -379,33 +380,22 @@ func (ctx *Context) TraceID() string {
// recordMessageMetadata records metadata for a sent message // recordMessageMetadata records metadata for a sent message
// Used to inherit BlockID and ThreadID in subsequent delta operations // Used to inherit BlockID and ThreadID in subsequent delta operations
func (ctx *Context) recordMessageMetadata(msg *message.Message) { func (ctx *Context) recordMessageMetadata(msg *message.Message) {
if msg.MessageID == "" { if msg.MessageID == "" || ctx.messageMetadata == nil {
return return
} }
ctx.metadataMu.Lock() ctx.messageMetadata.set(msg.MessageID, &MessageMetadata{
defer ctx.metadataMu.Unlock()
if ctx.messageMetadata == nil {
ctx.messageMetadata = make(map[string]*MessageMetadata)
}
ctx.messageMetadata[msg.MessageID] = &MessageMetadata{
MessageID: msg.MessageID, MessageID: msg.MessageID,
BlockID: msg.BlockID, BlockID: msg.BlockID,
ThreadID: msg.ThreadID, ThreadID: msg.ThreadID,
} })
} }
// getMessageMetadata retrieves metadata for a message by ID // getMessageMetadata retrieves metadata for a message by ID
// Returns nil if message metadata is not found // Returns nil if message metadata is not found
func (ctx *Context) getMessageMetadata(messageID string) *MessageMetadata { func (ctx *Context) getMessageMetadata(messageID string) *MessageMetadata {
ctx.metadataMu.RLock()
defer ctx.metadataMu.RUnlock()
if ctx.messageMetadata == nil { if ctx.messageMetadata == nil {
return nil return nil
} }
return ctx.messageMetadata.get(messageID)
return ctx.messageMetadata[messageID]
} }

10
agent/context/jsapi_test.go
View file

@ -26,7 +26,6 @@ func TestJsValue(t *testing.T) {
cxt := &context.Context{ cxt := &context.Context{
ChatID: "ChatID-123456", ChatID: "ChatID-123456",
AssistantID: "AssistantID-1234", AssistantID: "AssistantID-1234",
Sid: "Sid-1234",
Context: stdContext.Background(), Context: stdContext.Background(),
IDGenerator: message.NewIDGenerator(), IDGenerator: message.NewIDGenerator(),
} }
@ -91,12 +90,10 @@ func TestJsValueConcurrent(t *testing.T) {
for j := 0; j < iterationsPerGoroutine; j++ { for j := 0; j < iterationsPerGoroutine; j++ {
chatID := fmt.Sprintf("ChatID-%d-%d", routineID, j) chatID := fmt.Sprintf("ChatID-%d-%d", routineID, j)
assistantID := fmt.Sprintf("AssistantID-%d-%d", routineID, j) assistantID := fmt.Sprintf("AssistantID-%d-%d", routineID, j)
sid := fmt.Sprintf("Sid-%d-%d", routineID, j)
cxt := &context.Context{ cxt := &context.Context{
ChatID: chatID, ChatID: chatID,
AssistantID: assistantID, AssistantID: assistantID,
Sid: sid,
Context: stdContext.Background(), Context: stdContext.Background(),
IDGenerator: message.NewIDGenerator(), IDGenerator: message.NewIDGenerator(),
} }
@ -156,7 +153,6 @@ func TestJsValueRegistrationAndCleanup(t *testing.T) {
cxt := &context.Context{ cxt := &context.Context{
ChatID: fmt.Sprintf("ChatID-%d", i), ChatID: fmt.Sprintf("ChatID-%d", i),
AssistantID: fmt.Sprintf("AssistantID-%d", i), AssistantID: fmt.Sprintf("AssistantID-%d", i),
Sid: fmt.Sprintf("Sid-%d", i),
Context: stdContext.Background(), Context: stdContext.Background(),
IDGenerator: message.NewIDGenerator(), IDGenerator: message.NewIDGenerator(),
} }
@ -337,12 +333,6 @@ func TestJsValueAllFields(t *testing.T) {
assert.Equal(t, "engineering", extra["department"], "constraints.extra.department mismatch") assert.Equal(t, "engineering", extra["department"], "constraints.extra.department mismatch")
assert.Equal(t, "us-west", extra["region"], "constraints.extra.region mismatch") assert.Equal(t, "us-west", extra["region"], "constraints.extra.region mismatch")
// Verify deprecated fields are NOT exported
_, hasSid := result["sid"]
assert.False(t, hasSid, "sid (deprecated) should not be exported")
_, hasSilent := result["silent"]
assert.False(t, hasSilent, "silent (deprecated) should not be exported")
// Note: We can't directly check goMaps cleanup as it's in the bridge package // Note: We can't directly check goMaps cleanup as it's in the bridge package
} }

43
agent/context/types.go
View file

@ -202,23 +202,55 @@ type MessageMetadata struct {
ThreadID string // Thread ID ThreadID string // Thread ID
} }
// messageMetadataStore provides thread-safe storage for message metadata
type messageMetadataStore struct {
data map[string]*MessageMetadata
mu sync.RWMutex
}
// newMessageMetadataStore creates a new message metadata store
func newMessageMetadataStore() *messageMetadataStore {
return &messageMetadataStore{
data: make(map[string]*MessageMetadata),
}
}
// set stores metadata for a message (thread-safe)
func (s *messageMetadataStore) set(messageID string, metadata *MessageMetadata) {
s.mu.Lock()
defer s.mu.Unlock()
s.data[messageID] = metadata
}
// get retrieves metadata for a message (thread-safe)
func (s *messageMetadataStore) get(messageID string) *MessageMetadata {
s.mu.RLock()
defer s.mu.RUnlock()
return s.data[messageID]
}
// Context the context // Context the context
type Context struct { type Context struct {
// Context // Context
context.Context context.Context
// External
ID string `json:"id"` // Context ID for external interrupt identification ID string `json:"id"` // Context ID for external interrupt identification
Space plan.Space `json:"-"` // Shared data space, it will be used to share data between the request and the call Space plan.Space `json:"-"` // Shared data space, it will be used to share data between the request and the call
Cache store.Store `json:"-"` // Cache store, it will be used to store the message cache, default is "__yao.agent.cache" Cache store.Store `json:"-"` // Cache store, it will be used to store the message cache, default is "__yao.agent.cache"
Stack *Stack `json:"-"` // Stack, current active stack of the request Stack *Stack `json:"-"` // Stack, current active stack of the request
Stacks map[string]*Stack `json:"-"` // Stacks, all stacks in this request (for trace logging) Stacks map[string]*Stack `json:"-"` // Stacks, all stacks in this request (for trace logging)
Writer Writer `json:"-"` // Writer, it will be used to write response data to the client Writer Writer `json:"-"` // Writer, it will be used to write response data to the client
Skip *Skip `json:"skip,omitempty"` // Skip configuration (history, trace, etc.), nil means don't skip anything IDGenerator *message.IDGenerator `json:"-"` // ID generator for this context (chunk, message, block, thread IDs)
// Internal
trace traceTypes.Manager `json:"-"` // Trace manager, lazy initialized on first access trace traceTypes.Manager `json:"-"` // Trace manager, lazy initialized on first access
output *output.Output `json:"-"` // Output, it will be used to write response data to the client output *output.Output `json:"-"` // Output, it will be used to write response data to the client
IDGenerator *message.IDGenerator `json:"-"` // ID generator for this context (chunk, message, block, thread IDs) messageMetadata *messageMetadataStore `json:"-"` // Thread-safe message metadata store for delta operations
messageMetadata map[string]*MessageMetadata `json:"-"` // Message metadata cache for delta operations (inheriting BlockID/ThreadID)
metadataMu sync.RWMutex `json:"-"` // Mutex for concurrent access to messageMetadata // Skip configuration (history, trace, etc.), nil means don't skip anything
Skip *Skip `json:"skip,omitempty"` // Skip configuration (history, trace, etc.), nil means don't skip anything
// Model capabilities (set by assistant, used by output adapters) // Model capabilities (set by assistant, used by output adapters)
Capabilities *ModelCapabilities `json:"-"` // Model capabilities for the current connector Capabilities *ModelCapabilities `json:"-"` // Model capabilities for the current connector
@ -230,7 +262,6 @@ type Context struct {
Authorized *types.AuthorizedInfo `json:"authorized,omitempty"` // Authorized information Authorized *types.AuthorizedInfo `json:"authorized,omitempty"` // Authorized information
ChatID string `json:"chat_id,omitempty"` // Chat ID, use to select chat ChatID string `json:"chat_id,omitempty"` // Chat ID, use to select chat
AssistantID string `json:"assistant_id,omitempty"` // Assistant ID, use to select assistant AssistantID string `json:"assistant_id,omitempty"` // Assistant ID, use to select assistant
Sid string `json:"sid" yaml:"-"` // Session ID (Deprecated, use Authorized instead)
Connector string `json:"connector,omitempty"` // Connector, use to select the connector of the LLM Model, Default is Assistant.Connector Connector string `json:"connector,omitempty"` // Connector, use to select the connector of the LLM Model, Default is Assistant.Connector
Search *bool `json:"search,omitempty"` // Search mode, default is true Search *bool `json:"search,omitempty"` // Search mode, default is true
@ -251,8 +282,6 @@ type Context struct {
// CUI Context information // CUI Context information
Route string `json:"route,omitempty"` // The route of the request, it will be used to identify the route of the request Route string `json:"route,omitempty"` // The route of the request, it will be used to identify the route of the request
Metadata map[string]interface{} `json:"metadata,omitempty"` // The metadata of the request, it will be used to pass data to the page Metadata map[string]interface{} `json:"metadata,omitempty"` // The metadata of the request, it will be used to pass data to the page
Silent bool `json:"silent,omitempty"` // Silent mode (Deprecated, use Referer instead)
} }
// Stack represents the call stack node for tracing agent-to-agent calls // Stack represents the call stack node for tracing agent-to-agent calls

4
agent/jsapi/jsapi.go
View file

@ -1,4 +0,0 @@
package jsapi
// JSAPI Register the JavaScript API
// Agent API will be registered as a third party object

65
agent/load.go
View file

@ -6,19 +6,19 @@ import (
"github.com/yaoapp/gou/application" "github.com/yaoapp/gou/application"
"github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/connector"
"github.com/yaoapp/kun/exception"
"github.com/yaoapp/yao/agent/api"
"github.com/yaoapp/yao/agent/assistant" "github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n" "github.com/yaoapp/yao/agent/i18n"
mongoStore "github.com/yaoapp/yao/agent/store/mongo" storeMongo "github.com/yaoapp/yao/agent/store/mongo"
redisStore "github.com/yaoapp/yao/agent/store/redis" storeRedis "github.com/yaoapp/yao/agent/store/redis"
store "github.com/yaoapp/yao/agent/store/types" store "github.com/yaoapp/yao/agent/store/types"
xunStore "github.com/yaoapp/yao/agent/store/xun" "github.com/yaoapp/yao/agent/store/xun"
"github.com/yaoapp/yao/agent/types" "github.com/yaoapp/yao/agent/types"
"github.com/yaoapp/yao/config" "github.com/yaoapp/yao/config"
) )
var agentDSL *types.DSL
// Load load AIGC // Load load AIGC
func Load(cfg config.Config) error { func Load(cfg config.Config) error {
@ -59,8 +59,7 @@ func Load(cfg config.Config) error {
setting.Uses.Prompt = setting.Uses.Default setting.Uses.Prompt = setting.Uses.Default
} }
// Initialize Agent API agentDSL = &setting
api.Agent = &api.API{DSL: &setting}
// Store Setting // Store Setting
err = initStore() err = initStore()
@ -89,12 +88,9 @@ func Load(cfg config.Config) error {
return nil return nil
} }
// GetAgent returns the Agent instance // GetAgent returns the Agent settings
func GetAgent() *api.API { func GetAgent() *types.DSL {
if api.Agent == nil { return agentDSL
exception.New("Agent is not initialized", 500).Throw()
}
return api.Agent
} }
// initGlobalI18n initialize the global i18n // initGlobalI18n initialize the global i18n
@ -126,7 +122,7 @@ func initModelCapabilities() error {
return err return err
} }
api.Agent.DSL.Models = models agentDSL.Models = models
return nil return nil
} }
@ -134,57 +130,52 @@ func initModelCapabilities() error {
func initStore() error { func initStore() error {
var err error var err error
if api.Agent.DSL.StoreSetting.Connector == "default" || api.Agent.DSL.StoreSetting.Connector == "" { if agentDSL.StoreSetting.Connector == "default" || agentDSL.StoreSetting.Connector == "" {
api.Agent.DSL.Store, err = xunStore.NewXun(api.Agent.DSL.StoreSetting) agentDSL.Store, err = xun.NewXun(agentDSL.StoreSetting)
return err return err
} }
// other connector // other connector
conn, err := connector.Select(api.Agent.DSL.StoreSetting.Connector) conn, err := connector.Select(agentDSL.StoreSetting.Connector)
if err != nil { if err != nil {
return fmt.Errorf("load connectors error: %s", err.Error()) return fmt.Errorf("load connectors error: %s", err.Error())
} }
if conn.Is(connector.DATABASE) { if conn.Is(connector.DATABASE) {
api.Agent.DSL.Store, err = xunStore.NewXun(api.Agent.DSL.StoreSetting) agentDSL.Store, err = xun.NewXun(agentDSL.StoreSetting)
return err return err
} else if conn.Is(connector.REDIS) { } else if conn.Is(connector.REDIS) {
api.Agent.DSL.Store = redisStore.NewRedis() agentDSL.Store = storeRedis.NewRedis()
return nil return nil
} else if conn.Is(connector.MONGO) { } else if conn.Is(connector.MONGO) {
api.Agent.DSL.Store = mongoStore.NewMongo() agentDSL.Store = storeMongo.NewMongo()
return nil return nil
} }
return fmt.Errorf("Agent store connector %s not support", api.Agent.DSL.StoreSetting.Connector) return fmt.Errorf("Agent store connector %s not support", agentDSL.StoreSetting.Connector)
} }
// initAssistant initialize the assistant // initAssistant initialize the assistant
func initAssistant() error { func initAssistant() error {
// Set Storage // Set Storage
assistant.SetStorage(api.Agent.DSL.Store) assistant.SetStorage(agentDSL.Store)
// Assistant Vision
if api.Agent.DSL.Vision != nil {
assistant.SetVision(api.Agent.DSL.Vision)
}
// Set global Uses configuration // Set global Uses configuration
if api.Agent.DSL.Uses != nil { if agentDSL.Uses != nil {
globalUses := &context.Uses{ globalUses := &context.Uses{
Vision: api.Agent.DSL.Uses.Vision, Vision: agentDSL.Uses.Vision,
Audio: api.Agent.DSL.Uses.Audio, Audio: agentDSL.Uses.Audio,
Search: api.Agent.DSL.Uses.Search, Search: agentDSL.Uses.Search,
Fetch: api.Agent.DSL.Uses.Fetch, Fetch: agentDSL.Uses.Fetch,
} }
assistant.SetGlobalUses(globalUses) assistant.SetGlobalUses(globalUses)
} }
if api.Agent.DSL.Models != nil { if agentDSL.Models != nil {
assistant.SetModelCapabilities(api.Agent.DSL.Models) assistant.SetModelCapabilities(agentDSL.Models)
} }
// Load Built-in Assistants // Load Built-in Assistants
@ -199,14 +190,14 @@ func initAssistant() error {
return err return err
} }
api.Agent.DSL.Assistant = defaultAssistant agentDSL.Assistant = defaultAssistant
return nil return nil
} }
// defaultAssistant get the default assistant // defaultAssistant get the default assistant
func defaultAssistant() (*assistant.Assistant, error) { func defaultAssistant() (*assistant.Assistant, error) {
if api.Agent.DSL.Uses == nil || api.Agent.DSL.Uses.Default == "" { if agentDSL.Uses == nil || agentDSL.Uses.Default == "" {
return nil, fmt.Errorf("default assistant not found") return nil, fmt.Errorf("default assistant not found")
} }
return assistant.Get(api.Agent.DSL.Uses.Default) return assistant.Get(agentDSL.Uses.Default)
} }

411
agent/message/contents.go
View file

@ -1,411 +0,0 @@
package message
import (
"fmt"
"math/rand"
"strings"
"time"
"github.com/google/uuid"
jsoniter "github.com/json-iterator/go"
)
const (
// ContentStatusPending the content status pending
ContentStatusPending = iota
// ContentStatusDone the content status done
ContentStatusDone
// ContentStatusError the content status error
ContentStatusError
)
var tokens = map[string][2]string{
"think": {"<think>", "</think>"},
"tool": {"<tool>", "</tool>"},
}
// Contents the contents
type Contents struct {
Current int `json:"current"` // the current content index
Data []Data `json:"data"` // the data
token string // the current token
id string // the id of the contents
stack [][]string // the token stack
mapping map[string]string // the mapping of the token stack
}
// Data the data of the content
type Data struct {
Type string `json:"type"` // text, function, error, think, tool
ID string `json:"id"` // the id of the content
Bytes []byte `json:"bytes"` // the content bytes
Props map[string]interface{} `json:"props"` // the props
Begin int64 `json:"begin,omitempty"` // the begin time
End int64 `json:"end,omitempty"` // the end time
}
// Extra the extra of the content
type Extra struct {
ID string `json:"id,omitempty"` // the id of the content
Begin int64 `json:"begin,omitempty"` // the begin time
End int64 `json:"end,omitempty"` // the end time
}
// ScanCallbackParams the params of the scan callback
type ScanCallbackParams struct {
Token string
MessageID string
TokenID string
BeganAt int64
EndAt int64
Begin bool
End bool
Text string
Tails string
}
// NewContents create a new contents
func NewContents() *Contents {
return &Contents{
Current: -1,
Data: []Data{},
}
}
// ScanTokens scan the tokens
func (c *Contents) ScanTokens(messageID string, tokenID string, beganAt int64, cb func(params ScanCallbackParams)) {
text := strings.TrimSpace(c.Text())
// check the end of the token
if c.token != "" {
token := c.GetToken(c.token)
tokenType := c.GetTokenType(c.token)
// Check the end of the token
if index := strings.Index(text, token[1]); index >= 0 {
tails := ""
if index > 0 {
tails = text[index+len(token[1]):]
}
extra := Extra{
ID: c.id,
End: time.Now().UnixNano(),
}
c.UpdateType(tokenType, map[string]interface{}{"text": text}, extra)
c.NewText([]byte(tails), extra) // Create new text with the tails
cb(ScanCallbackParams{Token: tokenType, MessageID: c.id, TokenID: tokenID, BeganAt: beganAt, Begin: false, End: true, Text: text, Tails: tails, EndAt: extra.End})
c.ClearToken(c.token) // clear the token
return
}
// call the callback for the scanning of the token
cb(ScanCallbackParams{Token: tokenType, MessageID: c.id, TokenID: tokenID, BeganAt: beganAt, Begin: false, End: false, Text: text, Tails: "", EndAt: 0})
return
}
// scan the begin of the token
begin := false
for name, token := range tokens {
if index := strings.Index(text, token[0]); index >= 0 {
c.id = messageID
if c.id == "" {
c.id = GenerateNumericID("M")
}
tokenType := name
if tokenID != "" {
tokenType = c.GetTokenType(tokenID)
}
// First time scanning the token, generate the token ID and begin time
if tokenID == "" || tokenType != name {
tokenID = GenerateNumericID("T")
beganAt = time.Now().UnixNano()
begin = true
c.token = tokenID
c.AppendToken(tokenID, name)
c.UpdateType(name, map[string]interface{}{"text": text, "id": tokenID}, Extra{ID: c.id, Begin: beganAt, End: beganAt})
}
cb(ScanCallbackParams{Token: name, MessageID: c.id, TokenID: tokenID, BeganAt: beganAt, Begin: begin, End: false, Text: text, Tails: "", EndAt: 0}) // call the callback
}
}
}
// ClearToken clear the token
func (c *Contents) ClearToken(id string) {
c.token = ""
next := 0
if c.stack == nil {
c.stack = [][]string{}
}
if c.mapping == nil {
c.mapping = map[string]string{}
}
for i, node := range c.stack {
if node[0] == id {
next = i + 1
delete(c.mapping, id)
break
}
}
// Remove the token from the stack, and set the next token
if next > 0 && next < len(c.stack) {
c.stack = c.stack[next:]
c.token = c.stack[len(c.stack)-1][0]
}
}
// AppendToken append the token to the stack
func (c *Contents) AppendToken(id string, name string) {
if c.stack == nil {
c.stack = [][]string{}
}
if c.mapping == nil {
c.mapping = map[string]string{}
}
c.stack = append(c.stack, []string{id, name})
c.mapping[id] = name
c.token = id
}
// GetTokenType get the token type from the stack
func (c *Contents) GetTokenType(id string) string {
return c.mapping[id]
}
// GetToken get the token from the stack
func (c *Contents) GetToken(name string) [2]string {
typ, ok := c.mapping[name]
if !ok {
return [2]string{}
}
return tokens[typ]
}
// RemoveLastEmpty remove the last empty data
func (c *Contents) RemoveLastEmpty() {
if c.Current == -1 {
return
}
// Remove the last empty data
if len(c.Data[c.Current].Bytes) == 0 && c.Data[c.Current].Type == "text" {
c.Data = c.Data[:c.Current]
c.Current--
}
}
// NewText create a new text data and append to the contents
func (c *Contents) NewText(bytes []byte, extra ...Extra) *Contents {
data := Data{Type: "text", Bytes: bytes}
if len(extra) > 0 {
if extra[0].Begin != 0 {
data.Begin = extra[0].Begin
}
if extra[0].End != 0 {
data.End = extra[0].End
}
if extra[0].ID != "" {
data.ID = extra[0].ID
}
}
c.Data = append(c.Data, data)
c.Current++
return c
}
// NewType create a new type data and append to the contents
func (c *Contents) NewType(typ string, props map[string]interface{}, extra ...Extra) *Contents {
data := Data{
Type: typ,
Props: props,
}
if len(extra) > 0 {
if extra[0].Begin != 0 {
data.Begin = extra[0].Begin
}
if extra[0].End != 0 {
data.End = extra[0].End
}
if extra[0].ID != "" {
data.ID = extra[0].ID
}
}
c.Data = append(c.Data, data)
c.Current++
return c
}
// UpdateType update the type of the current content
func (c *Contents) UpdateType(typ string, props map[string]interface{}, extra ...Extra) *Contents {
if c.Current == -1 {
c.NewType(typ, props, extra...)
return c
}
if len(extra) > 0 {
if extra[0].Begin != 0 {
c.Data[c.Current].Begin = extra[0].Begin
}
if extra[0].End != 0 {
c.Data[c.Current].End = extra[0].End
}
if extra[0].ID != "" {
c.Data[c.Current].ID = extra[0].ID
}
}
c.Data[c.Current].Type = typ
if props != nil {
if c.Data[c.Current].Props == nil {
c.Data[c.Current].Props = map[string]interface{}{}
}
for k, v := range props {
c.Data[c.Current].Props[k] = v
}
}
return c
}
// NewError create a new error data and append to the contents
func (c *Contents) NewError(err []byte) *Contents {
c.Data = append(c.Data, Data{
Type: "error",
Bytes: err,
})
c.Current++
return c
}
// AppendText append the text to the current content
func (c *Contents) AppendText(bytes []byte, extra ...Extra) *Contents {
if c.Current == -1 {
c.NewText(bytes, extra...)
return c
}
if len(extra) > 0 {
if extra[0].ID != "" {
c.Data[c.Current].ID = extra[0].ID
}
if extra[0].Begin != 0 {
c.Data[c.Current].Begin = extra[0].Begin
}
if extra[0].End != 0 {
c.Data[c.Current].End = extra[0].End
}
}
c.Data[c.Current].Bytes = append(c.Data[c.Current].Bytes, bytes...)
return c
}
// AppendError append the error to the current content
func (c *Contents) AppendError(err []byte) *Contents {
if c.Current == -1 {
c.NewError(err)
return c
}
c.Data[c.Current].Bytes = append(c.Data[c.Current].Bytes, err...)
return c
}
// JSON returns the json representation
func (c *Contents) JSON() string {
raw, _ := jsoniter.MarshalToString(c.Data)
return raw
}
// Text returns the text of the current content
func (c *Contents) Text() string {
if c.Current == -1 {
return ""
}
return string(c.Data[c.Current].Bytes)
}
// CurrentType returns the type of the current content
func (c *Contents) CurrentType() string {
if c.Current == -1 {
return ""
}
return c.Data[c.Current].Type
}
// Map returns the map representation
func (data *Data) Map() (map[string]interface{}, error) {
v := map[string]interface{}{"type": data.Type}
if data.ID != "" {
v["id"] = data.ID
}
if data.Bytes != nil && data.Type == "text" {
v["text"] = string(data.Bytes)
}
if data.Props != nil && data.Type != "text" {
v["props"] = data.Props
}
return v, nil
}
// MarshalJSON returns the json representation
func (data *Data) MarshalJSON() ([]byte, error) {
v := map[string]interface{}{"type": data.Type}
if data.ID != "" {
v["id"] = data.ID
}
if data.Bytes != nil && data.Type == "text" {
v["text"] = string(data.Bytes)
}
if data.Props != nil && data.Type != "text" {
v["props"] = data.Props
}
// Add the begin and end time
if data.Begin != 0 {
v["begin"] = data.Begin
}
if data.End != 0 {
v["end"] = data.End
}
return jsoniter.Marshal(v)
}
// GenerateNumericID generates a 10-digit number using UUID as seed
func GenerateNumericID(prefix string) string {
// Generate UUID and use it as seed
id := uuid.New()
seed := int64(id[0])<<56 | int64(id[1])<<48 | int64(id[2])<<40 | int64(id[3])<<32 |
int64(id[4])<<24 | int64(id[5])<<16 | int64(id[6])<<8 | int64(id[7])
// Create a new random source using the seed
source := rand.NewSource(seed)
r := rand.New(source)
// Generate a number between 1000000000 and 9999999999 (10 digits)
num := r.Int63n(9000000000) + 1000000000
return fmt.Sprintf("%s%d", prefix, num)
}

650
agent/message/message.go
View file

@ -1,650 +0,0 @@
package message
import (
"fmt"
"os"
"strings"
"sync"
"github.com/fatih/color"
"github.com/gin-gonic/gin"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/helper"
"github.com/yaoapp/kun/exception"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/kun/maps"
"github.com/yaoapp/yao/attachment"
"github.com/yaoapp/yao/openai"
)
var locker = sync.Mutex{}
// New create a new message
func New() *Message {
return &Message{Actions: []Action{}, Props: map[string]interface{}{}}
}
// NewHistory create a new message from history
func NewHistory(history map[string]interface{}) ([]Message, error) {
if history == nil {
return []Message{}, nil
}
var copy map[string]interface{} = map[string]interface{}{}
for key, value := range history {
if key != "content" {
copy[key] = value
}
}
globalMessage := New().Map(copy)
messages := []Message{}
if content, ok := history["content"].(string); ok {
if strings.HasPrefix(content, "{") && strings.HasSuffix(content, "}") {
var msg Message = *globalMessage
if err := jsoniter.UnmarshalFromString(content, &msg); err != nil {
return nil, err
}
messages = append(messages, msg)
} else if strings.HasPrefix(content, "[") && strings.HasSuffix(content, "]") {
var msgs []Message
if err := jsoniter.UnmarshalFromString(content, &msgs); err != nil {
return nil, err
}
for _, msg := range msgs {
msg.AssistantID = globalMessage.AssistantID
msg.AssistantName = globalMessage.AssistantName
msg.AssistantAvatar = globalMessage.AssistantAvatar
msg.Role = globalMessage.Role
msg.Name = globalMessage.Name
msg.Mentions = globalMessage.Mentions
messages = append(messages, msg)
}
} else {
messages = append(messages, Message{Text: content})
}
}
return messages, nil
}
// NewContent create a new message from content
func NewContent(content string) ([]Message, error) {
messages := []Message{}
if strings.HasPrefix(content, "{") && strings.HasSuffix(content, "}") {
var msg Message
if err := jsoniter.UnmarshalFromString(content, &msg); err != nil {
return nil, err
}
messages = append(messages, msg)
} else if strings.HasPrefix(content, "[") && strings.HasSuffix(content, "]") {
var msgs []Message
if err := jsoniter.UnmarshalFromString(content, &msgs); err != nil {
return nil, err
}
for _, msg := range msgs {
messages = append(messages, msg)
}
} else {
messages = append(messages, Message{Text: content})
}
return messages, nil
}
// NewString create a new message from string
func NewString(content string, id ...string) (*Message, error) {
if strings.HasPrefix(content, "{") && strings.HasSuffix(content, "}") {
var msg Message
if err := jsoniter.UnmarshalFromString(content, &msg); err != nil {
return nil, err
}
return &msg, nil
}
if len(id) > 0 {
return &Message{ID: id[0], Text: content}, nil
}
return &Message{Text: content}, nil
}
// NewStringError create a new message from string error
func NewStringError(content string) (*Message, error) {
if strings.HasPrefix(content, "{") && strings.HasSuffix(content, "}") {
var msg = New()
var errorMessage openai.ErrorMessage
if err := jsoniter.UnmarshalFromString(content, &errorMessage); err != nil {
msg.Text = err.Error() + "\n" + content
return msg, nil
}
msg.Type = "error"
msg.Text = errorMessage.Error.Message
return msg, nil
}
return &Message{Text: content}, nil
}
// NewMap create a new message from map
func NewMap(content map[string]interface{}) (*Message, error) {
return New().Map(content), nil
}
// NewAny create a new message from any content
func NewAny(content interface{}) (*Message, error) {
switch v := content.(type) {
case string:
return NewString(v)
case map[string]interface{}:
return NewMap(v)
}
return nil, fmt.Errorf("unknown content type: %T", content)
}
// NewOpenAI create a new message from OpenAI response
func NewOpenAI(data []byte, isThinking bool) *Message {
// For debug environment, print the response data
if os.Getenv("YAO_AGENT_PRINT_RESPONSE_DATA") == "true" {
log.Trace("[Response Data] %s", string(data))
}
if data == nil || len(data) == 0 {
return nil
}
msg := New()
text := string(data)
data = []byte(strings.TrimPrefix(text, "data: "))
switch {
case strings.Contains(text, `"object":"chat.completion.chunk"`): // Delta content
var chunk openai.ChatCompletionChunk
err := jsoniter.Unmarshal(data, &chunk)
if err != nil {
color.Red("JSON parse error: %s", err.Error())
color.White(string(data))
msg.Text = "JSON parse error\n" + string(data)
msg.Type = "error"
msg.IsDone = true
}
// Empty content, then it is a pending message
if len(chunk.Choices) == 0 {
msg.Pending = true
return msg
}
// Tool calls
if len(chunk.Choices[0].Delta.ToolCalls) > 0 || chunk.Choices[0].FinishReason == "tool_calls" {
msg.Type = "tool_calls_native"
text := ""
if len(chunk.Choices[0].Delta.ToolCalls) > 0 {
id := chunk.Choices[0].Delta.ToolCalls[0].ID
function := chunk.Choices[0].Delta.ToolCalls[0].Function.Name
arguments := chunk.Choices[0].Delta.ToolCalls[0].Function.Arguments
text = arguments
if id != "" {
msg.IsBeginTool = true
msg.IsNew = true // mark as a new message
text = fmt.Sprintf(`{"id": "%s", "function": "%s", "arguments": %s`, id, function, arguments)
}
}
if chunk.Choices[0].FinishReason == "tool_calls" {
msg.IsEndTool = true
}
msg.Text = text
return msg
}
// Text content
if chunk.Choices[0].Delta.Content != "" {
msg.Type = "text"
msg.Text = chunk.Choices[0].Delta.Content
msg.IsDone = chunk.Choices[0].FinishReason == "stop" // is done when the content is finished
return msg
}
// Done messages
if chunk.Choices[0].FinishReason == "stop" || chunk.Choices[0].FinishReason == "tool_calls" {
msg.IsDone = true
return msg
}
// Reasoning content
if chunk.Choices[0].Delta.ReasoningContent != "" {
msg.Type = "think"
msg.Text = chunk.Choices[0].Delta.ReasoningContent
return msg
}
// Content is empty and is thinking, then it is a thinking message pending
if isThinking {
msg.Type = "think"
msg.Text = ""
return msg
}
msg.Text = ""
return msg
case strings.Contains(text, `"usage":`): // usage content
msg.IsDone = true
break
case strings.Contains(text, `[DONE]`):
msg.IsDone = true
return msg
case len(data) > 2 && data[0] == '{' && data[len(data)-1] == '}': // JSON content (error)
var error openai.Error
var errorMessage openai.ErrorMessage
if strings.Contains(string(data), `"error":`) {
if err := jsoniter.Unmarshal(data, &errorMessage); err != nil {
color.Red("JSON parse error: %s", err.Error())
color.White(string(data))
msg.Text = "JSON parse error\n" + string(data)
msg.Type = "error"
msg.IsDone = true
return msg
}
error = errorMessage.Error
} else {
err := jsoniter.Unmarshal(data, &error)
if err != nil {
color.Red("JSON parse error: %s", err.Error())
color.White(string(data))
msg.Text = "JSON parse error\n" + string(data)
msg.Type = "error"
msg.IsDone = true
return msg
}
}
message := error.Message
if message == "" {
message = "Unknown error occurred\n" + string(data)
}
msg.Type = "error"
msg.Text = message
msg.IsDone = true
return msg
case !strings.Contains(text, `data: `): // unknown message or uncompleted message
msg.Pending = true
msg.Text = text
return msg
default: // unknown message
str := strings.TrimPrefix(strings.Trim(string(data), "\""), "data: ")
msg.Type = "error"
msg.Text = str
return msg
}
return msg
}
// String returns the string representation
func (m *Message) String() string {
typ := m.Type
if typ == "" {
typ = "text"
}
switch typ {
case "text", "think", "tool", "tool_calls_native":
return m.Text
case "error":
return m.Text
default:
raw, _ := jsoniter.MarshalToString(map[string]interface{}{"type": m.Type, "props": m.Props})
return raw
}
}
// SetText set the text
func (m *Message) SetText(text string) *Message {
m.Text = text
if m.Data != nil {
if replaced := helper.Bind(text, m.Data); replaced != nil {
if replacedText, ok := replaced.(string); ok {
m.Text = replacedText
}
}
}
return m
}
// SetProps set the props
func (m *Message) SetProps(props map[string]interface{}) *Message {
m.Props = props
return m
}
// Error set the error
func (m *Message) Error(message interface{}) *Message {
m.Type = "error"
switch v := message.(type) {
case error:
m.Text = v.Error()
case string:
m.Text = v
default:
m.Text = fmt.Sprintf("%v", message)
}
return m
}
// SetResult set the result
func (m *Message) SetResult(result any) *Message {
m.Result = result
m.Type = "result" // set the type to result
return m
}
// SetContent set the content
func (m *Message) SetContent(content string) *Message {
if strings.HasPrefix(content, "{") && strings.HasSuffix(content, "}") {
var msg Message
if err := jsoniter.UnmarshalFromString(content, &msg); err != nil {
m.Text = err.Error() + "\n" + content
return m
}
*m = msg
} else {
m.Text = content
m.Type = "text"
}
return m
}
// AppendTo append the contents
func (m *Message) AppendTo(contents *Contents) *Message {
// Set type
if m.Type == "" {
m.Type = "text"
}
switch m.Type {
case "text", "think", "tool", "tool_calls_native":
if m.Text != "" {
if m.IsNew {
contents.NewText([]byte(m.Text), Extra{ID: m.ID, Begin: m.Begin, End: m.End})
return m
}
contents.AppendText([]byte(m.Text), Extra{ID: m.ID, Begin: m.Begin, End: m.End})
return m
}
return m
case "loading", "error", "action", "progress", "plan", "result": // Ignore progress, loading, plan and error messages
return m
default:
if m.IsNew {
contents.NewType(m.Type, m.Props)
return m
}
contents.UpdateType(m.Type, m.Props)
return m
}
}
// Content get the content
func (m *Message) Content() string {
content := map[string]interface{}{"text": m.Text}
if m.Attachments != nil {
content["attachments"] = m.Attachments
}
if m.Type != "" {
content["type"] = m.Type
}
contentRaw, _ := jsoniter.MarshalToString(content)
return contentRaw
}
// ToMap convert to map
func (m *Message) ToMap() map[string]interface{} {
return map[string]interface{}{
"content": m.Content(),
"role": m.Role,
"name": m.Name,
}
}
// Map set from map
func (m *Message) Map(msg map[string]interface{}) *Message {
if msg == nil {
return m
}
// Content {"text": "xxxx", "attachments": ... }
if content, ok := msg["content"].(string); ok {
if strings.HasPrefix(content, "{") && strings.HasSuffix(content, "}") {
var msg Message
if err := jsoniter.UnmarshalFromString(content, &msg); err != nil {
m.Text = err.Error() + "\n" + content
return m
}
*m = msg
} else {
m.Text = content
m.Type = "text"
}
}
// attachments
if attachments, has := msg["attachments"]; has {
raw, _ := jsoniter.Marshal(attachments)
m.Attachments = []attachment.Attachment{}
if err := jsoniter.Unmarshal(raw, &m.Attachments); err != nil {
color.Red("JSON parse error: %s", err.Error())
color.White(string(raw))
}
}
if role, ok := msg["role"].(string); ok {
m.Role = role
}
if name, ok := msg["name"].(string); ok {
m.Name = name
}
if text, ok := msg["text"].(string); ok {
m.Text = text
}
if typ, ok := msg["type"].(string); ok {
m.Type = typ
}
if done, ok := msg["done"].(bool); ok {
m.IsDone = done
}
if props, ok := msg["props"].(map[string]interface{}); ok {
m.Props = props
}
if isNew, ok := msg["new"].(bool); ok {
m.IsNew = isNew
}
if isDelta, ok := msg["delta"].(bool); ok {
m.IsDelta = isDelta
}
if assistantID, ok := msg["assistant_id"].(string); ok {
m.AssistantID = assistantID
// Set name
if m.Role == "assistant" {
m.Name = m.AssistantID
}
}
if assistantName, ok := msg["assistant_name"].(string); ok {
m.AssistantName = assistantName
}
if assistantAvatar, ok := msg["assistant_avatar"].(string); ok {
m.AssistantAvatar = assistantAvatar
}
if actions, ok := msg["actions"].([]interface{}); ok {
for _, action := range actions {
if v, ok := action.(map[string]interface{}); ok {
action := Action{}
if name, ok := v["name"].(string); ok {
action.Name = name
}
if t, ok := v["type"].(string); ok {
action.Type = t
}
if payload, ok := v["payload"].(map[string]interface{}); ok {
action.Payload = payload
}
m.Actions = append(m.Actions, action)
}
}
}
if data, ok := msg["data"].(map[string]interface{}); ok {
m.Data = data
}
return m
}
// Done set the done flag
func (m *Message) Done() *Message {
m.IsDone = true
return m
}
// Assistant set the assistant
func (m *Message) Assistant(id string, name string, avatar string) *Message {
m.AssistantID = id
m.AssistantName = name
m.AssistantAvatar = avatar
return m
}
// Action add an action
func (m *Message) Action(name string, t string, payload interface{}, next string) *Message {
if m.Data != nil {
payload = helper.Bind(payload, m.Data)
}
m.Actions = append(m.Actions, Action{
Name: name,
Type: t,
Payload: payload,
})
return m
}
// Bind replace with data
func (m *Message) Bind(data map[string]interface{}) *Message {
if data == nil {
return m
}
m.Data = maps.Of(data).Dot()
return m
}
// Callback callback the message
func (m *Message) Callback(fn interface{}) *Message {
if fn != nil {
switch v := fn.(type) {
case func(msg *Message):
if v == nil {
break
}
v(m)
break
case func():
if v == nil {
break
}
v()
break
default:
fmt.Println("no match callback")
break
}
}
return m
}
// WriteError writes an error message to response writer
func (m *Message) WriteError(w gin.ResponseWriter, message string) {
errMsg := strings.Trim(exception.New(message, 500).Message, "\"")
data := []byte(fmt.Sprintf(`{"text":"%s","type":"error"`, errMsg))
if m.IsDone {
data = []byte(fmt.Sprintf(`{"text":"%s","type":"error","done":true`, errMsg))
}
data = append([]byte("data: "), data...)
data = append(data, []byte("}\n\n")...)
if _, err := w.Write(data); err != nil {
color.Red("Write JSON Message Error: %s", message)
}
w.Flush()
}
// MarshalJSON implements json.Marshaler interface
func (m *Message) MarshalJSON() ([]byte, error) {
type Alias Message
return jsoniter.Marshal(&struct {
*Alias
}{
Alias: (*Alias)(m),
})
}
// UnmarshalJSON implements json.Unmarshaler interface
func (m *Message) UnmarshalJSON(data []byte) error {
type Alias Message
aux := &struct {
*Alias
}{
Alias: (*Alias)(m),
}
if err := jsoniter.Unmarshal(data, &aux); err != nil {
return err
}
return nil
}
// MarshalJSON implements json.Marshaler interface
func (a *Action) MarshalJSON() ([]byte, error) {
type Alias Action
return jsoniter.Marshal(&struct {
*Alias
}{
Alias: (*Alias)(a),
})
}
// UnmarshalJSON implements json.Unmarshaler interface
func (a *Action) UnmarshalJSON(data []byte) error {
type Alias Action
aux := &struct {
*Alias
}{
Alias: (*Alias)(a),
}
if err := jsoniter.Unmarshal(data, &aux); err != nil {
return err
}
return nil
}
// Write writes the message to response writer using the message queue
func (m *Message) Write(w gin.ResponseWriter) bool {
return WriteMessageAsync(m, w)
}

153
agent/message/queue.go
View file

@ -1,153 +0,0 @@
package message
import (
"sync"
"time"
"github.com/fatih/color"
"github.com/gin-gonic/gin"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/kun/log"
)
// AsyncMessageQueue represents a queue for handling message writes
type AsyncMessageQueue struct {
queue chan *AsyncTask
workers int
wg sync.WaitGroup
shutdown chan struct{}
}
// AsyncTask represents a task to write a message
type AsyncTask struct {
message *Message
writer gin.ResponseWriter
done chan bool
}
var (
defaultQueue *AsyncMessageQueue
queueOnce sync.Once
)
// GetQueue returns the default message queue instance
func GetQueue() *AsyncMessageQueue {
queueOnce.Do(func() {
defaultQueue = NewAsyncQueue(10) // Initialize with 10 workers
defaultQueue.Start()
})
return defaultQueue
}
// NewAsyncQueue creates a new message queue with the specified number of workers
func NewAsyncQueue(workers int) *AsyncMessageQueue {
return &AsyncMessageQueue{
queue: make(chan *AsyncTask, 1000), // Buffer size of 1000
workers: workers,
shutdown: make(chan struct{}),
}
}
// Start starts the message queue workers
func (mq *AsyncMessageQueue) Start() {
for i := 0; i < mq.workers; i++ {
mq.wg.Add(1)
go mq.worker()
}
}
// Stop stops the message queue workers
func (mq *AsyncMessageQueue) Stop() {
close(mq.shutdown)
mq.wg.Wait()
}
// worker processes messages from the queue
func (mq *AsyncMessageQueue) worker() {
defer mq.wg.Done()
for {
select {
case task := <-mq.queue:
if task == nil {
continue
}
success := writeMessageToResponse(task.message, task.writer)
if task.done != nil {
task.done <- success
}
case <-mq.shutdown:
return
}
}
}
// WriteMessageAsync writes the message to response writer using the message queue
func WriteMessageAsync(m *Message, w gin.ResponseWriter) bool {
done := make(chan bool, 1)
task := &AsyncTask{
message: m,
writer: w,
done: done,
}
// Try to send the task to the queue with a timeout
select {
case GetQueue().queue <- task:
// Wait for the message to be processed with a longer timeout
select {
case success := <-done:
return success
case <-time.After(5 * time.Second): // Increased timeout to 5 seconds
log.Error("Message processing timeout")
return false
}
case <-time.After(1 * time.Second): // Increased queue timeout to 1 second
log.Error("Queue is full, message dropped")
return false
}
}
// writeMessageToResponse writes the message directly to the response writer
func writeMessageToResponse(m *Message, w gin.ResponseWriter) bool {
// Sync write to response writer
locker.Lock()
defer locker.Unlock()
defer func() {
if r := recover(); r != nil {
// Ignore if done is true
if m.IsDone {
return
}
message := "Write Response Exception: (if client close the connection, it's normal) \n %s\n\n"
color.Red(message, r)
// Print the message
raw, _ := jsoniter.MarshalToString(m)
color.White("Message:\n %s", raw)
}
}()
// Ignore silent messages
if m.Silent {
return true
}
data, err := jsoniter.Marshal(m)
if err != nil {
log.Error("%s", err.Error())
return false
}
data = append([]byte("data: "), data...)
data = append(data, []byte("\n\n")...)
if _, err := w.Write(data); err != nil {
color.Red("Write JSON Message Error: %s", err.Error())
return false
}
w.Flush()
return true
}

48
agent/message/types.go
View file

@ -1,48 +0,0 @@
package message
import "github.com/yaoapp/yao/attachment"
// Message the message
type Message struct {
ID string `json:"id,omitempty"` // id for the message
ToolID string `json:"tool_id,omitempty"` // tool_id for the message
Text string `json:"text,omitempty"` // text content
Type string `json:"type,omitempty"` // error, text, plan, table, form, page, file, video, audio, image, markdown, json ...
Props map[string]interface{} `json:"props,omitempty"` // props for the types
IsDone bool `json:"done,omitempty"` // Mark as a done message from agent
IsNew bool `json:"new,omitempty"` // Mark as a new message from agent
IsDelta bool `json:"delta,omitempty"` // Mark as a delta message from agent
Actions []Action `json:"actions,omitempty"` // Conversation Actions for frontend
Attachments []attachment.Attachment `json:"attachments,omitempty"` // File attachments
Role string `json:"role,omitempty"` // user, assistant, system ...
Name string `json:"name,omitempty"` // name for the message
AssistantID string `json:"assistant_id,omitempty"` // assistant_id (for assistant role = assistant )
AssistantName string `json:"assistant_name,omitempty"` // assistant_name (for assistant role = assistant )
AssistantAvatar string `json:"assistant_avatar,omitempty"` // assistant_avatar (for assistant role = assistant )
Mentions []Mention `json:"menions,omitempty"` // Mentions for the message ( for user role = user )
Data map[string]interface{} `json:"-"` // data for the message
Pending bool `json:"-"` // pending for the message
Hidden bool `json:"hidden,omitempty"` // hidden for the message (not show in the UI and history)
Retry bool `json:"retry,omitempty"` // retry for the message
Silent bool `json:"silent,omitempty"` // silent for the message (not show in the UI and history)
IsTool bool `json:"-"` // is tool for the message for native tool_calls
IsBeginTool bool `json:"-"` // is new tool for the message for native tool_calls
IsEndTool bool `json:"-"` // is end tool for the message for native tool_calls
Result any `json:"result,omitempty"` // result for the message
Begin int64 `json:"begin,omitempty"` // begin at for the message // timestamp
End int64 `json:"end,omitempty"` // end at for the message // timestamp
}
// Mention represents a mention
type Mention struct {
ID string `json:"assistant_id"` // assistant_id
Name string `json:"name"` // name
Avatar string `json:"avatar,omitempty"` // avatar
}
// Action the action
type Action struct {
Name string `json:"name,omitempty"`
Type string `json:"type"`
Payload interface{} `json:"payload,omitempty"`
}

1
agent/plan/plan.go
View file

@ -1 +0,0 @@
package plan

332
agent/process.go
View file

@ -1,332 +0,0 @@
package agent
import (
"fmt"
"strconv"
"strings"
"github.com/gin-gonic/gin"
"github.com/yaoapp/gou/process"
"github.com/yaoapp/kun/exception"
"github.com/yaoapp/yao/agent/message"
store "github.com/yaoapp/yao/agent/store/types"
)
func init() {
process.RegisterGroup("agent", map[string]process.Handler{
"write": ProcessWrite,
"assistant.create": processAssistantCreate,
"assistant.save": processAssistantSave,
"assistant.delete": processAssistantDelete,
"assistant.search": processAssistantSearch,
"assistant.find": processAssistantFind,
"assistant.match": processAssistantMatch, // Match assistant by content and params
})
// Neo is deprecated, use agent instead (for backward compatibility, It will be removed in the future)
process.RegisterGroup("neo", map[string]process.Handler{
"write": ProcessWrite,
"assistant.create": processAssistantCreate,
"assistant.save": processAssistantSave,
"assistant.delete": processAssistantDelete,
"assistant.search": processAssistantSearch,
"assistant.find": processAssistantFind,
"assistant.match": processAssistantMatch, // Match assistant by content and params
})
}
// ProcessWrite process the write request
func ProcessWrite(process *process.Process) interface{} {
process.ValidateArgNums(2)
w, ok := process.Args[0].(gin.ResponseWriter)
if !ok {
exception.New("The first argument must be a io.Writer", 400).Throw()
return nil
}
data, ok := process.Args[1].([]interface{})
if !ok {
exception.New("The second argument must be a Array", 400).Throw()
return nil
}
for _, new := range data {
if v, ok := new.(map[string]interface{}); ok {
newMsg := message.New().Map(v)
newMsg.Write(w)
}
}
return nil
}
// processAssistantCreate process the assistant create request
func processAssistantCreate(process *process.Process) interface{} {
process.ValidateArgNums(1)
data := process.ArgsMap(0)
agent := GetAgent()
if agent.Store == nil {
exception.New("Agent store is not initialized", 500).Throw()
}
// Convert to AssistantModel
model, err := store.ToAssistantModel(data)
if err != nil {
exception.New("Invalid assistant data: %s", 400, err.Error()).Throw()
}
id, err := agent.Store.SaveAssistant(model)
if err != nil {
exception.New("Failed to create assistant: %s", 500, err.Error()).Throw()
}
return id
}
// processAssistantSave process the assistant save request
func processAssistantSave(process *process.Process) interface{} {
process.ValidateArgNums(1)
data := process.ArgsMap(0)
agent := GetAgent()
if agent.Store == nil {
exception.New("Agent store is not initialized", 500).Throw()
}
// Convert to AssistantModel
model, err := store.ToAssistantModel(data)
if err != nil {
exception.New("Invalid assistant data: %s", 400, err.Error()).Throw()
}
id, err := agent.Store.SaveAssistant(model)
if err != nil {
exception.New("Failed to save assistant: %s", 500, err.Error()).Throw()
}
return id
}
// processAssistantDelete process the assistant delete request
func processAssistantDelete(process *process.Process) interface{} {
process.ValidateArgNums(1)
assistantID := process.ArgsString(0)
agent := GetAgent()
if agent.Store == nil {
exception.New("Agent store is not initialized", 500).Throw()
}
err := agent.Store.DeleteAssistant(assistantID)
if err != nil {
exception.New("Failed to delete assistant: %s", 500, err.Error()).Throw()
}
return gin.H{"message": "ok"}
}
// processAssistantMatch process the assistant match request
func processAssistantMatch(process *process.Process) interface{} {
process.ValidateArgNums(1)
content := process.Args[0]
params := map[string]interface{}{}
if len(process.Args) > 1 {
params = process.ArgsMap(1)
}
// Limit default to 20
if _, has := params["limit"]; !has {
params["limit"] = 20
}
// Max limit to 100
if limit, has := params["limit"]; has {
switch v := limit.(type) {
case int:
if v > 100 {
params["limit"] = 100
}
case string:
limitInt, err := strconv.Atoi(v)
if err != nil {
exception.New("Invalid limit type: %T", 500, limit).Throw()
}
params["limit"] = limitInt
if limitInt > 100 {
params["limit"] = 100
}
default:
exception.New("Invalid limit type: %T", 500, limit).Throw()
}
}
// Match using Store
return assistantMatchStore(content, params)
}
// parseAssistantFilter parse common filter parameters
func parseAssistantFilter(params map[string]interface{}) store.AssistantFilter {
filter := store.AssistantFilter{}
// Parse page and pagesize
if page, ok := params["page"]; ok {
pageStr := fmt.Sprintf("%v", page)
if pageInt, err := strconv.Atoi(pageStr); err == nil {
filter.Page = pageInt
}
}
if pagesize, ok := params["pagesize"]; ok {
pagesizeStr := fmt.Sprintf("%v", pagesize)
if pagesizeInt, err := strconv.Atoi(pagesizeStr); err == nil {
filter.PageSize = pagesizeInt
}
}
// select
if sel, ok := params["select"]; ok {
switch v := sel.(type) {
case []interface{}:
filter.Select = []string{}
for _, field := range v {
switch v := field.(type) {
case string:
filter.Select = append(filter.Select, v)
case interface{}:
filter.Select = append(filter.Select, fmt.Sprintf("%v", v))
}
}
case []string:
filter.Select = v
case string:
fields := strings.Split(v, ",")
filter.Select = fields
}
}
// Parse tags
if tags, ok := params["tags"]; ok {
switch v := tags.(type) {
case []interface{}:
filter.Tags = make([]string, len(v))
for i, tag := range v {
filter.Tags[i] = fmt.Sprintf("%v", tag)
}
case []string:
filter.Tags = v
}
}
// Parse keywords
if keywords, ok := params["keywords"].(string); ok {
filter.Keywords = keywords
}
// Parse connector
if connector, ok := params["connector"].(string); ok {
filter.Connector = connector
}
// Parse mentionable
if mentionable, ok := params["mentionable"].(bool); ok {
filter.Mentionable = &mentionable
}
// Parse automated
if automated, ok := params["automated"].(bool); ok {
filter.Automated = &automated
}
return filter
}
func assistantMatchStore(content interface{}, params map[string]interface{}) interface{} {
agent := GetAgent()
if agent.Store == nil {
exception.New("Agent store is not initialized", 500).Throw()
}
// Convert limit to pagesize
if limit, has := params["limit"]; has {
params["pagesize"] = limit
}
params["page"] = 1
// Parse content to keywords if not empty
if content != nil {
contentStr := fmt.Sprintf("%v", content)
if contentStr != "" {
params["keywords"] = contentStr
}
}
filter := parseAssistantFilter(params)
res, err := agent.Store.GetAssistants(filter)
if err != nil {
exception.New("get assistants error: %s", 500, err).Throw()
}
return res.Data
}
// processAssistantSearch process the assistant search request
func processAssistantSearch(process *process.Process) interface{} {
params := process.ArgsMap(0)
filter := parseAssistantFilter(params)
// Get assistants
agent := GetAgent()
if agent.Store == nil {
exception.New("Agent store is not initialized", 500).Throw()
}
locale := "en"
if len(process.Args) > 1 {
locale = process.ArgsString(1)
}
res, err := agent.Store.GetAssistants(filter, locale)
if err != nil {
exception.New("get assistants error: %s", 500, err).Throw()
}
return res
}
// processAssistantFind process the assistant find request
func processAssistantFind(process *process.Process) interface{} {
process.ValidateArgNums(1)
assistantID := process.ArgsString(0)
agent := GetAgent()
if agent.Store == nil {
exception.New("Agent store is not initialized", 500).Throw()
}
filter := store.AssistantFilter{
AssistantID: assistantID,
Page: 1,
PageSize: 1,
}
locale := "en"
if len(process.Args) > 1 {
locale = process.ArgsString(1)
}
res, err := agent.Store.GetAssistants(filter, locale)
if err != nil {
exception.New("Failed to find assistant: %s", 500, err.Error()).Throw()
}
if len(res.Data) == 0 {
exception.New("Assistant not found: %s", 404, assistantID).Throw()
}
return res.Data[0]
}

513
agent/process_test.go
View file

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

3
agent/types/types.go
View file

@ -4,7 +4,6 @@ import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/yaoapp/yao/agent/assistant" "github.com/yaoapp/yao/agent/assistant"
store "github.com/yaoapp/yao/agent/store/types" store "github.com/yaoapp/yao/agent/store/types"
"github.com/yaoapp/yao/agent/vision"
) )
// DSL AI assistant // DSL AI assistant
@ -34,7 +33,7 @@ type DSL struct {
// 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
Vision *vision.Vision `json:"-" yaml:"-"` // Vision *vision.Vision `json:"-" yaml:"-"`
GuardHandlers []gin.HandlerFunc `json:"-" yaml:"-"` GuardHandlers []gin.HandlerFunc `json:"-" yaml:"-"`
} }

208
agent/vision/driver/local/storage.go
View file

@ -1,208 +0,0 @@
package local
import (
"bytes"
"context"
"crypto/sha256"
"fmt"
"image"
"image/jpeg"
"image/png"
"io"
"path/filepath"
"strings"
"time"
"github.com/yaoapp/gou/fs"
)
// MaxImageSize maximum image size (1920x1080)
const MaxImageSize = 1920
// Storage the local storage driver
type Storage struct {
Path string `json:"path" yaml:"path"`
Compression bool `json:"compression" yaml:"compression"`
BaseURL string `json:"base_url" yaml:"base_url"`
PreviewURL func(fileID string) string `json:"-" yaml:"-"`
}
// New create a new local storage
func New(options map[string]interface{}) (*Storage, error) {
storage := &Storage{
Compression: true,
}
if path, ok := options["path"].(string); ok {
storage.Path = path
}
if compression, ok := options["compression"].(bool); ok {
storage.Compression = compression
}
if baseURL, ok := options["base_url"].(string); ok {
storage.BaseURL = baseURL
}
if previewURL, ok := options["preview_url"].(func(string) string); ok {
storage.PreviewURL = previewURL
}
if storage.Path == "" {
return nil, fmt.Errorf("path is required")
}
return storage, nil
}
// Upload upload file to local storage
func (storage *Storage) Upload(ctx context.Context, filename string, reader io.Reader, contentType string) (string, error) {
data, err := fs.Get("data")
if err != nil {
return "", err
}
ext := filepath.Ext(filename)
id := storage.makeID(filename, ext)
path := filepath.Join(storage.Path, id)
// Create directory if not exists
dir := filepath.Dir(path)
if err := data.MkdirAll(dir, 0755); err != nil {
return "", err
}
// Check if compression is enabled and if it's an image
if storage.Compression && isImage(contentType) {
// Read the entire image into memory
content, err := io.ReadAll(reader)
if err != nil {
return "", fmt.Errorf("failed to read image: %w", err)
}
// Compress image
compressed, err := compressImage(content, contentType)
if err != nil {
return "", fmt.Errorf("failed to compress image: %w", err)
}
// Write compressed image
_, err = data.Write(path, bytes.NewReader(compressed), 0644)
if err != nil {
return "", err
}
} else {
// Write file without compression
_, err = data.Write(path, reader, 0644)
if err != nil {
return "", err
}
}
return id, nil
}
// Download download file from local storage
func (storage *Storage) Download(ctx context.Context, fileID string) (io.ReadCloser, string, error) {
data, err := fs.Get("data")
if err != nil {
return nil, "", err
}
path := filepath.Join(storage.Path, fileID)
reader, err := data.ReadCloser(path)
if err != nil {
return nil, "", err
}
contentType := "application/octet-stream"
if v, err := data.MimeType(path); err == nil {
contentType = v
}
return reader, contentType, nil
}
// URL get file url
func (storage *Storage) URL(ctx context.Context, fileID string) string {
if storage.PreviewURL != nil {
return storage.PreviewURL(fileID)
}
if storage.BaseURL != "" {
return fmt.Sprintf("%s/%s", strings.TrimRight(storage.BaseURL, "/"), fileID)
}
return fmt.Sprintf("%s/%s", storage.Path, fileID)
}
func (storage *Storage) makeID(filename string, ext string) string {
date := time.Now().Format("20060102")
hash := fmt.Sprintf("%x", sha256.Sum256([]byte(filename)))[:8]
name := strings.TrimSuffix(filepath.Base(filename), ext)
return fmt.Sprintf("%s/%s-%s%s", date, name, hash, ext)
}
// isImage checks if the content type is an image
func isImage(contentType string) bool {
return strings.HasPrefix(contentType, "image/")
}
// compressImage compresses the image while maintaining aspect ratio
func compressImage(data []byte, contentType string) ([]byte, error) {
// Decode image
img, _, err := image.Decode(bytes.NewReader(data))
if err != nil {
return nil, fmt.Errorf("failed to decode image: %w", err)
}
// Calculate new dimensions
bounds := img.Bounds()
width := bounds.Dx()
height := bounds.Dy()
var newWidth, newHeight int
if width > height {
if width > MaxImageSize {
newWidth = MaxImageSize
newHeight = int(float64(height) * (float64(MaxImageSize) / float64(width)))
} else {
return data, nil // No need to resize
}
} else {
if height > MaxImageSize {
newHeight = MaxImageSize
newWidth = int(float64(width) * (float64(MaxImageSize) / float64(height)))
} else {
return data, nil // No need to resize
}
}
// Create new image with new dimensions
newImg := image.NewRGBA(image.Rect(0, 0, newWidth, newHeight))
// Scale the image using bilinear interpolation
for y := 0; y < newHeight; y++ {
for x := 0; x < newWidth; x++ {
srcX := float64(x) * float64(width) / float64(newWidth)
srcY := float64(y) * float64(height) / float64(newHeight)
newImg.Set(x, y, img.At(int(srcX), int(srcY)))
}
}
// Encode image
var buf bytes.Buffer
switch contentType {
case "image/jpeg":
err = jpeg.Encode(&buf, newImg, &jpeg.Options{Quality: 85})
case "image/png":
err = png.Encode(&buf, newImg)
default:
return data, nil // Unsupported format, return original
}
if err != nil {
return nil, fmt.Errorf("failed to encode image: %w", err)
}
return buf.Bytes(), nil
}

150
agent/vision/driver/local/storage_test.go
View file

@ -1,150 +0,0 @@
package local
import (
"bytes"
"context"
"image"
"image/png"
"io"
"testing"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/test"
)
func TestLocalStorage(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
t.Run("Create Storage", func(t *testing.T) {
storage, err := New(map[string]interface{}{
"path": "/__vision_test",
"compression": true,
})
assert.NoError(t, err)
assert.NotNil(t, storage)
assert.Equal(t, "/__vision_test", storage.Path)
assert.True(t, storage.Compression)
})
t.Run("Upload and Download", func(t *testing.T) {
storage, err := New(map[string]interface{}{
"path": "/__vision_test",
"compression": true,
})
assert.NoError(t, err)
content := []byte("test content")
reader := bytes.NewReader(content)
fileID, err := storage.Upload(context.Background(), "test.txt", reader, "text/plain")
assert.NoError(t, err)
assert.NotEmpty(t, fileID)
// Download
reader2, contentType, err := storage.Download(context.Background(), fileID)
assert.NoError(t, err)
assert.Contains(t, contentType, "text/plain")
downloaded, err := io.ReadAll(reader2)
assert.NoError(t, err)
assert.Equal(t, content, downloaded)
})
t.Run("Upload and Download Image with Compression", func(t *testing.T) {
storage, err := New(map[string]interface{}{
"path": "/__vision_test",
"compression": true,
})
assert.NoError(t, err)
// Create a test image (2000x2000 pixels)
img := image.NewRGBA(image.Rect(0, 0, 2000, 2000))
var buf bytes.Buffer
err = png.Encode(&buf, img)
assert.NoError(t, err)
// Upload
reader := bytes.NewReader(buf.Bytes())
fileID, err := storage.Upload(context.Background(), "test.png", reader, "image/png")
assert.NoError(t, err)
assert.NotEmpty(t, fileID)
// Download and verify size
reader2, contentType, err := storage.Download(context.Background(), fileID)
assert.NoError(t, err)
assert.Equal(t, "image/png", contentType)
downloaded, err := io.ReadAll(reader2)
assert.NoError(t, err)
// Decode the downloaded image
downloadedImg, _, err := image.Decode(bytes.NewReader(downloaded))
assert.NoError(t, err)
// Verify dimensions
bounds := downloadedImg.Bounds()
assert.LessOrEqual(t, bounds.Dx(), MaxImageSize)
assert.LessOrEqual(t, bounds.Dy(), MaxImageSize)
})
t.Run("Upload Image without Compression", func(t *testing.T) {
storage, err := New(map[string]interface{}{
"path": "/__vision_test",
"compression": false,
})
assert.NoError(t, err)
// Create a test image (2000x2000 pixels)
img := image.NewRGBA(image.Rect(0, 0, 2000, 2000))
var buf bytes.Buffer
err = png.Encode(&buf, img)
assert.NoError(t, err)
// Upload
reader := bytes.NewReader(buf.Bytes())
fileID, err := storage.Upload(context.Background(), "test.png", reader, "image/png")
assert.NoError(t, err)
assert.NotEmpty(t, fileID)
// Download and verify size
reader2, contentType, err := storage.Download(context.Background(), fileID)
assert.NoError(t, err)
assert.Equal(t, "image/png", contentType)
downloaded, err := io.ReadAll(reader2)
assert.NoError(t, err)
// Decode the downloaded image
downloadedImg, _, err := image.Decode(bytes.NewReader(downloaded))
assert.NoError(t, err)
// Verify dimensions are unchanged
bounds := downloadedImg.Bounds()
assert.Equal(t, 2000, bounds.Dx())
assert.Equal(t, 2000, bounds.Dy())
})
t.Run("URL Generation", func(t *testing.T) {
storage, err := New(map[string]interface{}{
"path": "/__vision_test",
"compression": true,
})
assert.NoError(t, err)
fileID := "20240101/test-12345678.txt"
url := storage.URL(context.Background(), fileID)
assert.Equal(t, "/__vision_test/20240101/test-12345678.txt", url)
})
t.Run("Download Non-existent File", func(t *testing.T) {
storage, err := New(map[string]interface{}{
"path": "/__vision_test",
"compression": true,
})
assert.NoError(t, err)
_, _, err = storage.Download(context.Background(), "non-existent.txt")
assert.Error(t, err)
})
}

190
agent/vision/driver/openai/model.go
View file

@ -1,190 +0,0 @@
package openai
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"github.com/yaoapp/gou/fs"
)
// Model the OpenAI vision model
type Model struct {
APIKey string `json:"api_key" yaml:"api_key"`
Model string `json:"model" yaml:"model"`
Compression bool `json:"compression" yaml:"compression"`
Prompt string `json:"prompt" yaml:"prompt"`
}
// New create a new OpenAI vision model
func New(options map[string]interface{}) (*Model, error) {
model := &Model{
Model: "gpt-4-vision-preview",
Compression: true,
}
if apiKey, ok := options["api_key"].(string); ok {
model.APIKey = apiKey
}
if modelName, ok := options["model"].(string); ok {
model.Model = modelName
}
if compression, ok := options["compression"].(bool); ok {
model.Compression = compression
}
if prompt, ok := options["prompt"].(string); ok {
model.Prompt = prompt
}
if model.APIKey == "" {
return nil, fmt.Errorf("api_key is required")
}
return model, nil
}
// Analyze analyze image using OpenAI vision model
func (model *Model) Analyze(ctx context.Context, fileID string, prompt ...string) (map[string]interface{}, error) {
if model.APIKey == "" {
return nil, fmt.Errorf("api_key is required")
}
// Use default prompt if none provided
userPrompt := model.Prompt
if len(prompt) > 0 && prompt[0] != "" {
userPrompt = prompt[0]
}
// Check if fileID is a URL or base64 data
var imageURL string
if strings.HasPrefix(fileID, "data:image/") {
// Already a base64 data URL
imageURL = fileID
} else if strings.HasPrefix(fileID, "http://") || strings.HasPrefix(fileID, "https://") {
// Already a URL
imageURL = fileID
} else {
// Try to read the file and convert to base64
data, err := fs.Get("data")
if err != nil {
return nil, fmt.Errorf("failed to get data fs: %w", err)
}
reader, err := data.ReadCloser(fileID)
if err != nil {
return nil, fmt.Errorf("failed to read file: %w", err)
}
defer reader.Close()
content, err := io.ReadAll(reader)
if err != nil {
return nil, fmt.Errorf("failed to read content: %w", err)
}
// Get content type
contentType := "image/png" // default
if v, err := data.MimeType(fileID); err == nil {
contentType = v
}
// Convert to base64
base64Data := base64.StdEncoding.EncodeToString(content)
imageURL = fmt.Sprintf("data:%s;base64,%s", contentType, base64Data)
}
// Prepare the request body
reqBody := map[string]interface{}{
"model": model.Model,
"messages": []map[string]interface{}{
{
"role": "user",
"content": []map[string]interface{}{
{
"type": "text",
"text": userPrompt,
},
{
"type": "image_url",
"image_url": map[string]interface{}{
"url": imageURL,
},
},
},
},
},
"max_tokens": 1000,
}
jsonBody, err := json.Marshal(reqBody)
if err != nil {
return nil, fmt.Errorf("failed to marshal request body: %w", err)
}
// Create request
req, err := http.NewRequestWithContext(ctx, "POST", "https://api.openai.com/v1/chat/completions", bytes.NewBuffer(jsonBody))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", model.APIKey))
// Send request
client := &http.Client{}
resp, err := client.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
// Read response
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("OpenAI API error: %s", string(body))
}
// Parse response
var result map[string]interface{}
if err := json.Unmarshal(body, &result); err != nil {
return nil, fmt.Errorf("failed to parse response: %w", err)
}
// Extract content
choices, ok := result["choices"].([]interface{})
if !ok || len(choices) == 0 {
return nil, fmt.Errorf("invalid response format")
}
message, ok := choices[0].(map[string]interface{})["message"].(map[string]interface{})
if !ok {
return nil, fmt.Errorf("invalid response format")
}
content, ok := message["content"].(string)
if !ok {
return nil, fmt.Errorf("invalid response format")
}
// Try to parse content as JSON
var description map[string]interface{}
if err := json.Unmarshal([]byte(content), &description); err != nil {
// If not JSON, use the content as description
description = map[string]interface{}{
"description": content,
}
}
return description, nil
}

194
agent/vision/driver/openai/model_test.go
View file

@ -1,194 +0,0 @@
package openai
import (
"bytes"
"context"
"encoding/base64"
"os"
"testing"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/gou/fs"
"github.com/yaoapp/yao/agent/vision/driver/s3"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/test"
)
var (
// 1x1 transparent PNG
testImageBase64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
)
func TestOpenAIModel(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
t.Run("Create Model", func(t *testing.T) {
model, err := New(map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
})
assert.NoError(t, err)
assert.NotNil(t, model)
if model != nil {
assert.Equal(t, os.Getenv("OPENAI_API_KEY"), model.APIKey)
assert.Equal(t, os.Getenv("VISION_MODEL"), model.Model)
assert.True(t, model.Compression)
}
})
t.Run("Create Model with Invalid API Key", func(t *testing.T) {
_, err := New(map[string]interface{}{})
assert.Error(t, err)
assert.Contains(t, err.Error(), "api_key is required")
})
t.Run("Analyze with Base64 Image", func(t *testing.T) {
model, err := New(map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
})
assert.NoError(t, err)
// Use base64 image data
result, err := model.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "Describe this image in detail")
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotEmpty(t, result["description"])
})
t.Run("Analyze with URL", func(t *testing.T) {
if os.Getenv("S3_API") == "" || os.Getenv("S3_ACCESS_KEY") == "" ||
os.Getenv("S3_SECRET_KEY") == "" || os.Getenv("S3_BUCKET") == "" {
t.Skip("S3 environment variables not set")
}
model, err := New(map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
})
assert.NoError(t, err)
// Create S3 client and upload test image
s3Client, err := s3.New(map[string]interface{}{
"endpoint": os.Getenv("S3_API"),
"region": "auto",
"key": os.Getenv("S3_ACCESS_KEY"),
"secret": os.Getenv("S3_SECRET_KEY"),
"bucket": os.Getenv("S3_BUCKET"),
"prefix": "vision-test",
"expiration": "5m",
})
assert.NoError(t, err)
// Upload test image
imgData, err := base64.StdEncoding.DecodeString(testImageBase64)
assert.NoError(t, err)
reader := bytes.NewReader(imgData)
fileID, err := s3Client.Upload(context.Background(), "test.png", reader, "image/png")
assert.NoError(t, err)
// Get URL from S3
url := s3Client.URL(context.Background(), fileID)
assert.NotEmpty(t, url)
// Use S3 URL for analysis
result, err := model.Analyze(context.Background(), url, "Describe this image in detail")
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotEmpty(t, result["description"])
})
t.Run("Analyze with File ID", func(t *testing.T) {
model, err := New(map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
})
assert.NoError(t, err)
// Create test file
data, err := fs.Get("data")
assert.NoError(t, err)
// Write test image data
imgData, err := base64.StdEncoding.DecodeString(testImageBase64)
assert.NoError(t, err)
_, err = data.WriteFile("/__vision_test/test.png", imgData, 0644)
assert.NoError(t, err)
// Analyze using file ID
result, err := model.Analyze(context.Background(), "/__vision_test/test.png", "Describe this image in detail")
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotEmpty(t, result["description"])
})
t.Run("Analyze with Invalid File ID", func(t *testing.T) {
model, err := New(map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
})
assert.NoError(t, err)
_, err = model.Analyze(context.Background(), "/non-existent.png", "Describe this image in detail")
assert.Error(t, err)
assert.Contains(t, err.Error(), "failed to read file")
})
t.Run("Analyze with Invalid API Key", func(t *testing.T) {
model, err := New(map[string]interface{}{
"api_key": "invalid-key",
"model": os.Getenv("VISION_MODEL"),
})
assert.NoError(t, err)
_, err = model.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "Describe this image in detail")
assert.Error(t, err)
assert.Contains(t, err.Error(), "OpenAI API error")
})
t.Run("Analyze with Default Prompt", func(t *testing.T) {
model, err := New(map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
"prompt": "Default test prompt",
})
assert.NoError(t, err)
// Use base64 image data without providing a prompt
result, err := model.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotEmpty(t, result["description"])
})
t.Run("Analyze with Custom Prompt Overriding Default", func(t *testing.T) {
model, err := New(map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
"prompt": "Default test prompt",
})
assert.NoError(t, err)
// Use base64 image data with custom prompt
result, err := model.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "Custom test prompt")
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotEmpty(t, result["description"])
})
t.Run("Analyze with Empty Custom Prompt", func(t *testing.T) {
model, err := New(map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
"prompt": "Default test prompt",
})
assert.NoError(t, err)
// Use base64 image data with empty prompt (should use default)
result, err := model.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "")
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotEmpty(t, result["description"])
})
}

266
agent/vision/driver/s3/storage.go
View file

@ -1,266 +0,0 @@
package s3
import (
"bytes"
"context"
"fmt"
"image"
"image/jpeg"
"image/png"
"io"
"path/filepath"
"strings"
"time"
"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/credentials"
"github.com/aws/aws-sdk-go-v2/service/s3"
)
// DefaultExpiration default expiration time for presigned URLs (5 minutes)
const DefaultExpiration = 5 * time.Minute
// MaxImageSize maximum image size (1920x1080)
const MaxImageSize = 1920
// Storage the S3 storage driver
type Storage struct {
Endpoint string `json:"endpoint" yaml:"endpoint"`
Region string `json:"region" yaml:"region"`
Key string `json:"key" yaml:"key"`
Secret string `json:"secret" yaml:"secret"`
Bucket string `json:"bucket" yaml:"bucket"`
Expiration time.Duration `json:"expiration" yaml:"expiration"`
client *s3.Client
prefix string
compression bool
}
// New create a new S3 storage
func New(options map[string]interface{}) (*Storage, error) {
storage := &Storage{
Region: "auto",
Expiration: DefaultExpiration,
compression: true,
}
if endpoint, ok := options["endpoint"].(string); ok {
storage.Endpoint = endpoint
}
if region, ok := options["region"].(string); ok {
storage.Region = region
}
if key, ok := options["key"].(string); ok {
storage.Key = key
}
if secret, ok := options["secret"].(string); ok {
storage.Secret = secret
}
if bucket, ok := options["bucket"].(string); ok {
storage.Bucket = bucket
}
if prefix, ok := options["prefix"].(string); ok {
storage.prefix = prefix
}
if exp, ok := options["expiration"].(time.Duration); ok {
storage.Expiration = exp
}
if compression, ok := options["compression"].(bool); ok {
storage.compression = compression
}
// Validate required fields
if storage.Key == "" || storage.Secret == "" {
return nil, fmt.Errorf("key and secret are required")
}
if storage.Bucket == "" {
return nil, fmt.Errorf("bucket is required")
}
// Create S3 client
opts := s3.Options{
Region: storage.Region,
Credentials: credentials.NewStaticCredentialsProvider(storage.Key, storage.Secret, ""),
UsePathStyle: true,
}
if storage.Endpoint != "" {
// Remove bucket name from endpoint if present
endpoint := storage.Endpoint
if strings.Contains(endpoint, "/"+storage.Bucket) {
endpoint = strings.TrimSuffix(endpoint, "/"+storage.Bucket)
}
opts.BaseEndpoint = aws.String(endpoint)
}
storage.client = s3.New(opts)
return storage, nil
}
// Upload upload file to S3
func (storage *Storage) Upload(ctx context.Context, filename string, reader io.Reader, contentType string) (string, error) {
if storage.client == nil {
return "", fmt.Errorf("s3 client not initialized")
}
// Generate file ID
fileID := storage.makeID(filename, filepath.Ext(filename))
key := filepath.Join(storage.prefix, fileID)
// Check if compression is enabled and if it's an image
var body io.Reader
if storage.compression && isImage(contentType) {
// Read the entire image into memory
content, err := io.ReadAll(reader)
if err != nil {
return "", fmt.Errorf("failed to read image: %w", err)
}
// Compress image
compressed, err := compressImage(content, contentType)
if err != nil {
return "", fmt.Errorf("failed to compress image: %w", err)
}
body = bytes.NewReader(compressed)
} else {
body = reader
}
// Upload file
_, err := storage.client.PutObject(ctx, &s3.PutObjectInput{
Bucket: aws.String(storage.Bucket),
Key: aws.String(key),
Body: body,
ContentType: aws.String(contentType),
})
if err != nil {
return "", fmt.Errorf("failed to upload file: %w", err)
}
return fileID, nil
}
// Download download file from S3
func (storage *Storage) Download(ctx context.Context, fileID string) (io.ReadCloser, string, error) {
if storage.client == nil {
return nil, "", fmt.Errorf("s3 client not initialized")
}
key := filepath.Join(storage.prefix, fileID)
// Get object
result, err := storage.client.GetObject(ctx, &s3.GetObjectInput{
Bucket: aws.String(storage.Bucket),
Key: aws.String(key),
})
if err != nil {
return nil, "", fmt.Errorf("failed to download file: %w", err)
}
contentType := "application/octet-stream"
if result.ContentType != nil {
contentType = *result.ContentType
}
return result.Body, contentType, nil
}
// URL get file url with expiration
func (storage *Storage) URL(ctx context.Context, fileID string) string {
if storage.client == nil {
return ""
}
key := filepath.Join(storage.prefix, fileID)
presignClient := s3.NewPresignClient(storage.client)
request, err := presignClient.PresignGetObject(ctx, &s3.GetObjectInput{
Bucket: aws.String(storage.Bucket),
Key: aws.String(key),
}, s3.WithPresignExpires(storage.Expiration))
if err != nil {
return ""
}
return request.URL
}
func (storage *Storage) makeID(filename string, ext string) string {
date := time.Now().Format("20060102")
name := strings.TrimSuffix(filepath.Base(filename), ext)
return fmt.Sprintf("%s/%s-%d%s", date, name, time.Now().UnixNano(), ext)
}
// isImage checks if the content type is an image
func isImage(contentType string) bool {
return strings.HasPrefix(contentType, "image/")
}
// compressImage compresses the image while maintaining aspect ratio
func compressImage(data []byte, contentType string) ([]byte, error) {
// Decode image
img, _, err := image.Decode(bytes.NewReader(data))
if err != nil {
return nil, fmt.Errorf("failed to decode image: %w", err)
}
// Calculate new dimensions
bounds := img.Bounds()
width := bounds.Dx()
height := bounds.Dy()
var newWidth, newHeight int
if width > height {
if width > MaxImageSize {
newWidth = MaxImageSize
newHeight = int(float64(height) * (float64(MaxImageSize) / float64(width)))
} else {
return data, nil // No need to resize
}
} else {
if height > MaxImageSize {
newHeight = MaxImageSize
newWidth = int(float64(width) * (float64(MaxImageSize) / float64(height)))
} else {
return data, nil // No need to resize
}
}
// Create new image with new dimensions
newImg := image.NewRGBA(image.Rect(0, 0, newWidth, newHeight))
// Scale the image using bilinear interpolation
for y := 0; y < newHeight; y++ {
for x := 0; x < newWidth; x++ {
srcX := float64(x) * float64(width) / float64(newWidth)
srcY := float64(y) * float64(height) / float64(newHeight)
newImg.Set(x, y, img.At(int(srcX), int(srcY)))
}
}
// Encode image
var buf bytes.Buffer
switch contentType {
case "image/jpeg":
err = jpeg.Encode(&buf, newImg, &jpeg.Options{Quality: 85})
case "image/png":
err = png.Encode(&buf, newImg)
default:
return data, nil // Unsupported format, return original
}
if err != nil {
return nil, fmt.Errorf("failed to encode image: %w", err)
}
return buf.Bytes(), nil
}

204
agent/vision/driver/s3/storage_test.go
View file

@ -1,204 +0,0 @@
package s3
import (
"bytes"
"context"
"image"
"image/png"
"io"
"os"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/test"
)
func TestS3Storage(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
t.Run("Create Storage", func(t *testing.T) {
options := map[string]interface{}{
"endpoint": os.Getenv("S3_API"),
"region": "auto",
"key": os.Getenv("S3_ACCESS_KEY"),
"secret": os.Getenv("S3_SECRET_KEY"),
"bucket": os.Getenv("S3_BUCKET"),
"prefix": "vision-test",
"expiration": 10 * time.Minute,
"compression": true,
}
storage, err := New(options)
if err != nil {
t.Logf("Error creating storage: %v", err)
}
assert.NoError(t, err)
assert.NotNil(t, storage)
if storage != nil {
assert.Equal(t, os.Getenv("S3_API"), storage.Endpoint)
assert.Equal(t, "auto", storage.Region)
assert.Equal(t, os.Getenv("S3_ACCESS_KEY"), storage.Key)
assert.Equal(t, os.Getenv("S3_SECRET_KEY"), storage.Secret)
assert.Equal(t, os.Getenv("S3_BUCKET"), storage.Bucket)
assert.Equal(t, "vision-test", storage.prefix)
assert.Equal(t, 10*time.Minute, storage.Expiration)
assert.True(t, storage.compression)
}
})
t.Run("Upload and Download Image with Compression", func(t *testing.T) {
storage, err := New(map[string]interface{}{
"endpoint": os.Getenv("S3_API"),
"region": "auto",
"key": os.Getenv("S3_ACCESS_KEY"),
"secret": os.Getenv("S3_SECRET_KEY"),
"bucket": os.Getenv("S3_BUCKET"),
"prefix": "vision-test",
"expiration": 5 * time.Minute,
"compression": true,
})
if err != nil {
t.Skip("S3 configuration not available")
}
// Create a test image (2000x2000 pixels)
img := image.NewRGBA(image.Rect(0, 0, 2000, 2000))
var buf bytes.Buffer
err = png.Encode(&buf, img)
assert.NoError(t, err)
// Upload
reader := bytes.NewReader(buf.Bytes())
fileID, err := storage.Upload(context.Background(), "test.png", reader, "image/png")
assert.NoError(t, err)
assert.NotEmpty(t, fileID)
// Download and verify size
reader2, contentType, err := storage.Download(context.Background(), fileID)
assert.NoError(t, err)
assert.Equal(t, "image/png", contentType)
downloaded, err := io.ReadAll(reader2)
assert.NoError(t, err)
// Decode the downloaded image
downloadedImg, _, err := image.Decode(bytes.NewReader(downloaded))
assert.NoError(t, err)
// Verify dimensions
bounds := downloadedImg.Bounds()
assert.LessOrEqual(t, bounds.Dx(), MaxImageSize)
assert.LessOrEqual(t, bounds.Dy(), MaxImageSize)
})
t.Run("Upload Image without Compression", func(t *testing.T) {
storage, err := New(map[string]interface{}{
"endpoint": os.Getenv("S3_API"),
"region": "auto",
"key": os.Getenv("S3_ACCESS_KEY"),
"secret": os.Getenv("S3_SECRET_KEY"),
"bucket": os.Getenv("S3_BUCKET"),
"prefix": "vision-test",
"expiration": 5 * time.Minute,
"compression": false,
})
if err != nil {
t.Skip("S3 configuration not available")
}
// Create a test image (2000x2000 pixels)
img := image.NewRGBA(image.Rect(0, 0, 2000, 2000))
var buf bytes.Buffer
err = png.Encode(&buf, img)
assert.NoError(t, err)
// Upload
reader := bytes.NewReader(buf.Bytes())
fileID, err := storage.Upload(context.Background(), "test.png", reader, "image/png")
assert.NoError(t, err)
assert.NotEmpty(t, fileID)
// Download and verify size
reader2, contentType, err := storage.Download(context.Background(), fileID)
assert.NoError(t, err)
assert.Equal(t, "image/png", contentType)
downloaded, err := io.ReadAll(reader2)
assert.NoError(t, err)
// Decode the downloaded image
downloadedImg, _, err := image.Decode(bytes.NewReader(downloaded))
assert.NoError(t, err)
// Verify dimensions are unchanged
bounds := downloadedImg.Bounds()
assert.Equal(t, 2000, bounds.Dx())
assert.Equal(t, 2000, bounds.Dy())
})
t.Run("Upload and Download Text File", func(t *testing.T) {
storage, err := New(map[string]interface{}{
"endpoint": os.Getenv("S3_API"),
"region": "auto",
"key": os.Getenv("S3_ACCESS_KEY"),
"secret": os.Getenv("S3_SECRET_KEY"),
"bucket": os.Getenv("S3_BUCKET"),
"prefix": "vision-test",
"expiration": 5 * time.Minute,
"compression": true,
})
if err != nil {
t.Skip("S3 configuration not available")
}
content := []byte("test content")
reader := bytes.NewReader(content)
fileID, err := storage.Upload(context.Background(), "test.txt", reader, "text/plain")
assert.NoError(t, err)
assert.NotEmpty(t, fileID)
// Get presigned URL
url := storage.URL(context.Background(), fileID)
assert.NotEmpty(t, url)
assert.Contains(t, url, "X-Amz-Signature")
assert.Contains(t, url, "X-Amz-Expires")
// Download
reader2, contentType, err := storage.Download(context.Background(), fileID)
if err != nil {
t.Logf("Download error: %v", err)
t.FailNow()
}
assert.NoError(t, err)
assert.Contains(t, contentType, "text/plain")
if reader2 != nil {
downloaded, err := io.ReadAll(reader2)
assert.NoError(t, err)
assert.Equal(t, content, downloaded)
reader2.Close()
}
})
t.Run("Download Non-existent File", func(t *testing.T) {
storage, err := New(map[string]interface{}{
"endpoint": os.Getenv("S3_API"),
"region": "auto",
"key": os.Getenv("S3_ACCESS_KEY"),
"secret": os.Getenv("S3_SECRET_KEY"),
"bucket": os.Getenv("S3_BUCKET"),
"prefix": "vision-test",
"expiration": 5 * time.Minute,
"compression": true,
})
if err != nil {
t.Skip("S3 configuration not available")
}
_, _, err = storage.Download(context.Background(), "non-existent.txt")
assert.Error(t, err)
})
}

45
agent/vision/driver/types.go
View file

@ -1,45 +0,0 @@
package driver
import (
"context"
"io"
)
// Config the vision configuration
type Config struct {
Storage StorageConfig `json:"storage" yaml:"storage"`
Model ModelConfig `json:"model" yaml:"model"`
}
// StorageConfig the storage configuration
type StorageConfig struct {
Driver string `json:"driver" yaml:"driver"`
Options map[string]interface{} `json:"options" yaml:"options"`
}
// ModelConfig the model configuration
type ModelConfig struct {
Driver string `json:"driver" yaml:"driver"`
Options map[string]interface{} `json:"options" yaml:"options"`
}
// Storage the storage interface
type Storage interface {
Upload(ctx context.Context, filename string, reader io.Reader, contentType string) (string, error)
Download(ctx context.Context, fileID string) (io.ReadCloser, string, error)
URL(ctx context.Context, fileID string) string
}
// Model the vision model interface
type Model interface {
// Analyze analyzes an image file
// If prompt is empty, it will use the default prompt from model.options.prompt
Analyze(ctx context.Context, fileID string, prompt ...string) (map[string]interface{}, error)
}
// Response the vision response
type Response struct {
FileID string `json:"file_id" yaml:"file_id"`
URL string `json:"url" yaml:"url"`
Description map[string]interface{} `json:"description" yaml:"description"`
}

139
agent/vision/vision.go
View file

@ -1,139 +0,0 @@
package vision
import (
"context"
"fmt"
"io"
"os"
"strings"
"time"
"github.com/yaoapp/yao/agent/vision/driver"
"github.com/yaoapp/yao/agent/vision/driver/local"
"github.com/yaoapp/yao/agent/vision/driver/openai"
"github.com/yaoapp/yao/agent/vision/driver/s3"
)
// 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]interface{} {
converted := make(map[string]interface{})
for k, v := range options {
if str, ok := v.(string); ok {
converted[k] = parseEnvValue(str)
} else {
converted[k] = v
}
}
return converted
}
// Vision the vision service
type Vision struct {
storage driver.Storage
model driver.Model
}
// New create a new vision service
func New(cfg *driver.Config) (*Vision, error) {
// Parse environment variables in options
storageOptions := convertOptions(cfg.Storage.Options)
modelOptions := convertOptions(cfg.Model.Options)
// Create storage driver
var storage driver.Storage
var err error
switch cfg.Storage.Driver {
case "local":
storage, err = local.New(storageOptions)
case "s3":
// Convert expiration string to duration if present
if exp, ok := storageOptions["expiration"].(string); ok {
if duration, err := time.ParseDuration(exp); err == nil {
storageOptions["expiration"] = duration
}
}
storage, err = s3.New(storageOptions)
default:
return nil, fmt.Errorf("storage driver %s not supported", cfg.Storage.Driver)
}
if err != nil {
return nil, fmt.Errorf("create storage driver error: %s", err.Error())
}
// Create model driver
var model driver.Model
switch cfg.Model.Driver {
case "openai":
model, err = openai.New(modelOptions)
default:
return nil, fmt.Errorf("model driver %s not supported", cfg.Model.Driver)
}
if err != nil {
return nil, fmt.Errorf("create model driver error: %s", err.Error())
}
return &Vision{
storage: storage,
model: model,
}, nil
}
// Upload upload file
func (v *Vision) Upload(ctx context.Context, filename string, reader io.Reader, contentType string) (*driver.Response, error) {
fileID, err := v.storage.Upload(ctx, filename, reader, contentType)
if err != nil {
return nil, err
}
return &driver.Response{
FileID: fileID,
URL: v.storage.URL(ctx, fileID),
}, nil
}
// Analyze analyze image using vision model
func (v *Vision) Analyze(ctx context.Context, fileID string, prompt ...string) (*driver.Response, error) {
if v.model == nil {
return nil, fmt.Errorf("model is required")
}
var url string
// If the input is already a base64 data URL or a HTTP(S) URL, use it directly
if strings.HasPrefix(fileID, "data:image/") || strings.HasPrefix(fileID, "http://") || strings.HasPrefix(fileID, "https://") {
url = fileID
} else {
// Otherwise, try to get the URL from storage
url = v.storage.URL(ctx, fileID)
if url == "" {
return nil, fmt.Errorf("failed to get URL for file %s", fileID)
}
}
result, err := v.model.Analyze(ctx, url, prompt...)
if err != nil {
return nil, err
}
return &driver.Response{
FileID: fileID,
URL: url,
Description: result,
}, nil
}
// Download download file
func (v *Vision) Download(ctx context.Context, fileID string) (io.ReadCloser, string, error) {
return v.storage.Download(ctx, fileID)
}

502
agent/vision/vision_test.go
View file

@ -1,502 +0,0 @@
package vision
// import (
// "bytes"
// "context"
// "encoding/base64"
// "fmt"
// "image"
// "image/png"
// "io"
// "net/http"
// "net/http/httptest"
// "os"
// "testing"
// "github.com/stretchr/testify/assert"
// "github.com/yaoapp/gou/fs"
// "github.com/yaoapp/yao/agent/vision/driver"
// "github.com/yaoapp/yao/agent/vision/driver/local"
// "github.com/yaoapp/yao/config"
// "github.com/yaoapp/yao/test"
// )
// var (
// // 1x1 transparent PNG
// testImageBase64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
// )
// // MaxImageSize maximum image size (1920x1080)
// const MaxImageSize = local.MaxImageSize
// func TestVision(t *testing.T) {
// test.Prepare(t, config.Conf)
// defer test.Clean()
// // Setup test server for image hosting
// imgServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// // Log request for debugging
// t.Logf("Received request for: %s", r.URL.Path)
// // Always return the test image
// imgData, _ := base64.StdEncoding.DecodeString(testImageBase64)
// w.Header().Set("Content-Type", "image/png")
// w.Write(imgData)
// }))
// defer imgServer.Close()
// t.Logf("Test server running at: %s", imgServer.URL)
// t.Run("Create Vision Service", func(t *testing.T) {
// vision, err := createTestVision(imgServer.URL)
// assert.NoError(t, err)
// assert.NotNil(t, vision)
// })
// t.Run("Upload and Download with Local Storage", func(t *testing.T) {
// vision, err := createTestVision(imgServer.URL)
// assert.NoError(t, err)
// // Test with text file
// content := []byte("test content")
// reader := bytes.NewReader(content)
// resp, err := vision.Upload(context.Background(), "test.txt", reader, "text/plain")
// assert.NoError(t, err)
// assert.NotEmpty(t, resp.FileID)
// assert.NotEmpty(t, resp.URL)
// // Download
// reader2, contentType, err := vision.Download(context.Background(), resp.FileID)
// assert.NoError(t, err)
// assert.Contains(t, contentType, "text/plain")
// if reader2 != nil {
// downloaded, err := io.ReadAll(reader2)
// assert.NoError(t, err)
// assert.Equal(t, content, downloaded)
// reader2.Close()
// }
// })
// t.Run("Upload and Download with S3 Storage", func(t *testing.T) {
// vision, err := createTestVisionWithS3()
// if err != nil {
// t.Skip("S3 configuration not available")
// }
// // Test with text file
// content := []byte("test content")
// reader := bytes.NewReader(content)
// resp, err := vision.Upload(context.Background(), "test.txt", reader, "text/plain")
// assert.NoError(t, err)
// assert.NotEmpty(t, resp.FileID)
// assert.NotEmpty(t, resp.URL)
// // Download
// reader2, contentType, err := vision.Download(context.Background(), resp.FileID)
// assert.NoError(t, err)
// assert.Contains(t, contentType, "text/plain")
// if reader2 != nil {
// downloaded, err := io.ReadAll(reader2)
// assert.NoError(t, err)
// assert.Equal(t, content, downloaded)
// reader2.Close()
// }
// })
// t.Run("Analyze Image with Base64", func(t *testing.T) {
// // Create vision service
// cfg := &driver.Config{
// Storage: driver.StorageConfig{
// Driver: "local",
// Options: map[string]interface{}{
// "path": "/__vision_test",
// "compression": true,
// },
// },
// Model: driver.ModelConfig{
// Driver: "openai",
// Options: map[string]interface{}{
// "api_key": os.Getenv("OPENAI_API_KEY"),
// "model": os.Getenv("VISION_MODEL"),
// },
// },
// }
// vision, err := New(cfg)
// assert.NoError(t, err)
// // Use base64 data directly
// result, err := vision.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "Describe this image in detail")
// assert.NoError(t, err)
// assert.NotNil(t, result)
// assert.NotEmpty(t, result.Description)
// })
// t.Run("Analyze Image with File", func(t *testing.T) {
// // Create vision service
// cfg := &driver.Config{
// Storage: driver.StorageConfig{
// Driver: "local",
// Options: map[string]interface{}{
// "path": "/__vision_test",
// "compression": true,
// },
// },
// Model: driver.ModelConfig{
// Driver: "openai",
// Options: map[string]interface{}{
// "api_key": os.Getenv("OPENAI_API_KEY"),
// "model": os.Getenv("VISION_MODEL"),
// },
// },
// }
// vision, err := New(cfg)
// assert.NoError(t, err)
// // Create test file
// data, err := fs.Get("data")
// assert.NoError(t, err)
// // Write test image data
// imgData, err := base64.StdEncoding.DecodeString(testImageBase64)
// assert.NoError(t, err)
// _, err = data.WriteFile("/test.png", imgData, 0644)
// assert.NoError(t, err)
// // Analyze using file path
// result, err := vision.Analyze(context.Background(), "/test.png", "Describe this image in detail")
// assert.NoError(t, err)
// assert.NotNil(t, result)
// assert.NotEmpty(t, result.Description)
// })
// t.Run("Analyze Image with S3 URL", func(t *testing.T) {
// if os.Getenv("S3_API") == "" || os.Getenv("S3_ACCESS_KEY") == "" ||
// os.Getenv("S3_SECRET_KEY") == "" || os.Getenv("S3_BUCKET") == "" {
// t.Skip("S3 environment variables not set")
// }
// // Create vision service
// cfg := &driver.Config{
// Storage: driver.StorageConfig{
// Driver: "s3",
// Options: map[string]interface{}{
// "endpoint": os.Getenv("S3_API"),
// "region": "auto",
// "key": os.Getenv("S3_ACCESS_KEY"),
// "secret": os.Getenv("S3_SECRET_KEY"),
// "bucket": os.Getenv("S3_BUCKET"),
// "prefix": "vision-test",
// "expiration": "5m",
// },
// },
// Model: driver.ModelConfig{
// Driver: "openai",
// Options: map[string]interface{}{
// "api_key": os.Getenv("OPENAI_API_KEY"),
// "model": os.Getenv("VISION_MODEL"),
// },
// },
// }
// vision, err := New(cfg)
// assert.NoError(t, err)
// // Upload test image
// imgData, err := base64.StdEncoding.DecodeString(testImageBase64)
// assert.NoError(t, err)
// reader := bytes.NewReader(imgData)
// resp, err := vision.Upload(context.Background(), "test.png", reader, "image/png")
// assert.NoError(t, err)
// assert.NotEmpty(t, resp.FileID)
// assert.NotEmpty(t, resp.URL)
// // Analyze using S3 URL
// result, err := vision.Analyze(context.Background(), resp.URL, "Describe this image in detail")
// assert.NoError(t, err)
// assert.NotNil(t, result)
// assert.NotEmpty(t, result.Description)
// })
// t.Run("Invalid Model", func(t *testing.T) {
// cfg := &driver.Config{
// Storage: driver.StorageConfig{
// Driver: "local",
// Options: map[string]interface{}{
// "path": "/__vision_test",
// "compression": true,
// },
// },
// Model: driver.ModelConfig{
// Driver: "invalid",
// Options: map[string]interface{}{},
// },
// }
// _, err := New(cfg)
// assert.Error(t, err)
// assert.Contains(t, err.Error(), "model driver invalid not supported")
// })
// t.Run("Invalid Storage", func(t *testing.T) {
// cfg := &driver.Config{
// Storage: driver.StorageConfig{
// Driver: "invalid",
// Options: map[string]interface{}{},
// },
// Model: driver.ModelConfig{
// Driver: "openai",
// Options: map[string]interface{}{
// "api_key": "test",
// },
// },
// }
// _, err := New(cfg)
// assert.Error(t, err)
// assert.Contains(t, err.Error(), "storage driver invalid not supported")
// })
// t.Run("Upload and Download Image with Local Storage", func(t *testing.T) {
// vision, err := createTestVision(imgServer.URL)
// assert.NoError(t, err)
// // Create test image (2000x2000 pixels)
// img := image.NewRGBA(image.Rect(0, 0, 2000, 2000))
// var buf bytes.Buffer
// err = png.Encode(&buf, img)
// assert.NoError(t, err)
// // Upload
// reader := bytes.NewReader(buf.Bytes())
// resp, err := vision.Upload(context.Background(), "test.png", reader, "image/png")
// assert.NoError(t, err)
// assert.NotEmpty(t, resp.FileID)
// assert.NotEmpty(t, resp.URL)
// // Download and verify size
// reader2, contentType, err := vision.Download(context.Background(), resp.FileID)
// assert.NoError(t, err)
// assert.Equal(t, "image/png", contentType)
// downloaded, err := io.ReadAll(reader2)
// assert.NoError(t, err)
// // Decode the downloaded image
// downloadedImg, _, err := image.Decode(bytes.NewReader(downloaded))
// assert.NoError(t, err)
// // Verify dimensions
// bounds := downloadedImg.Bounds()
// assert.LessOrEqual(t, bounds.Dx(), MaxImageSize)
// assert.LessOrEqual(t, bounds.Dy(), MaxImageSize)
// })
// t.Run("Upload and Download Image with S3 Storage", func(t *testing.T) {
// vision, err := createTestVisionWithS3()
// if err != nil {
// t.Skip("S3 configuration not available")
// }
// // Create test image (2000x2000 pixels)
// img := image.NewRGBA(image.Rect(0, 0, 2000, 2000))
// var buf bytes.Buffer
// err = png.Encode(&buf, img)
// assert.NoError(t, err)
// // Upload
// reader := bytes.NewReader(buf.Bytes())
// resp, err := vision.Upload(context.Background(), "test.png", reader, "image/png")
// assert.NoError(t, err)
// assert.NotEmpty(t, resp.FileID)
// assert.NotEmpty(t, resp.URL)
// // Download and verify size
// reader2, contentType, err := vision.Download(context.Background(), resp.FileID)
// assert.NoError(t, err)
// assert.Equal(t, "image/png", contentType)
// downloaded, err := io.ReadAll(reader2)
// assert.NoError(t, err)
// // Decode the downloaded image
// downloadedImg, _, err := image.Decode(bytes.NewReader(downloaded))
// assert.NoError(t, err)
// // Verify dimensions
// bounds := downloadedImg.Bounds()
// assert.LessOrEqual(t, bounds.Dx(), MaxImageSize)
// assert.LessOrEqual(t, bounds.Dy(), MaxImageSize)
// })
// t.Run("Analyze Image with Default Prompt", func(t *testing.T) {
// // Create vision service with default prompt
// cfg := &driver.Config{
// Storage: driver.StorageConfig{
// Driver: "local",
// Options: map[string]interface{}{
// "path": "/__vision_test",
// "compression": true,
// },
// },
// Model: driver.ModelConfig{
// Driver: "openai",
// Options: map[string]interface{}{
// "api_key": os.Getenv("OPENAI_API_KEY"),
// "model": os.Getenv("VISION_MODEL"),
// "prompt": "Default test prompt",
// },
// },
// }
// vision, err := New(cfg)
// assert.NoError(t, err)
// // Use base64 data without providing a prompt
// result, err := vision.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64)
// assert.NoError(t, err)
// assert.NotNil(t, result)
// assert.NotEmpty(t, result.Description)
// })
// t.Run("Analyze Image with Custom Prompt", func(t *testing.T) {
// // Create vision service with default prompt
// cfg := &driver.Config{
// Storage: driver.StorageConfig{
// Driver: "local",
// Options: map[string]interface{}{
// "path": "/__vision_test",
// "compression": true,
// },
// },
// Model: driver.ModelConfig{
// Driver: "openai",
// Options: map[string]interface{}{
// "api_key": os.Getenv("OPENAI_API_KEY"),
// "model": os.Getenv("VISION_MODEL"),
// "prompt": "Default test prompt",
// },
// },
// }
// vision, err := New(cfg)
// assert.NoError(t, err)
// // Use base64 data with custom prompt
// result, err := vision.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "Custom test prompt")
// assert.NoError(t, err)
// assert.NotNil(t, result)
// assert.NotEmpty(t, result.Description)
// })
// t.Run("Analyze Image with Empty Custom Prompt", func(t *testing.T) {
// // Create vision service with default prompt
// cfg := &driver.Config{
// Storage: driver.StorageConfig{
// Driver: "local",
// Options: map[string]interface{}{
// "path": "/__vision_test",
// "compression": true,
// },
// },
// Model: driver.ModelConfig{
// Driver: "openai",
// Options: map[string]interface{}{
// "api_key": os.Getenv("OPENAI_API_KEY"),
// "model": os.Getenv("VISION_MODEL"),
// "prompt": "Default test prompt",
// },
// },
// }
// vision, err := New(cfg)
// assert.NoError(t, err)
// // Use base64 data with empty prompt (should use default)
// result, err := vision.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "")
// assert.NoError(t, err)
// assert.NotNil(t, result)
// assert.NotEmpty(t, result.Description)
// })
// }
// func createTestVision(baseURL string) (*Vision, error) {
// cfg := &driver.Config{
// Storage: driver.StorageConfig{
// Driver: "local",
// Options: map[string]interface{}{
// "path": "/__vision_test",
// "compression": true,
// "base_url": baseURL,
// },
// },
// Model: driver.ModelConfig{
// Driver: "openai",
// Options: map[string]interface{}{
// "api_key": os.Getenv("OPENAI_API_KEY"),
// "model": os.Getenv("VISION_MODEL"),
// "prompt": `# Objective
// You are a vision assistant, you can help the user to understand the image and describe it.
// ## Task Execution Steps
// 1. Understand the image/video and describe it.
// 2. Describe the image/video in detail.
// ## Result Format
// {
// "description": "The description of the image/video",
// "content": "The content of the image/video"
// }`,
// },
// },
// }
// return New(cfg)
// }
// func createTestVisionWithS3() (*Vision, error) {
// // Check required S3 environment variables
// if os.Getenv("S3_API") == "" || os.Getenv("S3_ACCESS_KEY") == "" ||
// os.Getenv("S3_SECRET_KEY") == "" || os.Getenv("S3_BUCKET") == "" {
// return nil, fmt.Errorf("S3 environment variables not set")
// }
// cfg := &driver.Config{
// Storage: driver.StorageConfig{
// Driver: "s3",
// Options: map[string]interface{}{
// "endpoint": os.Getenv("S3_API"),
// "region": "auto",
// "key": os.Getenv("S3_ACCESS_KEY"),
// "secret": os.Getenv("S3_SECRET_KEY"),
// "bucket": os.Getenv("S3_BUCKET"),
// "prefix": "vision-test",
// "expiration": "5m",
// },
// },
// Model: driver.ModelConfig{
// Driver: "openai",
// Options: map[string]interface{}{
// "api_key": os.Getenv("OPENAI_API_KEY"),
// "model": os.Getenv("VISION_MODEL"),
// "prompt": `# Objective
// You are a vision assistant, you can help the user to understand the image and describe it.
// ## Task Execution Steps
// 1. Understand the image/video and describe it.
// 2. Describe the image/video in detail.
// ## Result Format
// {
// "description": "The description of the image/video",
// "content": "The content of the image/video"
// }`,
// },
// },
// }
// return New(cfg)
// }

6
service/service.go
View file

@ -6,7 +6,6 @@ import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/yaoapp/gou/api" "github.com/yaoapp/gou/api"
"github.com/yaoapp/gou/server/http" "github.com/yaoapp/gou/server/http"
agent "github.com/yaoapp/yao/agent/api"
"github.com/yaoapp/yao/config" "github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/openapi" "github.com/yaoapp/yao/openapi"
"github.com/yaoapp/yao/share" "github.com/yaoapp/yao/share"
@ -36,11 +35,6 @@ func Start(cfg config.Config) (*http.Server, error) {
Timeout: 5 * time.Second, Timeout: 5 * time.Second,
}) })
// Agent API
if agent.Agent != nil {
agent.Agent.API(router, "/api/__yao/agent")
}
// OpenAPI Server // OpenAPI Server
if openapi.Server != nil { if openapi.Server != nil {
openapi.Server.Attach(router) openapi.Server.Attach(router)

11
widgets/app/app.go
View file

@ -17,7 +17,7 @@ import (
"github.com/yaoapp/gou/session" "github.com/yaoapp/gou/session"
"github.com/yaoapp/kun/exception" "github.com/yaoapp/kun/exception"
"github.com/yaoapp/kun/log" "github.com/yaoapp/kun/log"
agent "github.com/yaoapp/yao/agent/api" "github.com/yaoapp/yao/agent"
"github.com/yaoapp/yao/agent/assistant" "github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/config" "github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/data" "github.com/yaoapp/yao/data"
@ -550,15 +550,16 @@ func processXgen(process *process.Process) interface{} {
// The default assistant // The default assistant
agentConfig := map[string]interface{}{} agentConfig := map[string]interface{}{}
if agent.Agent != nil { agent := agent.GetAgent()
if agent != nil {
// Add Uses Settings // Add Uses Settings
if agent.Agent.DSL != nil && agent.Agent.DSL.Uses != nil { if agent.Uses != nil {
agentConfig["uses"] = agent.Agent.DSL.Uses agentConfig["uses"] = agent.Uses
} }
// Add Default Assistant Settings ( Will be removed later ) // Add Default Assistant Settings ( Will be removed later )
if ast, ok := agent.Agent.Assistant.(*assistant.Assistant); ok { if ast, ok := agent.Assistant.(*assistant.Assistant); ok {
agentConfig["default"] = map[string]interface{}{ agentConfig["default"] = map[string]interface{}{
"assistant_id": ast.ID, "assistant_id": ast.ID,
"assistant_name": ast.Name, "assistant_name": ast.Name,