feat(distill): P2 离线整合(UPDATE_PROMPT) + P3 双缓冲(token积累) + JSON健壮剥离修复
- P2: consolidate.go 新增 ConsolidateMemory + TextSimilarity，LLM三选一(update/delete/ignore)合并相似记忆
- P2: server.go 新增 POST /api/v1/consolidate/memory 手动触发端点
- P3: engine.go Enqueue 按 token 积累触发 flush(阈值2000)，batchTimeout兜底
- fix: LLM JSON 剥离增强(找首个{和最后}截取)，修复模型返回markdown/注释导致的parse error
This commit is contained in:
小唯
parent
0734ffaa5a
commit
82c3d25423
|
|
@ -0,0 +1,68 @@
|
|||
# P2+P3: 织忆记忆离线整合 + 双缓冲触发改造
|
||||
|
||||
> 2026-08-11 | 小唯 | 借鉴 zjunlp/LightMem（ICLR 2026）
|
||||
> 前置：P1 逐条事实提取已完成（commit 0734ffa）
|
||||
|
||||
## P2: 离线整合 UPDATE_PROMPT（记忆合并/冲突消解）
|
||||
|
||||
### 目标
|
||||
对相似记忆做 LLM 三选一决策（update 合并细节 / delete 冲突删旧 / ignore 不相关），解决记忆冗余和冲突。
|
||||
|
||||
### 实现（新增 `go/internal/distill/consolidate.go`）
|
||||
|
||||
1. **`ConsolidateMemory(ldb, llmConfig, namespace string)`** — 离线整合入口：
|
||||
- `ldb.Search("memories", zeroVec, 200, namespace)` 取全部记忆
|
||||
- 两两计算相似度（复用 bge 向量？简单方案：用 recall 端点向量检索找候选）
|
||||
- 对高相似候选对（score ≥ 0.85）调 LLM 三选一
|
||||
2. **UPDATE_PROMPT**（移植 LightMem 原文精髓）：
|
||||
- update：目标与候选描述同一事实但不完全一致 → 合并额外信息
|
||||
- delete：直接冲突且候选更新 → 删目标
|
||||
- ignore：不相关 → 跳过
|
||||
- 输出 JSON `{"action": "update"|"delete"|"ignore", "new_memory": "..."}`
|
||||
3. **执行**：
|
||||
- action=update → `UpdateMemoryContent(id, new_memory)`
|
||||
- action=delete → `DeleteMemory(id)`
|
||||
- action=ignore → 跳过
|
||||
4. **触发**：新增 API `POST /api/v1/consolidate/memory`（手动触发）+ 每日 cron 自动触发
|
||||
|
||||
### 依赖
|
||||
- `ldb.Search` / `ldb.UpdateMemoryContent` / `ldb.Delete`（需确认 Delete 存在）
|
||||
|
||||
## P3: 双缓冲触发（token 积累批量 distill）
|
||||
|
||||
### 目标
|
||||
LightMem 的 Sensory(512) → Short-term(2000) 双缓冲思想：织忆 distill 按 token 积累触发，而非按条数。
|
||||
|
||||
### 实现（改 `go/internal/distill/engine.go`）
|
||||
|
||||
1. **Engine 新增字段**：
|
||||
- `pendingTokens int` — 当前缓冲的累计 token 数
|
||||
- `flushTokenThreshold int` — 触发阈值（默认 2000，对应 LightMem short-term）
|
||||
- `maxBatchTokens int` — 单批上限（防止超大 batch）
|
||||
2. **Enqueue 改造**：
|
||||
- 入队时累加 `pendingTokens += estimateTokens(content)`（rune count / 2 中文近似）
|
||||
- `shouldFlush = pendingTokens >= flushTokenThreshold || len(queue) >= batchSize`
|
||||
- flush 后 `pendingTokens = 0`
|
||||
3. **token 估算**：简单函数 `estimateTokens(s) = len([]rune(s))/2`（中文≈1 token/字符，英文≈1 token/4字符，取折中）
|
||||
|
||||
### 配置
|
||||
- 通过环境变量 `DISTILL_FLUSH_TOKENS`（默认 2000）可调，避免硬编码
|
||||
|
||||
## 测试命令
|
||||
|
||||
```bash
|
||||
# P2 测试 — 触发手动整合
|
||||
curl -s -X POST -H "X-API-Key: zhiyi-dev-key-2026" -H "Content-Type: application/json" \
|
||||
-d '{"namespace":"hermes-main"}' \
|
||||
http://localhost:7821/api/v1/consolidate/memory
|
||||
|
||||
# P3 测试 — 提交 3 条小内容（累计 <2000 token），验证不立即 flush
|
||||
# 再提交大内容触发 flush，看日志 flush START 时机
|
||||
```
|
||||
|
||||
## 验收标准
|
||||
- [ ] go build 通过
|
||||
- [ ] P2: 手动触发后日志显示 update/delete/ignore 决策
|
||||
- [ ] P2: 相似记忆被合并（recall 不再返回重复内容）
|
||||
- [ ] P3: 小内容入队不立即 flush，达阈值才 flush
|
||||
- [ ] 部署后全链路健康
|
||||
|
|
@ -825,6 +825,92 @@ func NewServer() http.Handler {
|
|||
http.NotFound(w, r)
|
||||
}
|
||||
})
|
||||
// P2: 记忆离线整合（LightMem UPDATE_PROMPT）— 手动触发
|
||||
// POST /api/v1/consolidate/memory {"namespace":"hermes-main"}
|
||||
mux.HandleFunc("/api/v1/consolidate/memory", func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != "POST" {
|
||||
http.Error(w, "POST only", 405)
|
||||
return
|
||||
}
|
||||
if routes.DistillEngineRef == nil {
|
||||
respondJSON(w, 400, map[string]string{"error": "distill engine not initialized"})
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Namespace string `json:"namespace"`
|
||||
Limit int `json:"limit"`
|
||||
}
|
||||
_ = json.NewDecoder(r.Body).Decode(&req)
|
||||
if req.Namespace == "" {
|
||||
req.Namespace = "hermes-main"
|
||||
}
|
||||
if req.Limit <= 0 {
|
||||
req.Limit = 50
|
||||
}
|
||||
|
||||
// 取全部记忆（零向量搜索）
|
||||
zeroVec := make([]float32, 1024)
|
||||
mems, err := ldb.Search("memories", zeroVec, req.Limit, req.Namespace)
|
||||
if err != nil || len(mems) == 0 {
|
||||
respondJSON(w, 200, map[string]interface{}{"status": "ok", "processed": 0, "message": "no memories to consolidate"})
|
||||
return
|
||||
}
|
||||
|
||||
// 两两比较相似度（简单 Jaccard/词重叠启发式，避免调用嵌入）
|
||||
// 真实场景：应使用向量相似度，此处用词重叠近似
|
||||
processed := 0
|
||||
updated := 0
|
||||
deleted := 0
|
||||
ignored := 0
|
||||
errors_ := 0
|
||||
|
||||
for i := 0; i < len(mems); i++ {
|
||||
for j := i + 1; j < len(mems); j++ {
|
||||
sim := distill.TextSimilarity(mems[i].Content, mems[j].Content)
|
||||
if sim < 0.5 {
|
||||
continue
|
||||
}
|
||||
// 高相似 → LLM 决策（以 j 为目标，i 为候选）
|
||||
result, err := routes.DistillEngineRef.ConsolidateMemory(
|
||||
distill.MemoryCandidate{ID: mems[j].ID, Content: mems[j].Content},
|
||||
[]distill.MemoryCandidate{{ID: mems[i].ID, Content: mems[i].Content}},
|
||||
)
|
||||
if err != nil {
|
||||
errors_++
|
||||
continue
|
||||
}
|
||||
processed++
|
||||
switch result.Action {
|
||||
case "update":
|
||||
if result.NewMemory != "" {
|
||||
if err := ldb.UpdateMemoryContent(mems[j].ID, result.NewMemory, "consolidate"); err != nil {
|
||||
log.Printf("[consolidate] update failed %s: %v", mems[j].ID, err)
|
||||
} else {
|
||||
updated++
|
||||
}
|
||||
}
|
||||
case "delete":
|
||||
if err := ldb.SoftDelete(mems[j].ID, "consolidate: conflict"); err != nil {
|
||||
log.Printf("[consolidate] delete failed %s: %v", mems[j].ID, err)
|
||||
} else {
|
||||
deleted++
|
||||
}
|
||||
default:
|
||||
ignored++
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
respondJSON(w, 200, map[string]interface{}{
|
||||
"status": "ok",
|
||||
"processed": processed,
|
||||
"updated": updated,
|
||||
"deleted": deleted,
|
||||
"ignored": ignored,
|
||||
"errors": errors_,
|
||||
"total_memories": len(mems),
|
||||
})
|
||||
})
|
||||
// GET /api/v1/memories — list all memories (zero-vector search, for plugin compat)
|
||||
mux.HandleFunc("/api/v1/memories", adminAPI.ListMemories)
|
||||
// /api/v1/memory/{id}/versions — existing version history endpoint
|
||||
|
|
|
|||
|
|
@ -0,0 +1,254 @@
|
|||
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
|
||||
}
|
||||
|
|
@ -106,6 +106,10 @@ type Engine struct {
|
|||
batchTimeout time.Duration
|
||||
lastDistill time.Time
|
||||
|
||||
// P3 双缓冲：token 积累触发（LightMem 移植）
|
||||
pendingTokens int
|
||||
flushTokenThreshold int
|
||||
|
||||
// 成本控制
|
||||
dailyLimit int
|
||||
dailyUsed int
|
||||
|
|
@ -117,16 +121,29 @@ type Engine struct {
|
|||
|
||||
func NewEngine(llmEndpoint, llmModel, apiKey string) *Engine {
|
||||
return &Engine{
|
||||
LLMEndpoint: llmEndpoint,
|
||||
LLMModel: llmModel,
|
||||
APIKey: apiKey,
|
||||
batchSize: 10,
|
||||
batchTimeout: 5 * time.Minute,
|
||||
dailyLimit: 5000,
|
||||
client: &http.Client{Timeout: 120 * time.Second},
|
||||
LLMEndpoint: llmEndpoint,
|
||||
LLMModel: llmModel,
|
||||
APIKey: apiKey,
|
||||
batchSize: 10,
|
||||
batchTimeout: 5 * time.Minute,
|
||||
dailyLimit: 5000,
|
||||
flushTokenThreshold: 2000, // LightMem short-term buffer
|
||||
client: &http.Client{Timeout: 120 * time.Second},
|
||||
}
|
||||
}
|
||||
|
||||
// estimateTokens 粗略 token 估算（中文≈1 token/字符，英文≈1 token/4字符）
|
||||
// 用于 P3 双缓冲触发，不需要精确（只影响批量时机）
|
||||
func estimateTokens(s string) int {
|
||||
runes := len([]rune(s))
|
||||
if runes == 0 {
|
||||
return 0
|
||||
}
|
||||
// 中文按 1 token/字符，非中文按 1 token/4 字符近似
|
||||
// 简单折中：rune 数 / 2
|
||||
return (runes + 1) / 2
|
||||
}
|
||||
|
||||
// Enqueue 入队
|
||||
func (e *Engine) Enqueue(input DistillInput) {
|
||||
e.mu.Lock()
|
||||
|
|
@ -137,8 +154,10 @@ func (e *Engine) Enqueue(input DistillInput) {
|
|||
}
|
||||
|
||||
e.queue = append(e.queue, input)
|
||||
e.pendingTokens += estimateTokens(input.Content)
|
||||
|
||||
shouldFlush := len(e.queue) >= e.batchSize
|
||||
// P3 双缓冲触发：token 积累达阈值 或 条数达 batchSize 或 超时
|
||||
shouldFlush := e.pendingTokens >= e.flushTokenThreshold || len(e.queue) >= e.batchSize
|
||||
timeout := time.Since(e.lastDistill) > e.batchTimeout
|
||||
|
||||
if shouldFlush || (timeout && len(e.queue) > 0) {
|
||||
|
|
@ -156,6 +175,7 @@ func (e *Engine) flush() {
|
|||
|
||||
batch := e.queue
|
||||
e.queue = nil
|
||||
e.pendingTokens = 0
|
||||
e.lastDistill = time.Now()
|
||||
e.mu.Unlock()
|
||||
|
||||
|
|
@ -369,6 +389,14 @@ func (e *Engine) callLLM5D(content string) (LLMResponse, error) {
|
|||
llmContent = strings.TrimPrefix(llmContent, "```")
|
||||
llmContent = strings.TrimSuffix(llmContent, "```")
|
||||
llmContent = strings.TrimSpace(llmContent)
|
||||
// 健壮剥离：找第一个 { 和最后一个 } 截取（模型可能在 JSON 后加 markdown/注释）
|
||||
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)
|
||||
if err := json.Unmarshal([]byte(llmContent), &llmResp); err != nil {
|
||||
log.Printf("[distill] LLM JSON parse error: %v | content=%q", err, truncate(llmContent, 200))
|
||||
return LLMResponse{}, fmt.Errorf("parse score: %w", err)
|
||||
|
|
|
|||
Loading…
Reference in New Issue