Agent-System/internal/ctxmgr/chat_memory.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)
		}
		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)
	}

	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", " ")
}