Phase 11: test cues on facts, and close the phase

bench/README.md

@@ -85,10 +85,11 @@ The options are these:
 | `--sample` | `a`, `b`, or `all` for every question of the data file | `a` |
 | `--arms` | a comma-separated list of `full`, `retrieval`, and `memory` | every arm |
 | `--retrieval-ranking` | `vector` or `fused`, the ranking of the retrieval arm | `vector` |
-| `--memory-ranking` | `vector` or `fused`, the ranking of `capture` and `recall` | `vector` |
+| `--memory-ranking` | `vector` or `fused`, the ranking of `capture`, and of `recall` without `--recall-ranking` | `vector` |
+| `--recall-ranking` | `vector` or `fused`, the ranking of `recall` if it differs from `--memory-ranking`. With `--reuse`, a run recalls from the kept memory of a `--memory-ranking` run with another ranking. | the `--memory-ranking` |
 | `--recall-k` | the `k` of the memory recall | `30` |
 | `--recall-budget` | the `budget` of the memory recall, in characters | `8000` |
-| `--extractor` | a variant name in `VARIANTS` of `bench/prompts.py` | `baseline` |
+| `--extractor` | a variant name in `VARIANTS` of `bench/prompts.py`: `baseline`, `detail`, or `detail-cues`, which is `detail` with one to three cues for each fact | `baseline` |
 | `--no-fact-sources` | write the facts with no session source, as phases 5 and 6 did | off |
 | `--session-episodes` | capture each session as an episode, with the facts citing it | off |
 | `--chunk-chars` | store each session as chunks of at most this many characters, with each fact citing every chunk; needs `--session-episodes` | off |
@@ -210,7 +211,8 @@ with `--keep` or `--reuse`, the command reads the memory of each question from i
 database and does not run `capture`. If the result file names a background database, from
 a run with `--background`, the command captures again in that database, recalls over the
 background as the run did, and records whether the block has the SHA-256 of the run.
-`--ids` limits the diagnosis to the listed misses. A result file written before the runner
+If the run recalled with `--recall-ranking`, the command decides the cause of each miss
+with that ranking. `--ids` limits the diagnosis to the listed misses. A result file written before the runner
 recorded `chunk_chars` needs `--chunk-chars`. GPT-4o marks the stored facts that hold each answer.
 Each miss then gets one of four causes:
 

bench/diagnose.py

