wizardchen
45eab4a95f
Emit search_source metadata from RAG progress events so quick-answer steps, summaries, and IM labels no longer mislabel web-only results as knowledge base search.
104 lines
2.7 KiB
Go
104 lines
2.7 KiB
Go
package im
|
||
|
||
import (
|
||
"strings"
|
||
"testing"
|
||
)
|
||
|
||
func TestFormatIMToolLine_pendingWithQuery(t *testing.T) {
|
||
line := FormatIMToolLine(IMToolStep{
|
||
ToolName: "knowledge_search",
|
||
Pending: true,
|
||
Arguments: map[string]any{"query": "文明6"},
|
||
})
|
||
if line != "正在调用 知识库检索..." {
|
||
t.Fatalf("pending line = %q", line)
|
||
}
|
||
}
|
||
|
||
func TestFormatIMToolLine_searchDoneWithQueryAndSummary(t *testing.T) {
|
||
line := FormatIMToolLine(IMToolStep{
|
||
ToolName: "knowledge_search",
|
||
Success: true,
|
||
Arguments: map[string]any{
|
||
"query": "文明6",
|
||
},
|
||
Data: map[string]interface{}{
|
||
"results": []interface{}{map[string]interface{}{}, map[string]interface{}{}, map[string]interface{}{}},
|
||
"kb_counts": map[string]interface{}{"a": 1, "b": 1},
|
||
},
|
||
})
|
||
if !strings.Contains(line, "检索知识库：「文明6」") {
|
||
t.Fatalf("title missing query: %q", line)
|
||
}
|
||
if !strings.Contains(line, "找到 3 个结果，来自 2 个文件") {
|
||
t.Fatalf("summary missing: %q", line)
|
||
}
|
||
}
|
||
|
||
func TestFormatIMToolLine_grepPatterns(t *testing.T) {
|
||
line := FormatIMToolLine(IMToolStep{
|
||
ToolName: "grep_chunks",
|
||
Success: true,
|
||
Arguments: map[string]any{
|
||
"patterns": []any{"文明", "策略"},
|
||
},
|
||
Data: map[string]interface{}{
|
||
"total_matches": float64(5),
|
||
"document_count": float64(2),
|
||
},
|
||
})
|
||
if line != "搜索关键词：「文明、策略」 · 找到 5 个匹配片段，来自 2 个文档" {
|
||
t.Fatalf("grep line = %q", line)
|
||
}
|
||
}
|
||
|
||
func TestFormatIMRagPipelineLine_queryUnderstand(t *testing.T) {
|
||
pending := FormatIMRagPipelineLine(IMToolStep{
|
||
ToolName: "query_understand",
|
||
Pending: true,
|
||
})
|
||
if pending != "正在理解问题..." {
|
||
t.Fatalf("pending = %q", pending)
|
||
}
|
||
done := FormatIMRagPipelineLine(IMToolStep{
|
||
ToolName: "query_understand",
|
||
Success: true,
|
||
})
|
||
if done != "已完成问题理解" {
|
||
t.Fatalf("done = %q", done)
|
||
}
|
||
}
|
||
|
||
func TestFormatIMRagPipelineLine_searchWithQuery(t *testing.T) {
|
||
line := FormatIMRagPipelineLine(IMToolStep{
|
||
ToolName: "knowledge_search",
|
||
Pending: true,
|
||
Arguments: map[string]any{"query": "讯飞开放平台"},
|
||
})
|
||
if line != "正在检索知识库：「讯飞开放平台」" {
|
||
t.Fatalf("line = %q", line)
|
||
}
|
||
}
|
||
|
||
func TestFormatIMRagPipelineLine_webSearchWithQuery(t *testing.T) {
|
||
line := FormatIMRagPipelineLine(IMToolStep{
|
||
ToolName: "knowledge_search",
|
||
Pending: true,
|
||
Arguments: map[string]any{"query": "任素汐演唱会", "search_source": "web"},
|
||
})
|
||
if line != "正在检索网络：「任素汐演唱会」" {
|
||
t.Fatalf("line = %q", line)
|
||
}
|
||
}
|
||
|
||
func TestIMGetQueryText_joinsUniqueQueries(t *testing.T) {
|
||
got := imGetQueryText(map[string]any{
|
||
"query": "foo",
|
||
"queries": []any{"foo", "bar"},
|
||
})
|
||
if got != "foo，bar" {
|
||
t.Fatalf("query text = %q", got)
|
||
}
|
||
}
|