52 lines
1.5 KiB
Go
52 lines
1.5 KiB
Go
|
package workbench
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"strings"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||
|
|
)
|
||
|
|
|
||
|
|
type snapshotMemoryStub struct {
|
||
|
|
interfaces.MemoryService
|
||
|
|
calls int
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *snapshotMemoryStub) Recall(context.Context, string) interfaces.MemoryRecall {
|
||
|
|
m.calls++
|
||
|
|
return interfaces.MemoryRecall{Prompt: "当前用户的私有记忆"}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestFrozenAgentCannotReenableDisabledMemory(t *testing.T) {
|
||
|
|
for _, state := range []string{"live_disabled", "snapshot_disabled", "both_enabled"} {
|
||
|
|
t.Run(state, func(t *testing.T) {
|
||
|
|
s, ctx, _, _ := generationHarness(t)
|
||
|
|
disabled, enabled := false, true
|
||
|
|
live := &enabled
|
||
|
|
saved := &enabled
|
||
|
|
if state == "live_disabled" {
|
||
|
|
live = &disabled
|
||
|
|
}
|
||
|
|
if state == "snapshot_disabled" {
|
||
|
|
saved = &disabled
|
||
|
|
}
|
||
|
|
s.agents = generationAgentStub{config: types.CustomAgentConfig{MemoryEnabled: live}}
|
||
|
|
memories := &snapshotMemoryStub{}
|
||
|
|
s.memories = memories
|
||
|
|
ctx = context.WithValue(ctx, agentSnapshotsKey{}, map[string]AgentSnapshot{
|
||
|
|
"author": {Name: "Author", ModelID: "model", MemoryEnabled: saved, OutputLimit: 2000},
|
||
|
|
})
|
||
|
|
completion, err := s.prepareCompletion(ctx, "author", "写作", "编写本段")
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
wantRecall := state == "both_enabled"
|
||
|
|
if (memories.calls == 1) != wantRecall || strings.Contains(completion.messages[0].Content, "当前用户的私有记忆") != wantRecall {
|
||
|
|
t.Fatalf("memory preference ignored: state=%s calls=%d", state, memories.calls)
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|