Memory v2: LLM-proposed tags, 0-100% confidence buckets, provenance/citations; clean UI errors

Author: dev-pd

backend/app/agents/agent_b.py

@@ -1,9 +1,10 @@
 """Agent B — the learning agent.
 
 B is Agent A plus memory: the same shared pipeline and tools, but before each run
-it tags the page, retrieves matching playbook lessons, and injects them into a
-reserved slot on top of the unchanged core prompt. While learning on the train
-split it also carries the memory tools to jot observations into the diary.
+it injects the top trusted playbook rules into a reserved slot on top of the
+unchanged core prompt. While learning on the train split it also carries the
+memory tools to jot observations into the diary — choosing its own tags, shown the
+existing vocabulary so the janitor's grouping stays coherent.
 
 Everything else is identical to A — that is what makes any score difference
 attributable to memory.
@@ -13,8 +14,7 @@
 
 from app.agents.blind_spot import scan_blind_spots
 from app.agents.pipeline import run_pipeline
-from app.agents.prompts import CORE_SYSTEM, LEARNING_NOTE, render_memory_slot
-from app.memory.tagging import tags_for
+from app.agents.prompts import CORE_SYSTEM, render_learning_note, render_memory_slot
 from app.models import AgentRunResult
 from app.repositories import CompletionRepository, RunRepository
 from app.repositories.diary import DiaryRepository
@@ -24,6 +24,8 @@
 from app.tools.memory_tools import ReadMemoryTool, WriteMemoryTool
 from app.tools.registry import build_default_registry
 
+INJECT_TOP_RULES = 5
+
 
 async def run_agent_b(
     url: str,
@@ -44,22 +46,24 @@ async def run_agent_b(
     diary = diary or DiaryRepository()
     playbook = playbook or PlaybookRepository()
 
-    # Tag the page (fetch once; a cheap repo read in replay) to drive retrieval.
+    # Fetch the page once (cheap repo read in replay) — needed for the blind-spot scan.
     fetch_tool = base.get("fetch_page")
     assert isinstance(fetch_tool, FetchPageTool)
     page = await fetch_tool.run(FetchPageInput(url=url))
-    tags = tags_for(keyword, page)
 
-    # Build B's prompt: core + (retrieved lessons) + (learning note).
     system = CORE_SYSTEM
     extra_tools = []
     if memory_on:
-        entries = await playbook.get_by_tags(tags)
-        system += render_memory_slot([e.rule for e in entries])
+        top = await playbook.top(INJECT_TOP_RULES)
+        system += render_memory_slot([e.rule for e in top])
         extra_tools.append(ReadMemoryTool(playbook=playbook))
     if learn:
-        system += LEARNING_NOTE
-        extra_tools.append(WriteMemoryTool(diary=diary, run=run_index, tags=tags))
+        known_tags = {e.tag for e in await playbook.all()}
+        known_tags.update(tag for entry in await diary.all() for tag in entry.tags)
+        system += render_learning_note(sorted(known_tags))
+        extra_tools.append(
+            WriteMemoryTool(diary=diary, run=run_index, url=url, keyword=keyword)
+        )
 
     b_registry = ToolRegistry([*base.tools(), *extra_tools])
 

backend/app/agents/prompts.py

@@ -51,15 +51,20 @@ def render_memory_slot(lessons: list[str]) -> str:
     )
 
 
-# Shown to B only while learning on the train split.
-LEARNING_NOTE = (
-    "\n\nYou are learning. When competitors reveal a pattern this page is missing, "
-    "call write_memory with an abstract, UNCONDITIONAL rule for this kind of page "
-    "that names the exact recommendation to make next time — phrased so it applies "
-    "even when a future page has no competitor data. Use the shape: 'On <page kind>, "
-    "always recommend <recommendation_type> (target <target>): <one-line why>.' "
-    "Never tie it to this specific URL or keyword."
-)
+# Shown to B only while learning on the train split. The known tag vocabulary is
+# passed in so the LLM reuses tags (keeping the janitor's grouping coherent).
+def render_learning_note(known_tags: list[str]) -> str:
+    vocab = ", ".join(known_tags) if known_tags else "(none yet — you may create tags)"
+    return (
+        "\n\nYou are learning. When competitors reveal a pattern this page is missing, "
+        "call write_memory with an abstract, UNCONDITIONAL rule for this kind of page "
+        "that names the exact recommendation to make next time — phrased so it applies "
+        "even when a future page has no competitor data. Use the shape: 'On <page kind>, "
+        "always recommend <recommendation_type> (target <target>): <one-line why>.' "
+        f"Choose 1-4 tags for the pattern. Existing tags: {vocab}. Reuse an existing "
+        "tag when it fits; only add a new tag for a genuinely new kind of pattern. "
+        "Never tie a lesson to this specific URL or keyword."
+    )
 
 
 # Per-stage focus. Each runs its own short ReAct loop.

backend/app/api/memory.py

@@ -1,8 +1,9 @@
-"""Memory endpoints — a peek at what Agent B has learned."""
+"""Memory endpoints — a peek at what Agent B has learned, and why."""
 
 from fastapi import APIRouter
 
-from app.models.memory import PlaybookEntry
+from app.models.memory import DiaryEntry, PlaybookEntry
+from app.repositories.diary import DiaryRepository
 from app.repositories.playbook import PlaybookRepository
 
 router = APIRouter(prefix="/memory", tags=["memory"])
@@ -11,3 +12,9 @@
 @router.get("/playbook", response_model=list[PlaybookEntry])
 async def playbook() -> list[PlaybookEntry]:
     return await PlaybookRepository().all()
+
+
+@router.get("/diary", response_model=list[DiaryEntry])
+async def diary() -> list[DiaryEntry]:
+    """The raw observations — used to show the provenance behind each rule."""
+    return await DiaryRepository().all()

backend/app/memory/janitor.py

@@ -1,45 +1,53 @@
 """The janitor — consolidate the diary into the playbook.
 
