memoryweave/go/internal/distill/consolidate.go

package distill

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"
	"time"
)

// TextSimilarity 简单的词重叠相似度（0.0-1.0）
// 用于 P2 记忆整合候选筛选（避免每次调嵌入向量）
// 导出供 api 包使用
func TextSimilarity(a, b string) float64 {
	tokensA := tokenizeWords(a)
	tokensB := tokenizeWords(b)
	if len(tokensA) == 0 || len(tokensB) == 0 {
		return 0
	}
	setB := make(map[string]bool, len(tokensB))
	for _, t := range tokensB {
		setB[t] = true
	}
	overlap := 0
	for _, t := range tokensA {
		if setB[t] {
			overlap++
		}
	}
	// Jaccard 变体：重叠 / 较小集合大小
	denom := len(tokensA)
	if len(tokensB) < denom {
		denom = len(tokensB)
	}
	if denom == 0 {
		return 0
	}
	return float64(overlap) / float64(denom)
}

// tokenizeWords 中英文分词（中文按 2-gram，英文按单词）
func tokenizeWords(s string) []string {
	runes := []rune(s)
	var tokens []string
	// 提取英文单词和数字
	var cur strings.Builder
	flush := func() {
		if cur.Len() >= 2 {
			tokens = append(tokens, strings.ToLower(cur.String()))
		}
		cur.Reset()
	}
	for _, r := range runes {
		if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') {
			cur.WriteRune(r)
		} else {
			flush()
		}
	}
	flush()
	// 中文 2-gram（连续的汉字）
	var cn []rune
	flushCN := func() {
		if len(cn) >= 2 {
			for i := 0; i <= len(cn)-2; i++ {
				tokens = append(tokens, string(cn[i:i+2]))
			}
		}
		cn = nil
	}
	for _, r := range runes {
		if r >= 0x4e00 && r <= 0x9fa5 {
			cn = append(cn, r)
		} else {
			flushCN()
		}
	}
	flushCN()
	return tokens
}

// ─── P2: 离线整合（LightMem UPDATE_PROMPT 移植）────────────────

// UpdateAction LLM 决策结果
type UpdateAction struct {
	Action    string `json:"action"` // update / delete / ignore
	NewMemory string `json:"new_memory"`
}

// UpdatePrompt 记忆整合 prompt（移植 LightMem UPDATE_PROMPT 精髓）
const UpdatePrompt = `你是一个记忆管理助手。
你的任务是判断目标记忆应该被更新、删除还是忽略，基于候选源记忆。

决策规则:
1. update: 如果目标记忆和候选记忆描述的是同一个事实/事件但不完全一致（候选提供了更多细节、修正或澄清），更新目标记忆，整合额外信息。
2. delete: 如果目标记忆和候选记忆存在直接冲突，且候选记忆更新（时间更近），删除目标记忆。
3. ignore: 如果目标记忆和候选记忆不相关，不做任何操作，忽略。

附加指导:
- 只使用提供的信息，不要编造细节。
- 操作始终作用于目标记忆。不要修改或纠正候选记忆的内容。

输出必须是 JSON 结构:
{"action": "update" | "delete" | "ignore", "new_memory": "..."}

示例1:
目标记忆: "用户喜欢咖啡。"
候选记忆:
- "用户早上喜欢卡布奇诺。"
- "用户有时加班时喝浓缩咖啡。"
- "用户不喝无咖啡因咖啡。"
输出:
{"action": "update", "new_memory": "用户喜欢咖啡，尤其喜欢早上喝卡布奇诺、加班时喝浓缩咖啡，并且不喝无咖啡因咖啡。"}

示例2:
目标记忆: "用户目前住在纽约。"
候选记忆:
- "用户2023年搬到了旧金山。"
- "他们提到喜欢湾区的天气。"
输出:
{"action": "delete"}

示例3:
目标记忆: "用户正在学做意大利菜。"
候选记忆:
- "用户最近开始练瑜伽。"
- "他们买了一辆新自行车通勤。"
输出:
{"action": "ignore"}

以下是新的目标记忆和候选记忆。请根据规则决定合适的操作（update、delete 或 ignore）。

目标记忆: %s

候选记忆:
%s
`

