- memory/extractor: async LLM-based memory extraction from chat history - Parses JSON lines from LLM response as memory entries - Wire into orchestrator pipeline via PostProcess hook - Non-blocking: runs in goroutine with 30s timeout - Automatic dedup via existing keyword overlap logic
118 lines
3.3 KiB
Go
118 lines
3.3 KiB
Go
package memory
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
|
|
"git.yeij.top/AskaEth/Cyrene_v2/services/ai-core/model"
|
|
)
|
|
|
|
// LLMClient is the interface for calling an LLM.
|
|
type LLMClient interface {
|
|
Chat(ctx context.Context, messages []model.LLMMessage) (*model.LLMResponse, error)
|
|
}
|
|
|
|
// Extractor uses LLM to extract memories from conversations.
|
|
type Extractor struct {
|
|
svc *Service
|
|
llm LLMClient
|
|
enabled bool
|
|
}
|
|
|
|
// NewExtractor creates a memory extractor.
|
|
func NewExtractor(svc *Service, llm LLMClient) *Extractor {
|
|
return &Extractor{svc: svc, llm: llm, enabled: llm != nil}
|
|
}
|
|
|
|
// ExtractedMemory represents a single memory extracted by the LLM.
|
|
type ExtractedMemory struct {
|
|
Content string `json:"content"`
|
|
Category string `json:"category"`
|
|
Priority float64 `json:"priority"`
|
|
}
|
|
|
|
// ExtractFromConversation analyzes a conversation and extracts memories asynchronously.
|
|
// Does not block — runs in a goroutine.
|
|
func (e *Extractor) ExtractFromConversation(ctx context.Context, userID string, conversation []model.Message) {
|
|
if !e.enabled || e.llm == nil || len(conversation) == 0 {
|
|
return
|
|
}
|
|
|
|
go func() {
|
|
bgCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
|
|
memories, err := e.extract(bgCtx, conversation)
|
|
if err != nil {
|
|
return // Silent fail — don't block main flow
|
|
}
|
|
|
|
for _, m := range memories {
|
|
if m.Content == "" { continue }
|
|
if m.Priority <= 0 { m.Priority = 0.7 }
|
|
if m.Category == "" { m.Category = CatFact }
|
|
e.svc.Add(bgCtx, userID, "", m.Content, m.Category, m.Priority)
|
|
}
|
|
}()
|
|
}
|
|
|
|
func (e *Extractor) extract(ctx context.Context, messages []model.Message) ([]ExtractedMemory, error) {
|
|
// Build conversation text
|
|
var convText strings.Builder
|
|
for _, msg := range messages {
|
|
role := "用户"
|
|
if msg.Role == "assistant" || msg.Role == "action" {
|
|
role = "AI"
|
|
}
|
|
convText.WriteString(fmt.Sprintf("%s: %s\n", role, msg.Content))
|
|
}
|
|
|
|
prompt := []model.LLMMessage{
|
|
{Role: "system", Content: `你是一个记忆提取助手。从对话中提取关于用户的重要信息。
|
|
每条记忆用一行 JSON 表示,格式:{"content":"记忆内容","category":"category","priority":0.8}
|
|
|
|
类别可选: fact(事实), preference(偏好), event(事件), plan(计划)
|
|
优先级 0-1,重要信息用高优先级。
|
|
只提取新信息,不重复已有内容。
|
|
只输出 JSON 行,不要其他文字。`},
|
|
{Role: "user", Content: fmt.Sprintf("从以下对话中提取记忆:\n%s", convText.String())},
|
|
}
|
|
|
|
resp, err := e.llm.Chat(ctx, prompt)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("llm extract: %w", err)
|
|
}
|
|
|
|
return parseExtractedMemories(resp.Content), nil
|
|
}
|
|
|
|
func parseExtractedMemories(text string) []ExtractedMemory {
|
|
var memories []ExtractedMemory
|
|
lines := strings.Split(text, "\n")
|
|
for _, line := range lines {
|
|
line = strings.TrimSpace(line)
|
|
if line == "" || !strings.HasPrefix(line, "{") {
|
|
continue
|
|
}
|
|
var m ExtractedMemory
|
|
if err := json.Unmarshal([]byte(line), &m); err != nil {
|
|
// Try to extract just the JSON part
|
|
start := strings.Index(line, "{")
|
|
end := strings.LastIndex(line, "}")
|
|
if start >= 0 && end > start {
|
|
if err := json.Unmarshal([]byte(line[start:end+1]), &m); err == nil {
|
|
if m.Content != "" { memories = append(memories, m) }
|
|
}
|
|
}
|
|
continue
|
|
}
|
|
if m.Content != "" {
|
|
memories = append(memories, m)
|
|
}
|
|
}
|
|
return memories
|
|
}
|