-Scores each tag's pattern by blending frequency, strength, and recency, and
-promotes the ones that clear a tunable bar from the messy diary into the clean
-playbook. The playbook is always *regenerated* from the diary (never edited in
-place), so it can be rebuilt if it ever goes bad, and it is bounded to a top-K by
-confidence. The outcome signal comes from B's own observations and their
-recurrence — never B grading itself.
+Scores each tag's pattern by blending frequency, strength, and recency into a
+0–100% confidence, and promotes the ones that clear a bar from the messy diary
+into the clean playbook. The playbook is always *regenerated* from the diary
+(never edited in place), bounded to a top-K by confidence. Each promoted rule
+cites the diary rows it was distilled from (provenance).
+
+Tags are proposed by the LLM (reusing a shown vocabulary), so grouping by exact
+tag stays coherent without any clustering.
 """
 
 from __future__ import annotations
 
 from collections import defaultdict
 
-from app.models.memory import DiaryEntry, PlaybookEntry
-from app.repositories.diary import DiaryRepository
-from app.repositories.playbook import PlaybookRepository
+from app.models.memory import ConfidenceBucket, DiaryEntry, PlaybookEntry, Strength
 
 RECENCY_DECAY = 0.8  # weight of an observation one run older
+STRENGTH_POINTS: dict[Strength, int] = {"weak": 15, "medium": 30, "strong": 50}
 MIN_RUNS = 2  # a pattern must recur across >= 2 training runs to be trusted
-PROMOTION_BAR = 3.0  # min total confidence for a tag to promote at all
-TOP_RULES_PER_TAG = 3  # keep the strongest few distinct lessons per tag
+PROMOTION_CONFIDENCE = 40  # min 0-100 confidence to promote (the low/medium boundary)
+TOP_RULES_PER_TAG = 3
 TOP_K = 10  # bounded active playbook
 
 
-def _weight(entry: DiaryEntry, now: int) -> float:
-    return entry.strength * (RECENCY_DECAY ** (now - entry.run))
+def _points(entry: DiaryEntry, now: int) -> float:
+    return STRENGTH_POINTS[entry.strength] * (RECENCY_DECAY ** (now - entry.run))
+
+
+def confidence_of(entries: list[DiaryEntry], now: int) -> int:
+    """0–100% confidence from frequency × strength × recency (capped)."""
+    return min(100, round(sum(_points(e, now) for e in entries)))
+
+
+def bucket_of(confidence: int) -> ConfidenceBucket:
+    if confidence >= 70:
+        return "high"
+    if confidence >= 40:
+        return "medium"
+    return "low"
 
 
 def consolidate_entries(
     diary: list[DiaryEntry],
     *,
-    bar: float = PROMOTION_BAR,
+    bar: int = PROMOTION_CONFIDENCE,
     top_k: int = TOP_K,
 ) -> list[PlaybookEntry]:
-    """Pure consolidation: diary entries -> the promoted, bounded playbook.
-
-    A tag is trusted when its page-kind recurred across >= MIN_RUNS training runs
-    (that's the "proven pattern" gate — robust to the LLM phrasing lessons
-    differently each time). Within a trusted tag we keep the strongest few
-    distinct lessons by confidence (frequency x strength x recency).
-    """
+    """Pure consolidation: diary entries -> the promoted, bounded playbook."""
     if not diary:
         return []
     now = max(e.run for e in diary)
@@ -51,33 +59,33 @@ def consolidate_entries(
 
     promoted: list[PlaybookEntry] = []
     for tag, entries in by_tag.items():
-        runs = {e.run for e in entries}
-        if len(runs) < MIN_RUNS:
+        if len({e.run for e in entries}) < MIN_RUNS:
             continue  # pattern hasn't recurred across runs yet
-        if sum(_weight(e, now) for e in entries) < bar:
-            continue
 
-        # Score each distinct lesson under this tag.
         by_lesson: dict[str, list[DiaryEntry]] = defaultdict(list)
         for entry in entries:
             by_lesson[entry.lesson.strip()].append(entry)
 
         scored = sorted(
             (
-                (sum(_weight(e, now) for e in les_entries), lesson, les_entries)
+                (confidence_of(les_entries, now), lesson, les_entries)
                 for lesson, les_entries in by_lesson.items()
             ),
             reverse=True,
             key=lambda triple: triple[0],
         )
         for confidence, lesson, les_entries in scored[:TOP_RULES_PER_TAG]:
+            if confidence < bar:
+                continue
             promoted.append(
                 PlaybookEntry(
                     tag=tag,
                     rule=lesson,
-                    confidence=round(confidence, 3),
+                    confidence=confidence,
+                    confidence_bucket=bucket_of(confidence),
                     times_seen=len(les_entries),
                     last_run=max(e.run for e in les_entries),
+                    source_ids=[e.id for e in les_entries if e.id],
                 )
             )
 
@@ -86,13 +94,12 @@ def consolidate_entries(
 
 
 async def run_janitor(
-    diary_repo: DiaryRepository | None = None,
-    playbook_repo: PlaybookRepository | None = None,
-    *,
-    bar: float = PROMOTION_BAR,
-    top_k: int = TOP_K,
-) -> list[PlaybookEntry]:
+    diary_repo=None, playbook_repo=None, *, bar=PROMOTION_CONFIDENCE, top_k=TOP_K
+):
     """Read the diary, rebuild the playbook, write it atomically."""
+    from app.repositories.diary import DiaryRepository
+    from app.repositories.playbook import PlaybookRepository
+
     diary_repo = diary_repo or DiaryRepository()
     playbook_repo = playbook_repo or PlaybookRepository()
     diary = await diary_repo.all()

backend/app/memory/tagging.py

@@ -2,33 +2,44 @@
 
 `DiaryEntry` is the messy, append-only episodic log: raw tagged observations B
 jots during a run. `PlaybookEntry` is the clean, trusted, bounded semantic layer:
-abstract rules promoted from the diary, and the only thing injected back into B's
+abstract rules promoted from the diary, the only thing injected back into B's
 prompt.
 
-Both store *abstractions* (tagged patterns / rules), never answers glued to a
-specific URL or keyword — a lesson tied to specifics is a bug and would also trip
-the leak guard.
+Scoring is intuitive: each sighting has a **strength bucket** (weak/medium/strong),
+and a rule's **confidence is a 0–100%** score (frequency × strength × recency),
+itself bucketed low/medium/high. Both store *abstractions* (tagged patterns),
+never answers glued to a specific URL — a lesson tied to specifics is a bug.
 """
 
 from __future__ import annotations
 
