From b3cf5a09d98c8136a58de155ead451845b5cfafa Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 13 Dec 2025 14:05:26 +0800 Subject: [PATCH] Enhance Assistant and Search Modules with Contextual Logging and Refactorings - Added logging for hook start and completion in the Assistant's Stream method to improve traceability during execution. - Refactored the AgentGetterFunc to utilize the caller package, enhancing modularity and reducing circular dependencies. - Updated the CallAgent function to check for the initialized AgentGetterFunc from the caller package, ensuring proper agent loading. - Enhanced the Search handler to support an optional context parameter, improving flexibility for agent mode operations. - Refined the agentSearch function to delegate search requests to a new AgentProvider, streamlining the search process. - Updated DESIGN.md to reflect changes in search modes and the integration of the caller package, ensuring comprehensive documentation. --- agent/assistant/agent.go | 4 + agent/assistant/assistant.go | 6 +- agent/caller/caller.go | 17 ++ agent/content/tools.go | 13 +- agent/search/DESIGN.md | 4 +- agent/search/handlers/web/agent.go | 232 +++++++++++++++++++ agent/search/handlers/web/agent_test.go | 284 ++++++++++++++++++++++++ agent/search/handlers/web/handler.go | 57 +++-- agent/search/handlers/web/mcp.go | 203 +++++++++++++++++ agent/search/handlers/web/mcp_test.go | 241 ++++++++++++++++++++ 10 files changed, 1017 insertions(+), 44 deletions(-) create mode 100644 agent/caller/caller.go create mode 100644 agent/search/handlers/web/agent.go create mode 100644 agent/search/handlers/web/agent_test.go create mode 100644 agent/search/handlers/web/mcp.go create mode 100644 agent/search/handlers/web/mcp_test.go diff --git a/agent/assistant/agent.go b/agent/assistant/agent.go index 5250160a..12ed9669 100644 --- a/agent/assistant/agent.go +++ b/agent/assistant/agent.go @@ -367,6 +367,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa var nextResponse *context.NextHookResponse = nil if ast.HookScript != nil { + ctx.Logger.HookStart("Next") + // Begin step tracking for hook_next ast.BeginStep(ctx, context.StepTypeHookNext, map[string]interface{}{ "messages": fullMessages, @@ -393,6 +395,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa "response": nextResponse, }) + ctx.Logger.HookComplete("Next") + // Process Next hook response finalResponse, err = ast.processNextResponse(&NextProcessContext{ Context: ctx, diff --git a/agent/assistant/assistant.go b/agent/assistant/assistant.go index b048a92b..9ef9f589 100644 --- a/agent/assistant/assistant.go +++ b/agent/assistant/assistant.go @@ -5,7 +5,7 @@ import ( "path" "github.com/yaoapp/gou/fs" - "github.com/yaoapp/yao/agent/content" + "github.com/yaoapp/yao/agent/caller" agentContext "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/i18n" searchTypes "github.com/yaoapp/yao/agent/search/types" @@ -14,8 +14,8 @@ import ( ) func init() { - // Initialize AgentGetterFunc to allow content package to call agents - content.AgentGetterFunc = func(agentID string) (content.AgentCaller, error) { + // Initialize AgentGetterFunc to allow content and search packages to call agents + caller.AgentGetterFunc = func(agentID string) (caller.AgentCaller, error) { ast, err := Get(agentID) if err != nil { return nil, err diff --git a/agent/caller/caller.go b/agent/caller/caller.go new file mode 100644 index 00000000..fb4d2049 --- /dev/null +++ b/agent/caller/caller.go @@ -0,0 +1,17 @@ +// Package caller provides a shared interface for calling agents +// This package is used by both content and search packages to avoid circular dependencies +package caller + +import ( + agentContext "github.com/yaoapp/yao/agent/context" +) + +// AgentCaller interface for calling agents (to avoid circular dependency) +// Used by content handlers (vision, audio, etc.) and search handlers (agent mode) +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 +// This should be set by the assistant package during initialization +var AgentGetterFunc func(agentID string) (AgentCaller, error) diff --git a/agent/content/tools.go b/agent/content/tools.go index 9a75436e..cc9d7787 100644 --- a/agent/content/tools.go +++ b/agent/content/tools.go @@ -9,29 +9,22 @@ import ( jsoniter "github.com/json-iterator/go" "github.com/yaoapp/gou/mcp" "github.com/yaoapp/kun/log" + "github.com/yaoapp/yao/agent/caller" 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) - // fileInfoMutex protects concurrent access to files_info list in Space var fileInfoMutex sync.Mutex // 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 { + if caller.AgentGetterFunc == nil { return "", fmt.Errorf("AgentGetterFunc not initialized") } // Load the agent by ID using the injected function - agent, err := AgentGetterFunc(agentID) + agent, err := caller.AgentGetterFunc(agentID) if err != nil { return "", fmt.Errorf("failed to load agent %s: %w", agentID, err) } diff --git a/agent/search/DESIGN.md b/agent/search/DESIGN.md index 5f7afaac..a9812236 100644 --- a/agent/search/DESIGN.md +++ b/agent/search/DESIGN.md @@ -1110,7 +1110,7 @@ Tool format: `"builtin"`, `""` (Agent), `"mcp:."` (M | Mode | Example | Description | | --------- | ---------------------------- | -------------------------------------------------------------------------- | -| `builtin` | `"builtin"` | Use built-in providers (Tavily, Serper) | +| `builtin` | `"builtin"` | Use built-in providers (Tavily, Serper, SerpAPI) | | Agent | `"workers.search.web"` | AI-powered search: understand intent → optimize query → search → summarize | | MCP | `"mcp:my-server.web_search"` | External search tool via MCP protocol | @@ -1707,7 +1707,7 @@ Web search supports three modes via `uses.web`: | Mode | Value | Description | | ------- | ---------------------------- | ------------------------------------------- | -| Builtin | `"builtin"` | Direct API calls to Tavily/Serper | +| Builtin | `"builtin"` | Direct API calls to Tavily/Serper/SerpAPI | | Agent | `"workers.search.web"` | AI-powered search with intent understanding | | MCP | `"mcp:my-server.web_search"` | External search tool via MCP | diff --git a/agent/search/handlers/web/agent.go b/agent/search/handlers/web/agent.go new file mode 100644 index 00000000..e23c2bc6 --- /dev/null +++ b/agent/search/handlers/web/agent.go @@ -0,0 +1,232 @@ +package web + +import ( + "encoding/json" + "fmt" + "time" + + "github.com/yaoapp/yao/agent/caller" + agentContext "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/search/types" +) + +// AgentProvider implements web search using another agent (AI Search) +type AgentProvider struct { + agentID string // Agent/Assistant ID (e.g., "workers.search.web") +} + +// NewAgentProvider creates a new Agent provider +func NewAgentProvider(agentID string) *AgentProvider { + return &AgentProvider{ + agentID: agentID, + } +} + +// Search executes web search via agent delegation +// The agent can understand intent, generate optimized queries, and return structured results +func (p *AgentProvider) Search(ctx *agentContext.Context, req *types.Request) (*types.Result, error) { + startTime := time.Now() + + // Check if context is provided + if ctx == nil { + return &types.Result{ + Type: types.SearchTypeWeb, + Query: req.Query, + Source: req.Source, + Items: []*types.ResultItem{}, + Total: 0, + Duration: time.Since(startTime).Milliseconds(), + Error: "Agent mode requires context", + }, nil + } + + // Check if AgentGetterFunc is initialized + if caller.AgentGetterFunc == nil { + return &types.Result{ + Type: types.SearchTypeWeb, + Query: req.Query, + Source: req.Source, + Items: []*types.ResultItem{}, + Total: 0, + Duration: time.Since(startTime).Milliseconds(), + Error: "AgentGetterFunc not initialized", + }, nil + } + + // Get the agent + agent, err := caller.AgentGetterFunc(p.agentID) + if err != nil { + return &types.Result{ + Type: types.SearchTypeWeb, + Query: req.Query, + Source: req.Source, + Items: []*types.ResultItem{}, + Total: 0, + Duration: time.Since(startTime).Milliseconds(), + Error: fmt.Sprintf("Agent '%s' not found: %v", p.agentID, err), + }, nil + } + + // Build message for the agent + // Include search parameters in the message content + searchParams := map[string]interface{}{ + "query": req.Query, + "type": "web", + "source": string(req.Source), + } + + if req.Limit > 0 { + searchParams["limit"] = req.Limit + } + if len(req.Sites) > 0 { + searchParams["sites"] = req.Sites + } + if req.TimeRange != "" { + searchParams["time_range"] = req.TimeRange + } + + // Convert to JSON for the message + paramsJSON, err := json.Marshal(searchParams) + if err != nil { + return &types.Result{ + Type: types.SearchTypeWeb, + Query: req.Query, + Source: req.Source, + Items: []*types.ResultItem{}, + Total: 0, + Duration: time.Since(startTime).Milliseconds(), + Error: fmt.Sprintf("Failed to serialize search params: %v", err), + }, nil + } + + // Create message for the agent + message := agentContext.Message{ + Role: "user", + Content: string(paramsJSON), + } + + // Call the agent with skip options (no history, no output) + opts := &agentContext.Options{ + Skip: &agentContext.Skip{ + History: true, + Output: true, + }, + } + + response, err := agent.Stream(ctx, []agentContext.Message{message}, opts) + if err != nil { + return &types.Result{ + Type: types.SearchTypeWeb, + Query: req.Query, + Source: req.Source, + Items: []*types.ResultItem{}, + Total: 0, + Duration: time.Since(startTime).Milliseconds(), + Error: fmt.Sprintf("Agent call failed: %v", err), + }, nil + } + + // Parse the agent response + items, total, parseErr := p.parseAgentResponse(response, req.Source) + if parseErr != "" { + return &types.Result{ + Type: types.SearchTypeWeb, + Query: req.Query, + Source: req.Source, + Items: []*types.ResultItem{}, + Total: 0, + Duration: time.Since(startTime).Milliseconds(), + Error: parseErr, + }, nil + } + + return &types.Result{ + Type: types.SearchTypeWeb, + Query: req.Query, + Source: req.Source, + Items: items, + Total: total, + Duration: time.Since(startTime).Milliseconds(), + }, nil +} + +// parseAgentResponse parses the agent response into search result items +// The agent should return a JSON structure with search results +func (p *AgentProvider) parseAgentResponse(response interface{}, source types.SourceType) ([]*types.ResultItem, int, string) { + if response == nil { + return nil, 0, "Agent returned nil response" + } + + // Try to extract data from response + var data map[string]interface{} + + // Handle different response types + switch v := response.(type) { + case map[string]interface{}: + data = v + case string: + // Try to parse as JSON + if err := json.Unmarshal([]byte(v), &data); err != nil { + return nil, 0, fmt.Sprintf("Failed to parse agent response as JSON: %v", err) + } + default: + // Try to marshal and unmarshal + jsonBytes, err := json.Marshal(response) + if err != nil { + return nil, 0, fmt.Sprintf("Failed to serialize agent response: %v", err) + } + if err := json.Unmarshal(jsonBytes, &data); err != nil { + return nil, 0, fmt.Sprintf("Failed to parse agent response: %v", err) + } + } + + // Check for "next" field (custom hook data) + if next, hasNext := data["next"]; hasNext && next != nil { + if nextMap, ok := next.(map[string]interface{}); ok { + data = nextMap + } else if nextStr, ok := next.(string); ok { + // Try to parse as JSON + if err := json.Unmarshal([]byte(nextStr), &data); err != nil { + return nil, 0, fmt.Sprintf("Failed to parse next hook data: %v", err) + } + } + } + + // Extract items from data + items := []*types.ResultItem{} + total := 0 + + if itemsData, ok := data["items"].([]interface{}); ok { + for _, itemData := range itemsData { + if item, ok := itemData.(map[string]interface{}); ok { + resultItem := &types.ResultItem{ + Type: types.SearchTypeWeb, + Source: source, + } + + if title, ok := item["title"].(string); ok { + resultItem.Title = title + } + if content, ok := item["content"].(string); ok { + resultItem.Content = content + } + if url, ok := item["url"].(string); ok { + resultItem.URL = url + } + if score, ok := item["score"].(float64); ok { + resultItem.Score = score + } + + items = append(items, resultItem) + } + } + } + + if totalVal, ok := data["total"].(float64); ok { + total = int(totalVal) + } else { + total = len(items) + } + + return items, total, "" +} diff --git a/agent/search/handlers/web/agent_test.go b/agent/search/handlers/web/agent_test.go new file mode 100644 index 00000000..8438c06f --- /dev/null +++ b/agent/search/handlers/web/agent_test.go @@ -0,0 +1,284 @@ +package web_test + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/yaoapp/yao/agent/assistant" + agentContext "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/search/handlers/web" + "github.com/yaoapp/yao/agent/search/types" + "github.com/yaoapp/yao/agent/testutils" + oauthTypes "github.com/yaoapp/yao/openapi/oauth/types" +) + +// TestAgentProviderWithAssistantConfig tests AgentProvider using web-agent-caller assistant config +func TestAgentProviderWithAssistantConfig(t *testing.T) { + if testing.Short() { + t.Skip("Skipping integration test") + } + + testutils.Prepare(t) + defer testutils.Clean(t) + + // Load the web-agent-caller test assistant to get its config + ast, err := assistant.LoadPath("/assistants/tests/web-agent-caller") + require.NoError(t, err) + require.NotNil(t, ast) + require.NotNil(t, ast.Uses) + + // Verify assistant config + assert.Equal(t, "tests.web-agent-caller", ast.ID) + assert.Equal(t, "tests.web-agent", ast.Uses.Web) + + // Create AgentProvider from uses.web + provider := web.NewAgentProvider(ast.Uses.Web) + require.NotNil(t, provider) + + // Create a mock context + ctx := createTestContext(t) + + // Execute search + req := &types.Request{ + Query: "Yao App Engine", + Type: types.SearchTypeWeb, + Source: types.SourceAuto, + Limit: 5, + } + + result, err := provider.Search(ctx, req) + require.NoError(t, err) + require.NotNil(t, result) + + // Verify result structure + assert.Equal(t, types.SearchTypeWeb, result.Type) + assert.Equal(t, "Yao App Engine", result.Query) + assert.Equal(t, types.SourceAuto, result.Source) + + // Agent should return mock results from Next hook + if result.Error == "" { + assert.Greater(t, result.Total, 0) + assert.NotEmpty(t, result.Items) + assert.Greater(t, result.Duration, int64(0)) + + // Verify result item structure + for _, item := range result.Items { + assert.Equal(t, types.SearchTypeWeb, item.Type) + assert.Equal(t, types.SourceAuto, item.Source) + assert.NotEmpty(t, item.Title) + assert.NotEmpty(t, item.URL) + } + + t.Logf("Agent search returned %d results in %dms", result.Total, result.Duration) + } else { + t.Logf("Agent search returned error: %s", result.Error) + } +} + +// TestAgentProviderWithSiteRestriction tests AgentProvider with domain restriction +func TestAgentProviderWithSiteRestriction(t *testing.T) { + if testing.Short() { + t.Skip("Skipping integration test") + } + + testutils.Prepare(t) + defer testutils.Clean(t) + + // Create AgentProvider + provider := web.NewAgentProvider("tests.web-agent") + + // Create a mock context + ctx := createTestContext(t) + + // Execute search with site restriction + req := &types.Request{ + Query: "documentation", + Type: types.SearchTypeWeb, + Source: types.SourceHook, + Sites: []string{"github.com"}, + Limit: 3, + } + + result, err := provider.Search(ctx, req) + require.NoError(t, err) + require.NotNil(t, result) + + assert.Equal(t, types.SearchTypeWeb, result.Type) + assert.Equal(t, types.SourceHook, result.Source) + + if result.Error == "" { + // All results should be from github.com (mock data respects sites) + for _, item := range result.Items { + assert.Contains(t, item.URL, "github.com", "Result URL should be from github.com") + } + t.Logf("Site-restricted agent search returned %d results", result.Total) + } else { + t.Logf("Agent search returned error: %s", result.Error) + } +} + +// TestAgentProviderWithTimeRange tests AgentProvider with time range filter +func TestAgentProviderWithTimeRange(t *testing.T) { + if testing.Short() { + t.Skip("Skipping integration test") + } + + testutils.Prepare(t) + defer testutils.Clean(t) + + // Create AgentProvider + provider := web.NewAgentProvider("tests.web-agent") + + // Create a mock context + ctx := createTestContext(t) + + // Execute search with time range + req := &types.Request{ + Query: "artificial intelligence news", + Type: types.SearchTypeWeb, + Source: types.SourceAuto, + TimeRange: "week", + Limit: 5, + } + + result, err := provider.Search(ctx, req) + require.NoError(t, err) + require.NotNil(t, result) + + if result.Error == "" { + t.Logf("Time-ranged agent search (last week) returned %d results in %dms", result.Total, result.Duration) + } else { + t.Logf("Agent search returned error: %s", result.Error) + } +} + +// TestAgentProviderNotFound tests AgentProvider when agent is not found +func TestAgentProviderNotFound(t *testing.T) { + if testing.Short() { + t.Skip("Skipping integration test") + } + + testutils.Prepare(t) + defer testutils.Clean(t) + + // Create AgentProvider with non-existent agent + provider := web.NewAgentProvider("nonexistent.agent") + + // Create a mock context + ctx := createTestContext(t) + + req := &types.Request{ + Query: "test query", + Type: types.SearchTypeWeb, + Source: types.SourceAuto, + } + + result, err := provider.Search(ctx, req) + + // Should not return error, but result should have error message + require.NoError(t, err) + require.NotNil(t, result) + assert.NotEmpty(t, result.Error) + assert.Contains(t, result.Error, "not found") +} + +// TestAgentProviderWithoutContext tests AgentProvider without context +func TestAgentProviderWithoutContext(t *testing.T) { + if testing.Short() { + t.Skip("Skipping integration test") + } + + testutils.Prepare(t) + defer testutils.Clean(t) + + // Create AgentProvider + provider := web.NewAgentProvider("tests.web-agent") + + req := &types.Request{ + Query: "test query", + Type: types.SearchTypeWeb, + Source: types.SourceAuto, + } + + // Call without context (nil) + result, err := provider.Search(nil, req) + + // Should still work - agent provider handles nil context + require.NoError(t, err) + require.NotNil(t, result) + // May have error if context is required for agent call + t.Logf("Agent search without context: error=%s, total=%d", result.Error, result.Total) +} + +// TestWebHandlerAgentMode tests the web handler in agent mode +func TestWebHandlerAgentMode(t *testing.T) { + if testing.Short() { + t.Skip("Skipping integration test") + } + + testutils.Prepare(t) + defer testutils.Clean(t) + + // Create handler with agent mode + handler := web.NewHandler("tests.web-agent", nil) + require.NotNil(t, handler) + + // Verify type + assert.Equal(t, types.SearchTypeWeb, handler.Type()) + + // Create a mock context + ctx := createTestContext(t) + + // Execute search with context + req := &types.Request{ + Query: "Yao framework", + Type: types.SearchTypeWeb, + Source: types.SourceAuto, + Limit: 5, + } + + result, err := handler.SearchWithContext(ctx, req) + require.NoError(t, err) + require.NotNil(t, result) + + assert.Equal(t, types.SearchTypeWeb, result.Type) + assert.Equal(t, "Yao framework", result.Query) + + if result.Error == "" { + t.Logf("Handler agent mode returned %d results", result.Total) + } else { + t.Logf("Handler agent mode returned error: %s", result.Error) + } +} + +// TestWebHandlerAgentModeWithoutContext tests the web handler in agent mode without context +func TestWebHandlerAgentModeWithoutContext(t *testing.T) { + // Create handler with agent mode + handler := web.NewHandler("tests.web-agent", nil) + require.NotNil(t, handler) + + req := &types.Request{ + Query: "test", + Type: types.SearchTypeWeb, + Source: types.SourceAuto, + } + + // Call Search() without context (uses SearchWithContext with nil) + result, err := handler.Search(req) + require.NoError(t, err) + require.NotNil(t, result) + assert.NotEmpty(t, result.Error) + assert.Contains(t, result.Error, "requires context") +} + +// createTestContext creates a test context for agent calls +func createTestContext(t *testing.T) *agentContext.Context { + authorized := &oauthTypes.AuthorizedInfo{ + UserID: "test-user", + TenantID: "test-tenant", + } + ctx := agentContext.New(nil, authorized, "test-chat-id") + ctx.AssistantID = "tests.web-agent-caller" + return ctx +} diff --git a/agent/search/handlers/web/handler.go b/agent/search/handlers/web/handler.go index fd788f57..d45eaf54 100644 --- a/agent/search/handlers/web/handler.go +++ b/agent/search/handlers/web/handler.go @@ -4,6 +4,7 @@ import ( "fmt" "strings" + agentContext "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/search/types" ) @@ -24,7 +25,14 @@ func (h *Handler) Type() types.SearchType { } // Search executes web search based on uses.web mode +// ctx is optional and only required for agent mode func (h *Handler) Search(req *types.Request) (*types.Result, error) { + return h.SearchWithContext(nil, req) +} + +// SearchWithContext executes web search with optional agent context +// ctx is required for agent mode, optional for builtin and MCP modes +func (h *Handler) SearchWithContext(ctx *agentContext.Context, req *types.Request) (*types.Result, error) { switch { case h.usesWeb == "builtin" || h.usesWeb == "": return h.builtinSearch(req) @@ -32,7 +40,17 @@ func (h *Handler) Search(req *types.Request) (*types.Result, error) { return h.mcpSearch(req) default: // Agent mode: delegate to assistant for AI-powered search - return h.agentSearch(req) + if ctx == nil { + return &types.Result{ + Type: types.SearchTypeWeb, + Query: req.Query, + Source: req.Source, + Items: []*types.ResultItem{}, + Total: 0, + Error: "Agent mode requires context", + }, nil + } + return h.agentSearch(ctx, req) } } @@ -66,46 +84,27 @@ func (h *Handler) builtinSearch(req *types.Request) (*types.Result, error) { } // agentSearch delegates to an assistant for AI-powered search -func (h *Handler) agentSearch(req *types.Request) (*types.Result, error) { - // TODO: Implement agent mode - // 1. Call assistant with search request - // 2. Assistant understands intent, generates optimized queries - // 3. Assistant executes searches (may call builtin internally) - // 4. Assistant analyzes and returns structured results - return &types.Result{ - Type: types.SearchTypeWeb, - Query: req.Query, - Source: req.Source, - Items: []*types.ResultItem{}, - Total: 0, - Error: "Agent mode not yet implemented", - }, nil +func (h *Handler) agentSearch(ctx *agentContext.Context, req *types.Request) (*types.Result, error) { + provider := NewAgentProvider(h.usesWeb) + return provider.Search(ctx, req) } // mcpSearch calls external MCP tool func (h *Handler) mcpSearch(req *types.Request) (*types.Result, error) { - // TODO: Implement MCP mode // Parse "mcp:server.tool" mcpRef := strings.TrimPrefix(h.usesWeb, "mcp:") - parts := strings.SplitN(mcpRef, ".", 2) - if len(parts) != 2 { + + provider, err := NewMCPProvider(mcpRef) + if err != nil { return &types.Result{ Type: types.SearchTypeWeb, Query: req.Query, Source: req.Source, Items: []*types.ResultItem{}, Total: 0, - Error: fmt.Sprintf("Invalid MCP format, expected 'mcp:server.tool', got '%s'", h.usesWeb), + Error: fmt.Sprintf("Invalid MCP format: %v", err), }, nil } - // serverID, toolName := parts[0], parts[1] - // Call MCP tool - return &types.Result{ - Type: types.SearchTypeWeb, - Query: req.Query, - Source: req.Source, - Items: []*types.ResultItem{}, - Total: 0, - Error: "MCP mode not yet implemented", - }, nil + + return provider.Search(req) } diff --git a/agent/search/handlers/web/mcp.go b/agent/search/handlers/web/mcp.go new file mode 100644 index 00000000..82888707 --- /dev/null +++ b/agent/search/handlers/web/mcp.go @@ -0,0 +1,203 @@ +package web + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "time" + + "github.com/yaoapp/gou/mcp" + gouMCPTypes "github.com/yaoapp/gou/mcp/types" + "github.com/yaoapp/yao/agent/search/types" +) + +// MCPProvider implements web search using MCP tool +type MCPProvider struct { + serverID string // MCP server ID (e.g., "search") + toolName string // MCP tool name (e.g., "web_search") +} + +// NewMCPProvider creates a new MCP provider from "mcp:server.tool" format +func NewMCPProvider(mcpRef string) (*MCPProvider, error) { + // Parse "server.tool" format + parts := strings.SplitN(mcpRef, ".", 2) + if len(parts) != 2 { + return nil, fmt.Errorf("invalid MCP format, expected 'server.tool', got '%s'", mcpRef) + } + + return &MCPProvider{ + serverID: parts[0], + toolName: parts[1], + }, nil +} + +// Search executes web search via MCP tool +func (p *MCPProvider) Search(req *types.Request) (*types.Result, error) { + startTime := time.Now() + + // Select MCP client + client, err := mcp.Select(p.serverID) + if err != nil { + return &types.Result{ + Type: types.SearchTypeWeb, + Query: req.Query, + Source: req.Source, + Items: []*types.ResultItem{}, + Total: 0, + Duration: time.Since(startTime).Milliseconds(), + Error: fmt.Sprintf("MCP client '%s' not found: %v", p.serverID, err), + }, nil + } + + // Build MCP tool arguments + args := map[string]interface{}{ + "query": req.Query, + } + + if req.Limit > 0 { + args["limit"] = req.Limit + } + + if len(req.Sites) > 0 { + args["sites"] = req.Sites + } + + if req.TimeRange != "" { + args["time_range"] = req.TimeRange + } + + // Call MCP tool + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + result, err := client.CallTool(ctx, p.toolName, args) + if err != nil { + return &types.Result{ + Type: types.SearchTypeWeb, + Query: req.Query, + Source: req.Source, + Items: []*types.ResultItem{}, + Total: 0, + Duration: time.Since(startTime).Milliseconds(), + Error: fmt.Sprintf("MCP tool call failed: %v", err), + }, nil + } + + // Parse MCP result + items, total, parseErr := p.parseResult(result, req.Source) + if parseErr != "" { + return &types.Result{ + Type: types.SearchTypeWeb, + Query: req.Query, + Source: req.Source, + Items: []*types.ResultItem{}, + Total: 0, + Duration: time.Since(startTime).Milliseconds(), + Error: parseErr, + }, nil + } + + return &types.Result{ + Type: types.SearchTypeWeb, + Query: req.Query, + Source: req.Source, + Items: items, + Total: total, + Duration: time.Since(startTime).Milliseconds(), + }, nil +} + +// parseResult parses MCP tool result into search result items +func (p *MCPProvider) parseResult(result *gouMCPTypes.CallToolResponse, source types.SourceType) ([]*types.ResultItem, int, string) { + if result == nil { + return nil, 0, "MCP returned nil result" + } + + // Check for errors in result + if result.IsError { + errMsg := "MCP tool returned error" + if len(result.Content) > 0 && result.Content[0].Text != "" { + errMsg = result.Content[0].Text + } + return nil, 0, errMsg + } + + // Parse content - expect JSON data + if len(result.Content) == 0 { + return []*types.ResultItem{}, 0, "" + } + + // Try to extract data from content + var data map[string]interface{} + + for _, content := range result.Content { + // Check text content type + if content.Type == gouMCPTypes.ToolContentTypeText && content.Text != "" { + // Try to parse as JSON + if parsed, ok := parseJSON(content.Text); ok { + data = parsed + break + } + } + } + + if data == nil { + return []*types.ResultItem{}, 0, "" + } + + // Extract items from data + items := []*types.ResultItem{} + total := 0 + + if itemsData, ok := data["items"].([]interface{}); ok { + for _, itemData := range itemsData { + if item, ok := itemData.(map[string]interface{}); ok { + resultItem := &types.ResultItem{ + Type: types.SearchTypeWeb, + Source: source, + } + + if title, ok := item["title"].(string); ok { + resultItem.Title = title + } + if content, ok := item["content"].(string); ok { + resultItem.Content = content + } + if url, ok := item["url"].(string); ok { + resultItem.URL = url + } + if score, ok := item["score"].(float64); ok { + resultItem.Score = score + } + + items = append(items, resultItem) + } + } + } + + if totalVal, ok := data["total"].(float64); ok { + total = int(totalVal) + } else { + total = len(items) + } + + return items, total, "" +} + +// parseJSON attempts to parse a string as JSON +func parseJSON(s string) (map[string]interface{}, bool) { + // Simple JSON detection - if it starts with { and ends with } + s = strings.TrimSpace(s) + if !strings.HasPrefix(s, "{") || !strings.HasSuffix(s, "}") { + return nil, false + } + + // Use encoding/json for parsing + var result map[string]interface{} + if err := json.Unmarshal([]byte(s), &result); err != nil { + return nil, false + } + + return result, true +} diff --git a/agent/search/handlers/web/mcp_test.go b/agent/search/handlers/web/mcp_test.go new file mode 100644 index 00000000..3b9217d5 --- /dev/null +++ b/agent/search/handlers/web/mcp_test.go @@ -0,0 +1,241 @@ +package web_test + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/yaoapp/yao/agent/assistant" + "github.com/yaoapp/yao/agent/search/handlers/web" + "github.com/yaoapp/yao/agent/search/types" + "github.com/yaoapp/yao/agent/testutils" +) + +// TestMCPProviderWithAssistantConfig tests MCPProvider using web-mcp assistant config +func TestMCPProviderWithAssistantConfig(t *testing.T) { + if testing.Short() { + t.Skip("Skipping integration test") + } + + testutils.Prepare(t) + defer testutils.Clean(t) + + // Load the web-mcp test assistant to get its config + ast, err := assistant.LoadPath("/assistants/tests/web-mcp") + require.NoError(t, err) + require.NotNil(t, ast) + require.NotNil(t, ast.Uses) + + // Verify assistant config + assert.Equal(t, "tests.web-mcp", ast.ID) + assert.Equal(t, "mcp:search.web_search", ast.Uses.Web) + + // Create MCPProvider from uses.web + mcpRef := ast.Uses.Web[4:] // Remove "mcp:" prefix + provider, err := web.NewMCPProvider(mcpRef) + require.NoError(t, err) + require.NotNil(t, provider) + + // Execute search + req := &types.Request{ + Query: "Yao App Engine", + Type: types.SearchTypeWeb, + Source: types.SourceAuto, + Limit: 5, + } + + result, err := provider.Search(req) + require.NoError(t, err) + require.NotNil(t, result) + + // Verify result structure + assert.Equal(t, types.SearchTypeWeb, result.Type) + assert.Equal(t, "Yao App Engine", result.Query) + assert.Equal(t, types.SourceAuto, result.Source) + + // MCP should return mock results + if result.Error == "" { + assert.Greater(t, result.Total, 0) + assert.NotEmpty(t, result.Items) + assert.Greater(t, result.Duration, int64(0)) + + // Verify result item structure + for _, item := range result.Items { + assert.Equal(t, types.SearchTypeWeb, item.Type) + assert.Equal(t, types.SourceAuto, item.Source) + assert.NotEmpty(t, item.Title) + assert.NotEmpty(t, item.URL) + } + + t.Logf("MCP search returned %d results in %dms", result.Total, result.Duration) + } else { + t.Logf("MCP search returned error (expected if MCP not loaded): %s", result.Error) + } +} + +// TestMCPProviderWithSiteRestriction tests MCPProvider with domain restriction +func TestMCPProviderWithSiteRestriction(t *testing.T) { + if testing.Short() { + t.Skip("Skipping integration test") + } + + testutils.Prepare(t) + defer testutils.Clean(t) + + // Create MCPProvider + provider, err := web.NewMCPProvider("search.web_search") + require.NoError(t, err) + + // Execute search with site restriction + req := &types.Request{ + Query: "documentation", + Type: types.SearchTypeWeb, + Source: types.SourceHook, + Sites: []string{"github.com"}, + Limit: 3, + } + + result, err := provider.Search(req) + require.NoError(t, err) + require.NotNil(t, result) + + assert.Equal(t, types.SearchTypeWeb, result.Type) + assert.Equal(t, types.SourceHook, result.Source) + + if result.Error == "" { + t.Logf("Site-restricted MCP search returned %d results", result.Total) + } else { + t.Logf("MCP search returned error (expected if MCP not loaded): %s", result.Error) + } +} + +// TestMCPProviderWithTimeRange tests MCPProvider with time range filter +func TestMCPProviderWithTimeRange(t *testing.T) { + if testing.Short() { + t.Skip("Skipping integration test") + } + + testutils.Prepare(t) + defer testutils.Clean(t) + + // Create MCPProvider + provider, err := web.NewMCPProvider("search.web_search") + require.NoError(t, err) + + // Execute search with time range + req := &types.Request{ + Query: "artificial intelligence news", + Type: types.SearchTypeWeb, + Source: types.SourceAuto, + TimeRange: "week", + Limit: 5, + } + + result, err := provider.Search(req) + require.NoError(t, err) + require.NotNil(t, result) + + if result.Error == "" { + t.Logf("Time-ranged MCP search (last week) returned %d results in %dms", result.Total, result.Duration) + } else { + t.Logf("MCP search returned error (expected if MCP not loaded): %s", result.Error) + } +} + +// TestMCPProviderInvalidFormat tests MCPProvider with invalid format +func TestMCPProviderInvalidFormat(t *testing.T) { + // Test invalid format without dot + _, err := web.NewMCPProvider("invalid") + require.Error(t, err) + assert.Contains(t, err.Error(), "invalid MCP format") + + // Test empty string + _, err = web.NewMCPProvider("") + require.Error(t, err) + assert.Contains(t, err.Error(), "invalid MCP format") +} + +// TestMCPProviderNotFound tests MCPProvider when MCP server is not found +func TestMCPProviderNotFound(t *testing.T) { + if testing.Short() { + t.Skip("Skipping integration test") + } + + testutils.Prepare(t) + defer testutils.Clean(t) + + // Create MCPProvider with non-existent server + provider, err := web.NewMCPProvider("nonexistent.web_search") + require.NoError(t, err) + + req := &types.Request{ + Query: "test query", + Type: types.SearchTypeWeb, + Source: types.SourceAuto, + } + + result, err := provider.Search(req) + + // Should not return error, but result should have error message + require.NoError(t, err) + require.NotNil(t, result) + assert.NotEmpty(t, result.Error) + assert.Contains(t, result.Error, "not found") +} + +// TestWebHandlerMCPMode tests the web handler in MCP mode +func TestWebHandlerMCPMode(t *testing.T) { + if testing.Short() { + t.Skip("Skipping integration test") + } + + testutils.Prepare(t) + defer testutils.Clean(t) + + // Create handler with MCP mode + handler := web.NewHandler("mcp:search.web_search", nil) + require.NotNil(t, handler) + + // Verify type + assert.Equal(t, types.SearchTypeWeb, handler.Type()) + + // Execute search + req := &types.Request{ + Query: "Yao framework", + Type: types.SearchTypeWeb, + Source: types.SourceAuto, + Limit: 5, + } + + result, err := handler.Search(req) + require.NoError(t, err) + require.NotNil(t, result) + + assert.Equal(t, types.SearchTypeWeb, result.Type) + assert.Equal(t, "Yao framework", result.Query) + + if result.Error == "" { + t.Logf("Handler MCP mode returned %d results", result.Total) + } else { + t.Logf("Handler MCP mode returned error (expected if MCP not loaded): %s", result.Error) + } +} + +// TestWebHandlerInvalidMCPFormat tests the web handler with invalid MCP format +func TestWebHandlerInvalidMCPFormat(t *testing.T) { + // Create handler with invalid MCP format + handler := web.NewHandler("mcp:invalid", nil) + require.NotNil(t, handler) + + req := &types.Request{ + Query: "test", + Type: types.SearchTypeWeb, + Source: types.SourceAuto, + } + + result, err := handler.Search(req) + require.NoError(t, err) + require.NotNil(t, result) + assert.NotEmpty(t, result.Error) + assert.Contains(t, result.Error, "Invalid MCP format") +}