519 lines
14 KiB
Go
519 lines
14 KiB
Go
// 织忆 MemoryWeave — 蒸馏引擎核心
|
||
// 两阶段蒸馏：硬规则过滤 → LLM 5维度评估 → 批量蒸馏
|
||
|
||
package distill
|
||
|
||
import (
|
||
"bytes"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"log"
|
||
"net/http"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
// ─── 类型定义 ────────────────────────────────────────────────
|
||
|
||
// Category 记忆类别
|
||
type Category string
|
||
|
||
const (
|
||
CatSystemFact Category = "system_fact"
|
||
CatUserPref Category = "user_pref"
|
||
CatProjContext Category = "proj_context"
|
||
CatToolUsage Category = "tool_usage"
|
||
CatCodeSnippet Category = "code_snippet"
|
||
CatDecision Category = "decision"
|
||
)
|
||
|
||
// DistillInput 蒸馏输入
|
||
type DistillInput struct {
|
||
EpisodeID string
|
||
Content string
|
||
Category Category
|
||
Namespace string
|
||
AgentID string
|
||
}
|
||
|
||
// DistillResult 蒸馏输出
|
||
type DistillResult struct {
|
||
Facts []string
|
||
Entities []Entity
|
||
Score5D FiveDScore
|
||
Overall float64
|
||
}
|
||
|
||
// Entity 实体
|
||
type Entity struct {
|
||
Name string `json:"name"`
|
||
Type string `json:"type"` // entity / fact / decision / skill
|
||
Properties []string `json:"properties"`
|
||
}
|
||
|
||
// FiveDScore LLM 5维评估
|
||
type FiveDScore struct {
|
||
IS float64 `json:"is"` // Information Significance
|
||
SU float64 `json:"su"` // Strategic Utility
|
||
PA float64 `json:"pa"` // Practical Applicability
|
||
VD float64 `json:"vd"` // Validation Durability
|
||
RU float64 `json:"ru"` // Recall Usability
|
||
}
|
||
|
||
// LLMResponse LLM 完整响应（5D + 实体 + 事实）
|
||
type LLMResponse struct {
|
||
IS float64 `json:"is"`
|
||
SU float64 `json:"su"`
|
||
PA float64 `json:"pa"`
|
||
VD float64 `json:"vd"`
|
||
RU float64 `json:"ru"`
|
||
Entities []string `json:"entities"`
|
||
Facts []string `json:"facts"`
|
||
}
|
||
|
||
// LLM 5维权重
|
||
var Weights = FiveDScore{
|
||
IS: 0.20,
|
||
SU: 0.20,
|
||
PA: 0.15,
|
||
VD: 0.25,
|
||
RU: 0.20,
|
||
}
|
||
|
||
// ─── 蒸馏引擎 ────────────────────────────────────────────────
|
||
|
||
// Engine 蒸馏引擎
|
||
type Engine struct {
|
||
mu sync.Mutex
|
||
|
||
// LLM 配置
|
||
LLMEndpoint string
|
||
LLMModel string
|
||
APIKey string
|
||
|
||
// 队列
|
||
queue []DistillInput
|
||
batchSize int
|
||
batchTimeout time.Duration
|
||
lastDistill time.Time
|
||
|
||
// 成本控制
|
||
dailyLimit int
|
||
dailyUsed int
|
||
dailyReset time.Time
|
||
|
||
// HTTP客户端
|
||
client *http.Client
|
||
}
|
||
|
||
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: 30 * time.Second},
|
||
}
|
||
}
|
||
|
||
// Enqueue 入队
|
||
func (e *Engine) Enqueue(input DistillInput) {
|
||
e.mu.Lock()
|
||
defer e.mu.Unlock()
|
||
|
||
if !HardRulesPass(input.Content) {
|
||
return
|
||
}
|
||
|
||
e.queue = append(e.queue, input)
|
||
|
||
shouldFlush := len(e.queue) >= e.batchSize
|
||
timeout := time.Since(e.lastDistill) > e.batchTimeout
|
||
|
||
if shouldFlush || (timeout && len(e.queue) > 0) {
|
||
go e.flush()
|
||
}
|
||
}
|
||
|
||
// flush 批量蒸馏
|
||
func (e *Engine) flush() {
|
||
e.mu.Lock()
|
||
if len(e.queue) == 0 {
|
||
e.mu.Unlock()
|
||
return
|
||
}
|
||
|
||
batch := e.queue
|
||
e.queue = nil
|
||
e.lastDistill = time.Now()
|
||
e.mu.Unlock()
|
||
|
||
log.Printf("[distill] flush START: batch=%d items, dailyUsed=%d/%d, endpoint=%s, model=%s",
|
||
len(batch), e.dailyUsed, e.dailyLimit, e.LLMEndpoint, e.LLMModel)
|
||
|
||
// 成本控制检查
|
||
e.checkDailyLimit()
|
||
if e.dailyUsed >= e.dailyLimit {
|
||
log.Printf("[distill] flush FALLBACK: daily limit reached (%d/%d)", e.dailyUsed, e.dailyLimit)
|
||
results := fallbackDistill(batch)
|
||
e.emitResults(results)
|
||
return
|
||
}
|
||
|
||
// LLM 蒸馏
|
||
for _, input := range batch {
|
||
result := e.distillOne(input)
|
||
log.Printf("[distill] distillOne done: episode=%s facts=%d entities=%d",
|
||
input.EpisodeID, len(result.Facts), len(result.Entities))
|
||
e.dailyUsed++
|
||
e.emitResult(input, result)
|
||
}
|
||
log.Printf("[distill] flush DONE: processed %d items", len(batch))
|
||
}
|
||
|
||
// distillOne 蒸馏单条
|
||
func (e *Engine) distillOne(input DistillInput) DistillResult {
|
||
// 如果 LLM 端点不可用，降级
|
||
if e.LLMEndpoint == "" {
|
||
log.Printf("[distill] LLMEndpoint empty, fallback for %s", input.EpisodeID)
|
||
return fallbackSingle(input)
|
||
}
|
||
llmResp, err := e.callLLM5D(input.Content)
|
||
if err != nil {
|
||
log.Printf("[distill] callLLM5D err for %s: %v", input.EpisodeID, err)
|
||
return fallbackSingle(input)
|
||
}
|
||
|
||
score := FiveDScore{IS: llmResp.IS, SU: llmResp.SU, PA: llmResp.PA, VD: llmResp.VD, RU: llmResp.RU}
|
||
overall := score.IS*Weights.IS +
|
||
score.SU*Weights.SU +
|
||
score.PA*Weights.PA +
|
||
score.VD*Weights.VD +
|
||
score.RU*Weights.RU
|
||
|
||
log.Printf("[distill] score for %s: overall=%.3f is=%.2f su=%.2f pa=%.2f vd=%.2f ru=%.2f",
|
||
input.EpisodeID, overall, score.IS, score.SU, score.PA, score.VD, score.RU)
|
||
|
||
if overall < 0.7 && score.VD < 0.8 {
|
||
log.Printf("[distill] score below threshold for %s, skip", input.EpisodeID)
|
||
// 即使阈值未通过，仍然用 LLM 实体提取（如果有的话）
|
||
if len(llmResp.Entities) == 0 {
|
||
return DistillResult{}
|
||
}
|
||
}
|
||
|
||
// 优先使用 LLM 提取的实体和事实
|
||
var entities []Entity
|
||
var facts []string
|
||
|
||
if len(llmResp.Entities) > 0 {
|
||
for _, name := range llmResp.Entities {
|
||
entities = append(entities, Entity{
|
||
Name: name, Type: "entity", Properties: []string{"llm_extracted"},
|
||
})
|
||
}
|
||
} else {
|
||
// 降级：启发式提取
|
||
_, heuristicEntities := e.extractFacts(input.Content)
|
||
entities = heuristicEntities
|
||
}
|
||
|
||
// 事实
|
||
if len(llmResp.Facts) > 0 {
|
||
facts = llmResp.Facts
|
||
} else {
|
||
heuristicFacts, _ := e.extractFacts(input.Content)
|
||
facts = heuristicFacts
|
||
}
|
||
|
||
return DistillResult{
|
||
Facts: facts,
|
||
Entities: entities,
|
||
Score5D: score,
|
||
Overall: overall,
|
||
}
|
||
}
|
||
|
||
// callLLM5D 调用 LLM 进行 5维评估 + 实体/事实提取
|
||
func (e *Engine) callLLM5D(content string) (LLMResponse, error) {
|
||
prompt := fmt.Sprintf(`你是一个记忆质量评估器和信息提取器。分析以下内容，返回 JSON：
|
||
|
||
1. 5 维度评分（0-1）：
|
||
- is (Information Significance): 信息重要性
|
||
- su (Strategic Utility): 战略价值
|
||
- pa (Practical Applicability): 实用价值
|
||
- vd (Validation Durability): 验证耐久性
|
||
- ru (Recall Usability): 召回可用性
|
||
|
||
2. 提取命名实体和技术概念（entities）：重要的系统/工具/人名/技术名词
|
||
3. 提取核心事实陈述（facts）：具体的事实/决策/配置项
|
||
|
||
内容:
|
||
%s
|
||
|
||
只返回 JSON: {"is": 0.X, "su": 0.X, "pa": 0.X, "vd": 0.X, "ru": 0.X, "entities": ["entity1", "entity2"], "facts": ["fact1", "fact2"]}`, truncate(content, 500))
|
||
|
||
body := map[string]interface{}{
|
||
"model": e.LLMModel,
|
||
"messages": []map[string]string{
|
||
{"role": "user", "content": prompt},
|
||
},
|
||
"temperature": 0.2,
|
||
"max_tokens": 300,
|
||
}
|
||
|
||
jsonBody, err := json.Marshal(body)
|
||
if err != nil {
|
||
return LLMResponse{}, err
|
||
}
|
||
|
||
req, err := http.NewRequest("POST", e.LLMEndpoint, bytes.NewReader(jsonBody))
|
||
if err != nil {
|
||
return LLMResponse{}, err
|
||
}
|
||
req.Header.Set("Content-Type", "application/json")
|
||
if e.APIKey != "" {
|
||
req.Header.Set("Authorization", "Bearer "+e.APIKey)
|
||
}
|
||
|
||
resp, err := e.client.Do(req)
|
||
if err != nil {
|
||
return LLMResponse{}, err
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
respBody, _ := io.ReadAll(resp.Body)
|
||
|
||
var result struct {
|
||
Choices []struct {
|
||
Message struct {
|
||
Content string `json:"content"`
|
||
ReasoningContent string `json:"reasoning_content"`
|
||
} `json:"message"`
|
||
} `json:"choices"`
|
||
}
|
||
|
||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||
return LLMResponse{}, err
|
||
}
|
||
|
||
if len(result.Choices) == 0 {
|
||
return LLMResponse{}, fmt.Errorf("no choices in LLM response")
|
||
}
|
||
|
||
var llmResp LLMResponse
|
||
llmContent := result.Choices[0].Message.Content
|
||
if llmContent == "" {
|
||
llmContent = result.Choices[0].Message.ReasoningContent
|
||
}
|
||
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)
|
||
}
|
||
|
||
// 记录实体和事实数量
|
||
if len(llmResp.Entities) > 0 {
|
||
log.Printf("[distill] LLM entities for content: %d entities", len(llmResp.Entities))
|
||
}
|
||
if len(llmResp.Facts) > 0 {
|
||
log.Printf("[distill] LLM facts for content: %d facts", len(llmResp.Facts))
|
||
}
|
||
|
||
return llmResp, nil
|
||
}
|
||
|
||
// extractFacts 从内容中提取事实和实体（关键词 + 命名实体启发式）
|
||
func (e *Engine) extractFacts(content string) ([]string, []Entity) {
|
||
var facts []string
|
||
var entities []Entity
|
||
|
||
if len(content) > 20 {
|
||
facts = append(facts, truncate(content, 200))
|
||
}
|
||
|
||
// 启发式实体提取：提取大写单词、数字组合、中文命名实体
|
||
words := strings.Fields(content)
|
||
seen := make(map[string]bool)
|
||
for _, w := range words {
|
||
w = strings.Trim(w, ",.;:!?，。；：！？、\"'()（）[]【】")
|
||
if len(w) < 2 {
|
||
continue
|
||
}
|
||
// 大写字母开头（英文命名实体：Hermes, ComfyUI, Redis 等）
|
||
runes := []rune(w)
|
||
if len(runes) >= 2 && runes[0] >= 'A' && runes[0] <= 'Z' {
|
||
normalized := strings.ToLower(w)
|
||
if !seen[normalized] && !isStopWord(normalized) {
|
||
seen[normalized] = true
|
||
entities = append(entities, Entity{
|
||
Name: w, Type: "entity", Properties: []string{"extracted"},
|
||
})
|
||
}
|
||
}
|
||
// 中文实体（2-20 个纯中文字符，不含标点）
|
||
cleanChinese := stripNonChinese(w)
|
||
if len(cleanChinese) >= 2 && len(cleanChinese) <= 20 {
|
||
if !seen[cleanChinese] {
|
||
seen[cleanChinese] = true
|
||
isTech := containsAny(cleanChinese, []string{"端口", "配置", "系统", "内存", "显卡", "服务", "记忆", "织忆", "飞书", "版本", "模型", "工具", "项目", "安装", "部署", "目录", "路径", "开发", "语言"})
|
||
etype := "entity"
|
||
if isTech { etype = "fact" }
|
||
entities = append(entities, Entity{
|
||
Name: cleanChinese, Type: etype, Properties: []string{"extracted"},
|
||
})
|
||
}
|
||
}
|
||
// 数字/字母组合或纯数字（如 3050, 4GB, 8188, 7821）
|
||
if isTechToken(w) || isAllDigits(w) {
|
||
if !seen[w] {
|
||
seen[w] = true
|
||
entities = append(entities, Entity{
|
||
Name: w, Type: "entity", Properties: []string{"technical"},
|
||
})
|
||
}
|
||
}
|
||
}
|
||
|
||
return facts, entities
|
||
}
|
||
|
||
// isChineseSequence 判断是否为连续的中文字符串（长度在 minLen 到 maxLen 之间）
|
||
func isChineseSequence(s string, minLen, maxLen int) bool {
|
||
runes := []rune(s)
|
||
if len(runes) < minLen || len(runes) > maxLen {
|
||
return false
|
||
}
|
||
for _, r := range runes {
|
||
if r < 0x4E00 || r > 0x9FFF {
|
||
return false
|
||
}
|
||
}
|
||
return true
|
||
}
|
||
|
||
// isTechToken 判断是否为技术标记（包含数字+字母组合）
|
||
func isTechToken(s string) bool {
|
||
hasDigit := false
|
||
hasLetter := false
|
||
for _, r := range s {
|
||
if r >= '0' && r <= '9' { hasDigit = true }
|
||
if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') { hasLetter = true }
|
||
}
|
||
return hasDigit && hasLetter
|
||
}
|
||
|
||
// containsAny 判断字符串是否包含任意一个子串
|
||
func containsAny(s string, subs []string) bool {
|
||
for _, sub := range subs {
|
||
if strings.Contains(s, sub) {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
// isStopWord 判断是否为停用词
|
||
func isStopWord(w string) bool {
|
||
stops := map[string]bool{
|
||
"the": true, "a": true, "an": true, "is": true, "are": true,
|
||
"was": true, "were": true, "be": true, "been": true, "being": true,
|
||
"have": true, "has": true, "had": true, "do": true, "does": true,
|
||
"did": true, "will": true, "would": true, "could": true, "should": true,
|
||
"may": true, "might": true, "can": true, "shall": true,
|
||
"this": true, "that": true, "these": true, "those": true,
|
||
"it": true, "its": true, "they": true, "them": true, "their": true,
|
||
"we": true, "our": true, "you": true, "your": true, "he": true, "she": true,
|
||
"his": true, "her": true, "and": true, "or": true, "but": true,
|
||
"not": true, "no": true, "if": true, "then": true, "else": true,
|
||
"for": true, "with": true, "as": true, "at": true, "by": true,
|
||
"from": true, "in": true, "of": true, "on": true, "to": true, "up": true,
|
||
}
|
||
return stops[w]
|
||
}
|
||
|
||
// stripNonChinese 只保留中文字符
|
||
func stripNonChinese(s string) string {
|
||
var buf []rune
|
||
for _, r := range s {
|
||
if r >= 0x4E00 && r <= 0x9FFF {
|
||
buf = append(buf, r)
|
||
}
|
||
}
|
||
return string(buf)
|
||
}
|
||
|
||
// isAllDigits 判断是否为纯数字
|
||
func isAllDigits(s string) bool {
|
||
if len(s) == 0 {
|
||
return false
|
||
}
|
||
for _, r := range s {
|
||
if r < '0' || r > '9' {
|
||
return false
|
||
}
|
||
}
|
||
return true
|
||
}
|
||
|
||
// checkDailyLimit 每日限额检查
|
||
func (e *Engine) checkDailyLimit() {
|
||
now := time.Now()
|
||
if now.Sub(e.dailyReset) > 24*time.Hour {
|
||
e.dailyUsed = 0
|
||
e.dailyReset = now
|
||
}
|
||
}
|
||
|
||
// emitResult 发送蒸馏结果（回调）
|
||
func (e *Engine) emitResult(input DistillInput, result DistillResult) {
|
||
if len(result.Facts) == 0 {
|
||
return
|
||
}
|
||
// 由注册的回调处理
|
||
if OnDistillComplete != nil {
|
||
OnDistillComplete(input, result)
|
||
}
|
||
}
|
||
|
||
// emitResults 批量发送
|
||
func (e *Engine) emitResults(results map[DistillInput]DistillResult) {
|
||
for input, result := range results {
|
||
e.emitResult(input, result)
|
||
}
|
||
}
|
||
|
||
// OnDistillComplete 全局蒸馏完成回调
|
||
var OnDistillComplete func(DistillInput, DistillResult)
|
||
|
||
// ─── 辅助 ────────────────────────────────────────────────
|
||
|
||
func truncate(s string, maxLen int) string {
|
||
runes := []rune(s)
|
||
if len(runes) <= maxLen {
|
||
return s
|
||
}
|
||
return string(runes[:maxLen]) + "..."
|
||
}
|
||
|
||
// fallbackSingle 降级蒸馏（无 LLM）
|
||
func fallbackSingle(input DistillInput) DistillResult {
|
||
var facts []string
|
||
if len(input.Content) > 20 {
|
||
facts = append(facts, truncate(input.Content, 100))
|
||
}
|
||
return DistillResult{Facts: facts}
|
||
}
|
||
|
||
// fallbackDistill 批量降级蒸馏
|
||
func fallbackDistill(inputs []DistillInput) map[DistillInput]DistillResult {
|
||
results := make(map[DistillInput]DistillResult)
|
||
for _, input := range inputs {
|
||
results[input] = fallbackSingle(input)
|
||
}
|
||
return results
|
||
}
|