姚章浩
76 lines
2.3 KiB
Go
76 lines
2.3 KiB
Go
package qa
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"gitea.zheromiku.top/Weknora/dobotai-platform/internal/core"
|
|
"gitea.zheromiku.top/Weknora/dobotai-platform/internal/session"
|
|
"gitea.zheromiku.top/Weknora/dobotai-platform/internal/weknora"
|
|
)
|
|
|
|
func TestKnowledgeToolEventsMarkTurnAsSearched(t *testing.T) {
|
|
knowledgeTools := []string{
|
|
"knowledge_search",
|
|
"grep_chunks",
|
|
"list_knowledge_chunks",
|
|
"get_document_info",
|
|
"wiki_search",
|
|
"wiki_read_page",
|
|
"wiki_read_source_doc",
|
|
"query_knowledge_graph",
|
|
}
|
|
for index, toolName := range knowledgeTools {
|
|
t.Run(toolName, func(t *testing.T) {
|
|
repo := newMemoryRepository()
|
|
client := &fakeEventStreamingWeKnora{
|
|
fakeWeKnora: &fakeWeKnora{},
|
|
events: []weknora.StreamEvent{
|
|
{ResponseType: "tool_call", ToolName: toolName},
|
|
{ResponseType: "answer", FullAnswer: "answer", Done: true},
|
|
},
|
|
streamResult: weknora.ChatResult{RequestID: "request", Answer: "answer", References: []core.Reference{{KnowledgeID: "doc"}}},
|
|
}
|
|
service, err := New(repo, fixedResolver{}, client, session.GroupShared)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
messageID := "knowledge-tool-" + string(rune('a'+index))
|
|
if _, err = service.AskEvents(context.Background(), streamQARequest(messageID), func(weknora.StreamEvent) {}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
repo.mu.Lock()
|
|
stored := repo.turns["tenant:dingtalk:"+messageID]
|
|
repo.mu.Unlock()
|
|
if !stored.KnowledgeSearched {
|
|
t.Fatalf("tool %s did not mark knowledge searched", toolName)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestNonKnowledgeToolDoesNotMarkTurnAsSearched(t *testing.T) {
|
|
repo := newMemoryRepository()
|
|
client := &fakeEventStreamingWeKnora{
|
|
fakeWeKnora: &fakeWeKnora{},
|
|
events: []weknora.StreamEvent{
|
|
{ResponseType: "tool_call", ToolName: "query_understand"},
|
|
{ResponseType: "answer", FullAnswer: "answer", Done: true},
|
|
},
|
|
streamResult: weknora.ChatResult{RequestID: "request", Answer: "answer"},
|
|
}
|
|
service, err := New(repo, fixedResolver{}, client, session.GroupShared)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err = service.AskEvents(context.Background(), streamQARequest("non-knowledge-tool"), func(weknora.StreamEvent) {}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
repo.mu.Lock()
|
|
stored := repo.turns["tenant:dingtalk:non-knowledge-tool"]
|
|
repo.mu.Unlock()
|
|
if stored.KnowledgeSearched {
|
|
t.Fatal("non-knowledge tool marked turn as searched")
|
|
}
|
|
}
|