sundynix-agentix/sundynix-dispatcher/internal/eino/memory_extract.go

package eino

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"strings"

	"github.com/cloudwego/eino/schema"

	"github.com/sundynix/sundynix-dispatcher/internal/llm"
	"github.com/sundynix/sundynix-shared/contract"
	"github.com/sundynix/sundynix-shared/prompts"
)

// memOp 是 consolidate 的一条记忆操作（Mem0 式对账结果）。
type memOp struct {
	Op         string  `json:"op"` // ADD / UPDATE / DELETE / NOOP
	Key        string  `json:"key"`
	Value      string  `json:"value"`
	Importance float64 `json:"importance"` // 1~10，poignancy（读路径权重）
	Reason     string  `json:"reason"`
}

// consolidateMemory 写回阶段（异步、攒批）：把近期对话与已有画像交给 LLM 一次对账，
// 输出 ADD/UPDATE/DELETE/NOOP 操作并执行。ADD/UPDATE→upsert(带 importance)、DELETE→软删。
// 一次 LLM 调用同时做"抽取+对账"，从根上解决 exact-key 盲写导致的记忆腐烂。
func (o *Orchestrator) consolidateMemory(ctx context.Context, uid string, hist []*schema.Message) {
	if uid == "" || o.tools == nil || o.pool == nil || !o.pool.Ready() {
		return
	}
	dialogue := renderDialogue(hist)
	if strings.TrimSpace(dialogue) == "" {
		return
	}
	existing := parseProfile(o.fetchMemory(ctx, uid, ""))

	cctx, cancel := llmCtx(ctx)
	defer cancel()
	sys := prompts.Get(prompts.MemoryExtract)
	txt, err := o.pool.Chat(cctx, []llm.ChatMessage{
		{Role: "system", Content: sys},
		{Role: "user", Content: consolidatePrompt(dialogue, existing)},
	})
	if err != nil {
		log.Printf("[eino] (writeback) 记忆对账失败 user=%s: %v", uid, err)
		return
	}
	ops := sanitizeOps(parseOps(txt), existing)
	var add, del int
	for _, op := range ops {
		switch op.Op {
		case "DELETE":
			o.deleteMemory(ctx, uid, op.Key)
			del++
		default: // ADD / UPDATE
			o.upsertMemory(ctx, uid, op.Key, op.Value, op.Importance)
			add++
		}
	}
	if add+del > 0 {
		log.Printf("[eino] (writeback) 记忆对账 user=%s：写入/更新 %d，删除 %d", uid, add, del)
	}
}

// consolidatePrompt 拼对账提示词：给出已有记忆 + 近期对话，要求输出操作列表。
func consolidatePrompt(dialogue string, existing map[string]string) string {
	var ex strings.Builder
	if len(existing) == 0 {
		ex.WriteString("（暂无）")
	} else {
		for k, v := range existing {
			fmt.Fprintf(&ex, "- %s：%s\n", k, v)
		}
	}
	return fmt.Sprintf(`【已有记忆】
%s
【近期对话】
%s

请对账并只输出 JSON 数组，每项 {"op":"ADD|UPDATE|DELETE|NOOP","key":"维度","value":"值","importance":1到10的整数,"reason":"简述"}：
- ADD：对话中出现、已有记忆里没有的新长期偏好；
- UPDATE：与已有某条同维度但值变了（**key 必须复用已有的那个**）；
- DELETE：已有记忆被对话明确推翻/过时（key 用已有的）；
- NOOP：无需变更的不要输出。
importance 表示该偏好的重要程度（随口一提=低，强烈/反复强调=高）。只输出 JSON，无多余文字。`, strings.TrimRight(ex.String(), "\n"), dialogue)
}

