Cyrene_v2/internal/runtime/memory/llmextractor.go

package memory

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

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
)

// extractSystemPrompt tells the model.Adapter how to produce candidate memories:
// a JSON array of objects, one per memory worth persisting. Format is locked so
// parseCandidates can unmarshal it deterministically.
const extractSystemPrompt = `你是记忆提取器。从对话中提取值得长期记住的记忆项（用户偏好、习惯、约定、重要事件）。
只输出一个 JSON 数组，不要任何其他文字。每个元素形如：
{"content":"记忆内容","scope":"user|persona|project","kind":"relationship|world_event","tags":["标签"],"confidence":0.9,"project_id":"可选项目id","source_conversation":"溯源"}
去噪去重：只保留值得长期记住的，忽略寒暄与临时信息。`

// llmCandidate is the wire shape the extractor asks the LLM to emit: one JSON
// object per candidate memory. Scope/Kind default to the extractor's configured
// values when the model does not provide a valid one.
type llmCandidate struct {
	Content    string   `json:"content"`
	Scope      string   `json:"scope"`
	Kind       string   `json:"kind"`
	Tags       []string `json:"tags"`
	Confidence float64  `json:"confidence"`
	ProjectID  string   `json:"project_id"`
	Source     string   `json:"source_conversation"`
}

// LLMExtractor implements Extractor by asking a model.Adapter to derive candidate
// memory items from a conversation (design §1 LLM 异步自动提取). It runs the LLM
// call synchronously in Extract and on a background goroutine in ExtractAsync.
//
// It is idempotent and safe for concurrent use: candidate content is normalized
// and deduplicated against a mutex-guarded seen set, so re-extracting the same
// memory never yields a second copy, and the same content appearing twice inside
// one conversation collapses to a single item.
type LLMExtractor struct {
	adapter   model.Adapter
	personaID string
	roleID    string
	scope     Scope
	kind      Kind

	mu   sync.Mutex
	seen map[string]struct{}
}

// NewLLMExtractor builds an LLMExtractor whose extracted items are authored by
// personaID/roleID and default to ScopePersona + KindRelationship when the model
// does not specify otherwise. A nil adapter is rejected at use time via
// ErrNilAdapter, matching the lazy construction pattern of the rest of the
// runtime; callers may pass nil to build the struct and wire the adapter later.
func NewLLMExtractor(adapter model.Adapter, personaID, roleID string) *LLMExtractor {
	return &LLMExtractor{
		adapter:   adapter,
		personaID: personaID,
		roleID:    roleID,
		scope:     ScopePersona,
		kind:      KindRelationship,
		seen:      make(map[string]struct{}),
	}
}

