模型优化

main_test.go

@@ -129,6 +129,41 @@ func TestFilterToolSelections(t *testing.T) {
 	}
 }
 
+func TestEnsureTimeSelectionForRelativeSearch(t *testing.T) {
+	tools := map[string]ChatTool{
+		"time":   fakeChatTool{name: "time", enabled: true},
+		"search": fakeChatTool{name: "search", enabled: true},
+	}
+	selected := ensureTimeSelectionForRelativeQuery(
+		[]ToolSelection{{Name: "search", Reason: "查询历史事件"}},
+		tools,
+		[]ToolRouteConfig{{Name: "time"}, {Name: "search"}, {Name: "sql"}},
+		"历史上的今天都发生了什么？",
+	)
+	if len(selected) != 2 || selected[0].Name != "time" || selected[1].Name != "search" {
+		t.Fatalf("unexpected selected tools: %#v", selected)
+	}
+	if !strings.Contains(selected[0].Reason, "相对日期") {
+		t.Fatalf("unexpected time reason: %#v", selected[0])
+	}
+}
+
+func TestEnsureTimeSelectionSkipsOrdinarySearch(t *testing.T) {
+	tools := map[string]ChatTool{
+		"time":   fakeChatTool{name: "time", enabled: true},
+		"search": fakeChatTool{name: "search", enabled: true},
+	}
+	selected := ensureTimeSelectionForRelativeQuery(
+		[]ToolSelection{{Name: "search", Reason: "查询资料"}},
+		tools,
+		[]ToolRouteConfig{{Name: "time"}, {Name: "search"}},
+		"查一下 Go 语言官网",
+	)
+	if len(selected) != 1 || selected[0].Name != "search" {
+		t.Fatalf("unexpected selected tools: %#v", selected)
+	}
+}
+
 func TestRunTimeToolAddsHiddenDateRanges(t *testing.T) {
 	messages := []ChatMessage{{Role: "user", Content: "本月有什么日程安排"}}
 	withTime, err := runTimeTool(context.Background(), messages, "需要日期范围", func(chatSSEFrame) {})