@@ -142,6 +142,7 @@ def memory_config(result: dict[str, Any], chunk_chars: int | None = None) -> Con
         raise ValueError("the run stored chunks of a size that the file does not record")
     return Config(
         memory_ranking=found["ranking"],
+        recall_ranking=found.get("recall_ranking"),
         k=found["k"],
         budget=found["budget"],
         extractor=found["extractor"],
@@ -269,7 +270,7 @@ def diagnosed(
                 try:
                     for session in question.sessions:
                         capture_session(memory, namespace, session, replay, config, [])
-                    records, block = memory.recall(
+                    records, block = Memory(stores[config.recall_with]).recall(
                         question.question,
                         [namespace, BACKGROUND],
                         as_of=question.as_of,
@@ -282,7 +283,7 @@ def diagnosed(
                     sha = hashlib.sha256(block.encode()).hexdigest()
                     extra = {"block_reproduced": sha == saved[question.id]}
                     extra |= episode_counts(
-                        stores[config.memory_ranking], question, namespace, records, config.k
+                        stores[config.recall_with], question, namespace, records, config.k
                     )
                     diagnosed(question, stores, namespace, records, extra)
                 finally:
@@ -306,6 +307,7 @@ def diagnosed(
             "source": str(args.result.resolve().relative_to(Path.cwd())),
             "memory": {
                 "ranking": config.memory_ranking,
+                "recall_ranking": config.recall_with,
                 "k": config.k,
                 "budget": config.budget,
                 "extractor": config.extractor,
@@ -349,7 +351,7 @@ def diagnose(
     """
     sessions = {fact_source(session.id): session.id for session in question.sessions}
     sessions |= {episode_key(session.id): session.id for session in question.sessions}
-    store = stores[config.memory_ranking]
+    store = stores[config.recall_with]
     facts = store.list(namespace, Filter(kind=RECALLED_KINDS), limit=10_000)
     marked = _mark(marker, question.question, question.answer, facts)
     holding = {facts[n - 1].id for n in marked}
@@ -363,7 +365,7 @@ def diagnose(
         block = block_ids(top, question.as_of, config.k, config.budget)
     else:
         block = {record.id for record in recalled if record.kind != "episode"}
-    cause = verdict(holding, ranks[config.memory_ranking], block, config.k)
+    cause = verdict(holding, ranks[config.recall_with], block, config.k)
     present = None
     if cause == RECALLED:
         listed = [fact.content for fact in facts if fact.id in holding & block]

bench/e2e.py

@@ -91,7 +91,8 @@ class Config:
     """The configuration of the retrieval arm and the memory arm.
 
     `retrieval_ranking` ranks the sessions of the retrieval arm. The memory arm captures
-    and recalls with `memory_ranking`, recalls `k` records into a block of at most
+    with `memory_ranking`, and recalls with `recall_ranking`, or with `memory_ranking` if
+    it is None. It recalls `k` records into a block of at most
     `budget` characters, and extracts with the variant `extractor` of `prompts.VARIANTS`.
     With `fact_sources`, each fact carries the source of its session from `fact_source`.
     With `session_episodes`, `capture` stores each session as an episode under the key
@@ -109,9 +110,10 @@ class Config:
     session_episodes: bool = False
     episodes: str = "none"
     chunk_chars: int | None = None
+    recall_ranking: str | None = None
 
     def __post_init__(self) -> None:
-        for ranking in (self.retrieval_ranking, self.memory_ranking):
+        for ranking in (self.retrieval_ranking, self.memory_ranking, self.recall_with):
             if ranking not in RANKINGS:
                 raise ValueError(f"ranking must be one of {', '.join(RANKINGS)}, got {ranking!r}")
         if self.extractor not in prompts.VARIANTS:
@@ -125,6 +127,11 @@ def __post_init__(self) -> None:
         if self.chunk_chars is not None and (self.chunk_chars < 1 or not self.session_episodes):
             raise ValueError("chunk_chars must be at least 1, and needs session_episodes")
 
+    @property
+    def recall_with(self) -> str:
+        """Return the ranking of the memory recall."""
+        return self.memory_ranking if self.recall_ranking is None else self.recall_ranking
+
     @property
     def variant(self) -> prompts.Variant:
         """Return the extractor variant."""
@@ -243,16 +250,17 @@ def extractor(
 ) -> Callable[[str, datetime | None, list[Record]], list[Fact]]:
     """Return the extractor callable of `capture`: one call to the variant's model per session.
 
-    With `source`, each fact carries it in `sources`. A reply that reached the token limit
-    gives no facts, and `source` is appended to `at_limit`.
+    With `source`, each fact carries it in `sources`. A variant with cues gives each fact
+    the cues of its reply. A reply that reached the token limit gives no facts, and
+    `source` is appended to `at_limit`.
     """
     sources = (source,) if source else ()
 
     def extract(text: str, observed_at: datetime | None, current: list[Record]) -> list[Fact]:
         date = observed_at.strftime("%Y/%m/%d (%a) %H:%M") if observed_at else "unknown"
         known = [f"- {record.content}" for record in current]
         prompt = prompts.extract(text, date, known, variant.prompt)
-        reply = models(variant.model, prompt, EXTRACTOR_MAX_TOKENS, prompts.EXTRACT_SCHEMA)
+        reply = models(variant.model, prompt, EXTRACTOR_MAX_TOKENS, variant.schema)
         if reply.stopped:
             if at_limit is not None:
                 at_limit.append(source)
@@ -263,6 +271,7 @@ def extract(text: str, observed_at: datetime | None, current: list[Record]) -> l
                 kind="fact",
                 subject=item["subject"].strip(),
                 content=item["content"].strip(),
+                cues=tuple(cue.strip() for cue in item.get("cues", ()) if cue.strip()),
                 sources=sources,
             )
             for item in facts
@@ -348,7 +357,7 @@ def run_question(
                 written += len(
                     capture_session(memory, namespace, session, models, config, at_limit)
                 )
-        records, block = memory.recall(
+        records, block = Memory(stores[config.recall_with]).recall(
             question.question,
             [namespace] if background is None else [namespace, background],
             as_of=question.as_of,
@@ -896,6 +905,12 @@ def main(argv: Sequence[str] | None = None) -> Path | None:
     )
     parser.add_argument("--retrieval-ranking", choices=RANKINGS, default="vector")
     parser.add_argument("--memory-ranking", choices=RANKINGS, default="vector")
+    parser.add_argument(
+        "--recall-ranking",
+        choices=RANKINGS,
+        default=None,
+        help="the ranking of the memory recall, if not that of --memory-ranking",
+    )
     parser.add_argument("--recall-k", type=int, default=RECALL_K, help="k of the memory recall")
     parser.add_argument(
         "--recall-budget", type=int, default=RECALL_BUDGET, help="characters of the memory block"
@@ -978,11 +993,14 @@ def main(argv: Sequence[str] | None = None) -> Path | None:
             session_episodes=args.session_episodes,
             episodes=args.episodes,
             chunk_chars=args.chunk_chars,
+            recall_ranking=args.recall_ranking,
         )
     except ValueError as error:
         parser.error(str(error))
     if args.smoke:
-        config = replace(config, retrieval_ranking="vector", memory_ranking="vector")
+        config = replace(
+            config, retrieval_ranking="vector", memory_ranking="vector", recall_ranking=None
+        )
 
     all_questions = {question.id: question for question in load(DATA_DIR / DATA_FILE)}
     if args.smoke:
@@ -1156,6 +1174,7 @@ def main(argv: Sequence[str] | None = None) -> Path | None:
             "retrieval": {"ranking": config.retrieval_ranking, "sessions": RETRIEVED},
             "memory": {
                 "ranking": config.memory_ranking,
+                "recall_ranking": config.recall_with,
                 "k": config.k,
                 "budget": config.budget,
                 "extractor": config.extractor,
@@ -1184,7 +1203,7 @@ def main(argv: Sequence[str] | None = None) -> Path | None:
                         if config.chunk_chars is not None
                         else "extractor, then recall"
                     )
-                    + f" with the {config.memory_ranking} ranking, k "
+                    + f" with the {config.recall_with} ranking, k "
                     f"{config.k}, budget {config.budget} characters, episodes "
                     f"{config.episodes}, as_of the end of the question day"
                 ),

bench/prompts.py

@@ -120,6 +120,12 @@
 
 Session:
 {session}"""
+EXTRACT_DETAIL_CUES = EXTRACT_DETAIL.replace(
+    "Return an empty list if",
+    'Give `cues` as one to three short questions that the fact answers, in the words that \
+the user can use later to ask for it, such as "What car does the user drive?". Return an \
+empty list if',
+)
 EXTRACT_SCHEMA = {
     "type": "object",
     "properties": {
@@ -136,24 +142,55 @@
     "required": ["facts"],
     "additionalProperties": False,
 }
+EXTRACT_CUES_SCHEMA = {
+    "type": "object",
+    "properties": {
+        "facts": {
+            "type": "array",
+            "items": {
+                "type": "object",
+                "properties": {
+                    "subject": {"type": "string"},
+                    "content": {"type": "string"},
+                    "cues": {"type": "array", "items": {"type": "string"}},
+                },
+                "required": ["subject", "content", "cues"],
+                "additionalProperties": False,
+            },
+        }
+    },
+    "required": ["facts"],
+    "additionalProperties": False,
+}
 
 
 @dataclass(frozen=True)
 class Variant:
-    """An extractor variant: the prompt template and the model that reads it."""
+    """An extractor variant: the prompt template, the model that reads it, and its cues.
+
+    With `cues`, the reply schema is `EXTRACT_CUES_SCHEMA`, and each fact carries cues.
+    """
 
     prompt: str
     model: str
+    cues: bool = False
+
+    @property
+    def schema(self) -> dict[str, object]:
+        """Return the JSON schema of the reply."""
+        return EXTRACT_CUES_SCHEMA if self.cues else EXTRACT_SCHEMA
 
 
 VARIANTS = {
     "baseline": Variant(EXTRACT, EXTRACTOR),
     "detail": Variant(EXTRACT_DETAIL, EXTRACTOR),
+    "detail-cues": Variant(EXTRACT_DETAIL_CUES, EXTRACTOR, cues=True),
 }
 """The extractor variants by name. `baseline` is the extractor of phase 5.
 
 `detail` asks for each fact that the user states in passing, each item of a list that the
-assistant gives, and absolute dates.
+assistant gives, and absolute dates. `detail-cues` is `detail`, and also asks for one to
+three questions that each fact answers, as its cues.
 """
 
 

bench/tests/test_e2e.py

@@ -62,11 +62,14 @@ def __call__(
         self.prompts.append((model, prompt))
         if model == EXTRACTOR:
             session = prompt.split("Session:\n", 1)[1]
-            facts = [
+            facts: list[dict[str, Any]] = [
                 {"subject": "user", "content": line.removeprefix("user: ")}
                 for line in session.split("\n\n")
                 if line.startswith("user: ")
             ]
+            if schema is prompts.EXTRACT_CUES_SCHEMA:
+                for fact in facts:
+                    fact["cues"] = [f"What did the user say about {fact['content'][:20]}?", " "]
             return Reply(json.dumps({"facts": facts}), 10, 5)
         if prompt.startswith("I will give you a question") or "unanswerable" in prompt[:40]:
             return Reply("yes", 10, 1)
@@ -181,6 +184,35 @@ def test_config_refuses_unknown_values() -> None:
         Config(extractor="none")
     with pytest.raises(ValueError, match="at least 1"):
         Config(k=0)
+    with pytest.raises(ValueError, match="ranking"):
+        Config(recall_ranking="bm25")
+
+
+def test_recall_ranking_defaults_to_the_memory_ranking() -> None:
+    assert Config(memory_ranking="fused").recall_with == "fused"
+    assert Config(memory_ranking="fused", recall_ranking="vector").recall_with == "vector"
+
+
+def test_run_question_recalls_with_the_recall_ranking(questions: dict[str, Question]) -> None:
+    store = MemoryStore(embedder=fake_embedder, clock=FakeClock())
+    searched: list[str] = []
+
+    class Spy:
+        def __init__(self, name: str) -> None:
+            self.name = name
+
+        def __getattr__(self, attribute: str) -> Any:
+            if attribute == "search":
+                searched.append(self.name)
+            return getattr(store, attribute)
+
+    stores: Any = {"fused": Spy("fused"), "vector": Spy("vector")}
+    config = Config(memory_ranking="vector", recall_ranking="fused")
+    run_question(
+        questions["q_one"], FakeModels(), stores, lambda _: None, "r1", ("memory",), config
+    )
+    assert searched[-1] == "fused"
+    assert set(searched[:-1]) == {"vector"}
 
 
 def test_run_question_follows_the_config(
@@ -308,6 +340,21 @@ def empty(namespace: Sequence[str]) -> None:
     assert "answer_" not in block
 
 
+def test_the_cue_variant_puts_cues_on_each_fact(questions: dict[str, Question]) -> None:
+    session = questions["q_one"].sessions[0]
+    text = prompts.session_text(session)
+    variant = prompts.VARIANTS["detail-cues"]
+    assert variant.schema is prompts.EXTRACT_CUES_SCHEMA
+    assert prompts.VARIANTS["detail"].schema is prompts.EXTRACT_SCHEMA
+    assert "Give `cues` as one to three short questions" in variant.prompt
+    before, after = prompts.EXTRACT_DETAIL.split("Return an empty list if")
+    assert variant.prompt.startswith(before) and variant.prompt.endswith(after)
+    cued = extractor(FakeModels(), variant)(text, session.date, [])
+    plain = extractor(FakeModels(), prompts.VARIANTS["detail"])(text, session.date, [])
+    assert cued and all(len(fact.cues) == 1 and fact.cues[0].startswith("What") for fact in cued)
+    assert all(fact.cues == () for fact in plain)
+
+
 def test_a_reply_at_the_limit_gives_no_facts(questions: dict[str, Question]) -> None:
     session = questions["q_one"].sessions[0]