kevin/aichat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

拆分 main 功能模块

Browse Source 
Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
kevin
2026-06-17 11:29:51 +08:00
co-authored by Claude
parent 132ab2a1cb
commit ccc1260fe0
21 changed files with 2318 additions and 2074 deletions
Download Patch File Download Diff File Expand all files Collapse all files
main_test.go
+91 -83
View File
@@ -7,22 +7,40 @@ import (
"testing"
"time"
"aichat/config"
"aichat/llm"
"aichat/message"
"aichat/stream"
"aichat/toolrouter"
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
)
const testOpenAIBaseURL = "https://ark.cn-beijing.volces.com/api/v3"
func newTestAI(t *testing.T, configs []config.OpenAIConfig) *llm.State {
t.Helper()
ai, err := llm.NewState(configs)
if err != nil {
t.Fatal(err)
}
return ai
}
func TestNormalizeToolRouterConfigDefaults(t *testing.T) {
cfg := &Config{ToolRouter: ToolRouterConfig{Enabled: true}}
changed, err := normalizeToolRouterConfig(cfg)
cfg := &config.Config{ToolRouter: config.ToolRouterConfig{Enabled: true}}
changed, err := config.NormalizeToolRouterConfig(cfg)
if err != nil {
t.Fatal(err)
}
if !changed {
t.Fatal("expected defaults to change config")
}
if cfg.ToolRouter.Timeout != defaultToolRouterTimeout {
defaults := config.DefaultToolRouterConfig()
if cfg.ToolRouter.Timeout != defaults.Timeout {
t.Fatalf("timeout = %d", cfg.ToolRouter.Timeout)
}
if cfg.ToolRouter.MaxTokens != defaultToolRouterMaxTokens {
if cfg.ToolRouter.MaxTokens != defaults.MaxTokens {
t.Fatalf("max_tokens = %d", cfg.ToolRouter.MaxTokens)
}
if strings.TrimSpace(cfg.ToolRouter.SystemPrompt) == "" {
@@ -34,17 +52,17 @@ func TestNormalizeToolRouterConfigDefaults(t *testing.T) {
}
func TestNormalizeToolRouterConfigAddsTimeBeforeSQL(t *testing.T) {
cfg := &Config{ToolRouter: ToolRouterConfig{
cfg := &config.Config{ToolRouter: config.ToolRouterConfig{
Enabled: true,
Timeout: 1,
MaxTokens: 1,
SystemPrompt: "tools",
Tools: []ToolRouteConfig{
Tools: []config.ToolRouteConfig{
{Name: "search", Enabled: true},
{Name: "sql", Enabled: true},
},
}}
changed, err := normalizeToolRouterConfig(cfg)
changed, err := config.NormalizeToolRouterConfig(cfg)
if err != nil {
t.Fatal(err)
}
@@ -57,68 +75,59 @@ func TestNormalizeToolRouterConfigAddsTimeBeforeSQL(t *testing.T) {
}
func TestNormalizeToolRouterConfigDuplicateTools(t *testing.T) {
cfg := &Config{ToolRouter: ToolRouterConfig{
cfg := &config.Config{ToolRouter: config.ToolRouterConfig{
Enabled: true,
Timeout: 1,
MaxTokens: 1,
SystemPrompt: "tools",
Tools: []ToolRouteConfig{
Tools: []config.ToolRouteConfig{
{Name: "sql", Enabled: true},
{Name: " SQL ", Enabled: true},
},
}}
_, err := normalizeToolRouterConfig(cfg)
_, err := config.NormalizeToolRouterConfig(cfg)
if err == nil {
t.Fatal("expected duplicate tool error")
}
}
func TestAvailableAgentToolsUsesConfigOrderAndEnabled(t *testing.T) {
oldRouter := toolRouterState
oldSearch := searchState
oldSQL := sqlState
defer func() {
toolRouterState = oldRouter
searchState = oldSearch
sqlState = oldSQL
}()
toolRouterState = &ToolRouterState{cfg: &ToolRouterConfig{
ai := newTestAI(t, []config.OpenAIConfig{{Name: "test", APIKey: "key", BaseURL: testOpenAIBaseURL, Model: "test", Timeout: 1, Active: true}})
router, err := toolrouter.NewState(&config.ToolRouterConfig{
Enabled: true,
Tools: []ToolRouteConfig{
Tools: []config.ToolRouteConfig{
{Name: "search", Enabled: true},
{Name: "time", Enabled: true, Description: "custom time"},
{Name: "sql", Enabled: false},
},
}}
searchState = nil
sqlState = nil
}, ai)
if err != nil {
t.Fatal(err)
}
tools := availableAgentTools(&OpenAIProfile{}, nil)
tools := toolrouter.AvailableAgentTools(router, ai.ActiveProfile(), nil, nil, nil)
if len(tools) != 1 {
t.Fatalf("tools length = %d", len(tools))
}
if tools[0].name != "time" {
t.Fatalf("tool name = %s", tools[0].name)
if tools[0].Name() != "time" {
t.Fatalf("tool name = %s", tools[0].Name())
}
if tools[0].definition.Function == nil || tools[0].definition.Function.Description != "custom time" {
t.Fatalf("unexpected definition: %#v", tools[0].definition)
definition := tools[0].Definition()
if definition.Function == nil || definition.Function.Description != "custom time" {
t.Fatalf("unexpected definition: %#v", definition)
}
}
func TestRunAgentToolLoopAppendsToolMessages(t *testing.T) {
oldRouter := toolRouterState
defer func() { toolRouterState = oldRouter }()
ai := newTestAI(t, []config.OpenAIConfig{{Name: "test", APIKey: "key", BaseURL: testOpenAIBaseURL, Model: "test", Timeout: 1, Active: true}})
calls := 0
toolRouterState = &ToolRouterState{cfg: &ToolRouterConfig{
router, err := toolrouter.NewState(&config.ToolRouterConfig{
Enabled: true,
Timeout: 1,
MaxTokens: 128,
SystemPrompt: "use tools",
Tools: []ToolRouteConfig{{Name: "time", Enabled: true}},
}}
toolRouterState.complete = func(ctx context.Context, profile *OpenAIProfile, req model.CreateChatCompletionRequest, timeout time.Duration) (model.ChatCompletionResponse, error) {
Tools: []config.ToolRouteConfig{{Name: "time", Enabled: true}},
}, ai, toolrouter.WithCompleter(func(ctx context.Context, profile *llm.Profile, req model.CreateChatCompletionRequest, timeout time.Duration) (model.ChatCompletionResponse, error) {
calls++
if req.ToolChoice != model.ToolChoiceStringTypeAuto {
t.Fatalf("tool choice = %#v", req.ToolChoice)
@@ -129,10 +138,13 @@ func TestRunAgentToolLoopAppendsToolMessages(t *testing.T) {
if calls == 1 {
return model.ChatCompletionResponse{Choices: []*model.ChatCompletionChoice{{Message: model.ChatCompletionMessage{ToolCalls: []*model.ToolCall{{ID: "call_1", Type: model.ToolTypeFunction, Function: model.FunctionCall{Name: "time", Arguments: `{"reason":"需要当前日期"}`}}}}}}}, nil
}
return model.ChatCompletionResponse{Choices: []*model.ChatCompletionChoice{{Message: model.ChatCompletionMessage{Content: stringContent("done")}}}}, nil
return model.ChatCompletionResponse{Choices: []*model.ChatCompletionChoice{{Message: model.ChatCompletionMessage{Content: message.StringContent("done")}}}}, nil
}))
if err != nil {
t.Fatal(err)
}
messages, err := runAgentToolLoop(context.Background(), &OpenAIProfile{Config: OpenAIConfig{Model: "test"}}, []ChatMessage{{Role: "user", Content: "今天几号"}}, nil)
messages, err := toolrouter.RunAgentToolLoop(context.Background(), router, ai.ActiveProfile(), []message.ChatMessage{{Role: "user", Content: "今天几号"}}, nil, nil, nil)
if err != nil {
t.Fatal(err)
}
@@ -152,13 +164,13 @@ func TestRunAgentToolLoopAppendsToolMessages(t *testing.T) {
}
func TestExecuteAgentToolCallUnknownAndError(t *testing.T) {
unknown := executeAgentToolCall(context.Background(), &model.ToolCall{ID: "1", Type: model.ToolTypeFunction, Function: model.FunctionCall{Name: "missing"}}, map[string]agentTool{}, nil)
unknown := toolrouter.ExecuteAgentToolCall(context.Background(), &model.ToolCall{ID: "1", Type: model.ToolTypeFunction, Function: model.FunctionCall{Name: "missing"}}, map[string]toolrouter.AgentTool{}, nil)
if !strings.Contains(unknown, "未知工具") {
t.Fatalf("unknown result = %q", unknown)
}
failed := executeAgentToolCall(context.Background(), &model.ToolCall{ID: "2", Type: model.ToolTypeFunction, Function: model.FunctionCall{Name: "boom"}}, map[string]agentTool{
"boom": {name: "boom", execute: func(context.Context, string) (string, error) { return "", errors.New("bad args") }},
failed := toolrouter.ExecuteAgentToolCall(context.Background(), &model.ToolCall{ID: "2", Type: model.ToolTypeFunction, Function: model.FunctionCall{Name: "boom"}}, map[string]toolrouter.AgentTool{
"boom": toolrouter.NewAgentTool("boom", nil, func(context.Context, string) (string, error) { return "", errors.New("bad args") }),
}, nil)
if !strings.Contains(failed, "bad args") {
t.Fatalf("failed result = %q", failed)
@@ -166,21 +178,21 @@ func TestExecuteAgentToolCallUnknownAndError(t *testing.T) {
}
func TestRunAgentToolLoopMaxIterations(t *testing.T) {
oldRouter := toolRouterState
defer func() { toolRouterState = oldRouter }()
toolRouterState = &ToolRouterState{cfg: &ToolRouterConfig{
ai := newTestAI(t, []config.OpenAIConfig{{Name: "test", APIKey: "key", BaseURL: testOpenAIBaseURL, Model: "test", Timeout: 1, Active: true}})
router, err := toolrouter.NewState(&config.ToolRouterConfig{
Enabled: true,
Timeout: 1,
MaxTokens: 128,
SystemPrompt: "use tools",
Tools: []ToolRouteConfig{{Name: "time", Enabled: true}},
}}
toolRouterState.complete = func(context.Context, *OpenAIProfile, model.CreateChatCompletionRequest, time.Duration) (model.ChatCompletionResponse, error) {
Tools: []config.ToolRouteConfig{{Name: "time", Enabled: true}},
}, ai, toolrouter.WithCompleter(func(context.Context, *llm.Profile, model.CreateChatCompletionRequest, time.Duration) (model.ChatCompletionResponse, error) {
return model.ChatCompletionResponse{Choices: []*model.ChatCompletionChoice{{Message: model.ChatCompletionMessage{ToolCalls: []*model.ToolCall{{ID: "loop", Type: model.ToolTypeFunction, Function: model.FunctionCall{Name: "time", Arguments: `{}`}}}}}}}, nil
}))
if err != nil {
t.Fatal(err)
}
messages, err := runAgentToolLoop(context.Background(), &OpenAIProfile{Config: OpenAIConfig{Model: "test"}}, []ChatMessage{{Role: "user", Content: "今天"}}, nil)
messages, err := toolrouter.RunAgentToolLoop(context.Background(), router, ai.ActiveProfile(), []message.ChatMessage{{Role: "user", Content: "今天"}}, nil, nil, nil)
if err != nil {
t.Fatal(err)
}
@@ -191,7 +203,7 @@ func TestRunAgentToolLoopMaxIterations(t *testing.T) {
}
func TestBuildArkMessageImageTextOrder(t *testing.T) {
msg, err := buildArkMessage(ChatMessage{Role: "user", Content: "请描述图片", ImageURL: "data:image/png;base64,aGVsbG8="})
msg, err := message.BuildArkMessage(message.ChatMessage{Role: "user", Content: "请描述图片", ImageURL: "data:image/png;base64,aGVsbG8="})
if err != nil {
t.Fatal(err)
}
@@ -207,7 +219,7 @@ func TestBuildArkMessageImageTextOrder(t *testing.T) {
}
func TestBuildArkMessageImageOnly(t *testing.T) {
msg, err := buildArkMessage(ChatMessage{Role: "user", ImageURL: "data:image/png;base64,aGVsbG8="})
msg, err := message.BuildArkMessage(message.ChatMessage{Role: "user", ImageURL: "data:image/png;base64,aGVsbG8="})
if err != nil {
t.Fatal(err)
}
@@ -217,7 +229,7 @@ func TestBuildArkMessageImageOnly(t *testing.T) {
}
func TestThinkTagParserSingleChunk(t *testing.T) {
parser := &thinkTagParser{}
parser := &stream.Parser{}
visible, reasoning := parser.Accept("hello <think>abc</think> world")
flushVisible, flushReasoning := parser.Flush()
visible += flushVisible
@@ -228,7 +240,7 @@ func TestThinkTagParserSingleChunk(t *testing.T) {
}
func TestThinkTagParserAcrossChunks(t *testing.T) {
parser := &thinkTagParser{}
parser := &stream.Parser{}
var visible, reasoning string
for _, chunk := range []string{"hello <thi", "nk>abc</thi", "nk> world"} {
v, r := parser.Accept(chunk)
@@ -244,7 +256,7 @@ func TestThinkTagParserAcrossChunks(t *testing.T) {
}
func TestThinkTagParserUnclosedThink(t *testing.T) {
parser := &thinkTagParser{}
parser := &stream.Parser{}
visible, reasoning := parser.Accept("answer <think>still thinking")
v, r := parser.Flush()
visible += v
@@ -255,24 +267,24 @@ func TestThinkTagParserUnclosedThink(t *testing.T) {
}
func TestShouldParseThinkTags(t *testing.T) {
if !shouldParseThinkTags(&OpenAIProfile{Config: OpenAIConfig{BaseURL: "http://127.0.0.1:11434/v1"}}) {
if !llm.ShouldParseThinkTags(&llm.Profile{Config: config.OpenAIConfig{BaseURL: "http://127.0.0.1:11434/v1"}}) {
t.Fatal("expected local ollama to parse think tags")
}
if shouldParseThinkTags(&OpenAIProfile{Config: OpenAIConfig{BaseURL: defaultOpenAIBaseURL}}) {
if llm.ShouldParseThinkTags(&llm.Profile{Config: config.OpenAIConfig{BaseURL: testOpenAIBaseURL}}) {
t.Fatal("expected remote profile not to parse think tags by default")
}
falseValue := false
if shouldParseThinkTags(&OpenAIProfile{Config: OpenAIConfig{BaseURL: "http://127.0.0.1:11434/v1", ParseThinkTags: &falseValue}}) {
if llm.ShouldParseThinkTags(&llm.Profile{Config: config.OpenAIConfig{BaseURL: "http://127.0.0.1:11434/v1", ParseThinkTags: &falseValue}}) {
t.Fatal("explicit false should disable think parsing")
}
trueValue := true
if !shouldParseThinkTags(&OpenAIProfile{Config: OpenAIConfig{BaseURL: defaultOpenAIBaseURL, ParseThinkTags: &trueValue}}) {
if !llm.ShouldParseThinkTags(&llm.Profile{Config: config.OpenAIConfig{BaseURL: testOpenAIBaseURL, ParseThinkTags: &trueValue}}) {
t.Fatal("explicit true should enable think parsing")
}
}
func TestBuildToolDecisionMessagesRemovesImages(t *testing.T) {
messages, err := buildToolDecisionMessages([]ChatMessage{{Role: "user", Content: "描述这张图", ImageURL: "data:image/png;base64,aGVsbG8="}})
messages, err := message.BuildToolDecisionMessages([]message.ChatMessage{{Role: "user", Content: "描述这张图", ImageURL: "data:image/png;base64,aGVsbG8="}})
if err != nil {
t.Fatal(err)
}
@@ -288,17 +300,14 @@ func TestBuildToolDecisionMessagesRemovesImages(t *testing.T) {
}
func TestRunAgentToolLoopImageUsesTextOnlyDecisionMessages(t *testing.T) {
oldRouter := toolRouterState
defer func() { toolRouterState = oldRouter }()
toolRouterState = &ToolRouterState{cfg: &ToolRouterConfig{
ai := newTestAI(t, []config.OpenAIConfig{{Name: "chat", APIKey: "key", BaseURL: testOpenAIBaseURL, Model: "chat", Timeout: 1, Active: true}})
router, err := toolrouter.NewState(&config.ToolRouterConfig{
Enabled: true,
Timeout: 1,
MaxTokens: 128,
SystemPrompt: "use tools",
Tools: []ToolRouteConfig{{Name: "time", Enabled: true}},
}}
toolRouterState.complete = func(ctx context.Context, profile *OpenAIProfile, req model.CreateChatCompletionRequest, timeout time.Duration) (model.ChatCompletionResponse, error) {
Tools: []config.ToolRouteConfig{{Name: "time", Enabled: true}},
}, ai, toolrouter.WithCompleter(func(ctx context.Context, profile *llm.Profile, req model.CreateChatCompletionRequest, timeout time.Duration) (model.ChatCompletionResponse, error) {
for _, msg := range req.Messages {
if msg.Content != nil && len(msg.Content.ListValue) > 0 {
t.Fatalf("tool decision should not receive multimodal content: %#v", msg.Content)
@@ -313,10 +322,13 @@ func TestRunAgentToolLoopImageUsesTextOnlyDecisionMessages(t *testing.T) {
if !strings.Contains(joined, "工具判断阶段不读取图片内容") {
t.Fatalf("missing placeholder in decision messages: %q", joined)
}
return model.ChatCompletionResponse{Choices: []*model.ChatCompletionChoice{{Message: model.ChatCompletionMessage{Content: stringContent("no tool")}}}}, nil
return model.ChatCompletionResponse{Choices: []*model.ChatCompletionChoice{{Message: model.ChatCompletionMessage{Content: message.StringContent("no tool")}}}}, nil
}))
if err != nil {
t.Fatal(err)
}
messages, err := runAgentToolLoop(context.Background(), &OpenAIProfile{Config: OpenAIConfig{Model: "chat"}}, []ChatMessage{{Role: "user", Content: "描述这张图", ImageURL: "data:image/png;base64,aGVsbG8="}}, nil)
messages, err := toolrouter.RunAgentToolLoop(context.Background(), router, ai.ActiveProfile(), []message.ChatMessage{{Role: "user", Content: "描述这张图", ImageURL: "data:image/png;base64,aGVsbG8="}}, nil, nil, nil)
if err != nil {
t.Fatal(err)
}
@@ -332,32 +344,28 @@ func TestRunAgentToolLoopImageUsesTextOnlyDecisionMessages(t *testing.T) {
}
func TestRunAgentToolLoopUsesConfiguredRouterProfile(t *testing.T) {
oldRouter := toolRouterState
defer func() { toolRouterState = oldRouter }()
ai, err := NewOpenAIState([]OpenAIConfig{
{Name: "chat", APIKey: "key", BaseURL: defaultOpenAIBaseURL, Model: "chat-model", Timeout: 1, Active: true},
{Name: "router", APIKey: "key", BaseURL: defaultOpenAIBaseURL, Model: "router-model", Timeout: 1},
ai := newTestAI(t, []config.OpenAIConfig{
{Name: "chat", APIKey: "key", BaseURL: testOpenAIBaseURL, Model: "chat-model", Timeout: 1, Active: true},
{Name: "router", APIKey: "key", BaseURL: testOpenAIBaseURL, Model: "router-model", Timeout: 1},
})
if err != nil {
t.Fatal(err)
}
toolRouterState = &ToolRouterState{ai: ai, cfg: &ToolRouterConfig{
router, err := toolrouter.NewState(&config.ToolRouterConfig{
Enabled: true,
OpenAIName: "router",
Timeout: 1,
MaxTokens: 128,
SystemPrompt: "use tools",
Tools: []ToolRouteConfig{{Name: "time", Enabled: true}},
}}
toolRouterState.complete = func(ctx context.Context, profile *OpenAIProfile, req model.CreateChatCompletionRequest, timeout time.Duration) (model.ChatCompletionResponse, error) {
Tools: []config.ToolRouteConfig{{Name: "time", Enabled: true}},
}, ai, toolrouter.WithCompleter(func(ctx context.Context, profile *llm.Profile, req model.CreateChatCompletionRequest, timeout time.Duration) (model.ChatCompletionResponse, error) {
if profile.Config.Name != "router" || req.Model != "router-model" {
t.Fatalf("router profile not used: profile=%s model=%s", profile.Config.Name, req.Model)
}
return model.ChatCompletionResponse{Choices: []*model.ChatCompletionChoice{{Message: model.ChatCompletionMessage{Content: stringContent("no tool")}}}}, nil
return model.ChatCompletionResponse{Choices: []*model.ChatCompletionChoice{{Message: model.ChatCompletionMessage{Content: message.StringContent("no tool")}}}}, nil
}))
if err != nil {
t.Fatal(err)
}
_, err = runAgentToolLoop(context.Background(), &OpenAIProfile{Config: OpenAIConfig{Name: "chat", Model: "chat-model"}}, []ChatMessage{{Role: "user", Content: "今天"}}, nil)
_, err = toolrouter.RunAgentToolLoop(context.Background(), router, ai.ActiveProfile(), []message.ChatMessage{{Role: "user", Content: "今天"}}, nil, nil, nil)
if err != nil {
t.Fatal(err)
}