WeiCHEN
f12e2ed155
Close T20-B2/B6/B8/B9/B10 low-risk leaks: map-backed cancel reclaim, idempotency lock TTL in cleanupLoop, sync.Once information_schema probe, InvalidateLayer after chat memory upsert, and Subscribe before SSE headers.
447 lines
13 KiB
Go
447 lines
13 KiB
Go
package ctxmgr
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"sort"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
// ChatMemorySearchOptions 检索对话语义记忆选项
|
|
type ChatMemorySearchOptions struct {
|
|
Query string // 检索 query
|
|
UserID int64 // 必填：数据权限闸门（访问本人记忆）
|
|
ProjectID int64 // 可选：同项目公开记忆闸门
|
|
SessionID string // 可选：限定特定会话
|
|
ExcludeSessionID string // 可选：排除当前会话（避免与短期滑动窗口重复）
|
|
TopK int // 默认 3
|
|
MinScore float64 // 默认 0.25
|
|
}
|
|
|
|
// RecalledMemoryItem 召回的对话记忆项
|
|
type RecalledMemoryItem struct {
|
|
ID int64 `json:"id"`
|
|
Domain string `json:"domain"`
|
|
SessionID string `json:"session_id"`
|
|
ProjectID int64 `json:"project_id"`
|
|
UserID int64 `json:"user_id"`
|
|
TurnNo int `json:"turn_no"`
|
|
Title string `json:"title"`
|
|
Content string `json:"content"`
|
|
SimilarityScore float64 `json:"similarity_score"`
|
|
DecayedWeight float64 `json:"decayed_weight"`
|
|
FinalScore float64 `json:"final_score"`
|
|
CreatedAt time.Time `json:"created_at"`
|
|
IsOtherUser bool `json:"is_other_user"`
|
|
AuthorName string `json:"author_name,omitempty"`
|
|
Metadata map[string]any `json:"metadata,omitempty"`
|
|
}
|
|
|
|
// RecordTurnMemory 将一轮"问-答"对作为一条记忆切片向量化并持久化入库
|
|
// 触发点：session.sendMessageInternal 回复完成后异步执行（不阻塞回复）
|
|
// chunk_hash = sha256(session_id + ":" + turn_no)，天然幂等去重
|
|
func (m *ContextManager) RecordTurnMemory(
|
|
ctx context.Context,
|
|
sessionID string,
|
|
turnNo int,
|
|
userID, projectID int64,
|
|
userQuery, assistantReply, intentName, branch string,
|
|
isMilestone bool,
|
|
) (int64, error) {
|
|
return m.RecordTurnMemoryWithVisibility(ctx, sessionID, turnNo, userID, projectID, userQuery, assistantReply, intentName, branch, isMilestone, "private")
|
|
}
|
|
|
|
// RecordTurnMemoryWithVisibility 记录带显式可见性标签的对话切片 (private / project_public)
|
|
func (m *ContextManager) RecordTurnMemoryWithVisibility(
|
|
ctx context.Context,
|
|
sessionID string,
|
|
turnNo int,
|
|
userID, projectID int64,
|
|
userQuery, assistantReply, intentName, branch string,
|
|
isMilestone bool,
|
|
visibility string,
|
|
) (int64, error) {
|
|
if sessionID == "" {
|
|
return 0, errors.New("sessionID required")
|
|
}
|
|
cleanQ := strings.TrimSpace(userQuery)
|
|
cleanA := strings.TrimSpace(assistantReply)
|
|
if cleanQ == "" && cleanA == "" {
|
|
return 0, errors.New("query or assistant reply required")
|
|
}
|
|
if visibility == "" {
|
|
visibility = "private"
|
|
}
|
|
|
|
// 答复保留首 300 字符，控制切片信息密度
|
|
rA := []rune(cleanA)
|
|
if len(rA) > 300 {
|
|
cleanA = string(rA[:300]) + "..."
|
|
}
|
|
|
|
content := fmt.Sprintf("用户: %s\n助手: %s", cleanQ, cleanA)
|
|
title := fmt.Sprintf("对话·第%d轮", turnNo)
|
|
chunkHash := fmt.Sprintf("%x", sha256.Sum256([]byte(fmt.Sprintf("%s:%d", sessionID, turnNo))))
|
|
|
|
// 里程碑/槽位/决策轮次 raw_weight=1.20；普通问答 raw_weight=0.60
|
|
rawWeight := 0.60
|
|
if isMilestone {
|
|
rawWeight = 1.20
|
|
}
|
|
|
|
meta := map[string]any{
|
|
"session_id": sessionID,
|
|
"turn_no": turnNo,
|
|
"intent": intentName,
|
|
"branch": branch,
|
|
"visibility": visibility,
|
|
"is_milestone": isMilestone,
|
|
}
|
|
|
|
emb := m.embedder
|
|
if emb == nil {
|
|
emb = NewDeterministicEmbedder(1024)
|
|
}
|
|
vecs, err := emb.Embed(ctx, []string{content})
|
|
if err != nil || len(vecs) == 0 {
|
|
return 0, fmt.Errorf("embed turn content failed: %w", err)
|
|
}
|
|
vec := vecs[0]
|
|
|
|
if m.pool != nil {
|
|
metaJSON, _ := json.Marshal(meta)
|
|
vecLit := VectorLiteral(vec)
|
|
|
|
var uid *int64
|
|
if userID > 0 {
|
|
uid = &userID
|
|
}
|
|
var pid *int64
|
|
if projectID > 0 {
|
|
pid = &projectID
|
|
}
|
|
|
|
var chunkID int64
|
|
query := `
|
|
INSERT INTO context_memory_chunks (
|
|
domain, project_id, user_id, reference_id, title, content,
|
|
chunk_hash, embedding, importance_weight, metadata, created_at
|
|
)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8::vector, $9, $10, now())
|
|
ON CONFLICT (chunk_hash) DO UPDATE SET
|
|
title = EXCLUDED.title,
|
|
content = EXCLUDED.content,
|
|
embedding = EXCLUDED.embedding,
|
|
importance_weight = EXCLUDED.importance_weight,
|
|
metadata = EXCLUDED.metadata
|
|
RETURNING id
|
|
`
|
|
err := m.pool.QueryRow(ctx, query,
|
|
"chat_history", pid, uid, sessionID, title, content,
|
|
chunkHash, vecLit, rawWeight, metaJSON,
|
|
).Scan(&chunkID)
|
|
if err != nil {
|
|
return 0, fmt.Errorf("insert chat memory chunk to db failed: %w", err)
|
|
}
|
|
// C1 search cache keys include query text; drop the layer so new chunks are visible.
|
|
m.Cache().InvalidateLayer(LayerC1Search)
|
|
return chunkID, nil
|
|
}
|
|
|
|
// 离线内存兜底
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
|
|
chunkID := m.nextChunkID
|
|
m.nextChunkID++
|
|
item := MemoryChunkItem{
|
|
ID: chunkID,
|
|
Domain: "chat_history",
|
|
ProjectID: projectID,
|
|
UserID: userID,
|
|
ReferenceID: sessionID,
|
|
Title: title,
|
|
Content: content,
|
|
ChunkHash: chunkHash,
|
|
Embedding: vec,
|
|
ImportanceWeight: rawWeight,
|
|
Metadata: meta,
|
|
CreatedAt: time.Now(),
|
|
}
|
|
|
|
updated := false
|
|
for i, existing := range m.memoryChunks {
|
|
if existing.ChunkHash == chunkHash {
|
|
m.memoryChunks[i] = item
|
|
updated = true
|
|
chunkID = existing.ID
|
|
break
|
|
}
|
|
}
|
|
if !updated {
|
|
m.memoryChunks = append(m.memoryChunks, item)
|
|
}
|
|
|
|
m.Cache().InvalidateLayer(LayerC1Search)
|
|
return chunkID, nil
|
|
}
|
|
|
|
// SearchChatMemory 按语义余弦距离与时间半衰期衰减检索相关对话历史记忆
|
|
func (m *ContextManager) SearchChatMemory(ctx context.Context, opts ChatMemorySearchOptions) ([]RecalledMemoryItem, error) {
|
|
if strings.TrimSpace(opts.Query) == "" {
|
|
return nil, nil
|
|
}
|
|
if opts.TopK <= 0 {
|
|
opts.TopK = 3
|
|
}
|
|
if opts.MinScore <= 0 {
|
|
opts.MinScore = 0.25
|
|
}
|
|
|
|
cacheKey := HashKey(
|
|
"chat_memory",
|
|
fmt.Sprintf("%d", opts.UserID),
|
|
fmt.Sprintf("%d", opts.ProjectID),
|
|
opts.Query,
|
|
opts.SessionID,
|
|
opts.ExcludeSessionID,
|
|
fmt.Sprintf("%d", opts.TopK),
|
|
)
|
|
if cached, ok := m.Cache().Get(LayerC1Search, cacheKey); ok {
|
|
if items, ok := cached.([]RecalledMemoryItem); ok {
|
|
return items, nil
|
|
}
|
|
}
|
|
|
|
emb := m.embedder
|
|
if emb == nil {
|
|
emb = NewDeterministicEmbedder(1024)
|
|
}
|
|
qVecs, err := emb.Embed(ctx, []string{opts.Query})
|
|
if err != nil || len(qVecs) == 0 {
|
|
return nil, fmt.Errorf("embed search query failed: %w", err)
|
|
}
|
|
queryVec := qVecs[0]
|
|
|
|
if m.pool != nil {
|
|
vecLit := VectorLiteral(queryVec)
|
|
// SQL 物理权限第一闸门：
|
|
// 只能查本人的问答记录，或者同项目下明确标记为 project_public 的记录
|
|
query := `
|
|
SELECT c.id, c.domain, COALESCE(c.project_id, 0), COALESCE(c.user_id, 0),
|
|
c.reference_id, c.title, c.content, c.importance_weight,
|
|
c.metadata, c.created_at,
|
|
1 - (c.embedding <=> $1::vector) as cosine_sim,
|
|
COALESCE(u.name, '') as author_name
|
|
FROM context_memory_chunks c
|
|
LEFT JOIN users u ON c.user_id = u.id
|
|
WHERE c.domain = 'chat_history'
|
|
AND (
|
|
c.user_id = $2
|
|
OR (c.project_id = $3 AND c.project_id IS NOT NULL AND c.metadata->>'visibility' = 'project_public')
|
|
)
|
|
AND ($4 = '' OR c.reference_id != $4)
|
|
AND ($5 = '' OR c.reference_id = $5)
|
|
ORDER BY c.embedding <=> $1::vector ASC
|
|
LIMIT 15
|
|
`
|
|
rows, err := m.pool.Query(ctx, query, vecLit, opts.UserID, opts.ProjectID, opts.ExcludeSessionID, opts.SessionID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("search chat memory db failed: %w", err)
|
|
}
|
|
defer rows.Close()
|
|
|
|
var candidates []RecalledMemoryItem
|
|
for rows.Next() {
|
|
var it RecalledMemoryItem
|
|
var rawMeta []byte
|
|
var rawWeight float64
|
|
var authorName string
|
|
if err := rows.Scan(
|
|
&it.ID, &it.Domain, &it.ProjectID, &it.UserID,
|
|
&it.SessionID, &it.Title, &it.Content, &rawWeight,
|
|
&rawMeta, &it.CreatedAt, &it.SimilarityScore, &authorName,
|
|
); err == nil {
|
|
_ = json.Unmarshal(rawMeta, &it.Metadata)
|
|
it.AuthorName = authorName
|
|
if it.Metadata != nil {
|
|
if t, ok := it.Metadata["turn_no"].(float64); ok {
|
|
it.TurnNo = int(t)
|
|
}
|
|
}
|
|
it.IsOtherUser = opts.UserID > 0 && it.UserID != opts.UserID
|
|
|
|
isMs := false
|
|
if it.Metadata != nil {
|
|
if mVal, ok := it.Metadata["is_milestone"].(bool); ok {
|
|
isMs = mVal
|
|
}
|
|
}
|
|
it.DecayedWeight = CalculateDecayedWeight(rawWeight, it.CreatedAt, isMs, m.decayCfg)
|
|
it.FinalScore = it.SimilarityScore * it.DecayedWeight
|
|
if it.FinalScore >= opts.MinScore {
|
|
candidates = append(candidates, it)
|
|
}
|
|
}
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, fmt.Errorf("scan memory chunks failed: %w", err)
|
|
}
|
|
|
|
sort.Slice(candidates, func(i, j int) bool {
|
|
return candidates[i].FinalScore > candidates[j].FinalScore
|
|
})
|
|
if len(candidates) > opts.TopK {
|
|
candidates = candidates[:opts.TopK]
|
|
}
|
|
m.Cache().Set(LayerC1Search, cacheKey, candidates, 20)
|
|
return candidates, nil
|
|
}
|
|
|
|
// 离线内存兜底
|
|
m.mu.RLock()
|
|
defer m.mu.RUnlock()
|
|
|
|
var candidates []RecalledMemoryItem
|
|
for _, c := range m.memoryChunks {
|
|
if c.Domain != "chat_history" {
|
|
continue
|
|
}
|
|
|
|
// 隐私权限第一闸门
|
|
vis, _ := c.Metadata["visibility"].(string)
|
|
canAccess := (opts.UserID > 0 && c.UserID == opts.UserID) ||
|
|
(opts.ProjectID > 0 && c.ProjectID == opts.ProjectID && vis == "project_public")
|
|
if !canAccess {
|
|
continue
|
|
}
|
|
if opts.ExcludeSessionID != "" && c.ReferenceID == opts.ExcludeSessionID {
|
|
continue
|
|
}
|
|
if opts.SessionID != "" && c.ReferenceID != opts.SessionID {
|
|
continue
|
|
}
|
|
|
|
sim := cosineSimilarity(queryVec, c.Embedding)
|
|
isMs, _ := c.Metadata["is_milestone"].(bool)
|
|
decayed := CalculateDecayedWeight(c.ImportanceWeight, c.CreatedAt, isMs, m.decayCfg)
|
|
finalScore := sim * decayed
|
|
|
|
if finalScore >= opts.MinScore {
|
|
turnNo := 0
|
|
if t, ok := c.Metadata["turn_no"].(int); ok {
|
|
turnNo = t
|
|
} else if tFloat, ok := c.Metadata["turn_no"].(float64); ok {
|
|
turnNo = int(tFloat)
|
|
}
|
|
|
|
authorName := ""
|
|
if u, ok := m.users[c.UserID]; ok && u != nil {
|
|
authorName = u.Name
|
|
}
|
|
|
|
candidates = append(candidates, RecalledMemoryItem{
|
|
ID: c.ID,
|
|
Domain: c.Domain,
|
|
SessionID: c.ReferenceID,
|
|
ProjectID: c.ProjectID,
|
|
UserID: c.UserID,
|
|
TurnNo: turnNo,
|
|
Title: c.Title,
|
|
Content: c.Content,
|
|
SimilarityScore: sim,
|
|
DecayedWeight: decayed,
|
|
FinalScore: finalScore,
|
|
CreatedAt: c.CreatedAt,
|
|
IsOtherUser: opts.UserID > 0 && c.UserID != opts.UserID,
|
|
AuthorName: authorName,
|
|
Metadata: c.Metadata,
|
|
})
|
|
}
|
|
}
|
|
|
|
sort.Slice(candidates, func(i, j int) bool {
|
|
return candidates[i].FinalScore > candidates[j].FinalScore
|
|
})
|
|
if len(candidates) > opts.TopK {
|
|
candidates = candidates[:opts.TopK]
|
|
}
|
|
m.Cache().Set(LayerC1Search, cacheKey, candidates, 20)
|
|
return candidates, nil
|
|
}
|
|
|
|
// FormatRecalledMemorySnippet 将召回的历史切片格式化为标准 Prompt 片段（附来源与隐私免责声明）
|
|
// 每条片段 ≤ 150 字，TopK ≤ 3，总预算 ≤ 600 字
|
|
func FormatRecalledMemorySnippet(items []RecalledMemoryItem, currentSessionID string) string {
|
|
if len(items) == 0 {
|
|
return ""
|
|
}
|
|
|
|
var sb strings.Builder
|
|
sb.WriteString("【相关历史片段】（按语义相关度召回，仅供参考，其中标注他项的内容禁止复述细节）\n")
|
|
|
|
totalRunes := 0
|
|
const maxTotalRunes = 600
|
|
const maxItemRunes = 150
|
|
|
|
for _, it := range items {
|
|
sourceTag := ""
|
|
if currentSessionID != "" && it.SessionID == currentSessionID {
|
|
sourceTag = fmt.Sprintf("[本会话·第%d轮]", it.TurnNo)
|
|
} else if it.IsOtherUser {
|
|
author := it.AuthorName
|
|
if author == "" {
|
|
author = "项目成员"
|
|
}
|
|
dateStr := it.CreatedAt.Format("2006-01-02")
|
|
sourceTag = fmt.Sprintf("[项目公开·%s] (作者: %s)", dateStr, author)
|
|
} else {
|
|
dateStr := it.CreatedAt.Format("2006-01-02")
|
|
sourceTag = fmt.Sprintf("[历史会话·%s]", dateStr)
|
|
}
|
|
|
|
// 提取问答核心要点
|
|
summary := summarizeTurnContent(it.Content)
|
|
if it.IsOtherUser {
|
|
summary += " → 仅可引用客观结论，不得评价个人"
|
|
}
|
|
|
|
r := []rune(summary)
|
|
if len(r) > maxItemRunes {
|
|
summary = string(r[:maxItemRunes-3]) + "..."
|
|
}
|
|
|
|
line := fmt.Sprintf("- %s %s\n", sourceTag, summary)
|
|
lineRunes := len([]rune(line))
|
|
if totalRunes+lineRunes > maxTotalRunes {
|
|
break
|
|
}
|
|
sb.WriteString(line)
|
|
totalRunes += lineRunes
|
|
}
|
|
|
|
return sb.String()
|
|
}
|
|
|
|
func summarizeTurnContent(content string) string {
|
|
lines := strings.Split(content, "\n")
|
|
userPart := ""
|
|
asstPart := ""
|
|
for _, l := range lines {
|
|
lTrim := strings.TrimSpace(l)
|
|
if strings.HasPrefix(lTrim, "用户:") {
|
|
userPart = strings.TrimSpace(strings.TrimPrefix(lTrim, "用户:"))
|
|
} else if strings.HasPrefix(lTrim, "助手:") {
|
|
asstPart = strings.TrimSpace(strings.TrimPrefix(lTrim, "助手:"))
|
|
}
|
|
}
|
|
if userPart != "" && asstPart != "" {
|
|
return fmt.Sprintf("曾问\"%s\"，答复要点: %s", userPart, asstPart)
|
|
}
|
|
return strings.ReplaceAll(content, "\n", " ")
|
|
}
|