// renderDialogue 把近期消息渲染为"用户/助手"对话文本（取末尾若干条，限长）。
func renderDialogue(hist []*schema.Message) string {
	const maxMsgs = 8
	start := 0
	if len(hist) > maxMsgs {
		start = len(hist) - maxMsgs
	}
	var b strings.Builder
	for _, m := range hist[start:] {
		role := "用户"
		if m.Role == schema.Assistant {
			role = "助手"
		} else if m.Role == schema.System {
			continue
		}
		fmt.Fprintf(&b, "%s：%s\n", role, truncate(m.Content, 500))
	}
	return strings.TrimRight(b.String(), "\n")
}

// parseOps 解析 LLM 对账结果（容忍 ```json 围栏）为 []memOp。
func parseOps(txt string) []memOp {
	var ops []memOp
	if json.Unmarshal([]byte(stripFence(txt)), &ops) != nil {
		return nil
	}
	return ops
}

// sanitizeOps 校验/规整操作：规范 op 名；ADD/UPDATE 需 key+value；DELETE 需 key 且必须命中已有；
// importance 夹到 [1,10]；同 key 去重（保留末个）；过滤 NOOP。
func sanitizeOps(ops []memOp, existing map[string]string) []memOp {
	byKey := map[string]memOp{}
	order := []string{}
	for _, op := range ops {
		op.Op = strings.ToUpper(strings.TrimSpace(op.Op))
		op.Key = strings.TrimSpace(op.Key)
		op.Value = strings.TrimSpace(op.Value)
		if op.Key == "" {
			continue
		}
		switch op.Op {
		case "ADD", "UPDATE":
			if op.Value == "" {
				continue
			}
			if op.Importance < 1 {
				op.Importance = 1
			} else if op.Importance > 10 {
				op.Importance = 10
			}
		case "DELETE":
			if _, ok := existing[op.Key]; !ok { // 只删确实存在的，防误删/幻删
				continue
			}
		default: // NOOP 或未知 → 丢弃
			continue
		}
		if _, seen := byKey[op.Key]; !seen {
			order = append(order, op.Key)
		}
		byKey[op.Key] = op
	}
	out := make([]memOp, 0, len(order))
	for _, k := range order {
		out = append(out, byKey[k])
	}
	return out
}

// parseProfile 把 memory_get 渲染的画像（"- 维度：值" 多行）解析回 map，供对账。
func parseProfile(s string) map[string]string {
	m := map[string]string{}
	for _, line := range strings.Split(s, "\n") {
		line = strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(line), "-"))
		if line == "" {
			continue
		}
		for _, sep := range []string{"：", ":"} { // 兼容全角/半角冒号
			if i := strings.Index(line, sep); i > 0 {
				k := strings.TrimSpace(line[:i])
				if k != "" {
					m[k] = strings.TrimSpace(line[i+len(sep):])
				}
				break
			}
		}
	}
	return m
}

// upsertMemory 经 mcp-go memory_upsert 登记/更新一条偏好（带 importance）。
func (o *Orchestrator) upsertMemory(ctx context.Context, uid, key, value string, importance float64) {
	cctx, cancel := context.WithTimeout(ctx, toolCallTimeout)
	defer cancel()
	if _, err := o.tools.CallTool(cctx, contract.ToolSubjectGo("memory_upsert"),
		&contract.ToolCall{Tool: "memory_upsert", Args: map[string]any{"user_id": uid, "key": key, "value": value, "importance": importance}}); err != nil {
		log.Printf("[eino] memory_upsert 失败 %s=%s: %v", key, value, err)
	}
}

// deleteMemory 经 mcp-go memory_delete 软删一条偏好。
func (o *Orchestrator) deleteMemory(ctx context.Context, uid, key string) {
	cctx, cancel := context.WithTimeout(ctx, toolCallTimeout)
	defer cancel()
	if _, err := o.tools.CallTool(cctx, contract.ToolSubjectGo("memory_delete"),
		&contract.ToolCall{Tool: "memory_delete", Args: map[string]any{"user_id": uid, "key": key}}); err != nil {
		log.Printf("[eino] memory_delete 失败 %s: %v", key, err)
	}
}