59 lines
2.3 KiB
Go
59 lines
2.3 KiB
Go
|
package extraction
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"os"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestLocalModelContentValueDecisions(t *testing.T) {
|
||
|
|
baseURL := os.Getenv("LOCAL_MODEL_BASE_URL")
|
||
|
|
if baseURL == "" {
|
||
|
|
t.Skip("LOCAL_MODEL_BASE_URL is not set")
|
||
|
|
}
|
||
|
|
modelID := os.Getenv("LOCAL_MODEL_ID")
|
||
|
|
if modelID == "" {
|
||
|
|
modelID = "qwen2.5:7b"
|
||
|
|
}
|
||
|
|
generator, err := NewCompatibleGenerator("openai_compatible", baseURL, "ollama-local-key", modelID, 0.1, 1200, 2*time.Minute)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
tests := []struct {
|
||
|
|
name, content string
|
||
|
|
meaningful, storyWorth bool
|
||
|
|
}{
|
||
|
|
{
|
||
|
|
name: "discard small talk",
|
||
|
|
content: "豆包：您好，今天过得怎么样？\n长者：还行。\n豆包：吃饭了吗？\n长者：嗯，吃了。\n豆包：那您好好休息。\n长者：好，再见。",
|
||
|
|
},
|
||
|
|
{
|
||
|
|
name: "keep useful update without story",
|
||
|
|
content: "豆包：今天胃口怎么样？\n长者：今天胃口挺好，中午吃了一碗西红柿鸡蛋面，晚上想吃点清淡的。",
|
||
|
|
meaningful: true,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
name: "create story candidate from coherent memory",
|
||
|
|
content: "豆包：您还记得第一次当老师的情形吗？\n长者：记得。那是1978年秋天，我二十二岁，被分到河边小学。第一天是姐姐骑自行车送我去的，教室里有三十多个孩子。我紧张得把粉笔掉在地上，一个小姑娘帮我捡起来，全班都笑了，我也跟着笑。后来那个小姑娘一直读到大学，很多年后还回来看过我。",
|
||
|
|
meaningful: true, storyWorth: true,
|
||
|
|
},
|
||
|
|
}
|
||
|
|
for _, test := range tests {
|
||
|
|
t.Run(test.name, func(t *testing.T) {
|
||
|
|
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
|
||
|
|
defer cancel()
|
||
|
|
result, generateErr := generator.Generate(ctx, GenerationInput{Content: test.content})
|
||
|
|
if generateErr != nil {
|
||
|
|
t.Fatal(generateErr)
|
||
|
|
}
|
||
|
|
if result.Meaningful != test.meaningful || result.StoryWorthy != test.storyWorth || result.AllowStory != test.storyWorth {
|
||
|
|
t.Fatalf("decision mismatch: meaningful=%v story_worthy=%v allow_story=%v reason=%q summary=%q", result.Meaningful, result.StoryWorthy, result.AllowStory, result.DecisionReason, result.Summary)
|
||
|
|
}
|
||
|
|
if !result.Meaningful && (result.Summary != "" || result.StoryContent != "") {
|
||
|
|
t.Fatalf("discarded conversation produced content: %+v", result)
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|