// MemoryCandidate 候选记忆（用于 LLM 决策）
type MemoryCandidate struct {
	ID      string
	Content string
}

// ConsolidateInput 整合输入
type ConsolidateInput struct {
	Target   MemoryCandidate
	Candidates []MemoryCandidate
}

// ConsolidateResult 整合结果
type ConsolidateResult struct {
	Action      string
	NewMemory   string
	TargetID    string
	HasDecision bool
}

// ConsolidateMemory 对一对相似记忆做 LLM 决策
// 返回 true 表示 LLM 调用成功且给出了决策
func (e *Engine) ConsolidateMemory(target MemoryCandidate, candidates []MemoryCandidate) (ConsolidateResult, error) {
	if e.LLMEndpoint == "" {
		return ConsolidateResult{}, fmt.Errorf("LLM endpoint empty")
	}
	if len(candidates) == 0 {
		return ConsolidateResult{}, fmt.Errorf("no candidates")
	}

	// 组装候选列表
	var candBuilder strings.Builder
	for i, c := range candidates {
		candBuilder.WriteString(fmt.Sprintf("- %s", c.Content))
		if i < len(candidates)-1 {
			candBuilder.WriteString("\n")
		}
	}

	prompt := fmt.Sprintf(UpdatePrompt, target.Content, candBuilder.String())

	body := map[string]interface{}{
		"model":       e.LLMModel,
		"messages":    []map[string]string{{"role": "user", "content": prompt}},
		"temperature": 0.1,
		"max_tokens":  300,
	}
	jsonBody, err := json.Marshal(body)
	if err != nil {
		return ConsolidateResult{}, err
	}

	req, err := http.NewRequest("POST", e.LLMEndpoint, bytes.NewReader(jsonBody))
	if err != nil {
		return ConsolidateResult{}, err
	}
	req.Header.Set("Content-Type", "application/json")
	if e.APIKey != "" {
		req.Header.Set("Authorization", "Bearer "+e.APIKey)
	}

	client := &http.Client{Timeout: 60 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return ConsolidateResult{}, err
	}
	defer resp.Body.Close()

	respBody, _ := io.ReadAll(resp.Body)

	var result struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
	}
	if err := json.Unmarshal(respBody, &result); err != nil {
		return ConsolidateResult{}, err
	}
	if len(result.Choices) == 0 {
		return ConsolidateResult{}, fmt.Errorf("no choices")
	}

	llmContent := strings.TrimSpace(result.Choices[0].Message.Content)
	// 剥离 code fence
	llmContent = strings.TrimPrefix(llmContent, "```json")
	llmContent = strings.TrimPrefix(llmContent, "```")
	llmContent = strings.TrimSuffix(llmContent, "```")
	llmContent = strings.TrimSpace(llmContent)
	// 健壮剥离：找第一个 { 和最后一个 } 截取
	if idx := strings.Index(llmContent, "{"); idx > 0 {
		llmContent = llmContent[idx:]
	}
	if idx := strings.LastIndex(llmContent, "}"); idx >= 0 && idx < len(llmContent)-1 {
		llmContent = llmContent[:idx+1]
	}
	llmContent = strings.TrimSpace(llmContent)

	var action UpdateAction
	if err := json.Unmarshal([]byte(llmContent), &action); err != nil {
		log.Printf("[consolidate] LLM JSON parse error: %v | content=%q", err, truncate(llmContent, 200))
		return ConsolidateResult{}, err
	}

	action.Action = strings.ToLower(strings.TrimSpace(action.Action))
	return ConsolidateResult{
		Action:      action.Action,
		NewMemory:   action.NewMemory,
		TargetID:    target.ID,
		HasDecision: action.Action == "update" || action.Action == "delete",
	}, nil
}