-from pydantic import BaseModel, Field
+from typing import Literal
+
+from pydantic import BaseModel
+
+Strength = Literal["weak", "medium", "strong"]
+ConfidenceBucket = Literal["low", "medium", "high"]
 
 
 class DiaryEntry(BaseModel):
     """One raw observation B wrote during a run."""
 
+    id: str = ""  # Mongo row id (set on read) — used for provenance/citations
     tags: list[str]
     lesson: str
-    strength: int = Field(default=2, ge=1, le=3)  # weak=1, medium=2, strong=3
+    strength: Strength = "medium"
     run: int = 0  # training-run index, for recency
+    url: str = ""  # the page that prompted it (provenance only; never injected)
+    keyword: str = ""
 
 
 class PlaybookEntry(BaseModel):
     """A promoted, trusted rule. Tag-keyed; injected into B's prompt slot."""
 
     tag: str
     rule: str
-    confidence: float = 0.0
+    confidence: int = 0  # 0–100%
+    confidence_bucket: ConfidenceBucket = "low"
     times_seen: int = 0
     last_run: int = 0
+    source_ids: list[str] = []  # diary rows this rule was distilled from (citations)

backend/app/models/memory.py

@@ -10,11 +10,16 @@ class DiaryRepository(BaseRepository):
     collection_name = "diary"
 
     async def append(self, entry: DiaryEntry) -> None:
-        await self.collection.insert_one(entry.model_dump())
+        # Let Mongo assign the row id; we read it back as `id` for provenance.
+        await self.collection.insert_one(entry.model_dump(exclude={"id"}))
 
     async def all(self) -> list[DiaryEntry]:
         docs = await self.collection.find({}).to_list(length=None)
-        return [DiaryEntry.model_validate(doc) for doc in docs]
+        entries: list[DiaryEntry] = []
+        for doc in docs:
+            doc["id"] = str(doc.pop("_id"))
+            entries.append(DiaryEntry.model_validate(doc))
+        return entries
 
     async def clear(self) -> None:
         await self.collection.delete_many({})

backend/app/repositories/diary.py

@@ -27,6 +27,12 @@ async def get_by_tags(self, tags: list[str], top_k: int = 5) -> list[PlaybookEnt
         entries.sort(key=lambda e: e.confidence, reverse=True)
         return entries[:top_k]
 
+    async def top(self, k: int = 5) -> list[PlaybookEntry]:
+        """The highest-confidence rules overall (for prompt injection)."""
+        entries = await self.all()
+        entries.sort(key=lambda e: e.confidence, reverse=True)
+        return entries[:k]
+
     async def all(self) -> list[PlaybookEntry]:
         docs = await self.collection.find({}).to_list(length=None)
         return [PlaybookEntry.model_validate(doc) for doc in docs]