283 lines
7.6 KiB
Go
283 lines
7.6 KiB
Go
package ctxmgr_test
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"chat-agent-system/internal/ctxmgr"
|
|
)
|
|
|
|
func TestRecordTurnMemory_IdempotentAndWeight(t *testing.T) {
|
|
mgr := ctxmgr.NewContextManager(nil, nil, ctxmgr.DefaultTimeDecayConfig())
|
|
ctx := context.Background()
|
|
|
|
// 1. 记录普通闲聊轮次
|
|
id1, err := mgr.RecordTurnMemory(
|
|
ctx,
|
|
"sess_test_1",
|
|
1,
|
|
1001,
|
|
2001,
|
|
"你好，帮我看一下项目情况",
|
|
"好的，当前项目处于三角验证阶段...",
|
|
"CASUAL_CHAT",
|
|
"DIRECT_HIT",
|
|
false,
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("RecordTurnMemory failed: %v", err)
|
|
}
|
|
if id1 <= 0 {
|
|
t.Errorf("expected valid chunk ID, got %d", id1)
|
|
}
|
|
|
|
// 2. 幂等更新同一轮次（例如流式补充更新或重试）
|
|
id2, err := mgr.RecordTurnMemory(
|
|
ctx,
|
|
"sess_test_1",
|
|
1,
|
|
1001,
|
|
2001,
|
|
"你好，帮我看一下项目情况",
|
|
"好的，当前项目处于三角验证阶段，已为您整理待办事项。",
|
|
"CASUAL_CHAT",
|
|
"DIRECT_HIT",
|
|
false,
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("RecordTurnMemory retry failed: %v", err)
|
|
}
|
|
if id2 != id1 {
|
|
t.Errorf("expected idempotent chunk id %d, got %d", id1, id2)
|
|
}
|
|
|
|
// 3. 记录里程碑/决策轮次
|
|
idMilestone, err := mgr.RecordTurnMemory(
|
|
ctx,
|
|
"sess_test_1",
|
|
2,
|
|
1001,
|
|
2001,
|
|
"确认器械类型为超声手术设备，预算定为150万",
|
|
"已记录器械类型和预算约束，并同步到项目黑板。",
|
|
"CONFIRM_ACTION",
|
|
"DIRECT_HIT",
|
|
true,
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("RecordTurnMemory milestone failed: %v", err)
|
|
}
|
|
if idMilestone <= id1 {
|
|
t.Errorf("expected new chunk id > %d, got %d", id1, idMilestone)
|
|
}
|
|
}
|
|
|
|
func TestSearchChatMemory_PrivacyGate(t *testing.T) {
|
|
mgr := ctxmgr.NewContextManager(nil, nil, ctxmgr.DefaultTimeDecayConfig())
|
|
ctx := context.Background()
|
|
|
|
// 用户 A (ID: 101) 的私有会话
|
|
_, err := mgr.RecordTurnMemory(
|
|
ctx,
|
|
"sess_userA_priv",
|
|
1,
|
|
101,
|
|
2001,
|
|
"用户A私密提问：超声换能器设计工艺有缺陷怎么处理",
|
|
"建议重新核对压电陶瓷匹配层厚度与阻抗...",
|
|
"TECHNICAL_CONSULT",
|
|
"DIRECT_HIT",
|
|
false,
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("record user A private failed: %v", err)
|
|
}
|
|
|
|
// 用户 A (ID: 101) 的项目公开里程碑会话
|
|
_, err = mgr.RecordTurnMemoryWithVisibility(
|
|
ctx,
|
|
"sess_userA_pub",
|
|
2,
|
|
101,
|
|
2001,
|
|
"GB 9706.1 适用性确认完成，正式进入V0.5阶段",
|
|
"项目里程碑：标准适用性评估完成，全员可见。",
|
|
"EXECUTE_STAGE_ADVANCE",
|
|
"DIRECT_HIT",
|
|
true,
|
|
"project_public",
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("record user A public failed: %v", err)
|
|
}
|
|
|
|
// 用户 B (ID: 102) 的私有会话
|
|
_, err = mgr.RecordTurnMemory(
|
|
ctx,
|
|
"sess_userB_priv",
|
|
1,
|
|
102,
|
|
2001,
|
|
"用户B私有提问：GB 9706.1 绝缘耐压指标是多少",
|
|
"按照 GB 9706.1 标准要求，应用部分需耐受 4000V 试验电压...",
|
|
"TECHNICAL_CONSULT",
|
|
"DIRECT_HIT",
|
|
false,
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("record user B private failed: %v", err)
|
|
}
|
|
|
|
// Case 1: 用户 B 检索 "超声换能器" (用户 A 的私密话题)
|
|
resultsB, err := mgr.SearchChatMemory(ctx, ctxmgr.ChatMemorySearchOptions{
|
|
Query: "超声换能器设计工艺",
|
|
UserID: 102,
|
|
ProjectID: 2001,
|
|
TopK: 3,
|
|
MinScore: 0.1,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("SearchChatMemory for user B failed: %v", err)
|
|
}
|
|
// 绝对不能查出用户 A 的私有对话
|
|
for _, r := range resultsB {
|
|
if r.SessionID == "sess_userA_priv" || strings.Contains(r.Content, "用户A私密提问") {
|
|
t.Errorf("PRIVACY VIOLATION: User B recalled User A's private memory: %s", r.Content)
|
|
}
|
|
}
|
|
|
|
// Case 2: 用户 B 检索 "GB 9706.1"
|
|
// 用户 B 应当能查出：1) 自己的私有提问; 2) 用户 A 的公开里程碑
|
|
resultsBStandards, err := mgr.SearchChatMemory(ctx, ctxmgr.ChatMemorySearchOptions{
|
|
Query: "GB 9706.1 标准与适用性",
|
|
UserID: 102,
|
|
ProjectID: 2001,
|
|
TopK: 3,
|
|
MinScore: 0.1,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("SearchChatMemory failed: %v", err)
|
|
}
|
|
foundOwn := false
|
|
foundPublic := false
|
|
for _, r := range resultsBStandards {
|
|
if r.SessionID == "sess_userB_priv" {
|
|
foundOwn = true
|
|
if r.IsOtherUser {
|
|
t.Errorf("own turn should not have IsOtherUser=true")
|
|
}
|
|
}
|
|
if r.SessionID == "sess_userA_pub" {
|
|
foundPublic = true
|
|
if !r.IsOtherUser {
|
|
t.Errorf("other user's public turn should have IsOtherUser=true")
|
|
}
|
|
}
|
|
}
|
|
if !foundOwn {
|
|
t.Errorf("expected User B to recall their own chat memory")
|
|
}
|
|
if !foundPublic {
|
|
t.Errorf("expected User B to recall project_public memory from User A")
|
|
}
|
|
|
|
// Case 3: 用户 A 检索 "超声换能器"
|
|
// 用户 A 可以正常查出自己的私有提问
|
|
resultsA, err := mgr.SearchChatMemory(ctx, ctxmgr.ChatMemorySearchOptions{
|
|
Query: "超声换能器设计工艺",
|
|
UserID: 101,
|
|
ProjectID: 2001,
|
|
TopK: 3,
|
|
MinScore: 0.1,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("SearchChatMemory for user A failed: %v", err)
|
|
}
|
|
foundAOwn := false
|
|
for _, r := range resultsA {
|
|
if r.SessionID == "sess_userA_priv" {
|
|
foundAOwn = true
|
|
}
|
|
}
|
|
if !foundAOwn {
|
|
t.Errorf("expected User A to recall their own private memory")
|
|
}
|
|
}
|
|
|
|
func TestSearchChatMemory_ExcludeSessionID(t *testing.T) {
|
|
mgr := ctxmgr.NewContextManager(nil, nil, ctxmgr.DefaultTimeDecayConfig())
|
|
ctx := context.Background()
|
|
|
|
_, _ = mgr.RecordTurnMemory(ctx, "current_session", 1, 101, 2001, "测试同会话第一轮", "回复1", "CASUAL_CHAT", "", false)
|
|
_, _ = mgr.RecordTurnMemory(ctx, "past_session", 1, 101, 2001, "测试历史会话第一轮", "回复2", "CASUAL_CHAT", "", false)
|
|
|
|
// 排除当前会话
|
|
results, err := mgr.SearchChatMemory(ctx, ctxmgr.ChatMemorySearchOptions{
|
|
Query: "测试会话第一轮",
|
|
UserID: 101,
|
|
ProjectID: 2001,
|
|
ExcludeSessionID: "current_session",
|
|
TopK: 5,
|
|
MinScore: 0.1,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("SearchChatMemory failed: %v", err)
|
|
}
|
|
|
|
for _, r := range results {
|
|
if r.SessionID == "current_session" {
|
|
t.Errorf("expected current_session to be excluded, but found in results")
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestFormatRecalledMemorySnippet(t *testing.T) {
|
|
now := time.Date(2026, 9, 10, 10, 0, 0, 0, time.UTC)
|
|
items := []ctxmgr.RecalledMemoryItem{
|
|
{
|
|
SessionID: "sess_cur",
|
|
TurnNo: 9,
|
|
Content: "用户: GB 9706.1 对超声设备适用吗？\n助手: 适用，另需满足 YY 9706.102 EMC 规范。",
|
|
CreatedAt: now,
|
|
IsOtherUser: false,
|
|
SimilarityScore: 0.95,
|
|
},
|
|
{
|
|
SessionID: "sess_past",
|
|
TurnNo: 3,
|
|
Content: "用户: 客户预算通常是多少\n助手: 通常按千万级申报，但一期需严格控制在150万内。",
|
|
CreatedAt: now.Add(-2 * 24 * time.Hour),
|
|
IsOtherUser: false,
|
|
SimilarityScore: 0.88,
|
|
},
|
|
{
|
|
SessionID: "sess_other",
|
|
TurnNo: 5,
|
|
AuthorName: "李工",
|
|
Content: "用户: V0.1需求收集表终审结论\n助手: 需求收集表终审通过，进入框架阶段。",
|
|
CreatedAt: now.Add(-4 * 24 * time.Hour),
|
|
IsOtherUser: true,
|
|
SimilarityScore: 0.82,
|
|
},
|
|
}
|
|
|
|
snippet := ctxmgr.FormatRecalledMemorySnippet(items, "sess_cur")
|
|
if !strings.Contains(snippet, "【相关历史片段】") {
|
|
t.Fatalf("expected header not found")
|
|
}
|
|
if !strings.Contains(snippet, "[本会话·第9轮]") {
|
|
t.Errorf("expected [本会话·第9轮] tag")
|
|
}
|
|
if !strings.Contains(snippet, "[历史会话·2026-09-08]") {
|
|
t.Errorf("expected [历史会话·2026-09-08] tag")
|
|
}
|
|
if !strings.Contains(snippet, "[项目公开·2026-09-06] (作者: 李工)") {
|
|
t.Errorf("expected project public tag with author")
|
|
}
|
|
if !strings.Contains(snippet, "不得评价个人") {
|
|
t.Errorf("expected privacy disclaimer for other user's snippet")
|
|
}
|
|
}
|