77 lines
2.7 KiB
Go
77 lines
2.7 KiB
Go
|
package workbench
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"strings"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||
|
|
)
|
||
|
|
|
||
|
|
type snapshotInstructions struct {
|
||
|
|
text string
|
||
|
|
err error
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *snapshotInstructions) ResolveAgentInstructions(context.Context, *types.CustomAgent) (string, error) {
|
||
|
|
return s.text, s.err
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestFrozenAgentRetainsInstructionsButRechecksSkillAccess(t *testing.T) {
|
||
|
|
s, ctx, _, model := generationHarness(t)
|
||
|
|
s.agents = generationAgentStub{config: types.CustomAgentConfig{SystemPrompt: "原人格", InstructionSkillIDs: []string{"skill"}, Temperature: 0.3, MaxCompletionTokens: 1000}}
|
||
|
|
instructions := &snapshotInstructions{text: "原技能"}
|
||
|
|
s.SetInstructionResolver(instructions)
|
||
|
|
snapshots := map[string]AgentSnapshot{}
|
||
|
|
ctx = context.WithValue(ctx, agentSnapshotsKey{}, snapshots)
|
||
|
|
if _, err := s.complete(ctx, "agent", "system", "prompt"); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
s.agents = generationAgentStub{config: types.CustomAgentConfig{SystemPrompt: "新人格", Temperature: 0.8, MaxCompletionTokens: 4000}}
|
||
|
|
instructions.text = "新技能"
|
||
|
|
if _, err := s.complete(ctx, "agent", "system", "prompt"); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if !strings.Contains(model.messages[0].Content, "原人格") || !strings.Contains(model.messages[0].Content, "原技能") || strings.Contains(model.messages[0].Content, "新技能") || model.options.Temperature != 0.3 || model.options.MaxCompletionTokens != 1000 {
|
||
|
|
t.Fatal("a saved task did not preserve its Agent and skill inputs")
|
||
|
|
}
|
||
|
|
instructions.err = ErrForbidden
|
||
|
|
if _, err := s.complete(ctx, "agent", "system", "prompt"); !errors.Is(err, ErrForbidden) {
|
||
|
|
t.Fatalf("revoked skill must not be usable via snapshot: %v", err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
type snapshotMemory struct {
|
||
|
|
interfaces.MemoryService
|
||
|
|
user string
|
||
|
|
calls int
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *snapshotMemory) Recall(ctx context.Context, _ string) interfaces.MemoryRecall {
|
||
|
|
m.user, _ = types.UserIDFromContext(ctx)
|
||
|
|
m.calls++
|
||
|
|
return interfaces.MemoryRecall{Prompt: "偏好：简洁表达"}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTaskMemoryUsesCurrentPrincipalAndAgentOptOut(t *testing.T) {
|
||
|
|
s, ctx, _, model := generationHarness(t)
|
||
|
|
m := &snapshotMemory{}
|
||
|
|
s.SetMemoryService(m)
|
||
|
|
if _, err := s.complete(ctx, "agent", "system", "prompt"); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if m.user != "writer" || !strings.Contains(model.messages[0].Content, "简洁表达") {
|
||
|
|
t.Fatal("missing scoped task memory")
|
||
|
|
}
|
||
|
|
disabled := false
|
||
|
|
s.agents = generationAgentStub{config: types.CustomAgentConfig{MemoryEnabled: &disabled}}
|
||
|
|
if _, err := s.complete(ctx, "agent", "system", "prompt"); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if m.calls != 1 || strings.Contains(model.messages[0].Content, "简洁表达") {
|
||
|
|
t.Fatal("agent opt-out ignored")
|
||
|
|
}
|
||
|
|
}
|