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

Enhance Assistant's message building and completion options handling

Browse source 
- Refactored the BuildRequest method to improve clarity and maintainability by separating the building of completion options and final messages.
- Updated the buildMessages function to include MCP samples as system messages, enhancing the message prioritization logic.
- Introduced applyMCPTools method to integrate MCP tools into completion options, improving functionality and flexibility in tool usage.
- Enhanced error handling in completion options to ensure robust application flow.
This commit is contained in:
Max 2025-11-29 11:08:42 +08:00
parent cfde900ada
commit 151d7f19e0
5 changed files with 725 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
agent/assistant/agent.go
View file

@ -7,6 +7,7 @@ import (
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/connector"
"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"
@ -114,6 +115,10 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
// Execute tool calls
// ================================================
if completionResponse != nil && completionResponse.ToolCalls != nil {
fmt.Println("--- completionResponse ToolCalls --------------------------------")
utils.Dump(completionResponse.ToolCalls)
fmt.Println("--------------------------------")
}
// Request Done hook ( Optional )

74
agent/assistant/build.go
View file

@ -9,14 +9,14 @@ import (
// BuildRequest build the LLM request
func (ast *Assistant) BuildRequest(ctx *context.Context, messages []context.Message, createResponse *context.HookCreateResponse) ([]context.Message, *context.CompletionOptions, error) {
// Build final messages with proper priority
finalMessages, err := ast.buildMessages(ctx, messages, createResponse)
// Build completion options from createResponse and ctx (includes MCP tools)
options, mcpSamplesPrompt, err := ast.buildCompletionOptions(ctx, createResponse)
if err != nil {
return nil, nil, err
}
// Build completion options from createResponse and ctx
options, err := ast.buildCompletionOptions(ctx, createResponse)
// Build final messages with proper priority (includes MCP samples if available)
finalMessages, err := ast.buildMessages(ctx, messages, createResponse, mcpSamplesPrompt)
if err != nil {
return nil, nil, err
}
@ -25,9 +25,9 @@ func (ast *Assistant) BuildRequest(ctx *context.Context, messages []context.Mess
}
// buildMessages builds the final message list with proper priority
// Priority: Prompts > createResponse.Messages > input messages
// Priority: Prompts > MCP Samples > createResponse.Messages > input messages
// If createResponse is nil or has no messages, use input messages
func (ast *Assistant) buildMessages(ctx *context.Context, messages []context.Message, createResponse *context.HookCreateResponse) ([]context.Message, error) {
func (ast *Assistant) buildMessages(ctx *context.Context, messages []context.Message, createResponse *context.HookCreateResponse, mcpSamplesPrompt string) ([]context.Message, error) {
var finalMessages []context.Message
// If createResponse is nil or has no messages, use input messages
@ -38,6 +38,16 @@ func (ast *Assistant) buildMessages(ctx *context.Context, messages []context.Mes
finalMessages = createResponse.Messages
}
// Add MCP samples prompt as a system message (if available)
if mcpSamplesPrompt != "" {
mcpSamplesMsg := context.Message{
Role: context.RoleSystem,
Content: mcpSamplesPrompt,
}
// Prepend MCP samples before other messages
finalMessages = append([]context.Message{mcpSamplesMsg}, finalMessages...)
}
// ⚠️ Just for testing, will remove later
// If we have prompts, prepend them to the beginning
if len(ast.Prompts) > 0 {
@ -64,12 +74,13 @@ func (ast *Assistant) buildMessages(ctx *context.Context, messages []context.Mes
// buildCompletionOptions builds completion options from multiple sources
// Priority (lowest to highest, later overrides earlier): ast > ctx > createResponse
// The priority means: if createResponse has a value, use it; else use ctx; else use ast
func (ast *Assistant) buildCompletionOptions(ctx *context.Context, createResponse *context.HookCreateResponse) (*context.CompletionOptions, error) {
// Returns (options, mcpSamplesPrompt, error)
func (ast *Assistant) buildCompletionOptions(ctx *context.Context, createResponse *context.HookCreateResponse) (*context.CompletionOptions, string, error) {
options := &context.CompletionOptions{}
// Layer 1 (base): Apply ast - Assistant configuration
if err := ast.applyAssistantOptions(options); err != nil {
return nil, err
return nil, "", err
}
// Layer 2 (middle): Apply ctx - Context configuration (overrides ast)
@ -80,7 +91,13 @@ func (ast *Assistant) buildCompletionOptions(ctx *context.Context, createRespons
ast.applyCreateResponseOptions(options, createResponse)
}
return options, nil
// Add MCP tools if configured and get samples prompt
mcpSamplesPrompt, err := ast.applyMCPTools(ctx, options)
if err != nil {
return nil, "", fmt.Errorf("failed to apply MCP tools: %w", err)
}
return options, mcpSamplesPrompt, nil
}
// applyAssistantOptions applies options from ast.Options to CompletionOptions
@ -330,3 +347,42 @@ func (ast *Assistant) getUses() *context.Uses {
// Priority 2: Global settings only
return globalUses
}
// applyMCPTools adds MCP tools to completion options and returns samples prompt
// Returns (samplesPrompt, error)
func (ast *Assistant) applyMCPTools(ctx *context.Context, options *context.CompletionOptions) (string, error) {
if ast.MCP == nil || len(ast.MCP.Servers) == 0 {
return "", nil
}
// Build MCP tools and get samples prompt
mcpTools, samplesPrompt, err := ast.buildMCPTools(ctx)
if err != nil {
return "", fmt.Errorf("failed to build MCP tools: %w", err)
}
// Convert mcpTools to map format for CompletionOptions.Tools
if len(mcpTools) > 0 {
toolMaps := make([]map[string]interface{}, len(mcpTools))
for i, tool := range mcpTools {
toolMaps[i] = map[string]interface{}{
"type": "function",
"function": map[string]interface{}{
"name": tool.Name,
"description": tool.Description,
"parameters": tool.Parameters,
},
}
}
// Add MCP tools to existing tools (append to preserve existing tools)
if options.Tools == nil {
options.Tools = toolMaps
} else {
options.Tools = append(options.Tools, toolMaps...)
}
}
return samplesPrompt, nil
}

222
agent/assistant/build_mcp_test.go Normal file
View file

@ -0,0 +1,222 @@
package assistant_test
import (
"testing"
"github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/testutils"
)
// TestBuildRequest_MCP tests MCP tool integration in BuildRequest
func TestBuildRequest_MCP(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
agent, err := assistant.Get("tests.mcptest")
if err != nil {
t.Fatalf("Failed to get tests.mcptest assistant: %s", err.Error())
}
ctx := newTestContext("chat-test-mcp", "tests.mcptest")
t.Run("MCPToolsLoaded", func(t *testing.T) {
inputMessages := []context.Message{{Role: context.RoleUser, Content: "test mcp tools"}}
// Build LLM request
_, options, err := agent.BuildRequest(ctx, inputMessages, nil)
if err != nil {
t.Fatalf("Failed to build LLM request: %s", err.Error())
}
// Verify that tools are loaded
if options.Tools == nil {
t.Fatal("Expected tools to be loaded, got nil")
}
if len(options.Tools) == 0 {
t.Fatal("Expected at least some MCP tools, got empty list")
}
// Count MCP tools (should be filtered to only ping and echo)
mcpToolCount := 0
var toolNames []string
for _, toolMap := range options.Tools {
fn, ok := toolMap["function"].(map[string]interface{})
if !ok {
continue
}
name, ok := fn["name"].(string)
if ok {
toolNames = append(toolNames, name)
mcpToolCount++
}
}
t.Logf("Found %d MCP tools: %v", mcpToolCount, toolNames)
// Verify tool count (should be exactly 2: ping and echo)
if mcpToolCount != 2 {
t.Errorf("Expected 2 MCP tools (ping, echo), got %d: %v", mcpToolCount, toolNames)
}
// Verify specific tools exist
hasEchoPing := false
hasEchoEcho := false
for _, name := range toolNames {
if name == "echo__ping" {
hasEchoPing = true
}
if name == "echo__echo" {
hasEchoEcho = true
}
}
if !hasEchoPing {
t.Error("Expected 'echo__ping' tool to be present")
}
if !hasEchoEcho {
t.Error("Expected 'echo__echo' tool to be present")
}
// Verify that 'status' tool is NOT included (filtered out)
for _, name := range toolNames {
if name == "echo__status" {
t.Error("Tool 'echo__status' should be filtered out but was found")
}
}
t.Log("✓ MCP tools loaded and filtered correctly")
})
t.Run("MCPSamplesPrompt", func(t *testing.T) {
inputMessages := []context.Message{{Role: context.RoleUser, Content: "test mcp samples"}}
// Build LLM request
finalMessages, _, err := agent.BuildRequest(ctx, inputMessages, nil)
if err != nil {
t.Fatalf("Failed to build LLM request: %s", err.Error())
}
// Check if messages contain MCP samples prompt
// The samples prompt should be added as a system message
hasMCPSamples := false
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem {
if content, ok := msg.Content.(string); ok {
if len(content) > 50 &&
(contains(content, "MCP Tool Usage Examples") ||
contains(content, "echo.ping") ||
contains(content, "echo.echo")) {
hasMCPSamples = true
t.Logf("Found MCP samples prompt (length: %d chars)", len(content))
break
}
}
}
}
// Note: samples may not exist for echo tools, so this is informational
if hasMCPSamples {
t.Log("✓ MCP samples prompt included in messages")
} else {
t.Log("ℹ No MCP samples prompt found (may not have sample files)")
}
})
t.Run("MCPToolNameFormat", func(t *testing.T) {
inputMessages := []context.Message{{Role: context.RoleUser, Content: "test tool format"}}
// Build LLM request
_, options, err := agent.BuildRequest(ctx, inputMessages, nil)
if err != nil {
t.Fatalf("Failed to build LLM request: %s", err.Error())
}
// Verify tool name format: server_id.tool_name
for _, toolMap := range options.Tools {
fn, ok := toolMap["function"].(map[string]interface{})
if !ok {
continue
}
name, ok := fn["name"].(string)
if ok {
// Parse tool name
serverID, toolName, ok := assistant.ParseMCPToolName(name)
if !ok {
t.Errorf("Tool name '%s' is not in correct format (server_id.tool_name)", name)
continue
}
// Verify server ID
if serverID != "echo" {
t.Errorf("Expected server_id 'echo', got '%s' for tool '%s'", serverID, name)
}
// Verify tool name is either ping or echo
if toolName != "ping" && toolName != "echo" {
t.Errorf("Expected tool name 'ping' or 'echo', got '%s'", toolName)
}
t.Logf("✓ Tool name format correct: %s → (%s, %s)", name, serverID, toolName)
}
}
})
t.Run("MCPToolSchema", func(t *testing.T) {
inputMessages := []context.Message{{Role: context.RoleUser, Content: "test tool schema"}}
// Build LLM request
_, options, err := agent.BuildRequest(ctx, inputMessages, nil)
if err != nil {
t.Fatalf("Failed to build LLM request: %s", err.Error())
}
// Verify tool schema structure
for _, toolMap := range options.Tools {
// Verify type field
if toolType, ok := toolMap["type"].(string); !ok || toolType != "function" {
t.Errorf("Expected tool type 'function', got: %v", toolMap["type"])
}
// Verify function field exists
fn, ok := toolMap["function"].(map[string]interface{})
if !ok {
t.Error("Tool missing 'function' field or wrong type")
continue
}
// Verify required fields
if _, hasName := fn["name"]; !hasName {
t.Error("Tool function missing 'name' field")
}
if _, hasDesc := fn["description"]; !hasDesc {
t.Error("Tool function missing 'description' field")
}
if _, hasParams := fn["parameters"]; !hasParams {
t.Error("Tool function missing 'parameters' field")
}
t.Logf("✓ Tool schema valid: %v", fn["name"])
}
})
}
// Helper function to check if string contains substring
func contains(s, substr string) bool {
return len(s) >= len(substr) &&
(s == substr ||
len(s) > len(substr) &&
(s[:len(substr)] == substr ||
s[len(s)-len(substr):] == substr ||
findSubstring(s, substr)))
}
func findSubstring(s, substr string) bool {
for i := 0; i <= len(s)-len(substr); i++ {
if s[i:i+len(substr)] == substr {
return true
}
}
return false
}

197
agent/assistant/mcp.go Normal file
View file

@ -0,0 +1,197 @@
package assistant
import (
"context"
"fmt"
"strings"
"github.com/yaoapp/gou/mcp"
mcpTypes "github.com/yaoapp/gou/mcp/types"
"github.com/yaoapp/kun/log"
agentContext "github.com/yaoapp/yao/agent/context"
)
const (
// MaxMCPTools maximum number of MCP tools to include (to avoid overwhelming the LLM)
MaxMCPTools = 20
)
// MCPToolName formats a tool name with MCP server prefix
// Format: server_id__tool_name (double underscore separator)
// Dots in server_id are replaced with single underscores
// Examples:
// - ("echo", "ping") → "echo__ping"
// - ("github.enterprise", "search") → "github_enterprise__search"
//
// Naming constraint: MCP server_id MUST NOT contain underscores (_)
// Only dots (.), letters, numbers, and hyphens (-) are allowed in server_id
func MCPToolName(serverID, toolName string) string {
if serverID == "" || toolName == "" {
return ""
}
// Replace dots with single underscores in server_id
cleanServerID := strings.ReplaceAll(serverID, ".", "_")
// Use double underscore as separator
return fmt.Sprintf("%s__%s", cleanServerID, toolName)
}
// ParseMCPToolName parses a formatted MCP tool name into server ID and tool name
// Splits by double underscore (__), then restores dots in server_id
// Examples:
// - "echo__ping" → ("echo", "ping")
// - "github_enterprise__search" → ("github.enterprise", "search")
//
// Returns (serverID, toolName, true) if valid format, ("", "", false) otherwise
func ParseMCPToolName(formattedName string) (string, string, bool) {
if formattedName == "" {
return "", "", false
}
// Split by double underscore
parts := strings.Split(formattedName, "__")
if len(parts) != 2 {
return "", "", false
}
cleanServerID := parts[0]
toolName := parts[1]
// Validate that both parts are non-empty
if cleanServerID == "" || toolName == "" {
return "", "", false
}
// Restore dots in server_id (replace single underscores back to dots)
serverID := strings.ReplaceAll(cleanServerID, "_", ".")
return serverID, toolName, true
}
// MCPTool represents a simplified MCP tool for building LLM requests
type MCPTool struct {
Name string
Description string
Parameters interface{}
}
// buildMCPTools builds tool definitions and samples system prompt from MCP servers
// Returns (tools, samplesPrompt, error)
func (ast *Assistant) buildMCPTools(ctx *agentContext.Context) ([]MCPTool, string, error) {
if ast.MCP == nil || len(ast.MCP.Servers) == 0 {
return nil, "", nil
}
mcpCtx := context.Background()
allTools := make([]MCPTool, 0)
samplesBuilder := strings.Builder{}
hasSamples := false
// Process each MCP server in order
for _, serverConfig := range ast.MCP.Servers {
if len(allTools) >= MaxMCPTools {
log.Warn("[Assistant MCP] Reached maximum tool limit (%d), skipping remaining servers", MaxMCPTools)
break
}
// Get MCP client
client, err := mcp.Select(serverConfig.ServerID)
if err != nil {
log.Warn("[Assistant MCP] Failed to select MCP client '%s': %v", serverConfig.ServerID, err)
continue
}
// Get tools list (filter by serverConfig.Tools if specified)
toolsResponse, err := client.ListTools(mcpCtx, "")
if err != nil {
log.Warn("[Assistant MCP] Failed to list tools for '%s': %v", serverConfig.ServerID, err)
continue
}
// Build tool filter map if specified
toolFilter := make(map[string]bool)
if len(serverConfig.Tools) > 0 {
for _, toolName := range serverConfig.Tools {
toolFilter[toolName] = true
}
}
// Process each tool
for _, tool := range toolsResponse.Tools {
// Check tool limit
if len(allTools) >= MaxMCPTools {
break
}
// Apply tool filter if specified
if len(toolFilter) > 0 && !toolFilter[tool.Name] {
continue
}
// Format tool name with server prefix
formattedName := MCPToolName(serverConfig.ServerID, tool.Name)
// Convert MCP tool to MCPTool format
mcpTool := MCPTool{
Name: formattedName,
Description: tool.Description,
Parameters: tool.InputSchema,
}
allTools = append(allTools, mcpTool)
// Try to get samples for this tool
samples, err := client.ListSamples(mcpCtx, mcpTypes.SampleTool, tool.Name)
if err == nil && len(samples.Samples) > 0 {
if !hasSamples {
samplesBuilder.WriteString("\n\n## MCP Tool Usage Examples\n\n")
samplesBuilder.WriteString("The following examples demonstrate how to use MCP tools correctly:\n\n")
hasSamples = true
}
samplesBuilder.WriteString(fmt.Sprintf("### %s\n\n", formattedName))
if tool.Description != "" {
samplesBuilder.WriteString(fmt.Sprintf("**Description**: %s\n\n", tool.Description))
}
for i, sample := range samples.Samples {
if i >= 3 { // Limit to 3 examples per tool
break
}
samplesBuilder.WriteString(fmt.Sprintf("**Example %d", i+1))
if sample.Name != "" {
samplesBuilder.WriteString(fmt.Sprintf(" - %s", sample.Name))
}
samplesBuilder.WriteString("**:\n")
// Check metadata for description
if sample.Metadata != nil {
if desc, ok := sample.Metadata["description"].(string); ok && desc != "" {
samplesBuilder.WriteString(fmt.Sprintf("- Description: %s\n", desc))
}
}
if sample.Input != nil {
samplesBuilder.WriteString(fmt.Sprintf("- Input: `%v`\n", sample.Input))
}
if sample.Output != nil {
samplesBuilder.WriteString(fmt.Sprintf("- Output: `%v`\n", sample.Output))
}
samplesBuilder.WriteString("\n")
}
}
}
log.Trace("[Assistant MCP] Loaded %d tools from server '%s'", len(toolsResponse.Tools), serverConfig.ServerID)
}
samplesPrompt := ""
if hasSamples {
samplesPrompt = samplesBuilder.String()
}
log.Trace("[Assistant MCP] Total MCP tools loaded: %d", len(allTools))
return allTools, samplesPrompt, nil
}

236
agent/assistant/mcp_test.go Normal file
View file

@ -0,0 +1,236 @@
package assistant_test
import (
"testing"
"github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/agent/testutils"
)
func TestMCPToolName(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
tests := []struct {
name string
serverID string
toolName string
wantResult string
}{
{
name: "Simple tool name",
serverID: "github",
toolName: "search_repos",
wantResult: "github__search_repos",
},
{
name: "Server with dots",
serverID: "github.enterprise",
toolName: "search_repos",
wantResult: "github_enterprise__search_repos",
},
{
name: "Tool with underscores",
serverID: "customer-db",
toolName: "create_customer",
wantResult: "customer-db__create_customer",
},
{
name: "Complex server with multiple dots",
serverID: "com.example.mcp",
toolName: "tool_name",
wantResult: "com_example_mcp__tool_name",
},
{
name: "Empty server ID",
serverID: "",
toolName: "tool",
wantResult: "",
},
{
name: "Empty tool name",
serverID: "server",
toolName: "",
wantResult: "",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := assistant.MCPToolName(tt.serverID, tt.toolName)
if result != tt.wantResult {
t.Errorf("MCPToolName() = %v, want %v", result, tt.wantResult)
}
})
}
}
func TestParseMCPToolName(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
tests := []struct {
name string
formattedName string
wantServerID string
wantToolName string
wantOK bool
}{
{
name: "Valid simple format",
formattedName: "github__search_repos",
wantServerID: "github",
wantToolName: "search_repos",
wantOK: true,
},
{
name: "Server with dots restored",
formattedName: "github_enterprise__search_repos",
wantServerID: "github.enterprise",
wantToolName: "search_repos",
wantOK: true,
},
{
name: "Complex server ID with multiple dots",
formattedName: "com_example_mcp_server__tool_name",
wantServerID: "com.example.mcp.server",
wantToolName: "tool_name",
wantOK: true,
},
{
name: "Tool name with underscores",
formattedName: "server__create_new_user",
wantServerID: "server",
wantToolName: "create_new_user",
wantOK: true,
},
{
name: "Server with hyphens",
formattedName: "mcp-server__tool",
wantServerID: "mcp-server",
wantToolName: "tool",
wantOK: true,
},
{
name: "Invalid format - no double underscore",
formattedName: "invalid",
wantServerID: "",
wantToolName: "",
wantOK: false,
},
{
name: "Invalid format - empty string",
formattedName: "",
wantServerID: "",
wantToolName: "",
wantOK: false,
},
{
name: "Invalid format - only double underscore",
formattedName: "__",
wantServerID: "",
wantToolName: "",
wantOK: false,
},
{
name: "Invalid format - ends with double underscore",
formattedName: "server__",
wantServerID: "",
wantToolName: "",
wantOK: false,
},
{
name: "Invalid format - starts with double underscore",
formattedName: "__tool",
wantServerID: "",
wantToolName: "",
wantOK: false,
},
{
name: "Invalid format - multiple double underscores",
formattedName: "server__middle__tool",
wantServerID: "",
wantToolName: "",
wantOK: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
serverID, toolName, ok := assistant.ParseMCPToolName(tt.formattedName)
if serverID != tt.wantServerID {
t.Errorf("ParseMCPToolName() serverID = %v, want %v", serverID, tt.wantServerID)
}
if toolName != tt.wantToolName {
t.Errorf("ParseMCPToolName() toolName = %v, want %v", toolName, tt.wantToolName)
}
if ok != tt.wantOK {
t.Errorf("ParseMCPToolName() ok = %v, want %v", ok, tt.wantOK)
}
})
}
}
func TestMCPToolName_RoundTrip(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
tests := []struct {
name string
serverID string
toolName string
}{
{
name: "Simple IDs",
serverID: "github",
toolName: "search_repos",
},
{
name: "Server with dots",
serverID: "github.enterprise",
toolName: "search",
},
{
name: "Complex server ID",
serverID: "com.example.mcp.server",
toolName: "tool_name",
},
{
name: "Server with dashes",
serverID: "mcp-server-123",
toolName: "tool_with_underscores",
},
{
name: "Mixed dots and dashes",
serverID: "github.enterprise-prod",
toolName: "api_call",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// Format
formatted := assistant.MCPToolName(tt.serverID, tt.toolName)
if formatted == "" {
t.Fatal("MCPToolName() returned empty string")
}
// Parse
serverID, toolName, ok := assistant.ParseMCPToolName(formatted)
// Verify round-trip
if !ok {
t.Fatal("ParseMCPToolName() failed")
}
if serverID != tt.serverID {
t.Errorf("Round-trip failed: serverID = %v, want %v", serverID, tt.serverID)
}
if toolName != tt.toolName {
t.Errorf("Round-trip failed: toolName = %v, want %v", toolName, tt.toolName)
}
t.Logf("✓ Round-trip successful: (%s, %s) → %s → (%s, %s)",
tt.serverID, tt.toolName, formatted, serverID, toolName)
})
}
}