// Extract returns the candidate memory items the adapter derives from
// conversation. A nil adapter returns ErrNilAdapter; a canceled context aborts
// before the LLM call. The response is parsed as a JSON array and then
// deduplicated, so items already extracted by this extractor are dropped.
func (e *LLMExtractor) Extract(ctx context.Context, conversation []model.Message) ([]MemoryItem, error) {
	if e.adapter == nil {
		return nil, ErrNilAdapter
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	resp, err := e.adapter.Chat(ctx, model.Request{
		Messages: buildExtractionPrompt(conversation),
		Purpose:  "memory_extract",
	}, model.Options{})
	if err != nil {
		return nil, fmt.Errorf("memory: llm extract: %w", err)
	}
	cands, err := parseCandidates(resp.Content)
	if err != nil {
		return nil, err
	}
	return e.dedupeAndBuild(cands), nil
}

// AsyncResult is the outcome of an ExtractAsync call: either candidate items or
// the error that aborted extraction.
type AsyncResult struct {
	Items []MemoryItem
	Err   error
}

// ExtractAsync submits the conversation to a background goroutine and returns a
// channel that delivers exactly one AsyncResult when extraction finishes, closed
// by the goroutine. A nil adapter returns ErrNilAdapter immediately. If ctx is
// canceled before the goroutine runs, the result carries the cancellation error.
// Because Extract is idempotent, repeated async extraction of the same memory is
// safe and never double-records.
func (e *LLMExtractor) ExtractAsync(ctx context.Context, conversation []model.Message) (<-chan AsyncResult, error) {
	if e.adapter == nil {
		return nil, ErrNilAdapter
	}
	out := make(chan AsyncResult, 1)
	go func() {
		defer close(out)
		if err := ctx.Err(); err != nil {
			out <- AsyncResult{Err: err}
			return
		}
		items, err := e.Extract(ctx, conversation)
		out <- AsyncResult{Items: items, Err: err}
	}()
	return out, nil
}

// buildExtractionPrompt turns the conversation into the model request messages:
// the system instruction plus one user message carrying the transcript.
func buildExtractionPrompt(conversation []model.Message) []model.Message {
	var b strings.Builder
	for _, m := range conversation {
		b.WriteString(string(m.Role))
		b.WriteString(": ")
		b.WriteString(messageText(m))
		b.WriteByte('\n')
	}
	return []model.Message{
		model.TextMessage(model.RoleSystem, extractSystemPrompt),
		model.TextMessage(model.RoleUser, b.String()),
	}
}

// parseCandidates decodes the adapter's text response into candidates. A fenced
// ```json block is unwrapped so the model may wrap its array in markdown. Blank
// or non-array content yields no candidates; malformed JSON is an error so the
// caller can decide whether to retry or fall back.
func parseCandidates(s string) ([]llmCandidate, error) {
	s = strings.TrimSpace(s)
	if strings.HasPrefix(s, "```") {
		s = strings.TrimPrefix(s, "```")
		s = strings.TrimSuffix(s, "```")
		s = strings.TrimSpace(s)
		s = strings.TrimPrefix(s, "json")
		s = strings.TrimSpace(s)
	}
	if s == "" {
		return nil, nil
	}
	if strings.HasPrefix(s, "{") {
		return nil, fmt.Errorf("memory: llm candidates: expected a JSON array")
	}
	var cands []llmCandidate
	if err := json.Unmarshal([]byte(s), &cands); err != nil {
		return nil, fmt.Errorf("memory: decode llm candidates: %w", err)
	}
	if cands == nil {
		return nil, nil
	}
	return cands, nil
}

// dedupeAndBuild converts valid candidates to MemoryItems, dropping blank content
// and any content already emitted by this extractor (idempotent dedup). It
// requires e.mu; only valid items mark their dedup key so a malformed candidate
// does not suppress a later valid one with the same text.
func (e *LLMExtractor) dedupeAndBuild(cands []llmCandidate) []MemoryItem {
	e.mu.Lock()
	defer e.mu.Unlock()

	out := make([]MemoryItem, 0, len(cands))
	for _, c := range cands {
		content := strings.TrimSpace(c.Content)
		if content == "" {
			continue
		}
		key := dedupeKey(content)
		if _, ok := e.seen[key]; ok {
			continue
		}
		it := MemoryItem{
			PersonaID:          e.personaID,
			RoleID:             e.roleID,
			ProjectID:          strings.TrimSpace(c.ProjectID),
			Scope:              Scope(normalizeScope(c.Scope, string(e.scope))),
			Kind:               Kind(normalizeKind(c.Kind, string(e.kind))),
			Content:            content,
			Tags:               c.Tags,
			Confidence:         c.Confidence,
			SourceConversation: strings.TrimSpace(c.Source),
		}
		if err := it.validate(); err != nil {
			continue
		}
		e.seen[key] = struct{}{}
		out = append(out, it)
	}
	return out
}

// dedupeKey normalizes content for idempotent dedup: case-folded and whitespace-
// collapsed, so "记住 A" and "记住 a" collide into one memory.
func dedupeKey(content string) string {
	return strings.Join(strings.Fields(strings.ToLower(content)), " ")
}

// normalizeScope returns s when it names a defined scope, otherwise fallback.
func normalizeScope(s, fallback string) string {
	if Scope(s).valid() {
		return s
	}
	return fallback
}

// normalizeKind returns s when it names a defined kind, otherwise fallback.
func normalizeKind(s, fallback string) string {
	if Kind(s).valid() {
		return s
	}
	return fallback
}