first commit
This commit is contained in:
@@ -0,0 +1,434 @@
|
||||
package dream
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// This file implements the real, LLM-backed Consolidator (SPEC §5.1 step 5,
|
||||
// §5.1.1). It turns the deterministic Plan into a prompt for the main-session
|
||||
// model, asks it to confirm semantic merges and conservative prunes, and parses
|
||||
// the strict-JSON response back into a ConsolidateResult the Runner applies.
|
||||
//
|
||||
// The deterministic half (exact dedup, dead-path cleanup, MEMORY.md rewrite,
|
||||
// Reconcile, counters) stays in the Runner; this file is purely the semantic
|
||||
// LLM step. It never touches disk — it only produces decisions the Runner
|
||||
// path-guards and applies.
|
||||
|
||||
// completeFn performs a single LLM completion: given the system and user
|
||||
// prompts it returns the model's text response, or an error for a hard
|
||||
// transport/provider failure. The provider-backed implementation lives in
|
||||
// apply.go (newModelCompleter); tests inject a canned function so no live model
|
||||
// is ever called.
|
||||
type completeFn func(ctx context.Context, systemPrompt, userPrompt string) (string, error)
|
||||
|
||||
// defaultBodyBudget caps how many bytes of each entry body are shown to the
|
||||
// model, bounding prompt size on large memory libraries (SPEC §8.2 token
|
||||
// budget). Bodies longer than this are truncated with a marker so the model
|
||||
// still sees the leading, usually most salient, content.
|
||||
const defaultBodyBudget = 6000
|
||||
|
||||
// llmConsolidator is the production Consolidator: it drives the main-session
|
||||
// model through complete and parses the response. bodyBudget (0 → default)
|
||||
// bounds per-entry body size in the prompt.
|
||||
type llmConsolidator struct {
|
||||
complete completeFn
|
||||
bodyBudget int
|
||||
}
|
||||
|
||||
// Consolidate builds the prompt from the plan, runs one model completion, and
|
||||
// parses the response into merge/prune decisions. A hard model/transport error
|
||||
// is returned (the Runner maps it to a failed run, SPEC §5.5). A well-formed
|
||||
// call whose text cannot be parsed is NOT an error: it yields an empty result
|
||||
// with an explanatory note, so an unparseable response conservatively KEEPs
|
||||
// everything (PRD FR-14) and the deterministic pass still applies.
|
||||
func (c *llmConsolidator) Consolidate(ctx context.Context, in ConsolidateInput) (ConsolidateResult, error) {
|
||||
if c.complete == nil {
|
||||
return ConsolidateResult{}, fmt.Errorf("dream: llmConsolidator has no completion function")
|
||||
}
|
||||
|
||||
var res ConsolidateResult
|
||||
eligible := eligibleFiles(in.Plan)
|
||||
if len(eligible) > 0 {
|
||||
budget := c.bodyBudget
|
||||
if budget <= 0 {
|
||||
budget = defaultBodyBudget
|
||||
}
|
||||
prompt := buildConsolidatePrompt(in, eligible, budget)
|
||||
|
||||
raw, err := c.complete(ctx, dreamSystemPrompt, prompt)
|
||||
if err != nil {
|
||||
return ConsolidateResult{}, fmt.Errorf("dream: model completion: %w", err)
|
||||
}
|
||||
|
||||
allowed := make(map[string]struct{}, len(eligible))
|
||||
for _, f := range eligible {
|
||||
allowed[filepath.Clean(f.Path)] = struct{}{}
|
||||
}
|
||||
res = parseConsolidateResponse(raw, allowed)
|
||||
}
|
||||
|
||||
// Distillation pass (SPEC §5.3, PRD US-005/FR-13): a SEPARATE model call over
|
||||
// the recent-session transcripts the Runner collected. It runs even when the
|
||||
// library is empty (nothing to merge/prune) so a first-time distill can seed
|
||||
// memory from sessions. A hard model failure aborts the run; a well-formed
|
||||
// call yielding nothing simply adds no entries (Runner records the no-op).
|
||||
if strings.TrimSpace(in.Transcripts) != "" {
|
||||
if err := c.distill(ctx, in, &res); err != nil {
|
||||
return ConsolidateResult{}, err
|
||||
}
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
|
||||
// distill runs the JSONL distillation model call and folds the resulting new
|
||||
// entries into res: it prompts the distiller with the transcripts plus a summary
|
||||
// of the existing library (so the model avoids re-proposing known facts), parses
|
||||
// the response into path-guarded NewEntry writes deduped against the existing
|
||||
// memory, and bumps res.Distilled by the number added. A hard model/transport
|
||||
// error is returned so the Runner marks the run failed (SPEC §5.5); an
|
||||
// unparseable or empty response adds nothing and is not an error (conservative
|
||||
// KEEP, PRD FR-14).
|
||||
func (c *llmConsolidator) distill(ctx context.Context, in ConsolidateInput, res *ConsolidateResult) error {
|
||||
prompt := buildDistillPrompt(in)
|
||||
raw, err := c.complete(ctx, dreamDistillSystemPrompt, prompt)
|
||||
if err != nil {
|
||||
return fmt.Errorf("dream: distill completion: %w", err)
|
||||
}
|
||||
entries, notes := parseDistillResponse(raw, in.Plan.Files, in.MemoryRoot, in.ProjectDir)
|
||||
res.NewEntries = append(res.NewEntries, entries...)
|
||||
res.Distilled += len(entries)
|
||||
res.Notes = append(res.Notes, notes...)
|
||||
return nil
|
||||
}
|
||||
|
||||
// buildDistillPrompt renders the distiller's user prompt: the target scope, a
|
||||
// compact list of the titles/paths of existing memories (so the model does not
|
||||
// re-propose known facts), and the recent-session transcripts. Existing bodies
|
||||
// are summarized (path + leading text) rather than dumped in full to keep the
|
||||
// prompt bounded; the Go side still enforces near-duplicate rejection.
|
||||
func buildDistillPrompt(in ConsolidateInput) string {
|
||||
var b strings.Builder
|
||||
b.WriteString("# Memory distillation request\n\n")
|
||||
if in.ProjectDir != "" {
|
||||
b.WriteString("Current project: ")
|
||||
b.WriteString(in.ProjectDir)
|
||||
b.WriteByte('\n')
|
||||
} else {
|
||||
b.WriteString("Scope: global only\n")
|
||||
}
|
||||
|
||||
existing := eligibleFiles(in.Plan)
|
||||
if len(existing) > 0 {
|
||||
b.WriteString(fmt.Sprintf("\n## Existing memories (%d) — do NOT re-propose these\n\n", len(existing)))
|
||||
for _, f := range existing {
|
||||
b.WriteString("- ")
|
||||
b.WriteString(clip(strings.TrimSpace(firstNonEmptyLine(f.Body)), 120))
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
}
|
||||
|
||||
b.WriteString("\n## Recent session transcripts\n\n")
|
||||
b.WriteString(in.Transcripts)
|
||||
b.WriteString("\n\nReturn the JSON object described in your instructions. Extract only genuinely new, durable facts. When in doubt, return no entries.\n")
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// firstNonEmptyLine returns the first non-blank line of s (trimmed), or "" when
|
||||
// s is entirely blank. Used to label an existing memory in the distill prompt.
|
||||
func firstNonEmptyLine(s string) string {
|
||||
for _, line := range strings.Split(s, "\n") {
|
||||
if t := strings.TrimSpace(line); t != "" {
|
||||
return t
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// eligibleFiles is the subset of plan files the model may act on: it excludes
|
||||
// MEMORY.md index files (they are indexes, not entries — #521 NEW_WORK: we
|
||||
// special-case MEMORY.md out of merge/prune so the index is never folded into an
|
||||
// entry) and the redundant members of an exact-dedupe group (g.Paths[1:], which
|
||||
// the deterministic pass removes anyway — offering them would let the model
|
||||
// merge into a path about to be deleted). The representative g.Paths[0] stays.
|
||||
func eligibleFiles(plan Plan) []MemoryFile {
|
||||
drop := make(map[string]struct{})
|
||||
for _, g := range plan.DedupeGroups {
|
||||
for _, p := range g.Paths[1:] {
|
||||
drop[filepath.Clean(p)] = struct{}{}
|
||||
}
|
||||
}
|
||||
var out []MemoryFile
|
||||
for _, f := range plan.Files {
|
||||
if isMemoryIndex(f.Path) {
|
||||
continue
|
||||
}
|
||||
if _, dup := drop[filepath.Clean(f.Path)]; dup {
|
||||
continue
|
||||
}
|
||||
out = append(out, f)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// isMemoryIndex reports whether path is a scope MEMORY.md index file, which the
|
||||
// consolidation step must never merge, rewrite, or prune.
|
||||
func isMemoryIndex(path string) bool {
|
||||
return strings.EqualFold(filepath.Base(path), "MEMORY.md")
|
||||
}
|
||||
|
||||
// buildConsolidatePrompt renders the user prompt: the scope, the eligible
|
||||
// entries (path + scope/type + body, truncated to budget), and the deterministic
|
||||
// hints (near-dup candidate pairs, dead local-path references). It only lists
|
||||
// paths that are eligible, so the model is naturally steered away from MEMORY.md
|
||||
// and duplicate paths.
|
||||
func buildConsolidatePrompt(in ConsolidateInput, eligible []MemoryFile, budget int) string {
|
||||
var b strings.Builder
|
||||
b.WriteString("# Memory consolidation request\n\n")
|
||||
b.WriteString("Memory root: ")
|
||||
b.WriteString(in.MemoryRoot)
|
||||
b.WriteByte('\n')
|
||||
if in.ProjectDir != "" {
|
||||
b.WriteString("Active project scope: ")
|
||||
b.WriteString(in.ProjectDir)
|
||||
b.WriteByte('\n')
|
||||
} else {
|
||||
b.WriteString("Scope: global only\n")
|
||||
}
|
||||
b.WriteString(fmt.Sprintf("\n## Entries (%d)\n\n", len(eligible)))
|
||||
for i, f := range eligible {
|
||||
typ := f.Type
|
||||
if typ == "" {
|
||||
typ = "(root)"
|
||||
}
|
||||
b.WriteString(fmt.Sprintf("### [%d] %s\n", i+1, f.Path))
|
||||
b.WriteString(fmt.Sprintf("scope=%s type=%s bytes=%d\n\n", f.Scope, typ, f.Size))
|
||||
b.WriteString("```\n")
|
||||
b.WriteString(truncateBody(f.Body, budget))
|
||||
b.WriteString("\n```\n\n")
|
||||
}
|
||||
|
||||
if pairs := eligiblePairs(in.Plan, eligible); len(pairs) > 0 {
|
||||
b.WriteString("## Near-duplicate candidate pairs (merge only if truly overlapping)\n\n")
|
||||
for _, p := range pairs {
|
||||
b.WriteString(fmt.Sprintf("- %s <-> %s (similarity %.2f)\n", p.A, p.B, p.Similarity))
|
||||
}
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
|
||||
if refs := eligibleInvalidRefs(in.Plan, eligible); len(refs) > 0 {
|
||||
b.WriteString("## Entries referencing local files that no longer exist\n")
|
||||
b.WriteString("(the dead reference text is cleaned automatically; only PRUNE an entry if losing that reference leaves it meaningless)\n\n")
|
||||
for _, r := range refs {
|
||||
b.WriteString(fmt.Sprintf("- %s references missing %s\n", r.File, r.Ref))
|
||||
}
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
|
||||
b.WriteString("Return the JSON object described in your instructions. When in doubt, KEEP.\n")
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// truncateBody trims body to at most budget bytes on a rune boundary, appending
|
||||
// a marker when truncation occurred so the model knows content was elided.
|
||||
func truncateBody(body string, budget int) string {
|
||||
if budget <= 0 || len(body) <= budget {
|
||||
return body
|
||||
}
|
||||
cut := budget
|
||||
for cut > 0 && !isRuneStart(body[cut]) {
|
||||
cut--
|
||||
}
|
||||
return body[:cut] + "\n…[truncated]"
|
||||
}
|
||||
|
||||
// isRuneStart reports whether b is not a UTF-8 continuation byte, so a truncation
|
||||
// cut there does not split a multi-byte rune.
|
||||
func isRuneStart(b byte) bool { return b&0xC0 != 0x80 }
|
||||
|
||||
// eligiblePairs filters the plan's near-dup pairs to those whose BOTH members
|
||||
// are eligible (both are still offered to the model), so we never point the
|
||||
// model at a MEMORY.md or a to-be-deduped duplicate.
|
||||
func eligiblePairs(plan Plan, eligible []MemoryFile) []NearDupPair {
|
||||
ok := make(map[string]struct{}, len(eligible))
|
||||
for _, f := range eligible {
|
||||
ok[filepath.Clean(f.Path)] = struct{}{}
|
||||
}
|
||||
var out []NearDupPair
|
||||
for _, p := range plan.NearDupPairs {
|
||||
_, a := ok[filepath.Clean(p.A)]
|
||||
_, b := ok[filepath.Clean(p.B)]
|
||||
if a && b {
|
||||
out = append(out, p)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// eligibleInvalidRefs filters dead-path references to those on eligible files.
|
||||
func eligibleInvalidRefs(plan Plan, eligible []MemoryFile) []InvalidPathRef {
|
||||
ok := make(map[string]struct{}, len(eligible))
|
||||
for _, f := range eligible {
|
||||
ok[filepath.Clean(f.Path)] = struct{}{}
|
||||
}
|
||||
var out []InvalidPathRef
|
||||
for _, r := range plan.InvalidPathRefs {
|
||||
if _, found := ok[filepath.Clean(r.File)]; found {
|
||||
out = append(out, r)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// modelDecision mirrors the strict JSON output schema the dream prompt requires.
|
||||
type modelDecision struct {
|
||||
Merges []struct {
|
||||
Keep string `json:"keep"`
|
||||
Body string `json:"body"`
|
||||
Remove []string `json:"remove"`
|
||||
} `json:"merges"`
|
||||
Prunes []struct {
|
||||
Path string `json:"path"`
|
||||
Reason string `json:"reason"`
|
||||
} `json:"prunes"`
|
||||
Notes []string `json:"notes"`
|
||||
}
|
||||
|
||||
// parseConsolidateResponse turns the model's text into a ConsolidateResult,
|
||||
// validating every path against allowed (the eligible input paths). It is
|
||||
// conservative by construction: any decision that references an unknown path, a
|
||||
// MEMORY.md index, an empty merged body, or an empty prune reason is dropped
|
||||
// rather than obeyed, and an unparseable response yields an empty result with a
|
||||
// note (PRD FR-14 — default to KEEP on uncertainty). It never returns an error;
|
||||
// callers treat a hard model failure separately.
|
||||
func parseConsolidateResponse(raw string, allowed map[string]struct{}) ConsolidateResult {
|
||||
body := extractJSONObject(raw)
|
||||
if body == "" {
|
||||
return ConsolidateResult{Notes: []string{"dream: model response contained no JSON object; kept all entries"}}
|
||||
}
|
||||
var dec modelDecision
|
||||
if err := json.Unmarshal([]byte(body), &dec); err != nil {
|
||||
return ConsolidateResult{Notes: []string{"dream: model response was not valid JSON; kept all entries"}}
|
||||
}
|
||||
|
||||
var res ConsolidateResult
|
||||
res.MergedBodies = make(map[string]string)
|
||||
delSet := make(map[string]struct{})
|
||||
|
||||
valid := func(p string) (string, bool) {
|
||||
clean := filepath.Clean(strings.TrimSpace(p))
|
||||
if clean == "" || clean == "." {
|
||||
return "", false
|
||||
}
|
||||
if isMemoryIndex(clean) {
|
||||
return "", false
|
||||
}
|
||||
if _, ok := allowed[clean]; !ok {
|
||||
return "", false
|
||||
}
|
||||
return clean, true
|
||||
}
|
||||
|
||||
for _, m := range dec.Merges {
|
||||
keep, ok := valid(m.Keep)
|
||||
if !ok || strings.TrimSpace(m.Body) == "" {
|
||||
// Unknown/invalid target or an empty rewrite: skip, keep everything.
|
||||
continue
|
||||
}
|
||||
removed := 0
|
||||
for _, r := range m.Remove {
|
||||
rp, ok := valid(r)
|
||||
if !ok || rp == keep {
|
||||
continue
|
||||
}
|
||||
if _, dup := delSet[rp]; dup {
|
||||
continue
|
||||
}
|
||||
delSet[rp] = struct{}{}
|
||||
removed++
|
||||
}
|
||||
if removed == 0 {
|
||||
// A merge that removes nothing is a no-op rewrite; ignore it to avoid
|
||||
// gratuitously touching a file the model merely echoed back.
|
||||
continue
|
||||
}
|
||||
res.MergedBodies[keep] = m.Body
|
||||
res.Merged += removed
|
||||
}
|
||||
|
||||
for _, p := range dec.Prunes {
|
||||
pp, ok := valid(p.Path)
|
||||
if !ok || strings.TrimSpace(p.Reason) == "" {
|
||||
// No path or no stated reason → conservative KEEP.
|
||||
continue
|
||||
}
|
||||
if _, dup := delSet[pp]; dup {
|
||||
continue
|
||||
}
|
||||
// Never prune an entry we are simultaneously keeping as a merge target.
|
||||
if _, kept := res.MergedBodies[pp]; kept {
|
||||
continue
|
||||
}
|
||||
delSet[pp] = struct{}{}
|
||||
res.Pruned++
|
||||
res.Notes = append(res.Notes, fmt.Sprintf("pruned %s: %s", pp, strings.TrimSpace(p.Reason)))
|
||||
}
|
||||
|
||||
if len(res.MergedBodies) == 0 {
|
||||
res.MergedBodies = nil
|
||||
}
|
||||
for p := range delSet {
|
||||
res.Deletions = append(res.Deletions, p)
|
||||
}
|
||||
sort.Strings(res.Deletions)
|
||||
|
||||
for _, n := range dec.Notes {
|
||||
if s := strings.TrimSpace(n); s != "" {
|
||||
res.Notes = append(res.Notes, s)
|
||||
}
|
||||
}
|
||||
return res
|
||||
}
|
||||
|
||||
// extractJSONObject returns the outermost {...} span of s, tolerating models
|
||||
// that wrap the object in prose or Markdown code fences. It returns "" when no
|
||||
// balanced object is found.
|
||||
func extractJSONObject(s string) string {
|
||||
start := strings.IndexByte(s, '{')
|
||||
if start < 0 {
|
||||
return ""
|
||||
}
|
||||
depth := 0
|
||||
inStr := false
|
||||
esc := false
|
||||
for i := start; i < len(s); i++ {
|
||||
ch := s[i]
|
||||
if inStr {
|
||||
switch {
|
||||
case esc:
|
||||
esc = false
|
||||
case ch == '\\':
|
||||
esc = true
|
||||
case ch == '"':
|
||||
inStr = false
|
||||
}
|
||||
continue
|
||||
}
|
||||
switch ch {
|
||||
case '"':
|
||||
inStr = true
|
||||
case '{':
|
||||
depth++
|
||||
case '}':
|
||||
depth--
|
||||
if depth == 0 {
|
||||
return s[start : i+1]
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
Reference in New Issue
Block a user