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

Merge pull request #1366 from trheyi/main

Browse source 
Enhance A2A calls with multimodal support
This commit is contained in:
Max 2025-12-04 12:20:39 +08:00 • committed by GitHub
commit a19f276b36
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
71 changed files with 4139 additions and 497 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

64
agent/assistant/agent.go
View file

@ -8,6 +8,7 @@ import (
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/kun/utils"
"github.com/yaoapp/yao/agent/assistant/handlers"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
@ -17,7 +18,7 @@ import (
// Stream stream the agent
// handler is optional, if not provided, a default handler will be used
func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Message, handler ...message.StreamFunc) (interface{}, error) {
func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Message, options ...*context.Options) (interface{}, error) {
log.Trace("[AGENT] Stream started: assistant=%s, contextID=%s", ast.ID, ctx.ID)
defer log.Trace("[AGENT] Stream ended: assistant=%s, contextID=%s", ast.ID, ctx.ID)
@ -38,16 +39,24 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
// Initialize
// ================================================
// Get or create options
var opts *context.Options
if len(options) > 0 && options[0] != nil {
opts = options[0]
} else {
opts = &context.Options{}
}
// Initialize stack and auto-handle completion/failure/restore
_, _, done := context.EnterStack(ctx, ast.ID, ctx.Referer)
_, _, done := context.EnterStack(ctx, ast.ID, opts)
defer done()
// Determine stream handler
streamHandler := ast.getStreamHandler(ctx, handler...)
streamHandler := ast.getStreamHandler(ctx, opts)
// Get connector and capabilities early (before sending stream_start)
// so that output adapters can use them when converting stream_start event
err = ast.initializeCapabilities(ctx)
err = ast.initializeCapabilities(ctx, opts)
if err != nil {
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
@ -77,7 +86,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
var createResponse *context.HookCreateResponse
if ast.Script != nil {
var err error
createResponse, err = ast.Script.Create(ctx, fullMessages)
createResponse, opts, err = ast.Script.Create(ctx, fullMessages, opts)
if err != nil {
ast.traceAgentFail(agentNode, err)
// Send error stream_end for root stack
@ -106,8 +115,16 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
return nil, err
}
// Build content - convert extended types (file, data) to standard LLM types (text, image_url, input_audio)
completionMessages, err = ast.BuildContent(ctx, completionMessages, completionOptions, opts)
if err != nil {
ast.traceAgentFail(agentNode, err)
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
return nil, err
}
// Execute the LLM streaming call
completionResponse, err = ast.executeLLMStream(ctx, completionMessages, completionOptions, agentNode, streamHandler)
completionResponse, err = ast.executeLLMStream(ctx, completionMessages, completionOptions, agentNode, streamHandler, opts)
if err != nil {
ast.traceAgentFail(agentNode, err)
// Send error stream_end for root stack
@ -186,7 +203,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
// Retry LLM call (streaming to keep user informed)
log.Trace("[AGENT] Retrying LLM for tool call correction (attempt %d/%d)", attempt+1, maxToolRetries-1)
currentResponse, err = ast.executeLLMForToolRetry(ctx, retryMessages, completionOptions, agentNode, streamHandler)
currentResponse, err = ast.executeLLMForToolRetry(ctx, retryMessages, completionOptions, agentNode, streamHandler, opts)
if err != nil {
log.Error("[AGENT] LLM retry failed: %v", err)
ast.traceAgentFail(agentNode, err)
@ -219,11 +236,11 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
if ast.Script != nil {
var err error
nextResponse, err = ast.Script.Next(ctx, &context.NextHookPayload{
nextResponse, opts, err = ast.Script.Next(ctx, &context.NextHookPayload{
Messages: fullMessages,
Completion: completionResponse,
Tools: toolCallResponses,
})
}, opts)
if err != nil {
ast.traceAgentFail(agentNode, err)
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
@ -302,14 +319,14 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
return finalResponse, nil
}
// GetConnector get the connector object, capabilities, and error with priority: createResponse > ctx > ast
// Note: createResponse.Connector is already applied to ctx.Connector by applyContextAdjustments in create.go
// GetConnector get the connector object, capabilities, and error with priority: opts.Connector > ast.Connector
// Note: opts.Connector may be set by Create hook's applyOptionsAdjustments
// Returns: (connector, capabilities, error)
func (ast *Assistant) GetConnector(ctx *context.Context) (connector.Connector, *openai.Capabilities, error) {
// Determine connector ID with priority
func (ast *Assistant) GetConnector(ctx *context.Context, opts ...*context.Options) (connector.Connector, *openai.Capabilities, error) {
// Determine connector ID with priority: opts.Connector > ast.Connector
connectorID := ast.Connector
if ctx.Connector != "" {
connectorID = ctx.Connector
if len(opts) > 0 && opts[0] != nil && opts[0].Connector != "" {
connectorID = opts[0].Connector
}
// If empty, return error
@ -345,10 +362,11 @@ func (ast *Assistant) Info(locale ...string) *message.AssistantInfo {
}
}
// getStreamHandler returns the stream handler from the provided handlers or a default one
func (ast *Assistant) getStreamHandler(ctx *context.Context, handler ...message.StreamFunc) message.StreamFunc {
if len(handler) > 0 && handler[0] != nil {
return handler[0]
// getStreamHandler returns the stream handler from options or a default one
func (ast *Assistant) getStreamHandler(ctx *context.Context, opts ...*context.Options) message.StreamFunc {
// Check if handler is provided in options
if len(opts) > 0 && opts[0] != nil && opts[0].Writer != nil {
return handlers.DefaultStreamHandler(ctx)
}
return handlers.DefaultStreamHandler(ctx)
}
@ -444,16 +462,20 @@ func (ast *Assistant) handleInterrupt(ctx *context.Context, signal *context.Inte
// initializeCapabilities gets connector and capabilities, then sets them in context
// This should be called early (before sending stream_start) so that output adapters
// can use capabilities when converting stream_start event
func (ast *Assistant) initializeCapabilities(ctx *context.Context) error {
func (ast *Assistant) initializeCapabilities(ctx *context.Context, opts *context.Options) error {
if ast.Prompts == nil && ast.MCP == nil {
return nil
}
_, capabilities, err := ast.GetConnector(ctx)
_, capabilities, err := ast.GetConnector(ctx, opts)
if err != nil {
return err
}
fmt.Println("--- initializeCapabilities debug ---")
utils.Dump(capabilities)
fmt.Println("--- end initializeCapabilities debug ---")
// Set capabilities in context for output adapters to use
if capabilities != nil {
ctx.Capabilities = capabilities

1
agent/assistant/agent_interrupt_test.go
View file

@ -22,7 +22,6 @@ func newTestContextWithInterrupt(chatID, assistantID string) *context.Context {
Space: plan.NewMemorySharedSpace(),
ChatID: chatID,
AssistantID: assistantID,
Connector: "",
Locale: "en-us",
Theme: "light",
Client: context.Client{

23
agent/assistant/assistant.go
View file

@ -5,11 +5,34 @@ import (
"path"
"github.com/yaoapp/gou/fs"
"github.com/yaoapp/yao/agent/content"
agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
store "github.com/yaoapp/yao/agent/store/types"
sui "github.com/yaoapp/yao/sui/core"
)
func init() {
// Initialize AgentGetterFunc to allow content package to call agents
content.AgentGetterFunc = func(agentID string) (content.AgentCaller, error) {
ast, err := Get(agentID)
if err != nil {
return nil, err
}
// Return a wrapper that implements AgentCaller interface
return &agentCallerWrapper{ast: ast}, nil
}
}
// agentCallerWrapper wraps Assistant to implement AgentCaller interface
type agentCallerWrapper struct {
ast *Assistant
}
func (w *agentCallerWrapper) Stream(ctx *agentContext.Context, messages []agentContext.Message, options ...*agentContext.Options) (interface{}, error) {
return w.ast.Stream(ctx, messages, options...)
}
// Get get the assistant by id
func Get(id string) (*Assistant, error) {
return LoadStore(id)

31
agent/assistant/build_content.go Normal file
View file

@ -0,0 +1,31 @@
package assistant
import (
"fmt"
"github.com/yaoapp/yao/agent/content"
"github.com/yaoapp/yao/agent/context"
)
// BuildContent processes messages through Vision function to convert extended content types
// (file, data) to standard LLM-compatible types (text, image_url, input_audio)
//
// This should be called after BuildRequest and before executing LLM call
func (ast *Assistant) BuildContent(ctx *context.Context, messages []context.Message, options *context.CompletionOptions, opts *context.Options) ([]context.Message, error) {
// Get connector and capabilities
_, capabilities, err := ast.GetConnector(ctx, opts)
if err != nil {
return nil, fmt.Errorf("failed to get connector: %w", err)
}
// Get Uses configuration from options (already merged in BuildRequest)
uses := options.Uses
// Process content through Vision function
processedMessages, err := content.Vision(ctx, capabilities, messages, uses)
if err != nil {
return nil, fmt.Errorf("failed to process content: %w", err)
}
return processedMessages, nil
}

2
agent/assistant/build_mcp_test.go
View file

@ -215,7 +215,7 @@ func TestBuildRequest_MCP(t *testing.T) {
// Call create hook to get createResponse
var createResponse *context.HookCreateResponse
if hookAgent.Script != nil {
createResponse, err = hookAgent.Script.Create(hookCtx, inputMessages)
createResponse, _, err = hookAgent.Script.Create(hookCtx, inputMessages, &context.Options{})
if err != nil {
t.Fatalf("Failed to call create hook: %s", err.Error())
}

12
agent/assistant/build_prompts_test.go
View file

@ -650,7 +650,7 @@ func TestPromptPresetAssistant(t *testing.T) {
}
// Call Create hook
createResponse, err := ast.Script.Create(ctx, messages)
createResponse, _, err := ast.Script.Create(ctx, messages, &context.Options{})
require.NoError(t, err)
require.NotNil(t, createResponse)
assert.Equal(t, "mode.friendly", createResponse.PromptPreset)
@ -681,7 +681,7 @@ func TestPromptPresetAssistant(t *testing.T) {
}
// Call Create hook
createResponse, err := ast.Script.Create(ctx, messages)
createResponse, _, err := ast.Script.Create(ctx, messages, &context.Options{})
require.NoError(t, err)
require.NotNil(t, createResponse)
assert.Equal(t, "mode.professional", createResponse.PromptPreset)
@ -718,7 +718,7 @@ func TestPromptPresetAssistant(t *testing.T) {
}
// Call Create hook
createResponse, err := ast.Script.Create(ctx, messages)
createResponse, _, err := ast.Script.Create(ctx, messages, &context.Options{})
require.NoError(t, err)
require.NotNil(t, createResponse)
require.NotNil(t, createResponse.DisableGlobalPrompts)
@ -753,7 +753,7 @@ func TestPromptPresetAssistant(t *testing.T) {
}
// Call Create hook
createResponse, err := ast.Script.Create(ctx, messages)
createResponse, _, err := ast.Script.Create(ctx, messages, &context.Options{})
require.NoError(t, err)
require.NotNil(t, createResponse)
assert.Equal(t, "mode.friendly", createResponse.PromptPreset)
@ -788,7 +788,7 @@ func TestPromptPresetAssistant(t *testing.T) {
}
// Call Create hook
createResponse, err := ast.Script.Create(ctx, messages)
createResponse, _, err := ast.Script.Create(ctx, messages, &context.Options{})
require.NoError(t, err)
require.NotNil(t, createResponse)
assert.Equal(t, "non.existent.preset", createResponse.PromptPreset)
@ -819,7 +819,7 @@ func TestPromptPresetAssistant(t *testing.T) {
}
// Call Create hook - should return nil
createResponse, err := ast.Script.Create(ctx, messages)
createResponse, _, err := ast.Script.Create(ctx, messages, &context.Options{})
require.NoError(t, err)
assert.Nil(t, createResponse)

9
agent/assistant/build_test.go
View file

@ -18,7 +18,6 @@ func newTestContext(chatID, assistantID string) *context.Context {
Space: plan.NewMemorySharedSpace(),
ChatID: chatID,
AssistantID: assistantID,
Connector: "",
Locale: "en-us",
Theme: "light",
Client: context.Client{
@ -64,7 +63,7 @@ func TestBuildRequest(t *testing.T) {
inputMessages := []context.Message{{Role: "user", Content: "no_override"}}
// Call Create hook
createResponse, err := agent.Script.Create(ctx, inputMessages)
createResponse, _, err := agent.Script.Create(ctx, inputMessages, &context.Options{})
if err != nil {
t.Fatalf("Failed to call Create hook: %s", err.Error())
}
@ -112,7 +111,7 @@ func TestBuildRequest(t *testing.T) {
t.Run("OverrideTemperature", func(t *testing.T) {
inputMessages := []context.Message{{Role: "user", Content: "override_temperature"}}
createResponse, err := agent.Script.Create(ctx, inputMessages)
createResponse, _, err := agent.Script.Create(ctx, inputMessages, &context.Options{})
if err != nil {
t.Fatalf("Failed to call Create hook: %s", err.Error())
}
@ -143,7 +142,7 @@ func TestBuildRequest(t *testing.T) {
t.Run("OverrideAll", func(t *testing.T) {
inputMessages := []context.Message{{Role: "user", Content: "override_all"}}
createResponse, err := agent.Script.Create(ctx, inputMessages)
createResponse, _, err := agent.Script.Create(ctx, inputMessages, &context.Options{})
if err != nil {
t.Fatalf("Failed to call Create hook: %s", err.Error())
}
@ -196,7 +195,7 @@ func TestBuildRequest(t *testing.T) {
t.Run("OverrideRouteMetadata", func(t *testing.T) {
inputMessages := []context.Message{{Role: "user", Content: "override_route_metadata"}}
createResponse, err := agent.Script.Create(ctx, inputMessages)
createResponse, _, err := agent.Script.Create(ctx, inputMessages, &context.Options{})
if err != nil {
t.Fatalf("Failed to call Create hook: %s", err.Error())
}

12
agent/assistant/handlers/stream.go
View file

@ -107,6 +107,11 @@ func (s *streamState) handleMessageStart(data []byte) int {
startData.MessageID = messageID
}
// Auto-set ThreadID from Stack for nested agent calls
if startData.ThreadID == "" && s.ctx.Stack != nil && !s.ctx.Stack.IsRoot() {
startData.ThreadID = s.ctx.Stack.ID
}
// Initialize message state with the correct message ID
s.inGroup = true
s.currentGroupID = messageID
@ -312,11 +317,18 @@ func (s *streamState) handleMessageEnd(data []byte) int {
msgType = message.TypeText // Fallback to text if type not set
}
// Get ThreadID from Stack for nested agent calls
var threadID string
if s.ctx.Stack != nil && !s.ctx.Stack.IsRoot() {
threadID = s.ctx.Stack.ID
}
// Build EventMessageEndData with complete content
endData := message.EventMessageEndData{
MessageID: s.currentGroupID, // Use the message ID
Type: msgType,
Timestamp: time.Now().UnixMilli(),
ThreadID: threadID, // Include ThreadID for concurrent stream identification
DurationMs: durationMs,
ChunkCount: s.chunkCount,
Status: "completed",

52
agent/assistant/hook/create.go
View file

@ -9,53 +9,57 @@ import (
)
// Create create a new assistant
func (s *Script) Create(ctx *context.Context, messages []context.Message) (*context.HookCreateResponse, error) {
res, err := s.Execute(ctx, "Create", messages)
// opts is optional - if provided, will be adjusted based on hook response
func (s *Script) Create(ctx *context.Context, messages []context.Message, opts ...*context.Options) (*context.HookCreateResponse, *context.Options, error) {
// Get or create options
var options *context.Options
if len(opts) > 0 && opts[0] != nil {
options = opts[0]
} else {
options = &context.Options{}
}
// Execute hook with ctx, messages, and options (convert options to map for JS)
optionsMap := options.ToMap()
res, err := s.Execute(ctx, "Create", messages, optionsMap)
if err != nil {
return nil, err
return nil, nil, err
}
response, err := s.getHookCreateResponse(res)
if err != nil {
return nil, err
return nil, nil, err
}
// Apply context adjustments from the response back to the context
// Apply adjustments from the response
if response != nil {
s.applyContextAdjustments(ctx, response)
s.applyOptionsAdjustments(options, response)
}
return response, nil
return response, options, nil
}
// applyContextAdjustments applies context field overrides from the hook response back to the context
// applyContextAdjustments applies session-level field overrides from the hook response back to the context
func (s *Script) applyContextAdjustments(ctx *context.Context, response *context.HookCreateResponse) {
// Override assistant ID if provided
if response.AssistantID != "" {
ctx.AssistantID = response.AssistantID
}
// Note: AssistantID cannot be overridden - it's set at initialization and immutable
// Override connector if provided
if response.Connector != "" {
ctx.Connector = response.Connector
}
// Override locale if provided
// Override locale if provided (session-level)
if response.Locale != "" {
ctx.Locale = response.Locale
}
// Override theme if provided
// Override theme if provided (session-level)
if response.Theme != "" {
ctx.Theme = response.Theme
}
// Override route if provided
// Override route if provided (session-level)
if response.Route != "" {
ctx.Route = response.Route
}
// Merge or override metadata if provided
// Merge or override metadata if provided (session-level)
if len(response.Metadata) > 0 {
if ctx.Metadata == nil {
ctx.Metadata = make(map[string]interface{})
@ -67,6 +71,14 @@ func (s *Script) applyContextAdjustments(ctx *context.Context, response *context
}
}
// applyOptionsAdjustments applies call-level field overrides from the hook response to options
func (s *Script) applyOptionsAdjustments(opts *context.Options, response *context.HookCreateResponse) {
// Override connector if provided (call-level parameter)
if response.Connector != "" {
opts.Connector = response.Connector
}
}
// getHookCreateResponse convert the result to a HookCreateResponse
func (s *Script) getHookCreateResponse(res interface{}) (*context.HookCreateResponse, error) {
// Handle nil result

17
agent/assistant/hook/create_bench_test.go
View file

@ -34,7 +34,7 @@ func BenchmarkSimpleStandardMode(b *testing.B) {
b.ResetTimer()
for i := 0; i < b.N; i++ {
ctx := newBenchContext("bench-simple-standard", "tests.create")
_, err := agent.Script.Create(ctx, []context.Message{
_, _, err := agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "Hello"},
})
if err != nil {
@ -61,7 +61,7 @@ func BenchmarkSimplePerformanceMode(b *testing.B) {
b.ResetTimer()
for i := 0; i < b.N; i++ {
ctx := newBenchContext("bench-simple-performance", "tests.create")
_, err := agent.Script.Create(ctx, []context.Message{
_, _, err := agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "Hello"},
})
if err != nil {
@ -95,7 +95,7 @@ func BenchmarkBusinessStandardMode(b *testing.B) {
for i := 0; i < b.N; i++ {
scenario := scenarios[i%len(scenarios)]
ctx := newBenchContext("bench-business-standard", "tests.create")
_, err := agent.Script.Create(ctx, []context.Message{
_, _, err := agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: scenario.content},
})
if err != nil {
@ -125,7 +125,7 @@ func BenchmarkBusinessPerformanceMode(b *testing.B) {
for i := 0; i < b.N; i++ {
scenario := scenarios[i%len(scenarios)]
ctx := newBenchContext("bench-business-performance", "tests.create")
_, err := agent.Script.Create(ctx, []context.Message{
_, _, err := agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: scenario.content},
})
if err != nil {
@ -159,7 +159,7 @@ func BenchmarkConcurrentSimpleStandardMode(b *testing.B) {
i := 0
for pb.Next() {
ctx := newBenchContext("bench-concurrent-simple-standard", "tests.create")
_, err := agent.Script.Create(ctx, []context.Message{
_, _, err := agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "Hello"},
})
if err != nil {
@ -191,7 +191,7 @@ func BenchmarkConcurrentSimplePerformanceMode(b *testing.B) {
i := 0
for pb.Next() {
ctx := newBenchContext("bench-concurrent-simple", "tests.create")
_, err := agent.Script.Create(ctx, []context.Message{
_, _, err := agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "Hello"},
})
if err != nil {
@ -226,7 +226,7 @@ func BenchmarkConcurrentBusinessStandardMode(b *testing.B) {
for pb.Next() {
scenario := scenarios[i%len(scenarios)]
ctx := newBenchContext("bench-concurrent-business-standard", "tests.create")
_, err := agent.Script.Create(ctx, []context.Message{
_, _, err := agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: scenario.content},
})
if err != nil {
@ -261,7 +261,7 @@ func BenchmarkConcurrentBusinessPerformanceMode(b *testing.B) {
for pb.Next() {
scenario := scenarios[i%len(scenarios)]
ctx := newBenchContext("bench-concurrent-business", "tests.create")
_, err := agent.Script.Create(ctx, []context.Message{
_, _, err := agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: scenario.content},
})
if err != nil {
@ -301,7 +301,6 @@ func newBenchContext(chatID, assistantID string) *context.Context {
Space: plan.NewMemorySharedSpace(),
ChatID: chatID,
AssistantID: assistantID,
Connector: "",
Locale: "en-us",
Theme: "light",
Client: context.Client{

27
agent/assistant/hook/create_mem_test.go
View file

@ -36,7 +36,7 @@ func TestMemoryLeakStandardMode(t *testing.T) {
// Warm up - execute a few times to stabilize memory
for i := 0; i < 10; i++ {
ctx := newMemTestContext("warmup", "tests.create")
_, _ = agent.Script.Create(ctx, []context.Message{
_, _, _ = agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "Hello"},
})
ctx.Release()
@ -52,7 +52,7 @@ func TestMemoryLeakStandardMode(t *testing.T) {
iterations := 1000
for i := 0; i < iterations; i++ {
ctx := newMemTestContext("mem-test-standard", "tests.create")
_, err := agent.Script.Create(ctx, []context.Message{
_, _, err := agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "Hello"},
})
if err != nil {
@ -124,7 +124,7 @@ func TestMemoryLeakPerformanceMode(t *testing.T) {
// Warm up - execute a few times to stabilize memory and fill isolate pool
for i := 0; i < 20; i++ {
ctx := newMemTestContext("warmup", "tests.create")
_, _ = agent.Script.Create(ctx, []context.Message{
_, _, _ = agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "Hello"},
})
ctx.Release()
@ -140,7 +140,7 @@ func TestMemoryLeakPerformanceMode(t *testing.T) {
iterations := 1000
for i := 0; i < iterations; i++ {
ctx := newMemTestContext("mem-test-performance", "tests.create")
_, err := agent.Script.Create(ctx, []context.Message{
_, _, err := agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "Hello"},
})
if err != nil {
@ -221,7 +221,7 @@ func TestMemoryLeakBusinessScenarios(t *testing.T) {
// Warm up
for i := 0; i < 10; i++ {
ctx := newMemTestContext("warmup", "tests.create")
_, _ = agent.Script.Create(ctx, []context.Message{
_, _, _ = agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "return_full"},
})
ctx.Release()
@ -240,7 +240,7 @@ func TestMemoryLeakBusinessScenarios(t *testing.T) {
iterations := 200
for i := 0; i < iterations; i++ {
ctx := newMemTestContext("mem-test-business", "tests.create")
_, err := agent.Script.Create(ctx, []context.Message{
_, _, err := agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: scenario.content},
})
if err != nil {
@ -298,7 +298,7 @@ func TestMemoryLeakConcurrent(t *testing.T) {
// Warm up
for i := 0; i < 20; i++ {
ctx := newMemTestContext("warmup", "tests.create")
_, _ = agent.Script.Create(ctx, []context.Message{
_, _, _ = agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "Hello"},
})
ctx.Release()
@ -321,7 +321,7 @@ func TestMemoryLeakConcurrent(t *testing.T) {
defer func() { done <- true }()
for i := 0; i < iterPerGoroutine; i++ {
ctx := newMemTestContext("mem-test-concurrent", "tests.create")
_, err := agent.Script.Create(ctx, []context.Message{
_, _, err := agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "Hello"},
})
if err != nil {
@ -383,7 +383,7 @@ func TestMemoryLeakNestedCalls(t *testing.T) {
// Warm up
for i := 0; i < 10; i++ {
ctx := newMemTestContext("warmup", "tests.create")
_, _ = agent.Script.Create(ctx, []context.Message{
_, _, _ = agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "nested_script_call"},
})
ctx.Release()
@ -400,7 +400,7 @@ func TestMemoryLeakNestedCalls(t *testing.T) {
iterations := 200
for i := 0; i < iterations; i++ {
ctx := newMemTestContext("mem-test-nested", "tests.create")
_, err := agent.Script.Create(ctx, []context.Message{
_, _, err := agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "deep_nested_call"},
})
if err != nil {
@ -459,7 +459,7 @@ func TestMemoryLeakNestedConcurrent(t *testing.T) {
// Warm up
for i := 0; i < 20; i++ {
ctx := newMemTestContext("warmup", "tests.create")
_, _ = agent.Script.Create(ctx, []context.Message{
_, _, _ = agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "nested_script_call"},
})
ctx.Release()
@ -482,7 +482,7 @@ func TestMemoryLeakNestedConcurrent(t *testing.T) {
defer func() { done <- true }()
for i := 0; i < iterPerGoroutine; i++ {
ctx := newMemTestContext("mem-test-nested-concurrent", "tests.create")
_, err := agent.Script.Create(ctx, []context.Message{
_, _, err := agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "deep_nested_call"},
})
if err != nil {
@ -549,7 +549,7 @@ func TestIsolateDisposal(t *testing.T) {
iterations := 100
for i := 0; i < iterations; i++ {
ctx := newMemTestContext("disposal-test", "tests.create")
_, err := agent.Script.Create(ctx, []context.Message{
_, _, err := agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "Hello"},
})
if err != nil {
@ -615,7 +615,6 @@ func newMemTestContext(chatID, assistantID string) *context.Context {
Space: plan.NewMemorySharedSpace(),
ChatID: chatID,
AssistantID: assistantID,
Connector: "",
Locale: "en-us",
Theme: "light",
Client: context.Client{

4
agent/assistant/hook/create_nested_test.go
View file

@ -29,7 +29,7 @@ func TestNestedScriptCall(t *testing.T) {
// Call with deep_nested_call scenario
// This will: hook -> scripts.tests.create.NestedCall -> GetRoles -> model
res, err := agent.Script.Create(ctx, []context.Message{
res, _, err := agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "deep_nested_call"},
})
@ -86,7 +86,7 @@ func TestNestedScriptCallConcurrent(t *testing.T) {
for j := 0; j < iterations; j++ {
ctx := newTestContext("test-concurrent", "tests.create")
_, err := agent.Script.Create(ctx, []context.Message{
_, _, err := agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "deep_nested_call"},
})

31
agent/assistant/hook/create_test.go
View file

@ -19,7 +19,6 @@ func newTestContext(chatID, assistantID string) *context.Context {
Space: plan.NewMemorySharedSpace(),
ChatID: chatID,
AssistantID: assistantID,
Connector: "",
Locale: "en-us",
Theme: "light",
Client: context.Client{
@ -74,7 +73,7 @@ func TestCreate(t *testing.T) {
// Test scenario 1: Return null (should get nil response)
t.Run("ReturnNull", func(t *testing.T) {
res, err := agent.Script.Create(ctx, []context.Message{{Role: "user", Content: "return_null"}})
res, _, err := agent.Script.Create(ctx, []context.Message{{Role: "user", Content: "return_null"}})
if err != nil {
t.Fatalf("Failed to create with null return: %s", err.Error())
}
@ -85,7 +84,7 @@ func TestCreate(t *testing.T) {
// Test scenario 2: Return undefined (should get nil response)
t.Run("ReturnUndefined", func(t *testing.T) {
res, err := agent.Script.Create(ctx, []context.Message{{Role: "user", Content: "return_undefined"}})
res, _, err := agent.Script.Create(ctx, []context.Message{{Role: "user", Content: "return_undefined"}})
if err != nil {
t.Fatalf("Failed to create with undefined return: %s", err.Error())
}
@ -96,7 +95,7 @@ func TestCreate(t *testing.T) {
// Test scenario 3: Return empty object (should get empty HookCreateResponse)
t.Run("ReturnEmpty", func(t *testing.T) {
res, err := agent.Script.Create(ctx, []context.Message{{Role: "user", Content: "return_empty"}})
res, _, err := agent.Script.Create(ctx, []context.Message{{Role: "user", Content: "return_empty"}})
if err != nil {
t.Fatalf("Failed to create with empty return: %s", err.Error())
}
@ -110,7 +109,7 @@ func TestCreate(t *testing.T) {
// Test scenario 4: Return full response with all fields
t.Run("ReturnFull", func(t *testing.T) {
res, err := agent.Script.Create(ctx, []context.Message{{Role: "user", Content: "return_full"}})
res, _, err := agent.Script.Create(ctx, []context.Message{{Role: "user", Content: "return_full"}})
if err != nil {
t.Fatalf("Failed to create with full return: %s", err.Error())
}
@ -166,7 +165,7 @@ func TestCreate(t *testing.T) {
// Test scenario 5: Return partial response
t.Run("ReturnPartial", func(t *testing.T) {
res, err := agent.Script.Create(ctx, []context.Message{{Role: "user", Content: "return_partial"}})
res, _, err := agent.Script.Create(ctx, []context.Message{{Role: "user", Content: "return_partial"}})
if err != nil {
t.Fatalf("Failed to create with partial return: %s", err.Error())
}
@ -197,7 +196,7 @@ func TestCreate(t *testing.T) {
// Test scenario 6: Process call - calls models.__yao.role.Get and adds to messages
t.Run("ReturnProcess", func(t *testing.T) {
res, err := agent.Script.Create(ctx, []context.Message{{Role: "user", Content: "return_process"}})
res, _, err := agent.Script.Create(ctx, []context.Message{{Role: "user", Content: "return_process"}})
if err != nil {
t.Fatalf("Failed to create with process return: %s", err.Error())
}
@ -225,7 +224,7 @@ func TestCreate(t *testing.T) {
// Test scenario 7: Default response
t.Run("ReturnDefault", func(t *testing.T) {
testContent := "Hello, how are you?"
res, err := agent.Script.Create(ctx, []context.Message{{Role: "user", Content: testContent}})
res, _, err := agent.Script.Create(ctx, []context.Message{{Role: "user", Content: testContent}})
if err != nil {
t.Fatalf("Failed to create with default return: %s", err.Error())
}
@ -252,7 +251,7 @@ func TestCreate(t *testing.T) {
// Test scenario 8: Verify context fields - validates all context fields in JavaScript
t.Run("VerifyContext", func(t *testing.T) {
res, err := agent.Script.Create(ctx, []context.Message{{Role: "user", Content: "verify_context"}})
res, _, err := agent.Script.Create(ctx, []context.Message{{Role: "user", Content: "verify_context"}})
if err != nil {
t.Fatalf("Failed to create with verify_context: %s", err.Error())
}
@ -304,7 +303,7 @@ func TestCreate(t *testing.T) {
adjustCtx := newTestContext("chat-test-adjust", "tests.create")
// Call the hook which should adjust context fields
res, err := agent.Script.Create(adjustCtx, []context.Message{{Role: "user", Content: "adjust_context"}})
res, _, err := agent.Script.Create(adjustCtx, []context.Message{{Role: "user", Content: "adjust_context"}})
if err != nil {
t.Fatalf("Failed to create with adjust_context: %s", err.Error())
}
@ -313,9 +312,7 @@ func TestCreate(t *testing.T) {
}
// Verify the response contains adjusted fields
if res.AssistantID != "adjusted.assistant" {
t.Errorf("Expected adjusted assistant_id 'adjusted.assistant', got: %s", res.AssistantID)
}
// Note: AssistantID cannot be overridden by hooks, removed from HookCreateResponse
if res.Connector != "adjusted-connector" {
t.Errorf("Expected adjusted connector 'adjusted-connector', got: %s", res.Connector)
}
@ -338,12 +335,8 @@ func TestCreate(t *testing.T) {
}
// Verify context fields were actually updated
if adjustCtx.AssistantID != "adjusted.assistant" {
t.Errorf("Context assistant_id not updated. Expected 'adjusted.assistant', got: %s", adjustCtx.AssistantID)
}
if adjustCtx.Connector != "adjusted-connector" {
t.Errorf("Context connector not updated. Expected 'adjusted-connector', got: %s", adjustCtx.Connector)
}
// Note: AssistantID is immutable and cannot be overridden
// Note: Connector is now in Options, not in Context
if adjustCtx.Locale != "zh-cn" {
t.Errorf("Context locale not updated. Expected 'zh-cn', got: %s", adjustCtx.Locale)
}

9
agent/assistant/hook/goroutine_leak_test.go
View file

@ -48,7 +48,7 @@ func TestGoroutineLeakDetailed(t *testing.T) {
for i := 0; i < iterations; i++ {
ctx := newLeakTestContext(fmt.Sprintf("leak-test-%d", i), "tests.create")
_, err := agent.Script.Create(ctx, []context.Message{
_, _, err := agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "Hello"},
})
if err != nil {
@ -126,7 +126,7 @@ func TestGoroutineLeakByComponent(t *testing.T) {
for i := 0; i < 10; i++ {
ctx := newLeakTestContext(fmt.Sprintf("test-%d", i), "tests.create")
_, _ = agent.Script.Create(ctx, []context.Message{
_, _, _ = agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "Hello"},
})
ctx.Release()
@ -184,7 +184,7 @@ func TestGoroutineLeakWithoutRelease(t *testing.T) {
for i := 0; i < 10; i++ {
ctx := newLeakTestContext(fmt.Sprintf("no-release-%d", i), "tests.create")
_, _ = agent.Script.Create(ctx, []context.Message{
_, _, _ = agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "Hello"},
})
// Intentionally NOT calling ctx.Release()
@ -205,7 +205,7 @@ func TestGoroutineLeakWithoutRelease(t *testing.T) {
for i := 0; i < 10; i++ {
ctx := newLeakTestContext(fmt.Sprintf("with-release-%d", i), "tests.create")
_, _ = agent.Script.Create(ctx, []context.Message{
_, _, _ = agent.Script.Create(ctx, []context.Message{
{Role: "user", Content: "Hello"},
})
ctx.Release() // WITH Release
@ -299,7 +299,6 @@ func newLeakTestContext(chatID, assistantID string) *context.Context {
Space: plan.NewMemorySharedSpace(),
ChatID: chatID,
AssistantID: assistantID,
Connector: "",
Locale: "en-us",
Theme: "light",
Client: context.Client{

24
agent/assistant/hook/next.go
View file

@ -9,7 +9,16 @@ import (
)
// Next next hook for the next action after the completion
func (s *Script) Next(ctx *context.Context, payload *context.NextHookPayload) (*context.NextHookResponse, error) {
// opts is optional - if provided, will be passed to the hook
func (s *Script) Next(ctx *context.Context, payload *context.NextHookPayload, opts ...*context.Options) (*context.NextHookResponse, *context.Options, error) {
// Get or create options
var options *context.Options
if len(opts) > 0 && opts[0] != nil {
options = opts[0]
} else {
options = &context.Options{}
}
// Convert payload to map for JS (use JSON tag names)
payloadMap := map[string]interface{}{
"messages": payload.Messages,
@ -18,12 +27,19 @@ func (s *Script) Next(ctx *context.Context, payload *context.NextHookPayload) (*
"error": payload.Error,
}
res, err := s.Execute(ctx, "Next", payloadMap)
// Execute hook with ctx, payload, and options (convert options to map for JS)
optionsMap := options.ToMap()
res, err := s.Execute(ctx, "Next", payloadMap, optionsMap)
if err != nil {
return nil, err
return nil, nil, err
}
return s.getNextHookResponse(res)
response, err := s.getNextHookResponse(res)
if err != nil {
return nil, nil, err
}
return response, options, nil
}
// getNextHookResponse convert the result to a NextHookResponse

19
agent/assistant/hook/next_test.go
View file

@ -20,7 +20,6 @@ func newTestContextForNext(chatID, assistantID string) *context.Context {
Space: plan.NewMemorySharedSpace(),
ChatID: chatID,
AssistantID: assistantID,
Connector: "",
Locale: "en-us",
Theme: "light",
Client: context.Client{
@ -86,7 +85,7 @@ func TestNext(t *testing.T) {
Error: "",
}
res, err := agent.Script.Next(ctx, payload)
res, _, err := agent.Script.Next(ctx, payload)
if err != nil {
t.Fatalf("Failed to execute Next hook with null return: %s", err.Error())
}
@ -106,7 +105,7 @@ func TestNext(t *testing.T) {
},
}
res, err := agent.Script.Next(ctx, payload)
res, _, err := agent.Script.Next(ctx, payload)
if err != nil {
t.Fatalf("Failed to execute Next hook with undefined return: %s", err.Error())
}
@ -126,7 +125,7 @@ func TestNext(t *testing.T) {
},
}
res, err := agent.Script.Next(ctx, payload)
res, _, err := agent.Script.Next(ctx, payload)
if err != nil {
t.Fatalf("Failed to execute Next hook with empty return: %s", err.Error())
}
@ -152,7 +151,7 @@ func TestNext(t *testing.T) {
},
}
res, err := agent.Script.Next(ctx, payload)
res, _, err := agent.Script.Next(ctx, payload)
if err != nil {
t.Fatalf("Failed to execute Next hook with custom data: %s", err.Error())
}
@ -199,7 +198,7 @@ func TestNext(t *testing.T) {
},
}
res, err := agent.Script.Next(ctx, payload)
res, _, err := agent.Script.Next(ctx, payload)
if err != nil {
t.Fatalf("Failed to execute Next hook: %s", err.Error())
}
@ -245,7 +244,7 @@ func TestNext(t *testing.T) {
},
}
res, err := agent.Script.Next(ctx, payload)
res, _, err := agent.Script.Next(ctx, payload)
if err != nil {
t.Fatalf("Failed to execute Next hook with delegate: %s", err.Error())
}
@ -307,7 +306,7 @@ func TestNext(t *testing.T) {
Error: "",
}
res, err := agent.Script.Next(ctx, payload)
res, _, err := agent.Script.Next(ctx, payload)
if err != nil {
t.Fatalf("Failed to execute Next hook: %s", err.Error())
}
@ -366,7 +365,7 @@ func TestNext(t *testing.T) {
},
}
res, err := agent.Script.Next(ctx, payload)
res, _, err := agent.Script.Next(ctx, payload)
if err != nil {
t.Fatalf("Failed to execute Next hook: %s", err.Error())
}
@ -410,7 +409,7 @@ func TestNext(t *testing.T) {
Error: "Tool execution failed: timeout",
}
res, err := agent.Script.Next(ctx, payload)
res, _, err := agent.Script.Next(ctx, payload)
if err != nil {
t.Fatalf("Failed to execute Next hook: %s", err.Error())
}

17
agent/assistant/hook/realworld_next_test.go
View file

@ -19,7 +19,6 @@ func newRealWorldNextContext(chatID, assistantID string) *context.Context {
Space: plan.NewMemorySharedSpace(),
ChatID: chatID,
AssistantID: assistantID,
Connector: "",
Locale: "en-us",
Theme: "light",
Client: context.Client{
@ -76,7 +75,7 @@ func TestRealWorldNextStandard(t *testing.T) {
Error: "",
}
response, err := agent.Script.Next(ctx, payload)
response, _, err := agent.Script.Next(ctx, payload)
if err != nil {
t.Fatalf("Next hook failed: %v", err)
}
@ -118,7 +117,7 @@ func TestRealWorldNextCustomData(t *testing.T) {
Error: "",
}
response, err := agent.Script.Next(ctx, payload)
response, _, err := agent.Script.Next(ctx, payload)
if err != nil {
t.Fatalf("Next hook failed: %v", err)
}
@ -165,7 +164,7 @@ func TestRealWorldNextDelegate(t *testing.T) {
Error: "",
}
response, err := agent.Script.Next(ctx, payload)
response, _, err := agent.Script.Next(ctx, payload)
if err != nil {
t.Fatalf("Next hook failed: %v", err)
}
@ -227,7 +226,7 @@ func TestRealWorldNextProcessTools(t *testing.T) {
Error: "",
}
response, err := agent.Script.Next(ctx, payload)
response, _, err := agent.Script.Next(ctx, payload)
if err != nil {
t.Fatalf("Next hook failed: %v", err)
}
@ -280,7 +279,7 @@ func TestRealWorldNextErrorRecovery(t *testing.T) {
Error: "System error: Database connection timeout",
}
response, err := agent.Script.Next(ctx, payload)
response, _, err := agent.Script.Next(ctx, payload)
if err != nil {
t.Fatalf("Next hook failed: %v", err)
}
@ -329,7 +328,7 @@ func TestRealWorldNextConditional(t *testing.T) {
Error: "",
}
response, err := agent.Script.Next(ctx, payload)
response, _, err := agent.Script.Next(ctx, payload)
if err != nil {
t.Fatalf("Next hook failed: %v", err)
}
@ -362,7 +361,7 @@ func TestRealWorldNextConditional(t *testing.T) {
Error: "",
}
response, err := agent.Script.Next(ctx, payload)
response, _, err := agent.Script.Next(ctx, payload)
if err != nil {
t.Fatalf("Next hook failed: %v", err)
}
@ -406,7 +405,7 @@ func TestRealWorldNextDefault(t *testing.T) {
Error: "",
}
response, err := agent.Script.Next(ctx, payload)
response, _, err := agent.Script.Next(ctx, payload)
if err != nil {
t.Fatalf("Next hook failed: %v", err)
}

33
agent/assistant/hook/realworld_stress_test.go
View file

@ -42,7 +42,7 @@ func TestRealWorldSimpleScenario(t *testing.T) {
{Role: "user", Content: "simple"},
}
response, err := agent.Script.Create(ctx, messages)
response, _, err := agent.Script.Create(ctx, messages)
if err != nil {
t.Fatalf("Create failed: %v", err)
}
@ -73,7 +73,7 @@ func TestRealWorldMCPScenarios(t *testing.T) {
{Role: "user", Content: "mcp_health"},
}
response, err := agent.Script.Create(ctx, messages)
response, _, err := agent.Script.Create(ctx, messages)
if err != nil {
t.Fatalf("Create failed: %v", err)
}
@ -117,7 +117,7 @@ func TestRealWorldMCPScenarios(t *testing.T) {
{Role: "user", Content: "mcp_tools"},
}
response, err := agent.Script.Create(ctx, messages)
response, _, err := agent.Script.Create(ctx, messages)
if err != nil {
t.Fatalf("Create failed: %v", err)
}
@ -162,7 +162,7 @@ func TestRealWorldMCPScenarios(t *testing.T) {
ctx := newRealWorldContext("test-full-workflow", "tests.realworld")
// Initialize stack for trace
stack, _, done := context.EnterStack(ctx, "tests.realworld", context.RefererAPI)
stack, _, done := context.EnterStack(ctx, "tests.realworld", &context.Options{})
defer done()
ctx.Stack = stack
@ -170,7 +170,7 @@ func TestRealWorldMCPScenarios(t *testing.T) {
{Role: "user", Content: "full_workflow"},
}
response, err := agent.Script.Create(ctx, messages)
response, _, err := agent.Script.Create(ctx, messages)
if err != nil {
t.Fatalf("Create failed: %v", err)
}
@ -237,7 +237,7 @@ func TestRealWorldTraceIntensive(t *testing.T) {
}
ctx := newRealWorldContext("test-trace-intensive", "tests.realworld")
stack, _, done := context.EnterStack(ctx, "tests.realworld", context.RefererAPI)
stack, _, done := context.EnterStack(ctx, "tests.realworld", &context.Options{})
defer done()
ctx.Stack = stack
@ -245,7 +245,7 @@ func TestRealWorldTraceIntensive(t *testing.T) {
{Role: "user", Content: "trace_intensive"},
}
response, err := agent.Script.Create(ctx, messages)
response, _, err := agent.Script.Create(ctx, messages)
if err != nil {
t.Fatalf("Create failed: %v", err)
}
@ -279,7 +279,7 @@ func TestRealWorldStressSimple(t *testing.T) {
{Role: "user", Content: "simple"},
}
response, err := agent.Script.Create(ctx, messages)
response, _, err := agent.Script.Create(ctx, messages)
if err != nil {
t.Fatalf("Iteration %d failed: %v", i, err)
}
@ -338,14 +338,14 @@ func TestRealWorldStressMCP(t *testing.T) {
ctx := newRealWorldContext(fmt.Sprintf("stress-mcp-%d", i), "tests.realworld")
// Initialize stack for trace
stack, _, done := context.EnterStack(ctx, "tests.realworld", context.RefererAPI)
stack, _, done := context.EnterStack(ctx, "tests.realworld", &context.Options{})
ctx.Stack = stack
messages := []context.Message{
{Role: "user", Content: scenario},
}
response, err := agent.Script.Create(ctx, messages)
response, _, err := agent.Script.Create(ctx, messages)
if err != nil {
t.Fatalf("Iteration %d (%s) failed: %v", i, scenario, err)
}
@ -427,14 +427,14 @@ func TestRealWorldStressFullWorkflow(t *testing.T) {
ctx := newRealWorldContext(fmt.Sprintf("stress-workflow-%d", i), "tests.realworld")
// Initialize stack for trace
stack, _, done := context.EnterStack(ctx, "tests.realworld", context.RefererAPI)
stack, _, done := context.EnterStack(ctx, "tests.realworld", &context.Options{})
ctx.Stack = stack
messages := []context.Message{
{Role: "user", Content: "full_workflow"},
}
response, err := agent.Script.Create(ctx, messages)
response, _, err := agent.Script.Create(ctx, messages)
if err != nil {
t.Fatalf("Iteration %d failed: %v", i, err)
}
@ -531,14 +531,14 @@ func TestRealWorldStressConcurrent(t *testing.T) {
)
// Initialize stack for trace
stack, _, done := context.EnterStack(ctx, "tests.realworld", context.RefererAPI)
stack, _, done := context.EnterStack(ctx, "tests.realworld", &context.Options{})
ctx.Stack = stack
messages := []context.Message{
{Role: "user", Content: scenario},
}
response, err := agent.Script.Create(ctx, messages)
response, _, err := agent.Script.Create(ctx, messages)
if err != nil {
errors <- fmt.Errorf("goroutine %d iteration %d (%s): %v", goroutineID, i, scenario, err)
done()
@ -658,14 +658,14 @@ func TestRealWorldStressResourceHeavy(t *testing.T) {
ctx := newRealWorldContext(fmt.Sprintf("stress-heavy-%d", i), "tests.realworld")
// Initialize stack for trace
stack, _, done := context.EnterStack(ctx, "tests.realworld", context.RefererAPI)
stack, _, done := context.EnterStack(ctx, "tests.realworld", &context.Options{})
ctx.Stack = stack
messages := []context.Message{
{Role: "user", Content: "resource_heavy"},
}
response, err := agent.Script.Create(ctx, messages)
response, _, err := agent.Script.Create(ctx, messages)
if err != nil {
t.Fatalf("Iteration %d failed: %v", i, err)
}
@ -726,7 +726,6 @@ func newRealWorldContext(chatID, assistantID string) *context.Context {
Space: plan.NewMemorySharedSpace(),
ChatID: chatID,
AssistantID: assistantID,
Connector: "gpt-4o",
Locale: "en-us",
Theme: "light",
Client: context.Client{

6
agent/assistant/llm.go
View file

@ -18,6 +18,7 @@ func (ast *Assistant) executeLLMStream(
completionOptions *context.CompletionOptions,
agentNode types.Node,
streamHandler message.StreamFunc,
opts *context.Options,
) (*context.CompletionResponse, error) {
// === Debug LLM Stream Start ===
@ -27,7 +28,7 @@ func (ast *Assistant) executeLLMStream(
// === End Debug ===
// Get connector object (capabilities were already set above, before stream_start)
conn, capabilities, err := ast.GetConnector(ctx)
conn, capabilities, err := ast.GetConnector(ctx, opts)
if err != nil {
ast.traceAgentFail(agentNode, err)
return nil, err
@ -92,10 +93,11 @@ func (ast *Assistant) executeLLMForToolRetry(
completionOptions *context.CompletionOptions,
agentNode types.Node,
streamHandler message.StreamFunc,
opts *context.Options,
) (*context.CompletionResponse, error) {
// Get connector object
conn, capabilities, err := ast.GetConnector(ctx)
conn, capabilities, err := ast.GetConnector(ctx, opts)
if err != nil {
ast.traceAgentFail(agentNode, err)
return nil, err

42
agent/assistant/load.go
View file

@ -531,7 +531,38 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
// locales
if locales, ok := data["locales"].(i18n.Map); ok {
assistant.Locales = locales
i18n.Locales[id] = locales.FlattenWithGlobal()
flattened := locales.FlattenWithGlobal()
// Auto-inject assistant name and description into all locales
// so that {{name}} and {{description}} templates can be resolved
for locale, i18nObj := range flattened {
if i18nObj.Messages == nil {
i18nObj.Messages = make(map[string]any)
}
// Add name and description if not already present
if _, exists := i18nObj.Messages["name"]; !exists && assistant.Name != "" {
i18nObj.Messages["name"] = assistant.Name
}
if _, exists := i18nObj.Messages["description"]; !exists && assistant.Description != "" {
i18nObj.Messages["description"] = assistant.Description
}
flattened[locale] = i18nObj
}
i18n.Locales[id] = flattened
} else {
// No locales defined, create default with name and description
if assistant.Name != "" || assistant.Description != "" {
defaultLocales := make(map[string]i18n.I18n)
defaultLocales["en"] = i18n.I18n{
Locale: "en",
Messages: map[string]any{
"name": assistant.Name,
"description": assistant.Description,
},
}
i18n.Locales[id] = defaultLocales
}
}
// Search options
@ -751,15 +782,6 @@ func (ast *Assistant) initialize() error {
}
ast.openai = api
// Check if the assistant supports vision
model := api.Model()
if v, ok := ast.Options["model"].(string); ok {
model = strings.TrimLeft(v, "moapi:")
}
if _, ok := VisionCapableModels[model]; ok {
ast.vision = true
}
// Check if the assistant has an init hook
if ast.Script != nil {
scriptCtx, err := ast.Script.NewContext("", nil)

19
agent/assistant/load_store_test.go
View file

@ -180,7 +180,6 @@ func newStoreTestContext(chatID, assistantID string) *context.Context {
Context: stdContext.Background(),
ChatID: chatID,
AssistantID: assistantID,
Connector: "",
Locale: "en-us",
Theme: "light",
Client: context.Client{
@ -269,7 +268,7 @@ function Create(ctx: any, messages: any[]): any {
ctx := newStoreTestContext("test-chat-id", assistantID)
messages := []context.Message{{Role: "user", Content: "Hello"}}
res, err := loaded.Script.Create(ctx, messages)
res, _, err := loaded.Script.Create(ctx, messages, &context.Options{})
require.NoError(t, err, "Create hook should execute without error")
require.NotNil(t, res, "Create hook should return a response")
@ -576,7 +575,7 @@ function Create(ctx: any, messages: any[]): any {
ctx := newStoreTestContext("test-chat-all-fields", assistantID)
messages := []context.Message{{Role: "user", Content: "Test message"}}
res, err := loaded.Script.Create(ctx, messages)
res, _, err := loaded.Script.Create(ctx, messages, &context.Options{})
require.NoError(t, err, "Create hook should execute without error")
require.NotNil(t, res, "Create hook should return a response")
@ -691,7 +690,7 @@ function Create(ctx: CreateContext, messages: Message[]): CreateResponse | null
{Role: "user", Content: "How are you?"},
}
res, err := loaded.Script.Create(ctx, messages)
res, _, err := loaded.Script.Create(ctx, messages, &context.Options{})
require.NoError(t, err, "TypeScript Create hook should execute without error")
require.NotNil(t, res, "Create hook should return a response")
@ -759,7 +758,7 @@ function Create(ctx: any, messages: any[]): any {
ctx := newStoreTestContext("null-test-chat", assistantID)
messages := []context.Message{{Role: "user", Content: "Hello"}}
res, err := loaded.Script.Create(ctx, messages)
res, _, err := loaded.Script.Create(ctx, messages, &context.Options{})
require.NoError(t, err, "Hook returning null should not error")
assert.Nil(t, res, "Hook returning null should return nil response")
}
@ -832,7 +831,7 @@ function Create(ctx: any, messages: any[]): any {
ctx := newStoreTestContext("preset-test-1", assistantID)
messages := []context.Message{{Role: "user", Content: "Be friendly please"}}
res, err := loaded.Script.Create(ctx, messages)
res, _, err := loaded.Script.Create(ctx, messages, &context.Options{})
require.NoError(t, err)
require.NotNil(t, res)
assert.Equal(t, "friendly", res.PromptPreset)
@ -843,7 +842,7 @@ function Create(ctx: any, messages: any[]): any {
ctx := newStoreTestContext("preset-test-2", assistantID)
messages := []context.Message{{Role: "user", Content: "Be professional"}}
res, err := loaded.Script.Create(ctx, messages)
res, _, err := loaded.Script.Create(ctx, messages, &context.Options{})
require.NoError(t, err)
require.NotNil(t, res)
assert.Equal(t, "professional", res.PromptPreset)
@ -854,7 +853,7 @@ function Create(ctx: any, messages: any[]): any {
ctx := newStoreTestContext("preset-test-3", assistantID)
messages := []context.Message{{Role: "user", Content: "Hello"}}
res, err := loaded.Script.Create(ctx, messages)
res, _, err := loaded.Script.Create(ctx, messages, &context.Options{})
require.NoError(t, err)
assert.Nil(t, res)
})
@ -916,7 +915,7 @@ function Create(ctx: any, messages: any[]): any {
ctx := newStoreTestContext("disable-test-1", assistantID)
messages := []context.Message{{Role: "user", Content: "disable_global prompts"}}
res, err := loaded.Script.Create(ctx, messages)
res, _, err := loaded.Script.Create(ctx, messages, &context.Options{})
require.NoError(t, err)
require.NotNil(t, res)
require.NotNil(t, res.DisableGlobalPrompts)
@ -928,7 +927,7 @@ function Create(ctx: any, messages: any[]): any {
ctx := newStoreTestContext("disable-test-2", assistantID)
messages := []context.Message{{Role: "user", Content: "enable_global prompts"}}
res, err := loaded.Script.Create(ctx, messages)
res, _, err := loaded.Script.Create(ctx, messages, &context.Options{})
require.NoError(t, err)
require.NotNil(t, res)
require.NotNil(t, res.DisableGlobalPrompts)

5
agent/assistant/next.go
View file

@ -55,7 +55,10 @@ func (ast *Assistant) handleDelegation(
// 2. Execute with the same Context (preserving ID, Space, Writer, etc.)
// 3. Call done() to pop from Stack when finished
// This ensures proper Stack tracing: parent assistant -> delegated assistant
return targetAssistant.Stream(ctx, delegate.Messages, streamHandler)
// Convert options map from delegate config to Options struct
delegateOpts := agentContext.OptionsFromMap(delegate.Options)
return targetAssistant.Stream(ctx, delegate.Messages, delegateOpts)
}
// buildStandardResponse builds the standard agent response when no custom Next hook processing is needed

25
agent/assistant/types.go
View file

@ -40,31 +40,6 @@ type Assistant struct {
// toolCalls bool // Whether this assistant supports tool_calls
}
// VisionCapableModels list of LLM models that support vision capabilities
var VisionCapableModels = map[string]bool{
// OpenAI Models
"gpt-4-vision-preview": true,
"gpt-4v": true, // Alias for gpt-4-vision-preview
// Anthropic Models
"claude-3-opus": true, // Most capable Claude model
"claude-3-sonnet": true, // Balanced Claude model
"claude-3-haiku": true, // Fast and efficient Claude model
// Google Models
"gemini-pro-vision": true,
// Open Source Models
"llava-13b": true,
"cogvlm": true,
"qwen-vl": true,
"yi-vl": true,
// Custom Models
"gpt-4o": true, // Custom OpenAI compatible model
"gpt-4o-mini": true, // Custom OpenAI compatible model - mini version
}
// MCPTool represents a simplified MCP tool for building LLM requests
// This is an internal representation used when collecting tools from MCP servers
// and preparing them for the LLM's tool calling interface

326
agent/content/README.md Normal file
View file

@ -0,0 +1,326 @@
# Content Processing Package
This package handles content transformation for multimodal messages in agent conversations. It is called **BEFORE** sending messages to the LLM and converts extended content types into standard LLM-compatible formats.
## ⚠️ Critical Design Principle
**Input**: Messages with extended content types (`file`, `data`, etc.)
**Output**: Messages with ONLY standard LLM-compatible types (`text`, `image_url`, `input_audio`)
The LLM should NEVER receive `type="file"` or `type="data"` content parts. These MUST be converted to `text` (or `image_url` for images if model supports vision).
## Architecture
```
Vision (main entry)
↓
Initialize processedFiles cache (map[fileID]text)
↓
processMessage (for each message)
↓
processContentPart (for each content part)
↓
Is uploader wrapper?
├── Yes → Check cache
│ ├── In cache? → Use cached text ✅
│ └── Not in cache → Try GetText(fileID) preview
│ ├── Has preview? → Use preview + cache ✅
│ └── No preview → Proceed to full processing ↓
└── No (HTTP/other) → Proceed to full processing ↓
↓
├── Fetch content (if needed)
│ ├── HTTP URL
│ └── Uploader Wrapper (__uploader://fileid)
↓
├── Determine Processing Strategy
│ ├── Model supports? → Format for model
│ └── Model doesn't support? → Use agent/MCP
↓
ProcessorRegistry
↓
├── ImageProcessor
├── AudioProcessor
├── PDFProcessor
├── WordProcessor
├── ExcelProcessor
└── TextProcessor
↓
Cache result (if uploader wrapper)
```
## Content Type Transformation
### Input → Output Mapping
| Input Type | Model Supports? | Output Type | Processing |
|------------|-----------------|-------------|------------|
| `text` | - | `text` | Pass through |
| `image_url` | ✅ Yes | `image_url` | Convert format if needed (base64/URL) |
| `image_url` | ❌ No | `text` | Use vision agent/MCP to describe |
| `input_audio` | ✅ Yes | `input_audio` | Keep as audio |
| `input_audio` | ❌ No | `text` | Transcribe using audio agent/MCP |
| `file` (image) | ✅ Yes | `image_url` | Same as image_url processing |
| `file` (image) | ❌ No | `text` | Use vision tool to describe |
| `file` (document) | - | `text` | Extract text from PDF/Word/Excel/etc |
| `data` | - | `text` | Fetch and format data sources |
### 1. Images and Audio
**If model supports (vision/audio capability):**
- Keep as multimodal content:
- `image_url`: Convert to appropriate format (OpenAI URL vs Claude base64)
- `input_audio`: Convert to base64 format
**If model doesn't support:**
- Convert to text:
- Use agent/MCP specified in `uses.Vision` or `uses.Audio`
- Extract text description or transcription
- Return as `type="text"` content
**HTTP URLs:**
- Fetch content first
- Then process the same way as above
### 2. Files (type="file")
**Critical**: All `type="file"` content MUST be converted to `text` or `image_url` (if image and model supports).
**Processing Steps:**
1. **Fetch file content**:
- Uploader wrapper: `__uploader://fileid` → Parse and fetch from attachment manager
- HTTP URL: Download from URL
2. **Detect file type** from content-type and magic bytes
3. **Process based on file type**:
| File Type | Output Type | Processing Method |
| ------------ | ----------- | -------------------------------------------------------------------------------------------------------- |
| **Image** | `image_url` or `text` | If model supports vision → `image_url`<br>If not → use vision tool → `text` |
| **PDF** | `text` | If `uses.Vision` supports PDF → use vision tool<br>Otherwise → extract text directly |
| **Word** | `text` | Extract text using Word document parser |
| **Excel** | `text` | Extract and format as readable table/CSV |
| **PPT** | `text` | Extract text and slide content |
| **CSV** | `text` | Format as readable table |
| **Text** | `text` | Read directly (with encoding detection) |
| **JSON/XML** | `text` | Pretty print for readability |
### 3. Data Sources (type="data")
**Critical**: All `type="data"` content MUST be converted to `text`.
**Processing Steps:**
1. **Parse DataContent.Sources** array
2. **Fetch data** from each source:
- `model`: Query data model
- `kb_collection`: Search knowledge base collection
- `kb_document`: Get document content
- `table`: Query database table
- `api`: Call API endpoint
- `mcp_resource`: Fetch MCP resource
3. **Format as readable text**:
- Tables: Format as markdown tables or CSV
- Documents: Include title and content
- JSON: Pretty print
4. **Return as** `type="text"` content
## Components
### Core Files
- **content.go** - Main entry point (`Vision` function)
- **types.go** - Type definitions and constants
- **interfaces.go** - Interface definitions
### Fetching
- **fetch.go** - Fetch content from HTTP or uploader
### Processors
- **processor.go** - Processor registry and routing
- **image.go** - Image processing
- **audio.go** - Audio processing
- **pdf.go** - PDF document processing
- **word.go** - Word document processing
- **excel.go** - Excel spreadsheet processing
- **text.go** - Plain text and CSV processing
## Frontend Message Format
The frontend (InputArea) sends messages in the following format:
### Image Attachments
```json
{
"type": "image_url",
"image_url": {
"url": "__yao.attachment://file_id",
"detail": "auto"
}
}
```
### File Attachments
```json
{
"type": "file",
"file": {
"url": "__yao.attachment://file_id",
"filename": "document.pdf"
}
}
```
The `url` field contains an uploader wrapper in the format `__uploader://fileid`.
## Data Structures
### ContentInfo
Holds information about content to be processed:
```go
type ContentInfo struct {
Source ContentSource // http, uploader, base64, local
FileType FileType // image, audio, pdf, word, excel, etc.
ContentType string // MIME type
URL string // Original URL or file ID
Data []byte // File data
// For uploader wrapper
UploaderName string
FileID string
}
```
### ProcessedContent
Result of content processing:
```go
type ProcessedContent struct {
Text string // Extracted text
ContentPart *context.ContentPart // For model input
Metadata map[string]interface{}
Error error
}
```
## Usage Example
```go
import (
"github.com/yaoapp/yao/agent/content"
"github.com/yaoapp/yao/agent/context"
)
// Process messages before sending to LLM
processedMessages, err := content.Vision(
ctx,
capabilities, // Model capabilities
messages, // Original messages
uses, // Tool specifications (vision, audio, etc.)
)
```
## Performance Optimization
### File Processing Cache
**Problem**: Same file (uploader wrapper) might appear in multiple messages or be referenced multiple times.
**Solution**: Three-level caching strategy:
1. **In-memory cache** (`processedFiles` map):
- Caches processed text for the duration of the Vision() call
- Key: file ID from uploader wrapper
- Value: extracted text content
2. **Attachment preview** (attachment.GetText with preview):
- Tries to get preview (first 2000 chars) from attachment manager
- If file was previously processed and saved, preview is available immediately
- Much faster than full file processing
3. **Full processing** (only if needed):
- Falls back to complete file processing if no cache/preview available
- Result is cached in memory and optionally saved to attachment manager
### Cache Flow
```go
// For uploader://file_id
1. Check processedFiles[file_id]
└── Found? → Return cached text ⚡ (fastest)
2. Not in cache → Call attachment.GetText(file_id, false) // preview only
└── Has preview? → Cache and return ⚡ (fast)
3. No preview → Process file fully 🔄 (slower)
└── Cache result in processedFiles
└── Optional: Save to attachment using SaveText for future use
```
### Benefits
- **Avoid duplicate processing**: Same file processed only once per Vision() call
- **Fast preview access**: Leverage pre-processed content from attachment manager
- **Reduced latency**: Especially important for large documents (PDFs, Word, Excel)
- **Resource efficient**: Less CPU/memory usage for repeated file references
## Implementation Status
### ✅ Completed
- [x] Package structure
- [x] Type definitions
- [x] Interface definitions
- [x] Skeleton functions with TODO comments
- [x] File processing cache infrastructure
- [x] Cache helper functions (tryGetCachedText, cacheProcessedText)
### 🚧 To Implement
- [ ] tryGetCachedText implementation (attachment.GetText integration)
- [ ] cacheProcessedText implementation (attachment.SaveText integration)
- [ ] HTTP fetching logic
- [ ] Uploader wrapper parsing and fetching
- [ ] Image processing (base64, vision API)
- [ ] Audio processing (transcription)
- [ ] PDF text extraction
- [ ] Word document parsing
- [ ] Excel spreadsheet parsing
- [ ] Text/CSV formatting
- [ ] Content part processing logic
- [ ] Model capability detection
- [ ] Agent/MCP tool invocation
## Configuration
Content processing behavior is controlled by:
1. **Model Capabilities** (`openai.Capabilities`)
- Determines if model can handle images/audio directly
- Specifies vision format (OpenAI vs Claude)
2. **Uses** (`context.Uses`)
```go
type Uses struct {
Vision string // "agent" or "mcp:server_id"
Audio string // "agent" or "mcp:server_id"
Search string
Fetch string
}
```
## Error Handling
- Errors during processing are logged but don't stop the entire pipeline
- Original content is kept if processing fails
- Graceful degradation: if advanced processing fails, fall back to simpler methods

61
agent/content/audio.go Normal file
View file

@ -0,0 +1,61 @@
package content
import (
"fmt"
"strings"
"github.com/yaoapp/gou/connector/openai"
agentContext "github.com/yaoapp/yao/agent/context"
)
// AudioHandler handles audio content
type AudioHandler struct{}
// CanHandle checks if this handler can handle the content type
func (h *AudioHandler) CanHandle(contentType string, fileType FileType) bool {
return fileType == FileTypeAudio || strings.HasPrefix(contentType, "audio/")
}
// Handle processes audio content
// Logic similar to image:
// 1. If model supports audio input -> convert to base64 format
// 2. If model doesn't support audio -> use agent/MCP specified in uses.Audio
func (h *AudioHandler) Handle(ctx *agentContext.Context, info *Info, capabilities *openai.Capabilities, uses *agentContext.Uses) (*Result, error) {
// TODO: Implement audio handling
// 1. Check model audio capabilities
// 2. If supported:
// - Encode audio as base64 with proper format
// 3. If not supported:
// - Call audio agent/MCP to transcribe audio to text
// 4. Return Result with text or ContentPart
return nil, fmt.Errorf("not implemented")
}
// handleWithAudioModel processes audio using model's audio capability
func (h *AudioHandler) handleWithAudioModel(ctx *agentContext.Context, info *Info) (*Result, error) {
// TODO: Implement audio model processing
// Format audio according to model's audio input format
return nil, fmt.Errorf("not implemented")
}
// handleWithAudioAgent processes audio using audio agent or MCP
func (h *AudioHandler) handleWithAudioAgent(ctx *agentContext.Context, info *Info, audioTool string) (string, error) {
// TODO: Implement audio agent/MCP processing
// 1. Parse audioTool (format: "agent" or "mcp:server_id")
// 2. Call appropriate tool to transcribe audio
// 3. Return transcribed text
return "", fmt.Errorf("not implemented")
}
// encodeAudioBase64 encodes audio data to base64 with proper format
func encodeAudioBase64(data []byte, contentType string) string {
// TODO: Implement audio base64 encoding
return ""
}
// detectAudioFormat detects audio format from content type or data
func detectAudioFormat(contentType string, data []byte) string {
// TODO: Implement audio format detection
// Return format like "wav", "mp3", "flac", etc.
return ""
}

470
agent/content/content.go Normal file
View file

@ -0,0 +1,470 @@
package content
import (
"fmt"
"strings"
"github.com/yaoapp/gou/connector/openai"
agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/attachment"
)
// Vision transforms extended content types to LLM-compatible formats
// This is the main entry point for content preprocessing before sending to LLM
//
// IMPORTANT: This function is called BEFORE sending messages to LLM (in agent.executeLLMStream)
// It must convert all extended content types to standard LLM-compatible types.
//
// Input Content Types (Extended):
// - type="text" -> Pass through (already standard)
// - type="image_url" -> Process based on model capability (may need base64 conversion or vision tool)
// - type="input_audio" -> Process based on model capability (may need transcription)
// - type="file" -> Convert to text or image_url (MUST be converted)
// - type="data" -> Convert to text (MUST be converted)
//
// Output Content Types (LLM-compatible only):
// - type="text" -> Text content
// - type="image_url" -> Image (only if model supports vision)
// - type="input_audio" -> Audio (only if model supports audio)
//
// Processing Logic:
// 1. For images (image_url):
// - If model supports vision -> keep as image_url (may convert URL to base64)
// - If model doesn't support -> use vision agent/MCP to extract text -> convert to type="text"
//
// 2. For audio (input_audio):
// - If model supports audio -> keep as input_audio
// - If model doesn't support -> use audio agent/MCP to transcribe -> convert to type="text"
//
// 3. For files (type="file"):
// - Parse uploader wrapper (__uploader://fileid) or fetch HTTP URL
// - Detect file type (PDF, Word, Excel, Image, etc.)
// - Process based on file type:
// - Images: same as image processing above
// - PDF: use vision tool if available, otherwise extract text -> type="text"
// - Word/Excel/PPT/CSV: extract text -> type="text"
// - MUST convert to type="text" or type="image_url" (if image and model supports)
//
// 4. For data (type="data"):
// - Fetch data from sources (models, KB, MCP resources, etc.)
// - Format as readable text
// - MUST convert to type="text"
//
// Return: Messages with only standard LLM-compatible content types (text, image_url, input_audio)
func Vision(ctx *agentContext.Context, capabilities *openai.Capabilities, messages []agentContext.Message, uses *agentContext.Uses) ([]agentContext.Message, error) {
// Initialize handlers and fetcher
registry := NewRegistry()
fetcher := NewFetcher()
// Cache for processed files (uploader wrapper -> extracted text)
// Ensures each file is only processed once
processedFiles := make(map[string]string)
// Process each message
processedMessages := make([]agentContext.Message, 0, len(messages))
for _, msg := range messages {
processedMsg, err := processMessage(ctx, &msg, capabilities, uses, registry, fetcher, processedFiles)
if err != nil {
// Log error but continue processing other messages
// TODO: Add proper logging
fmt.Printf("Warning: failed to process message: %v\n", err)
processedMessages = append(processedMessages, msg) // Keep original on error
continue
}
processedMessages = append(processedMessages, processedMsg)
}
return processedMessages, nil
}
// processMessage processes a single message and its content parts
func processMessage(
ctx *agentContext.Context,
msg *agentContext.Message,
capabilities *openai.Capabilities,
uses *agentContext.Uses,
registry *Registry,
fetcher Fetcher,
processedFiles map[string]string,
) (agentContext.Message, error) {
// If content is simple string, no processing needed
if _, ok := msg.GetContentAsString(); ok {
return *msg, nil
}
// Get content parts
parts, ok := msg.GetContentAsParts()
if !ok {
return *msg, nil
}
// Process each content part
processedParts := make([]agentContext.ContentPart, 0, len(parts))
for _, part := range parts {
processedPart, err := processContentPart(ctx, &part, capabilities, uses, registry, fetcher, processedFiles)
if err != nil {
// Log error and handle gracefully
fmt.Printf("Warning: failed to process content part: %v\n", err)
// For image_url that failed to process, convert to text description
// This prevents sending unsupported multimodal content to non-vision models
if part.Type == agentContext.ContentImageURL {
processedParts = append(processedParts, agentContext.ContentPart{
Type: agentContext.ContentText,
Text: fmt.Sprintf("[Image processing failed: %s]", part.ImageURL.URL),
})
} else {
// For other types, keep original
processedParts = append(processedParts, part)
}
continue
}
// If handling returned text, convert to text part
if processedPart.Text != "" {
processedParts = append(processedParts, agentContext.ContentPart{
Type: agentContext.ContentText,
Text: processedPart.Text,
})
} else if processedPart.ContentPart != nil {
// Use the processed content part (e.g., base64 image)
processedParts = append(processedParts, *processedPart.ContentPart)
} else {
// Keep original if no handling result
processedParts = append(processedParts, part)
}
}
// Return new message with processed content
return agentContext.Message{
Role: msg.Role,
Content: processedParts,
Name: msg.Name,
ToolCallID: msg.ToolCallID,
ToolCalls: msg.ToolCalls,
Refusal: msg.Refusal,
}, nil
}
// processContentPart processes a single content part
// IMPORTANT: Must convert extended types (file, data) to standard types (text, image_url, input_audio)
func processContentPart(
ctx *agentContext.Context,
part *agentContext.ContentPart,
capabilities *openai.Capabilities,
uses *agentContext.Uses,
registry *Registry,
fetcher Fetcher,
processedFiles map[string]string,
) (*Result, error) {
// 1. Handle standard types - pass through
switch part.Type {
case agentContext.ContentText:
// Text is already standard, pass through
return &Result{
ContentPart: part,
}, nil
case agentContext.ContentImageURL:
// Image URL - check if it needs processing
return processImageURLContent(ctx, part, capabilities, uses, registry, fetcher, processedFiles)
case agentContext.ContentInputAudio:
// Audio - check if it needs processing
return processAudioContent(ctx, part, capabilities, uses, registry, fetcher, processedFiles)
}
// 2. Handle extended types - MUST convert to standard types
switch part.Type {
case agentContext.ContentFile:
return processFileContent(ctx, part, capabilities, uses, registry, fetcher, processedFiles)
case agentContext.ContentData:
return processDataContent(ctx, part)
default:
// Unknown type, return error
return nil, fmt.Errorf("unsupported content type: %s", part.Type)
}
}
// processFileContent processes file content with caching
func processFileContent(
ctx *agentContext.Context,
part *agentContext.ContentPart,
capabilities *openai.Capabilities,
uses *agentContext.Uses,
registry *Registry,
fetcher Fetcher,
processedFiles map[string]string,
) (*Result, error) {
if part.File == nil || part.File.URL == "" {
return nil, fmt.Errorf("file content part missing URL")
}
url := part.File.URL
// Step 1: Try to get cached text (three-tier cache)
cachedText, found, err := tryGetCachedText(ctx, url, processedFiles)
if err != nil {
return nil, fmt.Errorf("failed to check cache: %w", err)
}
if found {
// Cache hit! Return as text
return &Result{
Text: cachedText,
}, nil
}
// Step 2: No cache, need to process the file
// Determine content source
source, sourceURL, err := determineContentSource(part)
if err != nil {
return nil, fmt.Errorf("failed to determine content source: %w", err)
}
// Fetch content
info, err := fetcher.Fetch(ctx, source, sourceURL)
if err != nil {
return nil, fmt.Errorf("failed to fetch content: %w", err)
}
// Detect file type if not already set
if info.FileType == FileTypeUnknown {
info.FileType = DetectFileType(info.ContentType, part.File.Filename)
}
// Process with appropriate handler
result, err := registry.Handle(ctx, info, capabilities, uses)
if err != nil {
return nil, fmt.Errorf("failed to handle content: %w", err)
}
// Step 3: Cache the result if it's text
if result.Text != "" {
if cacheErr := cacheProcessedText(ctx, url, result.Text, processedFiles); cacheErr != nil {
// Log error but don't fail the request
fmt.Printf("Warning: failed to cache processed text: %v\n", cacheErr)
}
}
return result, nil
}
// processImageURLContent processes image_url content
// If URL is uploader wrapper or HTTP, fetch and process it
func processImageURLContent(
ctx *agentContext.Context,
part *agentContext.ContentPart,
capabilities *openai.Capabilities,
uses *agentContext.Uses,
registry *Registry,
fetcher Fetcher,
processedFiles map[string]string,
) (*Result, error) {
if part.ImageURL == nil || part.ImageURL.URL == "" {
return nil, fmt.Errorf("image_url content missing URL")
}
url := part.ImageURL.URL
// If it's a data URI (base64), pass through
if strings.HasPrefix(url, "data:") {
return &Result{
ContentPart: part,
}, nil
}
// If it's uploader wrapper or HTTP URL, need to process
// Check cache first
cachedText, found, err := tryGetCachedText(ctx, url, processedFiles)
if err != nil {
return nil, fmt.Errorf("failed to check cache: %w", err)
}
if found {
// Cache hit! Return as text
return &Result{
Text: cachedText,
}, nil
}
// Determine source
source, sourceURL, err := determineContentSource(part)
if err != nil {
return nil, fmt.Errorf("failed to determine content source: %w", err)
}
// Fetch content
info, err := fetcher.Fetch(ctx, source, sourceURL)
if err != nil {
return nil, fmt.Errorf("failed to fetch image: %w", err)
}
// Set file type as image
info.FileType = FileTypeImage
// Process with image handler
result, err := registry.Handle(ctx, info, capabilities, uses)
if err != nil {
return nil, fmt.Errorf("failed to handle image: %w", err)
}
// Cache if result is text
if result.Text != "" {
if cacheErr := cacheProcessedText(ctx, url, result.Text, processedFiles); cacheErr != nil {
fmt.Printf("Warning: failed to cache processed text: %v\n", cacheErr)
}
}
return result, nil
}
// processAudioContent processes input_audio content
func processAudioContent(
ctx *agentContext.Context,
part *agentContext.ContentPart,
capabilities *openai.Capabilities,
uses *agentContext.Uses,
registry *Registry,
fetcher Fetcher,
processedFiles map[string]string,
) (*Result, error) {
if part.InputAudio == nil || part.InputAudio.Data == "" {
return nil, fmt.Errorf("input_audio content missing data")
}
// For now, pass through audio as-is
// TODO: Implement audio processing (transcription, etc.)
return &Result{
ContentPart: part,
}, nil
}
// processDataContent processes data content (converts to text)
func processDataContent(ctx *agentContext.Context, part *agentContext.ContentPart) (*Result, error) {
if part.Data == nil {
return nil, fmt.Errorf("data content part missing data")
}
// TODO: Implement data processing
// For now, just return error
return nil, fmt.Errorf("data content processing not implemented yet")
}
// determineContentSource determines where the content comes from
func determineContentSource(part *agentContext.ContentPart) (Source, string, error) {
var url string
// Extract URL based on content type
switch part.Type {
case agentContext.ContentFile:
if part.File == nil || part.File.URL == "" {
return "", "", fmt.Errorf("file content missing URL")
}
url = part.File.URL
case agentContext.ContentImageURL:
if part.ImageURL == nil || part.ImageURL.URL == "" {
return "", "", fmt.Errorf("image_url content missing URL")
}
url = part.ImageURL.URL
case agentContext.ContentInputAudio:
if part.InputAudio == nil || part.InputAudio.Data == "" {
return "", "", fmt.Errorf("input_audio content missing data")
}
// Audio data is base64, treat as base64 source
return SourceBase64, part.InputAudio.Data, nil
default:
return "", "", fmt.Errorf("unsupported content type for source detection: %s", part.Type)
}
// Determine source type based on URL format
if strings.HasPrefix(url, "http://") || strings.HasPrefix(url, "https://") {
return SourceHTTP, url, nil
}
if strings.HasPrefix(url, "__") {
// Uploader wrapper format: __uploader://fileid
return SourceUploader, url, nil
}
if strings.HasPrefix(url, "data:") {
// Data URI (base64)
return SourceBase64, url, nil
}
// Default to treating as uploader if no prefix matches
return SourceUploader, url, nil
}
// shouldProcessWithModel checks if content should be processed by the model directly
func shouldProcessWithModel(capabilities *openai.Capabilities, fileType FileType) (bool, agentContext.VisionFormat) {
// TODO: Implement model capability check
// For images: check if model supports vision
// For audio: check if model supports audio input
// Return whether to use model and the format to use
return false, agentContext.VisionFormatNone
}
// getToolForProcessing gets the agent/MCP tool to use for processing
func getToolForProcessing(uses *agentContext.Uses, fileType FileType) string {
// TODO: Implement tool selection
// Based on file type, return the appropriate tool from uses
// - Images -> uses.Vision
// - Audio -> uses.Audio
// - PDF (if vision available) -> uses.Vision
return ""
}
// tryGetCachedText checks if the URL is an uploader wrapper and tries to get cached text
// Returns (text, found, error)
func tryGetCachedText(ctx *agentContext.Context, url string, processedFiles map[string]string) (string, bool, error) {
// Parse URL to check if it's an uploader wrapper
uploaderName, fileID, isWrapper := attachment.Parse(url)
if !isWrapper {
return "", false, nil // Not an uploader wrapper, no cache
}
// 1. Check in-memory cache for this Vision call
if text, ok := processedFiles[fileID]; ok {
return text, true, nil
}
// 2. Try attachment manager's content_preview (cross-call cache)
manager, exists := attachment.Managers[uploaderName]
if exists {
// GetText with fullContent=false to get preview (default)
text, err := manager.GetText(ctx.Context, fileID, false)
if err == nil && text != "" {
// Cache in-memory for this Vision call
processedFiles[fileID] = text
return text, true, nil
}
}
// No cache found
return "", false, nil
}
// cacheProcessedText caches the processed text for an uploader wrapper
func cacheProcessedText(ctx *agentContext.Context, url string, text string, processedFiles map[string]string) error {
// Parse URL to get uploader name and file ID
uploaderName, fileID, isWrapper := attachment.Parse(url)
if !isWrapper {
return nil // Not an uploader wrapper, nothing to cache
}
// 1. Cache in-memory for this Vision call
processedFiles[fileID] = text
// 2. Save to attachment manager for future Vision calls
manager, exists := attachment.Managers[uploaderName]
if exists {
return manager.SaveText(ctx.Context, fileID, text)
}
return nil
}

458
agent/content/content_vision_test.go Normal file
View file

@ -0,0 +1,458 @@
package content_test
import (
"bytes"
"context"
"image"
"image/color"
"image/png"
"mime/multipart"
"strings"
"testing"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/yao/agent/content"
agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/testutils"
"github.com/yaoapp/yao/attachment"
)
// setupTestUploader creates and registers a test uploader manager
// The manager will be registered with "__" prefix as required by attachment.Parse
func setupTestUploader(t *testing.T, name string) attachment.FileManager {
// Register with __ prefix to match Parse behavior
managerName := "__" + name
manager, err := attachment.Register(managerName, "local", attachment.ManagerOption{
Driver: "local",
MaxSize: "10M",
AllowedTypes: []string{"text/*", "image/*", "application/*"},
Options: map[string]interface{}{
"path": "/tmp/test_vision_attachments_" + name,
},
})
if err != nil {
t.Fatalf("Failed to register attachment manager '%s': %v", managerName, err)
}
return manager
}
// cleanupTestUploader removes the test uploader from registry
func cleanupTestUploader(name string) {
delete(attachment.Managers, "__"+name)
}
// generateTestImage creates a valid PNG image (100x100 red square)
func generateTestImage(t *testing.T) []byte {
img := image.NewRGBA(image.Rect(0, 0, 100, 100))
red := color.RGBA{255, 0, 0, 255}
for y := 0; y < 100; y++ {
for x := 0; x < 100; x++ {
img.Set(x, y, red)
}
}
var buf bytes.Buffer
if err := png.Encode(&buf, img); err != nil {
t.Fatalf("Failed to encode test image: %v", err)
}
return buf.Bytes()
}
// TestVision_TextFile tests Vision function with text/code file parsing
func TestVision_TextFile(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Setup test uploader
uploaderName := "test-vision-text"
manager := setupTestUploader(t, uploaderName)
defer cleanupTestUploader(uploaderName)
// 1. Create and upload a Go source file
testContent := `package main
import "fmt"
func main() {
fmt.Println("Hello, Vision Test!")
}
`
// Upload file
reader := strings.NewReader(testContent)
fileHeader := &attachment.FileHeader{
FileHeader: &multipart.FileHeader{
Filename: "main.go",
Size: int64(len(testContent)),
Header: make(map[string][]string),
},
}
fileHeader.Header.Set("Content-Type", "text/x-go")
uploadedFile, err := manager.Upload(context.Background(), fileHeader, reader, attachment.UploadOption{
Groups: []string{"vision", "test"},
})
if err != nil {
t.Fatalf("Failed to upload file: %v", err)
}
t.Logf("Uploaded file ID: %s", uploadedFile.ID)
// 2. Prepare Vision context (text files don't need special capabilities)
ctx := agentContext.New(context.Background(), nil, "test")
capabilities := &openai.Capabilities{}
messages := []agentContext.Message{
{
Role: "user",
Content: []agentContext.ContentPart{
{
Type: agentContext.ContentFile,
File: &agentContext.FileAttachment{
URL: "__" + uploaderName + "://" + uploadedFile.ID,
Filename: "main.go",
},
},
},
},
}
// 3. Call Vision function
result, err := content.Vision(ctx, capabilities, messages, nil)
if err != nil {
t.Fatalf("Vision function failed: %v", err)
}
if len(result) != 1 {
t.Fatalf("Expected 1 message, got %d", len(result))
}
// 4. Verify result
contentParts, ok := result[0].Content.([]agentContext.ContentPart)
if !ok {
t.Fatalf("Expected content to be []ContentPart, got %T", result[0].Content)
}
if len(contentParts) != 1 {
t.Fatalf("Expected 1 content part, got %d", len(contentParts))
}
// Should be converted to text
if contentParts[0].Type != agentContext.ContentText {
t.Errorf("Expected ContentText type, got %s", contentParts[0].Type)
}
if !strings.Contains(contentParts[0].Text, "package main") {
t.Errorf("Expected text to contain 'package main', got: %s", contentParts[0].Text)
}
if !strings.Contains(contentParts[0].Text, "Hello, Vision Test!") {
t.Errorf("Expected text to contain 'Hello, Vision Test!', got: %s", contentParts[0].Text)
}
t.Logf("✓ Text file successfully parsed: %d characters", len(contentParts[0].Text))
}
// TestVision_ImageWithVisionSupport tests image processing with vision-capable model
func TestVision_ImageWithVisionSupport(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Setup test uploader
uploaderName := "test-vision-image"
manager := setupTestUploader(t, uploaderName)
defer cleanupTestUploader(uploaderName)
// 1. Create and upload a test image (1x1 red PNG)
imageData := []byte{
0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A,
0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, 0x52,
0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01,
0x08, 0x02, 0x00, 0x00, 0x00, 0x90, 0x77, 0x53,
0xDE, 0x00, 0x00, 0x00, 0x0C, 0x49, 0x44, 0x41,
0x54, 0x08, 0xD7, 0x63, 0xF8, 0xCF, 0xC0, 0x00,
0x00, 0x03, 0x01, 0x01, 0x00, 0x18, 0xDD, 0x8D,
0xB4, 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4E,
0x44, 0xAE, 0x42, 0x60, 0x82,
}
// Upload image
reader := strings.NewReader(string(imageData))
fileHeader := &attachment.FileHeader{
FileHeader: &multipart.FileHeader{
Filename: "test.png",
Size: int64(len(imageData)),
Header: make(map[string][]string),
},
}
fileHeader.Header.Set("Content-Type", "image/png")
uploadedFile, err := manager.Upload(context.Background(), fileHeader, reader, attachment.UploadOption{
Groups: []string{"vision", "test"},
})
if err != nil {
t.Fatalf("Failed to upload image: %v", err)
}
// 2. Prepare Vision context with vision-capable model
ctx := agentContext.New(context.Background(), nil, "test")
// Construct capabilities with vision support (OpenAI format)
capabilities := &openai.Capabilities{
Vision: agentContext.VisionFormatOpenAI, // OpenAI vision format
}
messages := []agentContext.Message{
{
Role: "user",
Content: []agentContext.ContentPart{
{
Type: agentContext.ContentImageURL,
ImageURL: &agentContext.ImageURL{
URL: "__" + uploaderName + "://" + uploadedFile.ID,
},
},
},
},
}
// 3. Call Vision function (no uses needed for direct vision support)
result, err := content.Vision(ctx, capabilities, messages, nil)
if err != nil {
t.Fatalf("Vision function failed: %v", err)
}
if len(result) != 1 {
t.Fatalf("Expected 1 message, got %d", len(result))
}
// 4. Verify result
contentParts, ok := result[0].Content.([]agentContext.ContentPart)
if !ok {
t.Fatalf("Expected content to be []ContentPart, got %T", result[0].Content)
}
if len(contentParts) != 1 {
t.Fatalf("Expected 1 content part, got %d", len(contentParts))
}
// If model supports vision, should be image_url with base64
if capabilities.Vision != nil {
if contentParts[0].Type != agentContext.ContentImageURL {
t.Errorf("Expected ContentImageURL type, got %s", contentParts[0].Type)
}
if contentParts[0].ImageURL == nil {
t.Fatal("Expected ImageURL to be set")
}
if !strings.Contains(contentParts[0].ImageURL.URL, "data:image/png;base64,") {
t.Errorf("Expected base64 data URI, got: %s", contentParts[0].ImageURL.URL)
}
t.Logf("✓ Image processed with vision support: %d bytes (base64)", len(contentParts[0].ImageURL.URL))
} else {
// If no vision support, should fall back to text (via agent/MCP)
t.Logf("ℹ Model doesn't support vision, result type: %s", contentParts[0].Type)
}
}
// TestVision_ImageWithAgent tests image processing with vision agent when model doesn't support vision
// Note: This test demonstrates the agent fallback mechanism when the model doesn't support vision
func TestVision_ImageWithAgent(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Setup test uploader
uploaderName := "test-vision-agent"
manager := setupTestUploader(t, uploaderName)
defer cleanupTestUploader(uploaderName)
// 1. Generate and upload a valid test image (100x100 red PNG)
imageData := generateTestImage(t)
reader := strings.NewReader(string(imageData))
fileHeader := &attachment.FileHeader{
FileHeader: &multipart.FileHeader{
Filename: "test.png",
Size: int64(len(imageData)),
Header: make(map[string][]string),
},
}
fileHeader.Header.Set("Content-Type", "image/png")
uploadedFile, err := manager.Upload(context.Background(), fileHeader, reader, attachment.UploadOption{
Groups: []string{"vision", "test"},
})
if err != nil {
t.Fatalf("Failed to upload image: %v", err)
}
// 2. Prepare Vision context with proper setup
// Model does NOT support vision, but uses.Vision specifies a vision agent
ctx := agentContext.New(context.Background(), nil, "test")
// Capabilities without vision support
capabilities := &openai.Capabilities{
Vision: nil, // No vision support
}
// Uses configuration with vision agent
uses := &agentContext.Uses{
Vision: "tests.vision-helper", // Use vision-helper agent
}
messages := []agentContext.Message{
{
Role: "user",
Content: []agentContext.ContentPart{
{
Type: agentContext.ContentImageURL,
ImageURL: &agentContext.ImageURL{
URL: "__" + uploaderName + "://" + uploadedFile.ID,
},
},
},
},
}
// 3. Call Vision - should use agent since model doesn't support vision
result, err := content.Vision(ctx, capabilities, messages, uses)
if err != nil {
t.Fatalf("Vision function failed: %v", err)
}
if len(result) != 1 {
t.Fatalf("Expected 1 message, got %d", len(result))
}
// 4. Verify result is text (processed by vision agent)
contentParts, ok := result[0].Content.([]agentContext.ContentPart)
if !ok {
t.Fatalf("Expected content to be []ContentPart, got %T", result[0].Content)
}
if len(contentParts) != 1 {
t.Fatalf("Expected 1 content part, got %d", len(contentParts))
}
if contentParts[0].Type != agentContext.ContentText {
t.Errorf("Expected ContentText (from agent), got: %s", contentParts[0].Type)
}
if contentParts[0].Text == "" {
t.Error("Expected non-empty text from vision agent processing")
}
t.Logf("✓ Vision agent processed image to text: %d characters", len(contentParts[0].Text))
t.Logf("Agent response text:\n%s", contentParts[0].Text)
}
// TestVision_CachedContent tests that file content is cached and reused
func TestVision_CachedContent(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Setup test uploader
uploaderName := "test-vision-cache"
manager := setupTestUploader(t, uploaderName)
defer cleanupTestUploader(uploaderName)
// 1. Upload a text file
testContent := "Test content for caching verification"
reader := strings.NewReader(testContent)
fileHeader := &attachment.FileHeader{
FileHeader: &multipart.FileHeader{
Filename: "cache-test.txt",
Size: int64(len(testContent)),
Header: make(map[string][]string),
},
}
fileHeader.Header.Set("Content-Type", "text/plain")
uploadedFile, err := manager.Upload(context.Background(), fileHeader, reader, attachment.UploadOption{
Groups: []string{"vision", "test"},
})
if err != nil {
t.Fatalf("Failed to upload file: %v", err)
}
// 2. Prepare Vision context with same file referenced twice
ctx := agentContext.New(context.Background(), nil, "test")
// Construct simple capabilities (text files don't need vision)
capabilities := &openai.Capabilities{}
messages := []agentContext.Message{
{
Role: "user",
Content: []agentContext.ContentPart{
{
Type: agentContext.ContentFile,
File: &agentContext.FileAttachment{
URL: "__" + uploaderName + "://" + uploadedFile.ID,
Filename: "cache-test.txt",
},
},
{
Type: agentContext.ContentFile,
File: &agentContext.FileAttachment{
URL: "__" + uploaderName + "://" + uploadedFile.ID, // Same file
Filename: "cache-test.txt",
},
},
},
},
}
// 3. Call Vision
result, err := content.Vision(ctx, capabilities, messages, nil)
if err != nil {
t.Fatalf("Vision function failed: %v", err)
}
if len(result) != 1 {
t.Fatalf("Expected 1 message, got %d", len(result))
}
// 4. Verify both file references were processed
contentParts, ok := result[0].Content.([]agentContext.ContentPart)
if !ok {
t.Fatalf("Expected content to be []ContentPart")
}
if len(contentParts) != 2 {
t.Fatalf("Expected 2 content parts (both files), got %d", len(contentParts))
}
// Both should be text with same content
if contentParts[0].Type != agentContext.ContentText {
t.Errorf("First part: expected ContentText, got %s", contentParts[0].Type)
}
if contentParts[1].Type != agentContext.ContentText {
t.Errorf("Second part: expected ContentText, got %s", contentParts[1].Type)
}
if !strings.Contains(contentParts[0].Text, testContent) {
t.Errorf("First part text doesn't contain expected content")
}
if !strings.Contains(contentParts[1].Text, testContent) {
t.Errorf("Second part text doesn't contain expected content")
}
// Verify content was cached (check attachment manager)
cachedText, err := manager.GetText(context.Background(), uploadedFile.ID)
if err != nil {
t.Fatalf("Failed to get cached text: %v", err)
}
if cachedText == "" {
t.Error("Expected content to be cached in attachment manager")
}
t.Logf("✓ Content successfully cached and reused: %d characters", len(cachedText))
}

48
agent/content/excel.go Normal file
View file

@ -0,0 +1,48 @@
package content
import (
"fmt"
"strings"
"github.com/yaoapp/gou/connector/openai"
agentContext "github.com/yaoapp/yao/agent/context"
)
// ExcelHandler handles Microsoft Excel spreadsheets
type ExcelHandler struct{}
// CanHandle checks if this handler can handle the content type
func (h *ExcelHandler) CanHandle(contentType string, fileType FileType) bool {
return fileType == FileTypeExcel ||
contentType == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" ||
contentType == "application/vnd.ms-excel" ||
strings.Contains(contentType, "excel") ||
strings.Contains(contentType, "spreadsheet")
}
// Handle processes Excel spreadsheet content
func (h *ExcelHandler) Handle(ctx *agentContext.Context, info *Info, capabilities *openai.Capabilities, uses *agentContext.Uses) (*Result, error) {
// TODO: Implement Excel handling
// 1. Extract data from .xlsx or .xls file
// 2. Convert to text format (e.g., CSV-like or structured text)
// 3. Handle multiple sheets
// 4. Return Result with formatted text
return nil, fmt.Errorf("not implemented")
}
// extractExcelText extracts text from Excel file
func extractExcelText(data []byte, contentType string) (string, error) {
// TODO: Implement Excel text extraction
// Handle both .xls (old format) and .xlsx (new format)
// Consider using libraries like:
// - github.com/360EntSecGroup-Skylar/excelize for .xlsx
// Format output as readable text or CSV
return "", fmt.Errorf("not implemented")
}
// formatExcelAsText formats Excel data as readable text
func formatExcelAsText(sheets map[string][][]string) string {
// TODO: Format multiple sheets into readable text
// Include sheet names, headers, and data
return ""
}

90
agent/content/fetch.go Normal file
View file

@ -0,0 +1,90 @@
package content
import (
"fmt"
agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/attachment"
)
// DefaultFetcher implements the Fetcher interface
type DefaultFetcher struct{}
// NewFetcher creates a new default fetcher
func NewFetcher() Fetcher {
return &DefaultFetcher{}
}
// Fetch retrieves content from HTTP URL or uploader wrapper
func (f *DefaultFetcher) Fetch(ctx *agentContext.Context, source Source, url string) (*Info, error) {
switch source {
case SourceHTTP:
return f.fetchHTTP(ctx, url)
case SourceUploader:
return f.fetchUploader(ctx, url)
default:
return nil, fmt.Errorf("unsupported source: %s", source)
}
}
// fetchHTTP fetches content from an HTTP(S) URL
func (f *DefaultFetcher) fetchHTTP(ctx *agentContext.Context, url string) (*Info, error) {
// TODO: Implement HTTP fetch logic
// 1. Download file from URL
// 2. Detect content type
// 3. Detect file type based on content type and extension
// 4. Return Info with data
return nil, fmt.Errorf("not implemented")
}
// fetchUploader fetches content from uploader wrapper (__uploader://fileid)
func (f *DefaultFetcher) fetchUploader(ctx *agentContext.Context, wrapper string) (*Info, error) {
// 1. Parse wrapper to get uploader name and file ID
uploaderName, fileID, ok := attachment.Parse(wrapper)
if !ok {
return nil, fmt.Errorf("invalid uploader wrapper format: %s", wrapper)
}
// 2. Get attachment manager
var manager attachment.FileManager
var exists bool
// Try to get manager by name
manager, exists = attachment.Managers[uploaderName]
if !exists {
return nil, fmt.Errorf("uploader '%s' not found", uploaderName)
}
// 3. Get file info
file, err := manager.Info(ctx.Context, fileID)
if err != nil {
return nil, fmt.Errorf("failed to get file info: %w", err)
}
// 4. Read file content
data, err := manager.Read(ctx.Context, fileID)
if err != nil {
return nil, fmt.Errorf("failed to read file: %w", err)
}
// 5. Return Info with data
return &Info{
Data: data,
ContentType: file.ContentType,
FileType: DetectFileType(file.ContentType, file.Filename),
}, nil
}
// parseUploaderWrapper parses uploader wrapper format: __uploader://fileid
func parseUploaderWrapper(wrapper string) (uploaderName, fileID string, err error) {
// TODO: Implement wrapper parsing
// Format: __uploader://fileid
return "", "", fmt.Errorf("not implemented")
}
// detectFileType detects file type from content type and data
func detectFileType(contentType string, data []byte) FileType {
// TODO: Implement file type detection
// Based on content type and magic bytes
return FileTypeUnknown
}

173
agent/content/image.go Normal file
View file

@ -0,0 +1,173 @@
package content
import (
"fmt"
"strings"
"github.com/yaoapp/gou/connector/openai"
agentContext "github.com/yaoapp/yao/agent/context"
)
// ImageHandler handles image content
type ImageHandler struct{}
// CanHandle checks if this handler can handle the content type
func (h *ImageHandler) CanHandle(contentType string, fileType FileType) bool {
return fileType == FileTypeImage || strings.HasPrefix(contentType, "image/")
}
// Handle processes image content
// Logic:
// 1. If model supports vision -> convert to base64 or image_url format
// 2. If model doesn't support vision -> use agent/MCP specified in uses.Vision
func (h *ImageHandler) Handle(ctx *agentContext.Context, info *Info, capabilities *openai.Capabilities, uses *agentContext.Uses) (*Result, error) {
if len(info.Data) == 0 {
return nil, fmt.Errorf("no image data to process")
}
if capabilities == nil {
return nil, fmt.Errorf("no capabilities provided")
}
// Check if model supports vision
supportsVision, visionFormat := agentContext.GetVisionSupport(capabilities)
if supportsVision {
// Model supports vision - return as image_url ContentPart
contentPart, err := h.handleWithVisionModel(ctx, info, visionFormat)
if err != nil {
return nil, fmt.Errorf("failed to handle image with vision model: %w", err)
}
return &Result{
ContentPart: contentPart,
}, nil
}
// Model doesn't support vision - use vision agent/MCP
visionTool := ""
if uses != nil && uses.Vision != "" {
visionTool = uses.Vision
}
if visionTool == "" {
return nil, fmt.Errorf("model doesn't support vision and no vision tool specified in uses.Vision")
}
// Call vision agent/MCP to extract text
text, err := h.handleWithVisionAgent(ctx, info, visionTool)
if err != nil {
return nil, fmt.Errorf("failed to handle image with vision agent/MCP: %w", err)
}
return &Result{
Text: text,
}, nil
}
// handleWithVisionModel processes image using model's vision capability
func (h *ImageHandler) handleWithVisionModel(ctx *agentContext.Context, info *Info, format agentContext.VisionFormat) (*agentContext.ContentPart, error) {
// Encode image to base64
base64Data := encodeImageBase64(info.Data, info.ContentType)
// Format according to model's vision format
switch format {
case agentContext.VisionFormatOpenAI:
// OpenAI format: image_url with data URI
return &agentContext.ContentPart{
Type: agentContext.ContentImageURL,
ImageURL: &agentContext.ImageURL{
URL: base64Data,
Detail: agentContext.DetailAuto,
},
}, nil
case agentContext.VisionFormatClaude:
// Claude format: also uses image_url but may have different handling
// For now, use the same format as OpenAI
return &agentContext.ContentPart{
Type: agentContext.ContentImageURL,
ImageURL: &agentContext.ImageURL{
URL: base64Data,
Detail: agentContext.DetailAuto,
},
}, nil
case agentContext.VisionFormatDefault, "":
// Default format (when Vision: true) - use OpenAI format
return &agentContext.ContentPart{
Type: agentContext.ContentImageURL,
ImageURL: &agentContext.ImageURL{
URL: base64Data,
Detail: agentContext.DetailAuto,
},
}, nil
default:
return nil, fmt.Errorf("unsupported vision format: %s", format)
}
}
// handleWithVisionAgent processes image using vision agent or MCP
func (h *ImageHandler) handleWithVisionAgent(ctx *agentContext.Context, info *Info, visionTool string) (string, error) {
// Parse vision tool format
// Format can be:
// - "agent_id" (call agent)
// - "mcp:server_id" (call MCP tool)
if strings.HasPrefix(visionTool, "mcp:") {
// MCP tool
serverID := strings.TrimPrefix(visionTool, "mcp:")
return h.callMCPVisionTool(ctx, serverID, info)
}
// Agent call
return h.callVisionAgent(ctx, visionTool, info)
}
// callVisionAgent calls a vision agent to describe the image
func (h *ImageHandler) callVisionAgent(ctx *agentContext.Context, agentID string, info *Info) (string, error) {
// Prepare message with image
base64Data := EncodeToBase64DataURI(info.Data, info.ContentType)
message := agentContext.Message{
Role: agentContext.RoleUser,
Content: []agentContext.ContentPart{
{
Type: agentContext.ContentText,
Text: "Please describe this image in detail.",
},
{
Type: agentContext.ContentImageURL,
ImageURL: &agentContext.ImageURL{
URL: base64Data,
Detail: agentContext.DetailAuto,
},
},
},
}
return CallAgent(ctx, agentID, message)
}
// callMCPVisionTool calls an MCP vision tool to describe the image
func (h *ImageHandler) callMCPVisionTool(ctx *agentContext.Context, serverID string, info *Info) (string, error) {
// Prepare base64 encoded image for MCP tool
base64Data := EncodeToBase64DataURI(info.Data, info.ContentType)
// Prepare arguments for MCP tool
arguments := map[string]interface{}{
"image": base64Data,
"content_type": info.ContentType,
}
// Call MCP tool (typically "describe_image" or similar)
return CallMCPTool(ctx, serverID, "describe_image", arguments)
}
// encodeImageBase64 encodes image data to base64 with data URI prefix
func encodeImageBase64(data []byte, contentType string) string {
// Use the common function
if contentType == "" {
contentType = "image/png" // default for images
}
return EncodeToBase64DataURI(data, contentType)
}

264
agent/content/image_test.go Normal file
View file

@ -0,0 +1,264 @@
package content
import (
stdContext "context"
"encoding/base64"
"os"
"strings"
"testing"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/gou/plan"
agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/openapi/oauth/types"
"github.com/yaoapp/yao/test"
)
func TestMain(m *testing.M) {
// Setup test environment
test.Prepare(nil, config.Conf)
defer test.Clean()
// Run tests
code := m.Run()
os.Exit(code)
}
// newTestContext creates a Context for testing with commonly used fields pre-populated
func newTestContext(capabilities *openai.Capabilities) *agentContext.Context {
return &agentContext.Context{
Context: stdContext.Background(),
Space: plan.NewMemorySharedSpace(),
ChatID: "test-chat",
AssistantID: "test-assistant",
Locale: "en-us",
Theme: "light",
Client: agentContext.Client{
Type: "web",
UserAgent: "TestAgent/1.0",
IP: "127.0.0.1",
},
Referer: agentContext.RefererAPI,
Accept: agentContext.AcceptWebCUI,
Route: "",
Metadata: make(map[string]interface{}),
Capabilities: capabilities,
Authorized: &types.AuthorizedInfo{
Subject: "test-user",
ClientID: "test-client-id",
UserID: "test-user-123",
TeamID: "test-team-456",
TenantID: "test-tenant-789",
},
}
}
func TestImageHandler_CanHandle(t *testing.T) {
handler := &ImageHandler{}
tests := []struct {
name string
contentType string
fileType FileType
want bool
}{
{"PNG image", "image/png", FileTypeImage, true},
{"JPEG image", "image/jpeg", FileTypeImage, true},
{"GIF image", "image/gif", FileTypeImage, true},
{"WebP image", "image/webp", FileTypeImage, true},
{"Text (should not handle)", "text/plain", FileTypeText, false},
{"PDF (should not handle)", "application/pdf", FileTypePDF, false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := handler.CanHandle(tt.contentType, tt.fileType)
if got != tt.want {
t.Errorf("CanHandle(%q, %q) = %v, want %v", tt.contentType, tt.fileType, got, tt.want)
}
})
}
}
func TestImageHandler_Handle_WithVisionSupport(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
handler := &ImageHandler{}
// Create a simple test image (1x1 red PNG)
pngData := createTestPNG()
// Create capabilities with vision support
capabilities := &openai.Capabilities{
Vision: "openai", // Vision is enabled with OpenAI format
}
// Create test context
ctx := newTestContext(capabilities)
info := &Info{
FileType: FileTypeImage,
ContentType: "image/png",
Data: pngData,
}
result, err := handler.Handle(ctx, info, capabilities, nil)
if err != nil {
t.Fatalf("Handle() error = %v", err)
}
if result == nil {
t.Fatal("Expected non-nil result")
}
if result.ContentPart == nil {
t.Fatal("Expected ContentPart for vision-supported model")
}
if result.ContentPart.Type != agentContext.ContentImageURL {
t.Errorf("Expected ContentPart type = %v, got %v", agentContext.ContentImageURL, result.ContentPart.Type)
}
if result.ContentPart.ImageURL == nil {
t.Fatal("Expected ImageURL to be set")
}
// Verify base64 encoding
if result.ContentPart.ImageURL.URL == "" {
t.Error("Expected non-empty URL")
}
// Should be data URI format
if len(result.ContentPart.ImageURL.URL) < 20 {
t.Error("Expected data URI to be longer")
}
}
func TestImageHandler_Handle_WithoutVisionSupport(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
handler := &ImageHandler{}
// Create a simple test image
pngData := createTestPNG()
// Create capabilities WITHOUT vision support
capabilities := &openai.Capabilities{
Vision: nil, // No vision support
}
// Create test context
ctx := newTestContext(capabilities)
info := &Info{
FileType: FileTypeImage,
ContentType: "image/png",
Data: pngData,
}
// Should return error because no vision support and no tool
_, err := handler.Handle(ctx, info, capabilities, nil)
if err == nil {
t.Error("Expected error when no vision support and no tool specified")
}
}
func TestImageHandler_Handle_EmptyData(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
handler := &ImageHandler{}
capabilities := &openai.Capabilities{
Vision: "openai",
}
// Create test context
ctx := newTestContext(capabilities)
info := &Info{
FileType: FileTypeImage,
ContentType: "image/png",
Data: []byte{}, // Empty data
}
_, err := handler.Handle(ctx, info, capabilities, nil)
if err == nil {
t.Error("Expected error for empty image data")
}
}
func TestEncodeImageBase64(t *testing.T) {
tests := []struct {
name string
data []byte
contentType string
wantPrefix string
}{
{
name: "PNG image",
data: []byte{0x89, 0x50, 0x4E, 0x47}, // PNG magic number
contentType: "image/png",
wantPrefix: "data:image/png;base64,",
},
{
name: "JPEG image",
data: []byte{0xFF, 0xD8, 0xFF}, // JPEG magic number
contentType: "image/jpeg",
wantPrefix: "data:image/jpeg;base64,",
},
{
name: "Empty content type defaults to PNG",
data: []byte{0x01, 0x02, 0x03},
contentType: "",
wantPrefix: "data:image/png;base64,",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := encodeImageBase64(tt.data, tt.contentType)
// Check prefix
if !strings.HasPrefix(result, tt.wantPrefix) {
t.Errorf("Expected prefix %q, got %q", tt.wantPrefix, result[:len(tt.wantPrefix)])
}
// Verify base64 encoding by decoding
base64Part := result[len(tt.wantPrefix):]
decoded, err := base64.StdEncoding.DecodeString(base64Part)
if err != nil {
t.Errorf("Failed to decode base64: %v", err)
}
// Verify decoded data matches original
if len(decoded) != len(tt.data) {
t.Errorf("Decoded length = %d, want %d", len(decoded), len(tt.data))
}
for i := range decoded {
if decoded[i] != tt.data[i] {
t.Errorf("Decoded byte[%d] = %x, want %x", i, decoded[i], tt.data[i])
}
}
})
}
}
// createTestPNG creates a minimal valid PNG image (1x1 red pixel)
func createTestPNG() []byte {
// This is a minimal valid 1x1 red PNG image
return []byte{
0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, // PNG signature
0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, 0x52, // IHDR chunk
0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, // 1x1 dimensions
0x08, 0x02, 0x00, 0x00, 0x00, 0x90, 0x77, 0x53,
0xDE, 0x00, 0x00, 0x00, 0x0C, 0x49, 0x44, 0x41, // IDAT chunk
0x54, 0x08, 0xD7, 0x63, 0xF8, 0xCF, 0xC0, 0x00,
0x00, 0x03, 0x01, 0x01, 0x00, 0x18, 0xDD, 0x8D,
0xB4, 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4E, // IEND chunk
0x44, 0xAE, 0x42, 0x60, 0x82,
}
}

25
agent/content/interfaces.go Normal file
View file

@ -0,0 +1,25 @@
package content
import (
"github.com/yaoapp/gou/connector/openai"
agentContext "github.com/yaoapp/yao/agent/context"
)
// Handler defines the interface for handling different content types
// Converts content (images, documents, etc.) to text or standard formats
type Handler interface {
// CanHandle checks if this handler can handle the given content type
CanHandle(contentType string, fileType FileType) bool
// Handle converts the content and returns processed result
// ctx: agent context (passed from Vision function)
// capabilities: model capabilities (for vision/audio support detection)
// uses: configuration for external tools (agents/MCP servers)
Handle(ctx *agentContext.Context, info *Info, capabilities *openai.Capabilities, uses *agentContext.Uses) (*Result, error)
}
// Fetcher defines the interface for fetching content from different sources
type Fetcher interface {
// Fetch retrieves content from a URL or file ID
Fetch(ctx *agentContext.Context, source Source, url string) (*Info, error)
}

50
agent/content/pdf.go Normal file
View file

@ -0,0 +1,50 @@
package content
import (
"fmt"
"strings"
"github.com/yaoapp/gou/connector/openai"
agentContext "github.com/yaoapp/yao/agent/context"
)
// PDFHandler handles PDF documents
type PDFHandler struct{}
// CanHandle checks if this handler can handle the content type
func (h *PDFHandler) CanHandle(contentType string, fileType FileType) bool {
return fileType == FileTypePDF ||
contentType == "application/pdf" ||
strings.Contains(contentType, "pdf")
}
// Handle processes PDF content
// Logic:
// 1. Check if uses.Vision is specified and supports PDF
// 2. If yes, use vision tool to handle PDF (images + text)
// 3. If no, extract text directly from PDF
func (h *PDFHandler) Handle(ctx *agentContext.Context, info *Info, capabilities *openai.Capabilities, uses *agentContext.Uses) (*Result, error) {
// TODO: Implement PDF handling
// 1. Check if vision tool supports PDF
// 2. If yes:
// - Call vision tool to handle PDF (handles both text and images)
// 3. If no:
// - Extract text from PDF using default library
// 4. Return Result with extracted text
return nil, fmt.Errorf("not implemented")
}
// extractPDFText extracts text content from PDF
func extractPDFText(data []byte) (string, error) {
// TODO: Implement PDF text extraction
// Use a PDF library to extract text
// Consider preserving layout/structure
return "", fmt.Errorf("not implemented")
}
// handleWithVisionTool processes PDF using vision tool (for PDFs with images)
func handleWithVisionTool(ctx *agentContext.Context, data []byte, visionTool string) (string, error) {
// TODO: Implement vision tool PDF processing
// Some vision tools can handle PDF directly and extract both text and images
return "", fmt.Errorf("not implemented")
}

47
agent/content/registry.go Normal file
View file

@ -0,0 +1,47 @@
package content
import (
"fmt"
"github.com/yaoapp/gou/connector/openai"
agentContext "github.com/yaoapp/yao/agent/context"
)
// Registry holds all registered content handlers
type Registry struct {
handlers []Handler
}
// NewRegistry creates a new handler registry with default handlers
func NewRegistry() *Registry {
return &Registry{
handlers: []Handler{
&ImageHandler{},
&AudioHandler{},
&PDFHandler{},
&WordHandler{},
&ExcelHandler{},
&TextHandler{},
},
}
}
// GetHandler finds the appropriate handler for the given content
func (r *Registry) GetHandler(contentType string, fileType FileType) Handler {
for _, handler := range r.handlers {
if handler.CanHandle(contentType, fileType) {
return handler
}
}
return nil
}
// Handle processes content using the appropriate handler
func (r *Registry) Handle(ctx *agentContext.Context, info *Info, capabilities *openai.Capabilities, uses *agentContext.Uses) (*Result, error) {
handler := r.GetHandler(info.ContentType, info.FileType)
if handler == nil {
return nil, fmt.Errorf("no handler found for content type: %s, file type: %s", info.ContentType, info.FileType)
}
return handler.Handle(ctx, info, capabilities, uses)
}

123
agent/content/text.go Normal file
View file

@ -0,0 +1,123 @@
package content
import (
"encoding/json"
"fmt"
"strings"
"github.com/yaoapp/gou/connector/openai"
agentContext "github.com/yaoapp/yao/agent/context"
)
// TextHandler handles plain text, code files, CSV, JSON, XML, Markdown, etc.
type TextHandler struct{}
// CanHandle checks if this handler can handle the content type
func (h *TextHandler) CanHandle(contentType string, fileType FileType) bool {
// Handle explicit text file types
if fileType == FileTypeText || fileType == FileTypeCSV || fileType == FileTypeJSON {
return true
}
// Handle text-based MIME types
if strings.HasPrefix(contentType, "text/") {
return true
}
// Handle common text-based content types
textContentTypes := []string{
"application/json",
"application/xml",
"application/javascript",
"application/typescript",
"application/x-yaml",
"application/yaml",
"application/toml",
"application/x-sh",
"application/x-python",
"application/x-ruby",
"application/x-perl",
}
for _, ct := range textContentTypes {
if contentType == ct || strings.Contains(contentType, ct) {
return true
}
}
return false
}
// Handle processes text content
func (h *TextHandler) Handle(ctx *agentContext.Context, info *Info, capabilities *openai.Capabilities, uses *agentContext.Uses) (*Result, error) {
if len(info.Data) == 0 {
return nil, fmt.Errorf("no data to process")
}
var text string
var err error
// Handle different text formats
switch {
case info.FileType == FileTypeCSV || strings.Contains(info.ContentType, "csv"):
// Format CSV as readable text (for now, just return as-is, can enhance later)
text = string(info.Data)
case info.FileType == FileTypeJSON ||
info.ContentType == "application/json" ||
strings.Contains(info.ContentType, "json"):
// Pretty print JSON
text, err = formatJSONAsText(info.Data)
if err != nil {
// If JSON parsing fails, return raw text
text = string(info.Data)
}
case info.ContentType == "application/xml" ||
strings.Contains(info.ContentType, "xml"):
// For now, return XML as-is (can enhance formatting later)
text = string(info.Data)
default:
// Plain text, code files, markdown, etc.
text, err = readTextContent(info.Data, info.ContentType)
if err != nil {
return nil, fmt.Errorf("failed to read text content: %w", err)
}
}
return &Result{
Text: text,
}, nil
}
// readTextContent reads text content from data
func readTextContent(data []byte, contentType string) (string, error) {
// For now, assume UTF-8 encoding
// TODO: Add encoding detection if needed (e.g., using golang.org/x/text/encoding)
return string(data), nil
}
// formatCSVAsText formats CSV data as readable text
func formatCSVAsText(data []byte) (string, error) {
// TODO: Parse CSV and format as readable table
// Consider using encoding/csv package
// For now, just return as-is
return string(data), nil
}
// formatJSONAsText formats JSON data as readable text
func formatJSONAsText(data []byte) (string, error) {
// Pretty print JSON with indentation
var obj interface{}
if err := json.Unmarshal(data, &obj); err != nil {
return "", err
}
pretty, err := json.MarshalIndent(obj, "", " ")
if err != nil {
return "", err
}
return string(pretty), nil
}

152
agent/content/text_test.go Normal file
View file

@ -0,0 +1,152 @@
package content
import (
"testing"
)
func TestTextHandler_CanHandle(t *testing.T) {
handler := &TextHandler{}
tests := []struct {
name string
contentType string
fileType FileType
want bool
}{
{"Plain text", "text/plain", FileTypeText, true},
{"Markdown", "text/markdown", FileTypeText, true},
{"HTML", "text/html", FileTypeText, true},
{"JSON", "application/json", FileTypeJSON, true},
{"JavaScript", "application/javascript", FileTypeText, true},
{"TypeScript", "application/typescript", FileTypeText, true},
{"YAML", "application/yaml", FileTypeText, true},
{"CSV", "text/csv", FileTypeCSV, true},
{"XML", "application/xml", FileTypeText, true},
{"PDF (should not handle)", "application/pdf", FileTypePDF, false},
{"Image (should not handle)", "image/png", FileTypeImage, false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := handler.CanHandle(tt.contentType, tt.fileType)
if got != tt.want {
t.Errorf("CanHandle(%q, %q) = %v, want %v", tt.contentType, tt.fileType, got, tt.want)
}
})
}
}
func TestTextHandler_Handle(t *testing.T) {
handler := &TextHandler{}
tests := []struct {
name string
info *Info
wantErr bool
checkResult func(*testing.T, *Result)
}{
{
name: "Plain text",
info: &Info{
FileType: FileTypeText,
ContentType: "text/plain",
Data: []byte("Hello, World!"),
},
wantErr: false,
checkResult: func(t *testing.T, r *Result) {
if r.Text != "Hello, World!" {
t.Errorf("Expected 'Hello, World!', got %q", r.Text)
}
},
},
{
name: "JSON with pretty print",
info: &Info{
FileType: FileTypeJSON,
ContentType: "application/json",
Data: []byte(`{"name":"test","value":123}`),
},
wantErr: false,
checkResult: func(t *testing.T, r *Result) {
// Should be pretty printed
if len(r.Text) <= len(`{"name":"test","value":123}`) {
t.Errorf("JSON should be pretty printed, got: %q", r.Text)
}
},
},
{
name: "Code file (Go)",
info: &Info{
FileType: FileTypeText,
ContentType: "text/plain",
Data: []byte("package main\n\nfunc main() {\n\tprintln(\"Hello\")\n}"),
},
wantErr: false,
checkResult: func(t *testing.T, r *Result) {
if r.Text == "" {
t.Error("Expected non-empty text for Go code")
}
},
},
{
name: "Empty data",
info: &Info{
FileType: FileTypeText,
ContentType: "text/plain",
Data: []byte{},
},
wantErr: true,
},
}
// Create test context
testCtx := newTestContext(nil)
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result, err := handler.Handle(testCtx, tt.info, nil, nil)
if (err != nil) != tt.wantErr {
t.Errorf("Handle() error = %v, wantErr %v", err, tt.wantErr)
return
}
if !tt.wantErr && tt.checkResult != nil {
tt.checkResult(t, result)
}
})
}
}
func TestDetectFileType(t *testing.T) {
tests := []struct {
name string
contentType string
filename string
want FileType
}{
{"Go file", "text/plain", "main.go", FileTypeText},
{"Python file", "text/plain", "script.py", FileTypeText},
{"JavaScript file", "application/javascript", "app.js", FileTypeText},
{"TypeScript file", "text/plain", "index.ts", FileTypeText},
{"Markdown file", "text/markdown", "README.md", FileTypeText},
{"JSON file", "application/json", "config.json", FileTypeJSON},
{"YAML file", "text/plain", "config.yml", FileTypeText},
{"PDF file", "application/pdf", "document.pdf", FileTypePDF},
{"Image file", "image/png", "photo.png", FileTypeImage},
{"CSV file", "text/csv", "data.csv", FileTypeCSV},
{"XML file", "application/xml", "config.xml", FileTypeXML},
{"Shell script", "text/plain", "script.sh", FileTypeText},
{"Dockerfile", "text/plain", "Dockerfile", FileTypeText},
{"gitignore", "text/plain", ".gitignore", FileTypeText},
{"HTML", "text/html", "index.html", FileTypeText},
{"CSS", "text/css", "styles.css", FileTypeText},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := DetectFileType(tt.contentType, tt.filename)
if got != tt.want {
t.Errorf("DetectFileType(%q, %q) = %v, want %v", tt.contentType, tt.filename, got, tt.want)
}
})
}
}

188
agent/content/tools.go Normal file
View file

@ -0,0 +1,188 @@
package content
import (
"context"
"encoding/base64"
"fmt"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/mcp"
"github.com/yaoapp/kun/log"
agentContext "github.com/yaoapp/yao/agent/context"
)
// AgentCaller interface for calling agents (to avoid circular dependency)
type AgentCaller interface {
Stream(ctx *agentContext.Context, messages []agentContext.Message, options ...*agentContext.Options) (interface{}, error)
}
// AgentGetterFunc is a function type that gets an agent by ID
var AgentGetterFunc func(agentID string) (AgentCaller, error)
// CallAgent calls an agent to process content (vision, audio, etc.)
// This is a generic function that can be used by any handler
func CallAgent(ctx *agentContext.Context, agentID string, message agentContext.Message) (string, error) {
if AgentGetterFunc == nil {
return "", fmt.Errorf("AgentGetterFunc not initialized")
}
// Load the agent by ID using the injected function
agent, err := AgentGetterFunc(agentID)
if err != nil {
return "", fmt.Errorf("failed to load agent %s: %w", agentID, err)
}
// Call the agent with the message
messages := []agentContext.Message{message}
// Note: Connector is now in Options (call-level parameter), not Context
// For A2A calls, skip history and output (we only need the response data)
opts := &agentContext.Options{Skip: &agentContext.Skip{History: true, Output: true}} // Skip history and output
response, err := agent.Stream(ctx, messages, opts)
if err != nil {
return "", fmt.Errorf("failed to call agent %s: %w", agentID, err)
}
// Extract text from agent response
// Two formats are supported:
// 1. Custom Hook response (from Next hook)
// 2. Standard Agent Stream response (LLM completion)
return extractTextFromAgentResponse(response)
}
// extractTextFromAgentResponse extracts text from agent response
// Handles two response formats:
// 1. Custom Hook response: if it's a string, return directly; otherwise JSON stringify
// 2. Standard response: extract from completion.content
func extractTextFromAgentResponse(response interface{}) (string, error) {
if response == nil {
return "", fmt.Errorf("agent returned nil response")
}
// Try to parse as standard response format (has "completion" field with LLM result)
if responseMap, ok := response.(map[string]interface{}); ok {
// Check for completion field (standard LLM response)
if completion, hasCompletion := responseMap["completion"]; hasCompletion {
if completionMap, ok := completion.(map[string]interface{}); ok {
// Extract content from completion
if content, hasContent := completionMap["content"]; hasContent {
// Content can be string or structured
switch v := content.(type) {
case string:
return v, nil
case []interface{}:
// Handle multimodal content array
var text string
for _, part := range v {
if partMap, ok := part.(map[string]interface{}); ok {
if partType, _ := partMap["type"].(string); partType == "text" {
if textContent, ok := partMap["text"].(string); ok {
text += textContent
}
}
}
}
if text != "" {
return text, nil
}
}
}
}
}
// Check for data field (custom hook response with data wrapper)
if data, hasData := responseMap["data"]; hasData {
// If data is a string, return directly
if dataStr, ok := data.(string); ok {
return dataStr, nil
}
// Otherwise, JSON stringify
jsonBytes, err := jsoniter.Marshal(data)
if err != nil {
return "", fmt.Errorf("failed to serialize hook data response: %w", err)
}
return string(jsonBytes), nil
}
// If the map itself looks like content, try to extract
// This handles cases where the response is the content directly
if content, hasContent := responseMap["content"]; hasContent {
if contentStr, ok := content.(string); ok {
return contentStr, nil
}
}
}
// Custom Hook response: if it's a plain string, return directly
if responseStr, ok := response.(string); ok {
return responseStr, nil
}
// Otherwise, JSON stringify the response
jsonBytes, err := jsoniter.Marshal(response)
if err != nil {
return "", fmt.Errorf("failed to serialize agent response: %w", err)
}
return string(jsonBytes), nil
}
// CallMCPTool calls an MCP tool to process content
// This is a generic function that can be used by any handler
func CallMCPTool(ctx *agentContext.Context, serverID string, toolName string, arguments map[string]interface{}) (string, error) {
// Get MCP context for cancellation/timeout control
mcpCtx := ctx.Context
if mcpCtx == nil {
mcpCtx = context.Background()
}
// Get MCP client
client, err := mcp.Select(serverID)
if err != nil {
return "", fmt.Errorf("failed to select MCP client '%s': %w", serverID, err)
}
// Call the tool
log.Trace("[Content] Calling MCP tool: %s (server: %s)", toolName, serverID)
callResult, err := client.CallTool(mcpCtx, toolName, arguments)
if err != nil {
return "", fmt.Errorf("MCP tool call failed: %w", err)
}
// Check if result is an error
if callResult.IsError {
return "", fmt.Errorf("MCP tool returned error: %v", callResult.Content)
}
// Extract text content from result
// callResult.Content is []ToolContent
var text string
for _, content := range callResult.Content {
if content.Type == "text" {
text += content.Text
}
// Can also handle other types like image, resource if needed
}
if text == "" {
// If no text content found, return error
return "", fmt.Errorf("MCP tool returned no text content")
}
return text, nil
}
// EncodeToBase64DataURI encodes data to base64 with data URI prefix
// This is useful for encoding images, audio, or other binary data
func EncodeToBase64DataURI(data []byte, contentType string) string {
// Ensure we have a valid content type
if contentType == "" {
contentType = "application/octet-stream" // default
}
// Encode to base64
encoded := base64.StdEncoding.EncodeToString(data)
// Return data URI format
return fmt.Sprintf("data:%s;base64,%s", contentType, encoded)
}

261
agent/content/types.go Normal file
View file

@ -0,0 +1,261 @@
package content
import "github.com/yaoapp/yao/agent/context"
// FileType represents the type of file content
type FileType string
const (
// Image types
FileTypeImage FileType = "image"
// Audio types
FileTypeAudio FileType = "audio"
// Document types
FileTypeText FileType = "text"
FileTypePDF FileType = "pdf"
FileTypeWord FileType = "word"
FileTypeExcel FileType = "excel"
FileTypePPT FileType = "ppt"
FileTypeCSV FileType = "csv"
// Data types
FileTypeJSON FileType = "json"
FileTypeXML FileType = "xml"
// Binary
FileTypeBinary FileType = "binary"
// Other
FileTypeUnknown FileType = "unknown"
)
// Source represents where the content comes from
type Source string
const (
SourceHTTP Source = "http" // HTTP(S) URL
SourceUploader Source = "uploader" // Uploader wrapper: __uploader://fileid
SourceBase64 Source = "base64" // Base64 encoded data
SourceLocal Source = "local" // Local file path
)
// Result represents the result of content handling
type Result struct {
Text string // Extracted text content
ContentPart *context.ContentPart // Processed ContentPart (for model input)
Metadata map[string]interface{} // Additional metadata
Error error // Error if handling failed
}
// Info holds information about a content part to be handled
type Info struct {
Source Source // Where the content comes from
FileType FileType // Type of the file
ContentType string // MIME content type
URL string // Original URL or file ID
Data []byte // File data (if already fetched)
// For uploader wrapper
UploaderName string // Uploader name from wrapper
FileID string // File ID from wrapper
}
// DetectFileType detects file type from content type, filename, and file extension
func DetectFileType(contentType, filename string) FileType {
// Check by content type first
switch {
case contentType == "application/pdf":
return FileTypePDF
case contentType == "application/json":
return FileTypeJSON
case contentType == "application/xml" || contentType == "text/xml":
return FileTypeXML
case contentType == "text/csv":
return FileTypeCSV
case contentType == "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
contentType == "application/msword":
return FileTypeWord
case contentType == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
contentType == "application/vnd.ms-excel":
return FileTypeExcel
case contentType == "application/vnd.openxmlformats-officedocument.presentationml.presentation",
contentType == "application/vnd.ms-powerpoint":
return FileTypePPT
}
// Check image types
if isImageContentType(contentType) {
return FileTypeImage
}
// Check audio types
if isAudioContentType(contentType) {
return FileTypeAudio
}
// Check text types
if isTextContentType(contentType) {
return FileTypeText
}
// If content type doesn't help, check file extension
if filename != "" {
if ext := getFileExtension(filename); ext != "" {
return detectTypeByExtension(ext)
}
}
return FileTypeUnknown
}
// isImageContentType checks if content type is an image
func isImageContentType(contentType string) bool {
return contentType != "" &&
(contentType == "image/png" ||
contentType == "image/jpeg" ||
contentType == "image/jpg" ||
contentType == "image/gif" ||
contentType == "image/webp" ||
contentType == "image/svg+xml" ||
contentType == "image/bmp")
}
// isAudioContentType checks if content type is audio
func isAudioContentType(contentType string) bool {
return contentType != "" &&
(contentType == "audio/mpeg" ||
contentType == "audio/mp3" ||
contentType == "audio/wav" ||
contentType == "audio/ogg" ||
contentType == "audio/flac" ||
contentType == "audio/aac")
}
// isTextContentType checks if content type is text-based
func isTextContentType(contentType string) bool {
if contentType == "" {
return false
}
// Common text MIME types
textTypes := []string{
"text/plain",
"text/html",
"text/css",
"text/javascript",
"text/markdown",
"text/x-markdown",
"application/javascript",
"application/typescript",
"application/x-yaml",
"application/yaml",
"application/toml",
"application/x-sh",
"application/x-python",
"application/x-ruby",
"application/x-perl",
"application/x-php",
"application/x-go",
}
for _, t := range textTypes {
if contentType == t {
return true
}
}
return false
}
// getFileExtension extracts file extension from filename (without dot)
func getFileExtension(filename string) string {
for i := len(filename) - 1; i >= 0; i-- {
if filename[i] == '.' {
return filename[i+1:]
}
if filename[i] == '/' || filename[i] == '\\' {
break
}
}
return ""
}
// detectTypeByExtension detects file type by file extension
func detectTypeByExtension(ext string) FileType {
// Normalize to lowercase
ext = toLower(ext)
// Image extensions
imageExts := []string{"png", "jpg", "jpeg", "gif", "webp", "svg", "bmp", "ico"}
for _, e := range imageExts {
if ext == e {
return FileTypeImage
}
}
// Audio extensions
audioExts := []string{"mp3", "wav", "ogg", "flac", "aac", "m4a"}
for _, e := range audioExts {
if ext == e {
return FileTypeAudio
}
}
// Document extensions
switch ext {
case "pdf":
return FileTypePDF
case "doc", "docx":
return FileTypeWord
case "xls", "xlsx":
return FileTypeExcel
case "ppt", "pptx":
return FileTypePPT
case "csv":
return FileTypeCSV
case "json":
return FileTypeJSON
case "xml":
return FileTypeXML
}
// Code and text file extensions (very comprehensive list)
textExts := []string{
"txt", "text", "md", "markdown", "rst",
// Programming languages
"go", "py", "js", "ts", "jsx", "tsx", "java", "c", "cpp", "h", "hpp",
"cs", "rb", "php", "pl", "swift", "kt", "rs", "scala", "clj",
// Web
"html", "htm", "css", "scss", "sass", "less",
// Config
"yaml", "yml", "toml", "ini", "conf", "config",
// Shell
"sh", "bash", "zsh", "fish",
// Data
"sql", "graphql", "proto",
// Others
"log", "gitignore", "env", "dockerfile",
}
for _, e := range textExts {
if ext == e {
return FileTypeText
}
}
return FileTypeUnknown
}
// toLower converts ASCII string to lowercase (simple version)
func toLower(s string) string {
result := make([]byte, len(s))
for i := 0; i < len(s); i++ {
c := s[i]
if c >= 'A' && c <= 'Z' {
c += 'a' - 'A'
}
result[i] = c
}
return string(result)
}

39
agent/content/word.go Normal file
View file

@ -0,0 +1,39 @@
package content
import (
"fmt"
"strings"
"github.com/yaoapp/gou/connector/openai"
agentContext "github.com/yaoapp/yao/agent/context"
)
// WordHandler handles Microsoft Word documents
type WordHandler struct{}
// CanHandle checks if this handler can handle the content type
func (h *WordHandler) CanHandle(contentType string, fileType FileType) bool {
return fileType == FileTypeWord ||
contentType == "application/vnd.openxmlformats-officedocument.wordprocessingml.document" ||
contentType == "application/msword" ||
strings.Contains(contentType, "word")
}
// Handle processes Word document content
func (h *WordHandler) Handle(ctx *agentContext.Context, info *Info, capabilities *openai.Capabilities, uses *agentContext.Uses) (*Result, error) {
// TODO: Implement Word document handling
// 1. Extract text from .docx or .doc file
// 2. Preserve formatting information if needed
// 3. Return Result with extracted text
return nil, fmt.Errorf("not implemented")
}
// extractWordText extracts text from Word document
func extractWordText(data []byte, contentType string) (string, error) {
// TODO: Implement Word text extraction
// Handle both .doc (old format) and .docx (new format)
// Consider using libraries like:
// - github.com/unidoc/unioffice for .docx
// - Other libraries for .doc
return "", fmt.Errorf("not implemented")
}

17
agent/context/context.go
View file

@ -259,23 +259,6 @@ func (ctx *Context) Map() map[string]interface{} {
if ctx.AssistantID != "" {
data["assistant_id"] = ctx.AssistantID
}
if ctx.Connector != "" {
data["connector"] = ctx.Connector
}
if ctx.Search != nil {
data["search"] = *ctx.Search
}
// Arguments for call
if len(ctx.Args) > 0 {
data["args"] = ctx.Args
}
if ctx.Retry {
data["retry"] = ctx.Retry
}
if ctx.RetryTimes > 0 {
data["retry_times"] = ctx.RetryTimes
}
// Locale information
if ctx.Locale != "" {

3
agent/context/context_test.go
View file

@ -179,7 +179,7 @@ func TestGetCompletionRequest(t *testing.T) {
c.Request = req
// Call GetCompletionRequest
completionReq, ctx, err := GetCompletionRequest(c, cache)
completionReq, ctx, opts, err := GetCompletionRequest(c, cache)
if tt.expectError {
assert.Error(t, err)
@ -189,6 +189,7 @@ func TestGetCompletionRequest(t *testing.T) {
assert.NoError(t, err)
assert.NotNil(t, completionReq)
assert.NotNil(t, ctx)
assert.NotNil(t, opts)
// Verify CompletionRequest
assert.Equal(t, tt.expectedModel, completionReq.Model)

1
agent/context/interrupt_test.go
View file

@ -19,7 +19,6 @@ func newTestContextWithInterrupt(chatID, assistantID string) *Context {
Space: plan.NewMemorySharedSpace(),
ChatID: chatID,
AssistantID: assistantID,
Connector: "",
Locale: "en-us",
Theme: "light",
Client: Client{

16
agent/context/jsapi.go
View file

@ -36,13 +36,6 @@ func (ctx *Context) NewObject(v8ctx *v8go.Context) (*v8go.Value, error) {
// Set primitive fields in template
jsObject.Set("chat_id", ctx.ChatID)
jsObject.Set("assistant_id", ctx.AssistantID)
jsObject.Set("connector", ctx.Connector)
if ctx.Search != nil {
jsObject.Set("search", *ctx.Search)
}
jsObject.Set("retry", ctx.Retry)
jsObject.Set("retry_times", uint32(ctx.RetryTimes))
jsObject.Set("locale", ctx.Locale)
jsObject.Set("theme", ctx.Theme)
jsObject.Set("referer", ctx.Referer)
@ -97,15 +90,6 @@ func (ctx *Context) NewObject(v8ctx *v8go.Context) (*v8go.Value, error) {
}
// Set complex objects (maps, arrays) after instance creation using bridge
// Args array
if ctx.Args != nil {
argsVal, err := bridge.JsValue(v8ctx, ctx.Args)
if err == nil {
obj.Set("args", argsVal)
argsVal.Release() // Release Go-side Persistent handle, V8 internal reference remains
}
}
// Client object
clientData := map[string]interface{}{
"type": ctx.Client.Type,

33
agent/context/jsapi_mcp_test.go
View file

@ -22,8 +22,9 @@ func TestMCPListResources(t *testing.T) {
AssistantID: "test-assistant-id",
Locale: "en",
Context: stdContext.Background(),
Referer: context.RefererAPI,
}
stack, _, _ := context.EnterStack(ctx, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(ctx, "test-assistant", &context.Options{})
ctx.Stack = stack
res, err := v8.Call(v8.CallOptions{}, `
@ -66,8 +67,9 @@ func TestMCPReadResource(t *testing.T) {
AssistantID: "test-assistant-id",
Locale: "en",
Context: stdContext.Background(),
Referer: context.RefererAPI,
}
stack, _, _ := context.EnterStack(ctx, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(ctx, "test-assistant", &context.Options{})
ctx.Stack = stack
res, err := v8.Call(v8.CallOptions{}, `
@ -108,8 +110,9 @@ func TestMCPListTools(t *testing.T) {
AssistantID: "test-assistant-id",
Locale: "en",
Context: stdContext.Background(),
Referer: context.RefererAPI,
}
stack, _, _ := context.EnterStack(ctx, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(ctx, "test-assistant", &context.Options{})
ctx.Stack = stack
res, err := v8.Call(v8.CallOptions{}, `
@ -154,8 +157,9 @@ func TestMCPCallTool(t *testing.T) {
AssistantID: "test-assistant-id",
Locale: "en",
Context: stdContext.Background(),
Referer: context.RefererAPI,
}
stack, _, _ := context.EnterStack(ctx, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(ctx, "test-assistant", &context.Options{})
ctx.Stack = stack
res, err := v8.Call(v8.CallOptions{}, `
@ -196,8 +200,9 @@ func TestMCPCallTools(t *testing.T) {
AssistantID: "test-assistant-id",
Locale: "en",
Context: stdContext.Background(),
Referer: context.RefererAPI,
}
stack, _, _ := context.EnterStack(ctx, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(ctx, "test-assistant", &context.Options{})
ctx.Stack = stack
res, err := v8.Call(v8.CallOptions{}, `
@ -243,8 +248,9 @@ func TestMCPCallToolsParallel(t *testing.T) {
AssistantID: "test-assistant-id",
Locale: "en",
Context: stdContext.Background(),
Referer: context.RefererAPI,
}
stack, _, _ := context.EnterStack(ctx, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(ctx, "test-assistant", &context.Options{})
ctx.Stack = stack
res, err := v8.Call(v8.CallOptions{}, `
@ -290,8 +296,9 @@ func TestMCPListPrompts(t *testing.T) {
AssistantID: "test-assistant-id",
Locale: "en",
Context: stdContext.Background(),
Referer: context.RefererAPI,
}
stack, _, _ := context.EnterStack(ctx, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(ctx, "test-assistant", &context.Options{})
ctx.Stack = stack
res, err := v8.Call(v8.CallOptions{}, `
@ -334,8 +341,9 @@ func TestMCPGetPrompt(t *testing.T) {
AssistantID: "test-assistant-id",
Locale: "en",
Context: stdContext.Background(),
Referer: context.RefererAPI,
}
stack, _, _ := context.EnterStack(ctx, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(ctx, "test-assistant", &context.Options{})
ctx.Stack = stack
res, err := v8.Call(v8.CallOptions{}, `
@ -376,8 +384,9 @@ func TestMCPListSamples(t *testing.T) {
AssistantID: "test-assistant-id",
Locale: "en",
Context: stdContext.Background(),
Referer: context.RefererAPI,
}
stack, _, _ := context.EnterStack(ctx, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(ctx, "test-assistant", &context.Options{})
ctx.Stack = stack
res, err := v8.Call(v8.CallOptions{}, `
@ -418,8 +427,9 @@ func TestMCPGetSample(t *testing.T) {
AssistantID: "test-assistant-id",
Locale: "en",
Context: stdContext.Background(),
Referer: context.RefererAPI,
}
stack, _, _ := context.EnterStack(ctx, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(ctx, "test-assistant", &context.Options{})
ctx.Stack = stack
res, err := v8.Call(v8.CallOptions{}, `
@ -462,8 +472,9 @@ func TestMCPJsApiWithTrace(t *testing.T) {
AssistantID: "test-assistant-id",
Locale: "en",
Context: stdContext.Background(),
Referer: context.RefererAPI,
}
stack, _, _ := context.EnterStack(ctx, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(ctx, "test-assistant", &context.Options{})
ctx.Stack = stack
res, err := v8.Call(v8.CallOptions{}, `

15
agent/context/jsapi_release_test.go
View file

@ -22,10 +22,11 @@ func TestContextRelease(t *testing.T) {
AssistantID: "test-assistant-id",
Context: stdContext.Background(),
IDGenerator: message.NewIDGenerator(),
Referer: context.RefererAPI,
}
// Initialize stack and trace
stack, _, _ := context.EnterStack(cxt, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(cxt, "test-assistant", &context.Options{})
cxt.Stack = stack
res, err := v8.Call(v8.CallOptions{}, `
@ -75,10 +76,11 @@ func TestTraceRelease(t *testing.T) {
AssistantID: "test-assistant-id",
Context: stdContext.Background(),
IDGenerator: message.NewIDGenerator(),
Referer: context.RefererAPI,
}
// Initialize stack and trace
stack, _, _ := context.EnterStack(cxt, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(cxt, "test-assistant", &context.Options{})
cxt.Stack = stack
res, err := v8.Call(v8.CallOptions{}, `
@ -137,10 +139,11 @@ func TestContextReleaseWithTrace(t *testing.T) {
AssistantID: "test-assistant-id",
Context: stdContext.Background(),
IDGenerator: message.NewIDGenerator(),
Referer: context.RefererAPI,
}
// Initialize stack and trace
stack, _, _ := context.EnterStack(cxt, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(cxt, "test-assistant", &context.Options{})
cxt.Stack = stack
res, err := v8.Call(v8.CallOptions{}, `
@ -185,10 +188,11 @@ func TestTryFinallyPattern(t *testing.T) {
AssistantID: "test-assistant-id",
Context: stdContext.Background(),
IDGenerator: message.NewIDGenerator(),
Referer: context.RefererAPI,
}
// Initialize stack and trace
stack, _, _ := context.EnterStack(cxt, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(cxt, "test-assistant", &context.Options{})
cxt.Stack = stack
res, err := v8.Call(v8.CallOptions{}, `
@ -287,10 +291,11 @@ func TestTryFinallyPatternWithError(t *testing.T) {
AssistantID: "test-assistant-id",
Context: stdContext.Background(),
IDGenerator: message.NewIDGenerator(),
Referer: context.RefererAPI,
}
// Initialize stack and trace
stack, _, _ := context.EnterStack(cxt, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(cxt, "test-assistant", &context.Options{})
cxt.Stack = stack
_, err := v8.Call(v8.CallOptions{}, `

24
agent/context/jsapi_stress_test.go
View file

@ -38,7 +38,8 @@ func TestStressContextCreationAndRelease(t *testing.T) {
}
// Initialize stack and trace
stack, _, _ := context.EnterStack(cxt, "test-assistant", context.RefererAPI)
cxt.Referer = context.RefererAPI
stack, _, _ := context.EnterStack(cxt, "test-assistant", &context.Options{})
cxt.Stack = stack
_, err := v8.Call(v8.CallOptions{}, `
@ -111,10 +112,11 @@ func TestStressTraceOperations(t *testing.T) {
AssistantID: "test-assistant",
Context: stdContext.Background(),
IDGenerator: message.NewIDGenerator(),
Referer: context.RefererAPI,
}
// Initialize stack and trace
stack, _, _ := context.EnterStack(cxt, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(cxt, "test-assistant", &context.Options{})
cxt.Stack = stack
_, err := v8.Call(v8.CallOptions{}, fmt.Sprintf(`
function test(ctx) {
@ -188,9 +190,10 @@ func TestStressMCPOperations(t *testing.T) {
AssistantID: "test-assistant",
Context: stdContext.Background(),
IDGenerator: message.NewIDGenerator(),
Referer: context.RefererAPI,
}
stack, _, _ := context.EnterStack(cxt, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(cxt, "test-assistant", &context.Options{})
cxt.Stack = stack
startMemory := getMemStats()
@ -271,9 +274,10 @@ func TestStressConcurrentContexts(t *testing.T) {
AssistantID: "test-assistant",
Context: stdContext.Background(),
IDGenerator: message.NewIDGenerator(),
Referer: context.RefererAPI,
}
stack, _, _ := context.EnterStack(cxt, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(cxt, "test-assistant", &context.Options{})
cxt.Stack = stack
_, err := v8.Call(v8.CallOptions{}, `
@ -421,9 +425,10 @@ func TestStressReleasePatterns(t *testing.T) {
AssistantID: "test-assistant",
Context: stdContext.Background(),
IDGenerator: message.NewIDGenerator(),
Referer: context.RefererAPI,
}
stack, _, _ := context.EnterStack(cxt, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(cxt, "test-assistant", &context.Options{})
cxt.Stack = stack
_, err := v8.Call(v8.CallOptions{}, `
@ -459,9 +464,10 @@ func TestStressReleasePatterns(t *testing.T) {
AssistantID: "test-assistant",
Context: stdContext.Background(),
IDGenerator: message.NewIDGenerator(),
Referer: context.RefererAPI,
}
stack, _, _ := context.EnterStack(cxt, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(cxt, "test-assistant", &context.Options{})
cxt.Stack = stack
_, err := v8.Call(v8.CallOptions{}, `
@ -499,9 +505,10 @@ func TestStressReleasePatterns(t *testing.T) {
AssistantID: "test-assistant",
Context: stdContext.Background(),
IDGenerator: message.NewIDGenerator(),
Referer: context.RefererAPI,
}
stack, _, _ := context.EnterStack(cxt, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(cxt, "test-assistant", &context.Options{})
cxt.Stack = stack
_, err := v8.Call(v8.CallOptions{}, `
@ -544,9 +551,10 @@ func TestStressLongRunningTrace(t *testing.T) {
AssistantID: "test-assistant",
Context: stdContext.Background(),
IDGenerator: message.NewIDGenerator(),
Referer: context.RefererAPI,
}
stack, _, _ := context.EnterStack(cxt, "test-assistant", context.RefererAPI)
stack, _, _ := context.EnterStack(cxt, "test-assistant", &context.Options{})
cxt.Stack = stack
startMemory := getMemStats()

30
agent/context/jsapi_test.go
View file

@ -219,15 +219,9 @@ func TestJsValueAllFields(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
searchTrue := true
cxt := &context.Context{
ChatID: "test-chat-id",
AssistantID: "test-assistant-id",
Connector: "test-connector",
Search: &searchTrue,
Args: []interface{}{"arg1", "arg2", 123},
Retry: true,
RetryTimes: 3,
Locale: "zh-cn",
Theme: "dark",
Context: stdContext.Background(),
@ -279,21 +273,12 @@ func TestJsValueAllFields(t *testing.T) {
// Verify all fields
assert.Equal(t, "test-chat-id", result["chat_id"], "chat_id mismatch")
assert.Equal(t, "test-assistant-id", result["assistant_id"], "assistant_id mismatch")
assert.Equal(t, "test-connector", result["connector"], "connector mismatch")
assert.Equal(t, true, result["search"], "search mismatch")
assert.Equal(t, true, result["retry"], "retry mismatch")
assert.Equal(t, float64(3), result["retry_times"], "retry_times mismatch")
assert.Equal(t, "zh-cn", result["locale"], "locale mismatch")
assert.Equal(t, "dark", result["theme"], "theme mismatch")
assert.Equal(t, "api", result["referer"], "referer mismatch")
assert.Equal(t, "cui-web", result["accept"], "accept mismatch")
assert.Equal(t, "/dashboard/home", result["route"], "route mismatch")
// Verify args array
args, ok := result["args"].([]interface{})
assert.True(t, ok, "args should be an array")
assert.Equal(t, 3, len(args), "args length mismatch")
// Verify client object
client, ok := result["client"].(map[string]interface{})
assert.True(t, ok, "client should be an object")
@ -373,21 +358,6 @@ func testAllFieldsFunction(info *v8go.FunctionCallbackInfo) *v8go.Value {
if val, ok := getField("assistant_id"); ok {
result["assistant_id"] = val
}
if val, ok := getField("connector"); ok {
result["connector"] = val
}
if val, ok := getField("search"); ok {
result["search"] = val
}
if val, ok := getField("args"); ok {
result["args"] = val
}
if val, ok := getField("retry"); ok {
result["retry"] = val
}
if val, ok := getField("retry_times"); ok {
result["retry_times"] = val
}
if val, ok := getField("locale"); ok {
result["locale"] = val
}

3
agent/context/mcp_test.go
View file

@ -20,10 +20,11 @@ func newTestMCPContext() *context.Context {
ChatID: "test-chat",
AssistantID: "test-assistant",
Locale: "en",
Referer: context.RefererAPI,
}
// Initialize stack and trace
stack, traceID, _ := context.EnterStack(ctx, "test-assistant", context.RefererAPI)
stack, traceID, _ := context.EnterStack(ctx, "test-assistant", &context.Options{})
ctx.Stack = stack
_ = traceID // traceID is set in stack

37
agent/context/openapi.go
View file

@ -9,26 +9,27 @@ import (
"strings"
"github.com/gin-gonic/gin"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/store"
"github.com/yaoapp/yao/openapi/oauth/authorized"
)
// GetCompletionRequest parse completion request and create context from openapi request
// Returns: *CompletionRequest, *Context, error
func GetCompletionRequest(c *gin.Context, cache store.Store) (*CompletionRequest, *Context, error) {
// Returns: *CompletionRequest, *Context, *Options, error
func GetCompletionRequest(c *gin.Context, cache store.Store) (*CompletionRequest, *Context, *Options, error) {
// Get authorized information
authInfo := authorized.GetInfo(c)
// Parse completion request from payload or query first
completionReq, err := parseCompletionRequestData(c)
if err != nil {
return nil, nil, fmt.Errorf("failed to parse completion request: %w", err)
return nil, nil, nil, fmt.Errorf("failed to parse completion request: %w", err)
}
// Extract assistant ID using completionReq (can extract from model field)
assistantID, err := GetAssistantID(c, completionReq)
if err != nil {
return nil, nil, fmt.Errorf("failed to get assistant ID: %w", err)
return nil, nil, nil, fmt.Errorf("failed to get assistant ID: %w", err)
}
// Extract chat ID (may generate from messages if not provided)
@ -46,7 +47,7 @@ func GetCompletionRequest(c *gin.Context, cache store.Store) (*CompletionRequest
// Create context with unique ID using New() to ensure proper initialization
ctx := New(c.Request.Context(), authInfo, chatID)
// Set additional fields
// Set context fields (session-level state)
ctx.Cache = cache
ctx.Writer = c.Writer
ctx.AssistantID = assistantID
@ -61,14 +62,34 @@ func GetCompletionRequest(c *gin.Context, cache store.Store) (*CompletionRequest
}
ctx.Route = GetRoute(c, completionReq)
ctx.Metadata = GetMetadata(c, completionReq)
ctx.Skip = GetSkip(c, completionReq)
// Create Options (call-level parameters)
opts := &Options{
Context: c.Request.Context(),
Skip: GetSkip(c, completionReq),
}
// Try to extract custom connector from model field
// If model is a valid connector ID, set it to opts.Connector
// Otherwise, keep the standard OpenAI-compatible behavior (model as assistant ID)
if completionReq != nil && completionReq.Model != "" {
// Check if model is a valid connector (not containing "-yao_" which indicates assistant ID format)
if !strings.Contains(completionReq.Model, "-yao_") {
// Try to validate if it's a real connector
if _, err := connector.Select(completionReq.Model); err == nil {
// It's a valid connector, use it
opts.Connector = completionReq.Model
}
// If not a valid connector, ignore it (keep opts.Connector empty to use assistant's default)
}
}
// Initialize interrupt controller
ctx.Interrupt = NewInterruptController()
// Register context to global registry first (required for interrupt handler callback)
if err := Register(ctx); err != nil {
return nil, nil, fmt.Errorf("failed to register context: %w", err)
return nil, nil, nil, fmt.Errorf("failed to register context: %w", err)
}
// Start interrupt listener after registration
@ -76,7 +97,7 @@ func GetCompletionRequest(c *gin.Context, cache store.Store) (*CompletionRequest
// HTTP context cancellation is handled by LLM/Agent layers naturally
ctx.Interrupt.Start(ctx.ID)
return completionReq, ctx, nil
return completionReq, ctx, opts, nil
}
// getClientType parses the client type from User-Agent header

14
agent/context/openapi_test.go
View file

@ -851,7 +851,7 @@ func TestGetCompletionRequest_WriterInitialized(t *testing.T) {
c, _ := gin.CreateTestContext(w)
c.Request = req
completionReq, ctx, err := GetCompletionRequest(c, cache)
completionReq, ctx, opts, err := GetCompletionRequest(c, cache)
if err != nil {
t.Fatalf("Failed to get completion request: %v", err)
}
@ -867,6 +867,11 @@ func TestGetCompletionRequest_WriterInitialized(t *testing.T) {
t.Error("Expected ctx.Writer to be the same as gin context writer")
}
// Check that Options is initialized
if opts == nil {
t.Error("Expected opts to be initialized, got nil")
}
// Check other fields
if completionReq.Model != "gpt-4-yao_test" {
t.Errorf("Expected model 'gpt-4-yao_test', got '%s'", completionReq.Model)
@ -914,12 +919,17 @@ func TestGetCompletionRequest_ChatIDFallback(t *testing.T) {
c, _ := gin.CreateTestContext(w)
c.Request = req
_, ctx, err := GetCompletionRequest(c, cache)
_, ctx, opts, err := GetCompletionRequest(c, cache)
if err != nil {
t.Fatalf("Failed to get completion request: %v", err)
}
defer ctx.Release()
// Check that Options is initialized
if opts == nil {
t.Error("Expected opts to be initialized, got nil")
}
// ChatID should be generated (not empty)
if ctx.ChatID == "" {
t.Error("Expected ChatID to be generated via fallback, got empty string")

74
agent/context/options.go Normal file
View file

@ -0,0 +1,74 @@
package context
// ToMap converts Options struct to map for JSON serialization
func (opts *Options) ToMap() map[string]interface{} {
if opts == nil {
return nil
}
result := make(map[string]interface{})
// Add configurable fields (with json tags)
if opts.Connector != "" {
result["connector"] = opts.Connector
}
if opts.Mode != "" {
result["mode"] = opts.Mode
}
if opts.Search != nil {
result["search"] = *opts.Search
}
if opts.Skip != nil {
result["skip"] = opts.Skip
}
// Only add DisableGlobalPrompts if true (avoid false values in map)
if opts.DisableGlobalPrompts {
result["disable_global_prompts"] = opts.DisableGlobalPrompts
}
// Note: Runtime fields (Context, Writer) are not serialized (json:"-")
// They should not be included in the map
return result
}
// OptionsFromMap creates Options struct from map (e.g., from JS Hook)
func OptionsFromMap(m map[string]interface{}) *Options {
if m == nil {
return &Options{}
}
opts := &Options{}
// Extract configurable fields
if connector, ok := m["connector"].(string); ok {
opts.Connector = connector
}
if mode, ok := m["mode"].(string); ok {
opts.Mode = mode
}
if search, ok := m["search"].(bool); ok {
opts.Search = &search
}
if skipMap, ok := m["skip"].(map[string]interface{}); ok {
skip := &Skip{}
if history, ok := skipMap["history"].(bool); ok {
skip.History = history
}
if trace, ok := skipMap["trace"].(bool); ok {
skip.Trace = trace
}
if output, ok := skipMap["output"].(bool); ok {
skip.Output = output
}
opts.Skip = skip
}
if disableGlobalPrompts, ok := m["disable_global_prompts"].(bool); ok {
opts.DisableGlobalPrompts = disableGlobalPrompts
}
// Note: Context and Writer are runtime fields, not restored from map
// They should be set by the caller if needed
return opts
}

55
agent/context/output.go
View file

@ -23,6 +23,24 @@ func (ctx *Context) Send(msg *message.Message) error {
// Skip lifecycle events for event-type messages (prevent recursion)
isEventMessage := msg.Type == message.TypeEvent
// === Handle message_start event: record metadata for future delta chunks ===
if isEventMessage && msg.Props != nil {
if event, ok := msg.Props["event"].(string); ok && event == message.EventMessageStart {
if data, ok := msg.Props["data"].(message.EventMessageStartData); ok {
// Record metadata from message_start event
if data.MessageID != "" && ctx.messageMetadata != nil {
ctx.messageMetadata.setMessage(data.MessageID, &MessageMetadata{
MessageID: data.MessageID,
ThreadID: data.ThreadID,
Type: data.Type,
StartTime: time.Now(),
ChunkCount: 0, // Will be incremented by delta chunks
})
}
}
}
}
// === Delta operations: Auto-inherit and update metadata ===
if msg.Delta && msg.MessageID != "" && ctx.messageMetadata != nil {
if metadata := ctx.getMessageMetadata(msg.MessageID); metadata != nil {
@ -103,6 +121,7 @@ func (ctx *Context) Send(msg *message.Message) error {
MessageID: msg.MessageID,
Type: msg.Type,
Timestamp: time.Now().UnixMilli(),
ThreadID: msg.ThreadID, // Include ThreadID for concurrent stream identification
}
messageStartEvent := output.NewEventMessage(message.EventMessageStart, "Message started", messageStartData)
if err := ctx.sendRaw(messageStartEvent); err != nil {
@ -147,6 +166,7 @@ func (ctx *Context) Send(msg *message.Message) error {
MessageID: msg.MessageID,
Type: msg.Type,
Timestamp: time.Now().UnixMilli(),
ThreadID: metadata.ThreadID, // Include ThreadID for concurrent stream identification
DurationMs: durationMs,
ChunkCount: metadata.ChunkCount,
Status: "completed",
@ -191,6 +211,7 @@ func (ctx *Context) EndMessage(messageID string, content interface{}) error {
MessageID: messageID,
Type: metadata.Type,
Timestamp: time.Now().UnixMilli(),
ThreadID: metadata.ThreadID, // Include ThreadID for concurrent stream identification
DurationMs: durationMs,
ChunkCount: metadata.ChunkCount,
Status: "completed",
@ -276,16 +297,33 @@ func (ctx *Context) sendRaw(msg *message.Message) error {
return out.Send(msg)
}
// getWriter gets the effective Writer for the current context
// Priority: Skip.Output > Stack.Options.Writer > ctx.Writer
func (ctx *Context) getWriter() Writer {
// Check if output is explicitly skipped (for internal A2A calls)
if ctx.Stack != nil && ctx.Stack.Options != nil && ctx.Stack.Options.Skip != nil && ctx.Stack.Options.Skip.Output {
return nil // Explicitly disable output
}
// Check if current Stack has a Writer override
if ctx.Stack != nil && ctx.Stack.Options != nil && ctx.Stack.Options.Writer != nil {
return ctx.Stack.Options.Writer
}
return ctx.Writer
}
// getOutput gets the output writer for the context
func (ctx *Context) getOutput() (*output.Output, error) {
if ctx.output != nil {
return ctx.output, nil
// Check if current Stack has cached output
if ctx.Stack != nil && ctx.Stack.output != nil {
return ctx.Stack.output, nil
}
trace, _ := ctx.Trace()
var options message.Options = message.Options{
BaseURL: "/",
Writer: ctx.Writer,
Writer: ctx.getWriter(), // Use getWriter() to resolve Writer priority
Trace: trace,
Locale: ctx.Locale,
Accept: string(ctx.Accept),
@ -297,10 +335,15 @@ func (ctx *Context) getOutput() (*output.Output, error) {
options.Capabilities = &caps
}
var err error
ctx.output, err = output.NewOutput(options)
out, err := output.NewOutput(options)
if err != nil {
return nil, err
}
return ctx.output, nil
// Cache to current Stack (each Stack has its own output with its own Writer)
if ctx.Stack != nil {
ctx.Stack.output = out
}
return out, nil
}

18
agent/context/stack.go
View file

@ -9,7 +9,7 @@ import (
)
// NewStack creates a new root stack with the given trace ID and assistant ID
func NewStack(traceID, assistantID, referer string) *Stack {
func NewStack(traceID, assistantID, referer string, opts *Options) *Stack {
if traceID == "" {
traceID = uuid.New().String()
}
@ -25,13 +25,14 @@ func NewStack(traceID, assistantID, referer string) *Stack {
Depth: 0,
ParentID: "",
Path: []string{stackID},
Options: opts,
CreatedAt: now,
Status: StackStatusRunning,
}
}
// NewChildStack creates a child stack from the current stack
func (s *Stack) NewChildStack(assistantID, referer string) *Stack {
func (s *Stack) NewChildStack(assistantID, referer string, opts *Options) *Stack {
stackID := uuid.New().String()
now := time.Now().UnixMilli()
@ -48,6 +49,7 @@ func (s *Stack) NewChildStack(assistantID, referer string) *Stack {
Depth: s.Depth + 1,
ParentID: s.ID,
Path: path,
Options: opts,
CreatedAt: now,
Status: StackStatusRunning,
}
@ -134,6 +136,7 @@ func (s *Stack) Clone() *Stack {
Depth: s.Depth,
ParentID: s.ParentID,
Path: make([]string, len(s.Path)),
Options: s.Options, // Shallow copy of Options pointer
CreatedAt: s.CreatedAt,
Status: s.Status,
Error: s.Error,
@ -165,14 +168,17 @@ func (s *Stack) Clone() *Stack {
//
// Usage:
//
// stack, traceID, done := context.EnterStack(ctx, assistantID, referer)
// stack, traceID, done := context.EnterStack(ctx, assistantID, opts)
// defer done()
// // ... your code here ...
func EnterStack(ctx *Context, assistantID, referer string) (*Stack, string, func()) {
func EnterStack(ctx *Context, assistantID string, opts *Options) (*Stack, string, func()) {
var stack *Stack
var parentStack *Stack
var traceID string
// Get referer from ctx (request source)
referer := ctx.Referer
// Initialize Stacks map if not exists
if ctx.Stacks == nil {
ctx.Stacks = make(map[string]*Stack)
@ -182,14 +188,14 @@ func EnterStack(ctx *Context, assistantID, referer string) (*Stack, string, func
// Create root stack for this assistant call (entry point)
// Generate a new trace ID for root
traceID = trace.GenTraceID()
stack = NewStack(traceID, assistantID, referer)
stack = NewStack(traceID, assistantID, referer, opts)
ctx.Stack = stack
} else {
// Create child stack for nested agent call
// Inherit trace ID from parent
parentStack = ctx.Stack
traceID = parentStack.TraceID
stack = ctx.Stack.NewChildStack(assistantID, referer)
stack = ctx.Stack.NewChildStack(assistantID, referer, opts)
ctx.Stack = stack
}

50
agent/context/stack_test.go
View file

@ -16,8 +16,9 @@ func TestNewStack(t *testing.T) {
traceID := "12345678"
assistantID := "test-assistant"
referer := RefererAPI
opts := &Options{}
stack := NewStack(traceID, assistantID, referer)
stack := NewStack(traceID, assistantID, referer, opts)
if stack == nil {
t.Fatal("Expected stack to be created, got nil")
@ -57,7 +58,7 @@ func TestNewStack_GenerateTraceID(t *testing.T) {
defer test.Clean()
// Empty traceID should generate a UUID
stack := NewStack("", "test-assistant", RefererAPI)
stack := NewStack("", "test-assistant", RefererAPI, &Options{})
if stack.TraceID == "" {
t.Error("Expected TraceID to be generated, got empty string")
@ -74,10 +75,10 @@ func TestNewChildStack(t *testing.T) {
defer test.Clean()
// Create parent stack
parentStack := NewStack("12345678", "parent-assistant", RefererAPI)
parentStack := NewStack("12345678", "parent-assistant", RefererAPI, &Options{})
// Create child stack
childStack := parentStack.NewChildStack("child-assistant", RefererAgent)
childStack := parentStack.NewChildStack("child-assistant", RefererAgent, &Options{})
if childStack == nil {
t.Fatal("Expected child stack to be created, got nil")
@ -121,7 +122,7 @@ func TestStackComplete(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
stack := NewStack("12345678", "test-assistant", RefererAPI)
stack := NewStack("12345678", "test-assistant", RefererAPI, &Options{})
// Wait a bit to have measurable duration
time.Sleep(10 * time.Millisecond)
@ -157,7 +158,7 @@ func TestStackFail(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
stack := NewStack("12345678", "test-assistant", RefererAPI)
stack := NewStack("12345678", "test-assistant", RefererAPI, &Options{})
testError := "test error message"
stack.Fail(nil)
@ -180,7 +181,7 @@ func TestStackTimeout(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
stack := NewStack("12345678", "test-assistant", RefererAPI)
stack := NewStack("12345678", "test-assistant", RefererAPI, &Options{})
stack.Timeout()
@ -199,9 +200,10 @@ func TestEnterStack_RootCreation(t *testing.T) {
ctx := &Context{
IDGenerator: message.NewIDGenerator(),
Referer: RefererAPI,
}
stack, traceID, done := EnterStack(ctx, "test-assistant", RefererAPI)
stack, traceID, done := EnterStack(ctx, "test-assistant", &Options{})
defer done()
if stack == nil {
@ -244,10 +246,11 @@ func TestEnterStack_ChildCreation(t *testing.T) {
ctx := &Context{
IDGenerator: message.NewIDGenerator(),
Referer: RefererAPI,
}
// Create parent
parentStack, parentTraceID, parentDone := EnterStack(ctx, "parent-assistant", RefererAPI)
parentStack, parentTraceID, parentDone := EnterStack(ctx, "parent-assistant", &Options{})
defer parentDone()
if parentStack == nil {
@ -255,7 +258,7 @@ func TestEnterStack_ChildCreation(t *testing.T) {
}
// Create child
childStack, childTraceID, childDone := EnterStack(ctx, "child-assistant", RefererAgent)
childStack, childTraceID, childDone := EnterStack(ctx, "child-assistant", &Options{})
defer childDone()
if childStack == nil {
@ -289,13 +292,14 @@ func TestEnterStack_DoneCallback(t *testing.T) {
ctx := &Context{
IDGenerator: message.NewIDGenerator(),
Referer: RefererAPI,
}
// Create parent
parentStack, _, parentDone := EnterStack(ctx, "parent-assistant", RefererAPI)
parentStack, _, parentDone := EnterStack(ctx, "parent-assistant", &Options{})
// Create child
childStack, _, childDone := EnterStack(ctx, "child-assistant", RefererAgent)
childStack, _, childDone := EnterStack(ctx, "child-assistant", &Options{})
// Child should be current
if ctx.Stack != childStack {
@ -330,16 +334,17 @@ func TestContextGetAllStacks(t *testing.T) {
ctx := &Context{
IDGenerator: message.NewIDGenerator(),
Referer: RefererAPI,
}
// Create multiple stacks
_, _, done1 := EnterStack(ctx, "assistant1", RefererAPI)
_, _, done1 := EnterStack(ctx, "assistant1", &Options{})
defer done1()
_, _, done2 := EnterStack(ctx, "assistant2", RefererAgent)
_, _, done2 := EnterStack(ctx, "assistant2", &Options{})
defer done2()
_, _, done3 := EnterStack(ctx, "assistant3", RefererAgent)
_, _, done3 := EnterStack(ctx, "assistant3", &Options{})
defer done3()
// Get all stacks
@ -356,9 +361,10 @@ func TestContextGetStackByID(t *testing.T) {
ctx := &Context{
IDGenerator: message.NewIDGenerator(),
Referer: RefererAPI,
}
stack, _, done := EnterStack(ctx, "test-assistant", RefererAPI)
stack, _, done := EnterStack(ctx, "test-assistant", &Options{})
defer done()
// Get stack by ID
@ -385,13 +391,14 @@ func TestContextGetStacksByTraceID(t *testing.T) {
ctx := &Context{
IDGenerator: message.NewIDGenerator(),
Referer: RefererAPI,
}
// Create parent and child (same trace ID)
_, traceID, done1 := EnterStack(ctx, "parent-assistant", RefererAPI)
_, traceID, done1 := EnterStack(ctx, "parent-assistant", &Options{})
defer done1()
_, _, done2 := EnterStack(ctx, "child-assistant", RefererAgent)
_, _, done2 := EnterStack(ctx, "child-assistant", &Options{})
defer done2()
// Get stacks by trace ID
@ -415,14 +422,15 @@ func TestContextGetRootStack(t *testing.T) {
ctx := &Context{
IDGenerator: message.NewIDGenerator(),
Referer: RefererAPI,
}
// Create parent
parentStack, _, done1 := EnterStack(ctx, "parent-assistant", RefererAPI)
parentStack, _, done1 := EnterStack(ctx, "parent-assistant", &Options{})
defer done1()
// Create child
_, _, done2 := EnterStack(ctx, "child-assistant", RefererAgent)
_, _, done2 := EnterStack(ctx, "child-assistant", &Options{})
defer done2()
// Get root stack
@ -445,7 +453,7 @@ func TestStackClone(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
original := NewStack("12345678", "test-assistant", RefererAPI)
original := NewStack("12345678", "test-assistant", RefererAPI, &Options{})
original.Complete()
clone := original.Clone()

99
agent/context/types.go
View file

@ -194,6 +194,7 @@ type AssistantInfo struct {
type Skip struct {
History bool `json:"history"` // Skip saving chat history (for internal calls like title/prompt generation)
Trace bool `json:"trace"` // Skip trace logging
Output bool `json:"output"` // Skip output to client (for internal A2A calls that only need response data)
}
// MessageMetadata stores metadata for sent messages
@ -232,12 +233,8 @@ type Context struct {
// Internal
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
messageMetadata *messageMetadataStore `json:"-"` // Thread-safe message metadata store for delta operations
// 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)
Capabilities *openai.Capabilities `json:"-"` // Model capabilities for the current connector
@ -248,13 +245,6 @@ type Context struct {
Authorized *types.AuthorizedInfo `json:"authorized,omitempty"` // Authorized information
ChatID string `json:"chat_id,omitempty"` // Chat ID, use to select chat
AssistantID string `json:"assistant_id,omitempty"` // Assistant ID, use to select assistant
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
// Arguments for call
Args []interface{} `json:"args,omitempty"` // Arguments for call, it will be used to pass data to the call
Retry bool `json:"retry,omitempty"` // Retry mode
RetryTimes uint8 `json:"retry_times,omitempty"` // Retry times
// Locale information
Locale string `json:"locale,omitempty"` // Locale
@ -270,6 +260,31 @@ type Context struct {
Metadata map[string]interface{} `json:"metadata,omitempty"` // The metadata of the request, it will be used to pass data to the page
}
// Options represents the options for the context
type Options struct {
// Original context, override the default context
Context context.Context `json:"-"` // Context, it will be used to pass the context to the call
// Writer, use to write response data to the client (override the default writer)
Writer Writer `json:"writer,omitempty"` // Writer, use to write response data to the client
// 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
// 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
// Disable global prompts, default is false
DisableGlobalPrompts bool `json:"disable_global_prompts,omitempty"` // Temporarily disable global prompts for this request
// Search mode, default is true
Search *bool `json:"search,omitempty"` // Search mode, default is true
// Agent mode, use to select the mode of the request, default is "chat"
Mode string `json:"mode,omitempty"` // Agent mode, use to select the mode of the request, default is "chat"
}
// Stack represents the call stack node for tracing agent-to-agent calls
// Uses a flat structure to avoid circular references and memory overhead
type Stack struct {
@ -277,6 +292,9 @@ type Stack struct {
ID string `json:"id"` // Unique stack node ID, used to identify this specific call
TraceID string `json:"trace_id"` // Shared trace ID for entire call tree, inherited from root
// Options
Options *Options `json:"options,omitempty"` // Options for the call
// Call context
AssistantID string `json:"assistant_id"` // Assistant handling this call
Referer string `json:"referer,omitempty"` // Call source: api, agent, tool, process, etc.
@ -294,6 +312,9 @@ type Stack struct {
// Metrics
DurationMs *int64 `json:"duration_ms,omitempty"` // Duration in milliseconds (calculated when completed)
// Runtime cache (not serialized)
output *output.Output `json:"-"` // Cached output instance for this stack
}
// Response the response
@ -331,12 +352,11 @@ type HookCreateResponse struct {
DisableGlobalPrompts *bool `json:"disable_global_prompts,omitempty"` // Temporarily disable global prompts for this request
// Context adjustments - allow hook to modify context fields
AssistantID string `json:"assistant_id,omitempty"` // Override assistant ID
Connector string `json:"connector,omitempty"` // Override connector
Locale string `json:"locale,omitempty"` // Override locale
Theme string `json:"theme,omitempty"` // Override theme
Route string `json:"route,omitempty"` // Override route
Metadata map[string]interface{} `json:"metadata,omitempty"` // Override or merge metadata
Connector string `json:"connector,omitempty"` // Override connector (call-level)
Locale string `json:"locale,omitempty"` // Override locale (session-level)
Theme string `json:"theme,omitempty"` // Override theme (session-level)
Route string `json:"route,omitempty"` // Override route (session-level)
Metadata map[string]interface{} `json:"metadata,omitempty"` // Override or merge metadata (session-level)
}
// NextHookPayload payload for the next hook
@ -372,9 +392,9 @@ type NextHookResponse struct {
// DelegateConfig configuration for delegating to another agent
type DelegateConfig struct {
AgentID string `json:"agent_id"` // Required: target agent ID
Messages []Message `json:"messages"` // Messages to send to target agent
AgentID string `json:"agent_id"` // Required: target agent ID
Messages []Message `json:"messages"` // Messages to send to target agent
Options map[string]interface{} `json:"options,omitempty"` // Optional: call-level options for delegation
}
// NextAction defines the action determined by Next hook response
@ -464,6 +484,8 @@ const (
ContentText ContentPartType = "text" // Text content
ContentImageURL ContentPartType = "image_url" // Image URL content (Vision)
ContentInputAudio ContentPartType = "input_audio" // Input audio content (Audio)
ContentFile ContentPartType = "file" // File attachment (documents, etc.)
ContentData ContentPartType = "data" // Generic data content (base64, binary, etc.)
)
// ContentPart represents a part of the message content (for multimodal messages)
@ -473,6 +495,8 @@ type ContentPart struct {
Text string `json:"text,omitempty"` // For type="text": the text content
ImageURL *ImageURL `json:"image_url,omitempty"` // For type="image_url": the image URL
InputAudio *InputAudio `json:"input_audio,omitempty"` // For type="input_audio": the input audio data
File *FileAttachment `json:"file,omitempty"` // For type="file": file attachment
Data *DataContent `json:"data,omitempty"` // For type="data": generic data content
}
// ImageDetailLevel represents the detail level for image processing
@ -497,6 +521,41 @@ type InputAudio struct {
Format string `json:"format"` // Required: Audio format (e.g., "wav", "mp3")
}
// FileAttachment represents a file attachment in the message content
// Compatible with frontend InputArea format: { type: 'file', file: { url, filename } }
type FileAttachment struct {
URL string `json:"url"` // Required: URL of the file (http:// or __uploader://fileid wrapper)
Filename string `json:"filename,omitempty"` // Optional: original filename
}
// DataSourceType represents the type of data source
type DataSourceType string
// Data source type constants
const (
DataSourceModel DataSourceType = "model" // Data model
DataSourceKBCollection DataSourceType = "kb_collection" // Knowledge base collection
DataSourceKBDocument DataSourceType = "kb_document" // Knowledge base document/file
DataSourceTable DataSourceType = "table" // Database table
DataSourceAPI DataSourceType = "api" // API endpoint
DataSourceMCPResource DataSourceType = "mcp_resource" // MCP (Model Context Protocol) resource
)
// DataSource represents a single data source reference
type DataSource struct {
Type DataSourceType `json:"type"` // Required: type of data source
Name string `json:"name"` // Required: name/identifier of the data source
ID string `json:"id,omitempty"` // Optional: specific ID (e.g., document ID, record ID)
Filters map[string]interface{} `json:"filters,omitempty"` // Optional: filters to apply
Metadata map[string]interface{} `json:"metadata,omitempty"` // Optional: additional metadata
}
// DataContent represents data source references in the message
// Used to reference data models, knowledge base collections, KB documents, etc.
type DataContent struct {
Sources []DataSource `json:"sources"` // Required: array of data source references
}
// ToolCallType represents the type of tool call
type ToolCallType string

6
agent/i18n/builtin.go
View file

@ -21,7 +21,7 @@ func init() {
"assistant.agent.stream.skipping": "Skipping output close (nested call)",
"assistant.agent.stream.close_error": "Failed to close output",
"assistant.agent.completion.label": "Agent Completion",
"assistant.agent.completion.description": "Final output from assistant",
"assistant.agent.completion.description": "Final output from {{name}}",
// LLM: providers/openai/openai.go Stream() function
"llm.openai.stream.label": "LLM %s",
@ -111,7 +111,7 @@ func init() {
"assistant.agent.stream.skipping": "跳过输出关闭（嵌套调用）",
"assistant.agent.stream.close_error": "关闭输出失败",
"assistant.agent.completion.label": "智能体完成",
"assistant.agent.completion.description": "智能体最终输出",
"assistant.agent.completion.description": "{{name}} 最终输出",
// LLM: providers/openai/openai.go Stream() function
"llm.openai.stream.label": "LLM %s",
@ -173,7 +173,7 @@ func init() {
"assistant.agent.stream.skipping": "跳过输出关闭（嵌套调用）",
"assistant.agent.stream.close_error": "关闭输出失败",
"assistant.agent.completion.label": "智能体完成",
"assistant.agent.completion.description": "智能体最终输出",
"assistant.agent.completion.description": "{{name}} 最终输出",
// LLM: providers/openai/openai.go Stream() function
"llm.openai.stream.label": "LLM %s",

1
agent/llm/providers/openai/claude_test.go
View file

@ -22,7 +22,6 @@ func newClaudeTestContext(chatID, connectorID string) *context.Context {
Space: plan.NewMemorySharedSpace(),
ChatID: chatID,
AssistantID: "test-assistant",
Connector: connectorID,
Locale: "en-us",
Theme: "light",
Client: context.Client{

1
agent/llm/providers/openai/deepseek_r1_test.go
View file

@ -401,7 +401,6 @@ func newDeepSeekTestContext(chatID, connectorID string) *context.Context {
Space: plan.NewMemorySharedSpace(),
ChatID: chatID,
AssistantID: "test-assistant",
Connector: connectorID,
Locale: "en-us",
Theme: "light",
Client: context.Client{

1
agent/llm/providers/openai/deepseek_v3_test.go
View file

@ -373,7 +373,6 @@ func newDeepSeekV3TestContext(chatID, connectorID string) *context.Context {
Space: plan.NewMemorySharedSpace(),
ChatID: chatID,
AssistantID: "test-assistant",
Connector: connectorID,
Locale: "en-us",
Theme: "light",
Client: context.Client{

1
agent/llm/providers/openai/gpt5_test.go
View file

@ -388,7 +388,6 @@ func newGPT5TestContext(chatID, connectorID string) *context.Context {
Space: plan.NewMemorySharedSpace(),
ChatID: chatID,
AssistantID: "test-assistant",
Connector: connectorID,
Locale: "en-us",
Theme: "light",
Client: context.Client{

8
agent/llm/providers/openai/openai.go
View file

@ -211,7 +211,11 @@ func (p *Provider) Stream(ctx *context.Context, messages []context.Message, opti
var lastErr error
// Get Go context for cancellation support
// Read from Stack.Options if available (call-level override)
goCtx := ctx.Context
if ctx.Stack != nil && ctx.Stack.Options != nil && ctx.Stack.Options.Context != nil {
goCtx = ctx.Stack.Options.Context
}
if goCtx == nil {
goCtx = gocontext.Background()
}
@ -808,7 +812,11 @@ func (p *Provider) Post(ctx *context.Context, messages []context.Message, option
var lastErr error
// Get Go context for cancellation support
// Read from Stack.Options if available (call-level override)
goCtx := ctx.Context
if ctx.Stack != nil && ctx.Stack.Options != nil && ctx.Stack.Options.Context != nil {
goCtx = ctx.Stack.Options.Context
}
if goCtx == nil {
goCtx = gocontext.Background()
}

1
agent/llm/providers/openai/openai_test.go
View file

@ -1508,7 +1508,6 @@ func newTestContext(chatID, connectorID string) *context.Context {
Space: plan.NewMemorySharedSpace(),
ChatID: chatID,
AssistantID: "test-assistant",
Connector: connectorID,
Locale: "en-us",
Theme: "light",
Client: context.Client{

1
agent/llm/providers/openai/temperature_test.go
View file

@ -340,7 +340,6 @@ func newTemperatureTestContext(chatID, connectorID string) *context.Context {
Space: plan.NewMemorySharedSpace(),
ChatID: chatID,
AssistantID: "test-assistant",
Connector: connectorID,
Locale: "en-us",
Theme: "light",
Client: context.Client{

2
agent/output/message/types.go
View file

@ -318,6 +318,7 @@ type EventMessageStartData struct {
MessageID string `json:"message_id"` // Message ID (M1, M2, M3...)
Type string `json:"type"` // Message type: "text" | "thinking" | "tool_call" | "refusal"
Timestamp int64 `json:"timestamp"` // Unix timestamp when message started
ThreadID string `json:"thread_id,omitempty"` // Thread ID (optional; for concurrent streams)
ToolCall *EventToolCallInfo `json:"tool_call,omitempty"` // Tool call metadata (if type is "tool_call")
Extra map[string]interface{} `json:"extra,omitempty"` // Additional metadata (for custom providers or future extensions)
}
@ -329,6 +330,7 @@ type EventMessageEndData struct {
MessageID string `json:"message_id"` // Message ID (M1, M2, M3...)
Type string `json:"type"` // Message type (same as in message_start)
Timestamp int64 `json:"timestamp"` // Unix timestamp when message ended
ThreadID string `json:"thread_id,omitempty"` // Thread ID (optional; for concurrent streams)
DurationMs int64 `json:"duration_ms"` // Duration of this message in milliseconds
ChunkCount int `json:"chunk_count"` // Number of data chunks in this message
Status string `json:"status"` // "completed" | "partial" | "error"

128
attachment/README.md
View file

@ -659,6 +659,134 @@ case "upload_failed":
}
```
## Text Content Storage
The attachment package supports storing parsed text content extracted from files (e.g., from PDFs, Word documents, or image OCR). This is useful for building search indexes or providing text-based previews.
The system automatically maintains two versions of the text content:
- **Full content** (`content`): Complete text, stored as longText (up to 4GB)
- **Preview** (`content_preview`): First 2000 characters, stored as text for quick access
### Saving Parsed Text Content
Use `SaveText` to store the extracted text content. It automatically saves both full content and preview:
```go
// Upload a PDF file
file, err := manager.Upload(ctx, fileHeader, reader, option)
if err != nil {
return err
}
// Extract text from the PDF (using your preferred library)
parsedText := extractTextFromPDF(file.ID)
// Save the parsed text (automatically saves both full and preview)
err = manager.SaveText(ctx, file.ID, parsedText)
if err != nil {
return fmt.Errorf("failed to save text content: %w", err)
}
```
### Retrieving Parsed Text Content
Use `GetText` to retrieve text content. By default, it returns the preview for better performance:
```go
// Get preview (first 2000 characters) - Fast, suitable for UI display
preview, err := manager.GetText(ctx, file.ID)
if err != nil {
return fmt.Errorf("failed to get preview: %w", err)
}
if preview == "" {
fmt.Println("No text content available for this file")
} else {
fmt.Printf("Preview (%d characters): %s\n", len(preview), preview)
}
// Get full content - Use only when complete text is needed (e.g., for indexing)
fullText, err := manager.GetText(ctx, file.ID, true)
if err != nil {
return fmt.Errorf("failed to get full text: %w", err)
}
fmt.Printf("Full content (%d characters)\n", len(fullText))
```
### Performance Optimization
The text content fields are optimized for different use cases:
| Field | Size Limit | Use Case | Performance |
|-------|------------|----------|-------------|
| `content_preview` | 2000 chars | Quick preview, UI display, snippets | ⚡ Very Fast |
| `content` | 4GB | Full text search, complete content | 🐌 Slow for large files |
**Best Practices:**
1. Use preview by default: `GetText(ctx, fileID)`
2. Only request full content when necessary: `GetText(ctx, fileID, true)`
3. Both fields are excluded from `List()` by default for optimal performance
4. Preview uses character (rune) count, not bytes, for proper UTF-8 handling
### Example: Complete Text Processing Workflow
```go
// 1. Upload file
file, err := manager.Upload(ctx, fileHeader, reader, option)
if err != nil {
return err
}
// 2. Process file based on content type
var parsedText string
switch {
case strings.HasPrefix(file.ContentType, "image/"):
// Use OCR to extract text from image
parsedText, err = performOCR(file.ID)
case file.ContentType == "application/pdf":
// Extract text from PDF
parsedText, err = extractPDFText(file.ID)
case strings.Contains(file.ContentType, "wordprocessingml"):
// Extract text from Word document
parsedText, err = extractWordText(file.ID)
}
if err != nil {
return fmt.Errorf("failed to extract text: %w", err)
}
// 3. Save the extracted text
if parsedText != "" {
err = manager.SaveText(ctx, file.ID, parsedText)
if err != nil {
return fmt.Errorf("failed to save text: %w", err)
}
}
// 4. Later, retrieve the text for search or display
savedText, err := manager.GetText(ctx, file.ID)
if err != nil {
return err
}
fmt.Printf("Retrieved text: %s\n", savedText)
```
### Text Content Features
- **Dual Storage**: Automatically maintains both full content and preview (2000 chars)
- **Size Limits**:
- Preview: 2000 characters (text type)
- Full content: Up to 4GB (longText type)
- **Smart Retrieval**: Returns preview by default, full content on demand
- **Update**: Text content can be updated at any time using `SaveText`
- **Clear**: Set text to empty string to clear both fields
- **UTF-8 Safe**: Preview uses character (rune) count, not bytes, ensuring proper multi-byte character handling
- **Performance**: Both `content` and `content_preview` fields are excluded by default in `List()` and `Info()` operations to avoid loading text data. Use `GetText()` to explicitly retrieve text content when needed
#### `RegisterDefault(name string) (*Manager, error)`
Registers a default attachment manager with sensible defaults for common file types.

101
attachment/manager.go
View file

@ -723,6 +723,16 @@ func (manager Manager) List(ctx context.Context, option ListOption) (*ListResult
for _, field := range option.Select {
queryParam.Select = append(queryParam.Select, field)
}
} else {
// Default: exclude the 'content' field (which may contain large text data)
// Only include it if explicitly requested in Select
queryParam.Select = []interface{}{
"id", "file_id", "uploader", "content_type", "name", "url", "description",
"type", "user_path", "path", "groups", "gzip", "bytes", "status",
"progress", "error", "preset", "public", "share",
"created_at", "updated_at", "deleted_at",
"__yao_created_by", "__yao_updated_by", "__yao_team_id", "__yao_tenant_id",
}
}
// Add filters
@ -1322,3 +1332,94 @@ func (manager Manager) getStoragePathFromDatabase(ctx context.Context, fileID st
return "", fmt.Errorf("invalid storage path for file ID: %s", fileID)
}
// GetText retrieves the parsed text content for a file by its ID
// By default, returns the preview (first 2000 characters) from 'content_preview' field
// Set fullContent to true to retrieve the complete text from 'content' field
func (manager Manager) GetText(ctx context.Context, fileID string, fullContent ...bool) (string, error) {
m := model.Select("__yao.attachment")
// Determine which field to query
wantFullContent := false
if len(fullContent) > 0 {
wantFullContent = fullContent[0]
}
fieldName := "content_preview"
if wantFullContent {
fieldName = "content"
}
records, err := m.Get(model.QueryParam{
Select: []interface{}{fieldName},
Wheres: []model.QueryWhere{
{Column: "file_id", Value: fileID},
},
Limit: 1,
})
if err != nil {
return "", fmt.Errorf("failed to query text content: %w", err)
}
if len(records) == 0 {
return "", fmt.Errorf("file not found: %s", fileID)
}
// Handle content field - it may be nil, string, or other types
if content, ok := records[0][fieldName].(string); ok {
return content, nil
}
// If content is nil or not a string, return empty string
return "", nil
}
// SaveText saves the parsed text content for a file by its ID
// Automatically saves both full content and preview (first 2000 characters)
// Updates both 'content' and 'content_preview' fields in the attachment record
func (manager Manager) SaveText(ctx context.Context, fileID string, text string) error {
m := model.Select("__yao.attachment")
// Check if record exists first
records, err := m.Get(model.QueryParam{
Select: []interface{}{"file_id"},
Wheres: []model.QueryWhere{
{Column: "file_id", Value: fileID},
},
Limit: 1,
})
if err != nil {
return fmt.Errorf("failed to check file existence: %w", err)
}
if len(records) == 0 {
return fmt.Errorf("file not found: %s", fileID)
}
// Create preview: first 2000 characters (or runes for proper UTF-8 handling)
preview := text
const maxPreviewLength = 2000
if len([]rune(text)) > maxPreviewLength {
preview = string([]rune(text)[:maxPreviewLength])
}
// Update both content and content_preview fields
updateData := map[string]interface{}{
"content": text,
"content_preview": preview,
}
_, err = m.UpdateWhere(model.QueryParam{
Wheres: []model.QueryWhere{
{Column: "file_id", Value: fileID},
},
}, updateData)
if err != nil {
return fmt.Errorf("failed to save text content: %w", err)
}
return nil
}

327
attachment/manager_test.go
View file

@ -1477,3 +1477,330 @@ func TestManagerLocalPath_ValidationFlow(t *testing.T) {
t.Logf("Warning: Failed to delete test file: %v", err)
}
}
func TestGetTextAndSaveText(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
manager, err := RegisterDefault("test-text-content")
if err != nil {
t.Fatalf("Failed to register manager: %v", err)
}
// Upload a test file
content := "This is a test file for text content storage"
reader := strings.NewReader(content)
fileHeader := &FileHeader{
FileHeader: &multipart.FileHeader{
Filename: "test-text.txt",
Size: int64(len(content)),
Header: make(map[string][]string),
},
}
fileHeader.Header.Set("Content-Type", "text/plain")
option := UploadOption{
Groups: []string{"test"},
OriginalFilename: "test-text.txt",
}
file, err := manager.Upload(context.Background(), fileHeader, reader, option)
if err != nil {
t.Fatalf("Failed to upload file: %v", err)
}
// Test 1: GetText on file without saved text (should return empty)
t.Run("GetTextEmpty", func(t *testing.T) {
text, err := manager.GetText(context.Background(), file.ID)
if err != nil {
t.Fatalf("Failed to get text: %v", err)
}
if text != "" {
t.Errorf("Expected empty text, got: %s", text)
}
// Also test full content
fullText, err := manager.GetText(context.Background(), file.ID, true)
if err != nil {
t.Fatalf("Failed to get full text: %v", err)
}
if fullText != "" {
t.Errorf("Expected empty full text, got: %s", fullText)
}
})
// Test 2: SaveText and verify
t.Run("SaveTextAndVerify", func(t *testing.T) {
parsedText := "This is the parsed text content from the file. It could be extracted from PDF, Word, or image OCR."
err := manager.SaveText(context.Background(), file.ID, parsedText)
if err != nil {
t.Fatalf("Failed to save text: %v", err)
}
// Retrieve the saved text
retrievedText, err := manager.GetText(context.Background(), file.ID)
if err != nil {
t.Fatalf("Failed to get saved text: %v", err)
}
if retrievedText != parsedText {
t.Errorf("Text mismatch. Expected: %s, Got: %s", parsedText, retrievedText)
}
t.Logf("Successfully saved and retrieved text content (%d characters)", len(retrievedText))
})
// Test 3: Update existing text
t.Run("UpdateText", func(t *testing.T) {
updatedText := "This is the updated parsed text content with additional information."
err := manager.SaveText(context.Background(), file.ID, updatedText)
if err != nil {
t.Fatalf("Failed to update text: %v", err)
}
retrievedText, err := manager.GetText(context.Background(), file.ID)
if err != nil {
t.Fatalf("Failed to get updated text: %v", err)
}
if retrievedText != updatedText {
t.Errorf("Updated text mismatch. Expected: %s, Got: %s", updatedText, retrievedText)
}
})
// Test 4: Save long text content and verify preview vs full content
t.Run("SaveLongText", func(t *testing.T) {
// Generate a large text content (10KB)
longText := strings.Repeat("This is a long text content that simulates parsing from a large document like PDF or Word. ", 100)
err := manager.SaveText(context.Background(), file.ID, longText)
if err != nil {
t.Fatalf("Failed to save long text: %v", err)
}
// Get preview (default, should be limited to 2000 characters)
previewText, err := manager.GetText(context.Background(), file.ID)
if err != nil {
t.Fatalf("Failed to get preview text: %v", err)
}
// Preview should be exactly 2000 characters (runes)
previewRunes := []rune(previewText)
if len(previewRunes) != 2000 {
t.Errorf("Preview length mismatch. Expected: 2000 runes, Got: %d runes", len(previewRunes))
}
// Get full content
fullText, err := manager.GetText(context.Background(), file.ID, true)
if err != nil {
t.Fatalf("Failed to get full text: %v", err)
}
if fullText != longText {
t.Errorf("Full text mismatch. Expected length: %d, Got: %d", len(longText), len(fullText))
}
t.Logf("Successfully saved long text - Preview: %d chars, Full: %d chars", len(previewText), len(fullText))
})
// Test 5: Test UTF-8 character handling in preview
t.Run("UTF8PreviewHandling", func(t *testing.T) {
// Create text with multi-byte UTF-8 characters (Chinese, emoji, etc.)
// Each Chinese character is 3 bytes, emoji is 4 bytes
chineseText := strings.Repeat("这是一个测试文本，包含中文字符。", 150) // Should exceed 2000 chars
emojiText := strings.Repeat("Hello 👋 World 🌍 ", 150)
// Test with Chinese text
err := manager.SaveText(context.Background(), file.ID, chineseText)
if err != nil {
t.Fatalf("Failed to save Chinese text: %v", err)
}
previewChinese, err := manager.GetText(context.Background(), file.ID)
if err != nil {
t.Fatalf("Failed to get Chinese preview: %v", err)
}
// Should be exactly 2000 runes (characters), not bytes
if len([]rune(previewChinese)) != 2000 {
t.Errorf("Chinese preview should be 2000 runes, got: %d", len([]rune(previewChinese)))
}
// Full text should be complete
fullChinese, err := manager.GetText(context.Background(), file.ID, true)
if err != nil {
t.Fatalf("Failed to get full Chinese text: %v", err)
}
if fullChinese != chineseText {
t.Errorf("Chinese text mismatch")
}
// Test with emoji text
err = manager.SaveText(context.Background(), file.ID, emojiText)
if err != nil {
t.Fatalf("Failed to save emoji text: %v", err)
}
previewEmoji, err := manager.GetText(context.Background(), file.ID)
if err != nil {
t.Fatalf("Failed to get emoji preview: %v", err)
}
if len([]rune(previewEmoji)) != 2000 {
t.Errorf("Emoji preview should be 2000 runes, got: %d", len([]rune(previewEmoji)))
}
t.Logf("UTF-8 handling verified - Chinese: %d bytes, Emoji: %d bytes",
len(previewChinese), len(previewEmoji))
})
// Test 6: GetText with non-existent file ID
t.Run("GetTextNonExistent", func(t *testing.T) {
_, err := manager.GetText(context.Background(), "non-existent-id")
if err == nil {
t.Error("Expected error for non-existent file ID")
}
if !strings.Contains(err.Error(), "file not found") {
t.Errorf("Expected 'file not found' error, got: %s", err.Error())
}
})
// Test 7: SaveText with non-existent file ID
t.Run("SaveTextNonExistent", func(t *testing.T) {
err := manager.SaveText(context.Background(), "non-existent-id", "some text")
if err == nil {
t.Error("Expected error for non-existent file ID")
}
if !strings.Contains(err.Error(), "file not found") {
t.Errorf("Expected 'file not found' error, got: %s", err.Error())
}
})
// Test 8: Save empty text (clear content)
t.Run("SaveEmptyText", func(t *testing.T) {
err := manager.SaveText(context.Background(), file.ID, "")
if err != nil {
t.Fatalf("Failed to save empty text: %v", err)
}
retrievedText, err := manager.GetText(context.Background(), file.ID)
if err != nil {
t.Fatalf("Failed to get empty text: %v", err)
}
if retrievedText != "" {
t.Errorf("Expected empty text, got: %s", retrievedText)
}
})
// Test 9: Verify List doesn't include content fields by default
t.Run("ListExcludesContentByDefault", func(t *testing.T) {
// Save some text content
testText := "This text should not appear in list results by default"
err := manager.SaveText(context.Background(), file.ID, testText)
if err != nil {
t.Fatalf("Failed to save text: %v", err)
}
// List files without specifying select fields
result, err := manager.List(context.Background(), ListOption{
Filters: map[string]interface{}{
"file_id": file.ID,
},
})
if err != nil {
t.Fatalf("Failed to list files: %v", err)
}
if len(result.Files) == 0 {
t.Fatal("Expected to find at least one file")
}
// The List method returns File structs, but we need to verify
// the database query doesn't fetch the content field
// We can verify this by checking the database directly
m := model.Select("__yao.attachment")
records, err := m.Get(model.QueryParam{
Wheres: []model.QueryWhere{
{Column: "file_id", Value: file.ID},
},
})
if err != nil {
t.Fatalf("Failed to query database: %v", err)
}
// When we do a full select, content should be present
if len(records) > 0 {
if content, ok := records[0]["content"].(string); ok && content == testText {
t.Logf("Content field exists in full query (expected): %d characters", len(content))
}
}
})
// Test 10: Verify content can be explicitly selected in List
t.Run("ListIncludesContentWhenExplicitlySelected", func(t *testing.T) {
// Save some text content
testText := "This text SHOULD appear when explicitly selected"
err := manager.SaveText(context.Background(), file.ID, testText)
if err != nil {
t.Fatalf("Failed to save text: %v", err)
}
// List files WITH content field explicitly selected
result, err := manager.List(context.Background(), ListOption{
Select: []string{"file_id", "name", "content"},
Filters: map[string]interface{}{
"file_id": file.ID,
},
})
if err != nil {
t.Fatalf("Failed to list files with content: %v", err)
}
if len(result.Files) == 0 {
t.Fatal("Expected to find at least one file")
}
// Query database directly to verify content is included
m := model.Select("__yao.attachment")
records, err := m.Get(model.QueryParam{
Select: []interface{}{"file_id", "name", "content"},
Wheres: []model.QueryWhere{
{Column: "file_id", Value: file.ID},
},
})
if err != nil {
t.Fatalf("Failed to query database: %v", err)
}
if len(records) == 0 {
t.Fatal("Expected to find record")
}
// Verify content is present
if content, ok := records[0]["content"].(string); ok {
if content != testText {
t.Errorf("Expected content '%s', got '%s'", testText, content)
}
t.Logf("Content field correctly included when explicitly selected: %d characters", len(content))
} else {
t.Error("Content field not found when explicitly selected")
}
})
// Clean up
err = manager.Delete(context.Background(), file.ID)
if err != nil {
t.Logf("Warning: Failed to delete test file: %v", err)
}
}

8
attachment/types.go
View file

@ -47,6 +47,14 @@ type FileManager interface {
// LocalPath gets the local path of the file
LocalPath(ctx context.Context, fileID string) (string, string, error)
// GetText retrieves the parsed text content for a file
// By default returns preview (first 2000 chars), set fullContent=true for complete text
GetText(ctx context.Context, fileID string, fullContent ...bool) (string, error)
// SaveText saves the parsed text content for a file
// Automatically saves both full content and preview
SaveText(ctx context.Context, fileID string, text string) error
}
// File the file

284
data/bindata.go
View file

File diff suppressed because it is too large Load diff

4
openapi/chat/completions.go
View file

@ -25,7 +25,7 @@ func GinCreateCompletions(c *gin.Context) {
return
}
completionReq, ctx, err := context.GetCompletionRequest(c, cache)
completionReq, ctx, opts, err := context.GetCompletionRequest(c, cache)
if err != nil {
fmt.Println("-----------------------------------------------")
fmt.Println("Error: ", err.Error())
@ -61,7 +61,7 @@ func GinCreateCompletions(c *gin.Context) {
// Stream the completion (uses default handler which sends to ctx.Writer)
// The Stream method will automatically close the writer and send [DONE] marker
log.Trace("[HTTP] Calling ast.Stream()")
_, err = ast.Stream(ctx, completionReq.Messages)
_, err = ast.Stream(ctx, completionReq.Messages, opts)
log.Trace("[HTTP] ast.Stream() returned, err=%v", err)
if err != nil {
response.RespondWithError(c, response.StatusInternalServerError, &response.ErrorResponse{

14
yao/models/attachment.mod.yao
View file

@ -45,6 +45,20 @@
"nullable": false,
"index": true
},
{
"name": "content",
"type": "longText",
"label": "Content",
"comment": "Full parsed text content from image, pdf, word and other file types",
"nullable": true
},
{
"name": "content_preview",
"type": "text",
"label": "Content Preview",
"comment": "Preview of parsed text content (first 2000 characters)",
"nullable": true
},
{
"name": "name",
"type": "string",