Phase 10: the vector ranking by default, and episode chunks

CHANGELOG.md

@@ -66,3 +66,13 @@ each one, and its failure message points to this file.
 - The block of `recall` and `profile` shows a source that is the id of a record in the
   namespace of its line as `<key>@<revision>`. Before, it showed the full id. The stored
   sources do not change.
+- `HotdataStore` ranks by content vector alone by default, as `MemoryStore` does. Before,
+  the default was `ranking="fused"`, which stays as an option. On LongMemEval-S, the BM25
+  arm of the fused ranking lowered the retrieval score.
+- `Memory.capture` takes `chunk_chars`, with `episode`. It cuts the text at blank lines
+  into chunks and stores chunk n as an episode under `<episode>-<n>`. The extractor still
+  reads the whole text, and each fact names every chunk in `sources`.
+- The block of `recall` and `profile` shows a run of sources in the namespace of its line,
+  with keys `<prefix>-<n>`, n in four digits counting up by 1, and one revision, as
+  `<prefix>-<first>…<last>@<revision>`. The
+  stored sources do not change.

CONTRIBUTING.md

@@ -53,7 +53,7 @@ model. The tests pass a fixed clock and a fake embedder to the store, from
   `tests/conftest.py`.
 - `tests/test_contract.py` runs `Memory` over every driver through the same `store`
   fixture. It pins the rendered blocks of `recall` and `profile` as exact strings.
-  `tests/test_chunks.py` tests the chunking of `load` without a driver.
+  `tests/test_chunks.py` tests the chunking of `load` and `capture` without a driver.
 - The in-memory driver is the reference for the Hotdata driver. `tests/test_oracle.py`
   builds the same records in both drivers and compares the results of `search`. The
   ranking by content vector alone must return the same records in the same order, with

README.md

@@ -22,7 +22,7 @@ The library has two layers:
   text is optional, and it takes an extractor callable that the caller supplies. `load`
   cuts a Markdown document into chunks, stores each chunk as an episode, and extracts
   facts that name their episode as a source. `capture` with an `episode` key stores its
-  text as an episode in the same way. `profile` returns facts, profiles, and procedures,
+  text as an episode in the same way, or as chunks with `chunk_chars`. `profile` returns facts, profiles, and procedures,
   never episodes. `recall` returns them too, and adds episodes after the facts only when
   the caller asks for them.
 

bench/README.md

@@ -37,6 +37,7 @@ The options are these:
 | `--location` | `local` or `cloud` | `local` |
 | `--url` | the address of the local engine | `http://localhost:3000` |
 | `--limit` | the first N scored questions | every scored question |
+| `--ablate` | also rank with the fused query less one arm at a time: `bm25-content`, `content-cue`, and `bm25`; needs `--driver hotdata` | off |
 
 With `--location cloud`, the runner reads the connection from `HOTDATA_API_KEY`,
 `HOTDATA_WORKSPACE`, and `HOTDATA_API_URL`, so set them in `.env` first. It ignores `--url`.
@@ -81,13 +82,14 @@ The options are these:
 |---|---|---|
 | `--sample` | `a` or `b` | `a` |
 | `--arms` | a comma-separated list of `full`, `retrieval`, and `memory` | every arm |
-| `--retrieval-ranking` | `fused` or `vector`, the ranking of the retrieval arm | `fused` |
-| `--memory-ranking` | `fused` or `vector`, the ranking of `capture` and `recall` | `fused` |
+| `--retrieval-ranking` | `vector` or `fused`, the ranking of the retrieval arm | `vector` |
+| `--memory-ranking` | `vector` or `fused`, the ranking of `capture` and `recall` | `vector` |
 | `--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` |
 | `--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 |
 | `--episodes` | `none`, `cited`, or `search`, the `episodes` of the memory recall; a value other than `none` needs `--session-episodes` | `none` |
 | `--replay` | call no model, and stop at the first reply that the cache does not hold | off |
 | `--workers` | the number of processes that run the questions | `1` |
@@ -100,9 +102,9 @@ The options are these:
 
 Each worker gets the questions in turn. Without `--keep` or `--reuse`, it runs them over a
 new local database of its own, and empties the namespaces of each question after it. Each
-table then holds one question at a time, as in a run with one worker. BM25 ranks
-over the whole table before the scope filter, so the rows of a second question in a table
-can change the ranking. An extractor variant names a prompt and a model. Each fact carries
+table then holds one question at a time, as in a run with one worker. With
+`--memory-ranking fused`, BM25 ranks over the whole table before the scope filter, so the
+rows of a second question in a table can change the ranking. An extractor variant names a prompt and a model. Each fact carries
 the source `session:<hash>`, where the hash is the first 12 hex digits of the SHA-256 of
 the session id. The block shows the sources to the reader, and the id of an evidence
 session holds `answer`, so the source does not show the id. If an extractor reply reaches
@@ -112,7 +114,9 @@ so a replay skips the same session. The result file counts these sessions in
 
 With `--session-episodes`, `capture` stores each session as an episode under the key
 `session-<hash>`, with the same hash. Each fact then cites that episode, and carries no
-`session:<hash>` source. The block shows the source as `session-<hash>@1`. The source
+`session:<hash>` source. The block shows the source as `session-<hash>@1`. With `--chunk-chars` as well, `capture`
+stores the session as chunks under `session-<hash>-<n>`, each fact cites every chunk, and
+the block shows the sources as one range, `session-<hash>-0001…<last>@1`. The source
 holds no run id, so a second run with another `--episodes` or `--recall-budget` replays
 `capture` from the cache. The result file lists the episode keys of each block in
 `episodes_in_block`, and the SHA-256 of each block in `block_sha256`. The result file is
@@ -132,14 +136,15 @@ uv run --group bench --env-file .env python -m bench.e2e --run --arms memory --e
 ```
 
 With `--keep`, each question runs over a new local database of its own, and the database
-stays after the run. One database holds one question, because BM25 ranks the whole table
-before the namespace filter, and the memory of a second question changes the block. The
+stays after the run. One database holds one question. With `--memory-ranking fused`, BM25
+ranks the whole table before the namespace filter, so the memory of a second question
+would change the block. The
 name of the database is `hotmemory-kept-<sample>-<extractor><digest>-<cites>-<ranking>-<commit>-<question id>`:
 
 | Part | What it holds |
 |---|---|
 | `<extractor><digest>` | The variant name, and 6 hex digits of the SHA-256 of its model and prompt. |
-| `<cites>` | `episodes` with `--session-episodes`, `sources` with fact sources, or `plain`. |
+| `<cites>` | `episodes` with `--session-episodes`, `chunks<size>` with `--chunk-chars` as well, `sources` with fact sources, or `plain`. |
 | `<ranking>` | The `--memory-ranking`. |
 | `<commit>` | The last commit that changed `src/hotmemory`. |
 
@@ -161,10 +166,11 @@ uv run --group bench --env-file .env python -m bench.e2e --run --smoke --arms me
 It runs the 18 questions of `bench/sample-smoke.json`, 3 of each type drawn from sample A
 with the seed `20261016`. Each worker runs its questions over a `MemoryStore` of its own,
 in its own process, and empties the memory of each question after it. It decides no gain: 18
-questions can be off by about 23 points, and `MemoryStore` ranks by vector only. Because
-the ranking differs, `capture` shows the extractor another list of known facts than a run
-on the local engine. So the first smoke run of a configuration pays for its extraction,
-and a later one replays it from the cache. Pass `--workers 4` to make the first run
+questions can be off by about 23 points. `MemoryStore` ranks by vector only, as
+`HotdataStore` does by default. With `--memory-ranking fused`, the ranking differs, so
+`capture` shows the extractor another list of known facts than a run on the local engine.
+The first smoke run of a configuration pays for the extraction that the cache does not
+hold, and a later one replays it from the cache. Pass `--workers 4` to make the first run
 faster. `--smoke` refuses `--keep`, `--reuse`, and `--sample`.
 
 Each run writes the file `bench/results/<name>.rows.jsonl` as it goes, where `<name>` is

bench/ablation.py

@@ -0,0 +1,114 @@
+"""The fused ranking of `HotdataStore` with a chosen set of its arms, for the ablation.
+
+Each ablated ranking builds the query of `hotmemory.hotdata._fused_query` with only the
+named arms: `bm25`, `content`, and `cue`. With every arm, the query is the same as that of
+`ranking="fused"`.
+"""
+
+from __future__ import annotations
+
+import builtins
+from collections.abc import Sequence
+
+from hotmemory import Filter, Hit
+from hotmemory.hotdata import (
+    CUE_TABLE,
+    READ_COLUMNS,
+    RRF_CONSTANT,
+    HotdataStore,
+    _bm25_text,
+    _depth,
+    _record,
+    _ref,
+    _tables,
+    _text,
+    _union,
+    _vector,
+)
+
+ARMS = ("bm25", "content", "cue")
+ABLATIONS = {
+    "bm25-content": ("bm25", "content"),
+    "content-cue": ("content", "cue"),
+    "bm25": ("bm25",),
+}
+
+
+class ArmStore(HotdataStore):
+    """A `HotdataStore` whose search fuses only the arms in `arms`."""
+
+    def __init__(self, *args: object, arms: Sequence[str], **kwargs: object) -> None:
+        if not arms or set(arms) - set(ARMS):
+            raise ValueError(f"arms must be a non-empty subset of {ARMS}, got {arms!r}")
+        super().__init__(*args, ranking="fused", **kwargs)  # type: ignore[arg-type]
+        self._arms = tuple(arms)
+
+    def search(
+        self,
+        query: str | None,
+        prefixes: Sequence[Sequence[str]],
+        filter: Filter | None = None,
+        k: int = 10,
+    ) -> builtins.list[Hit]:
+        if query is None or k == 0:
+            return super().search(query, prefixes, filter, k)
+        where = self._listed(prefixes, filter)
+        tables = _tables(filter)
+        vector = _vector(self._embed([query])[0])
+        sql = arm_query(tables, where, vector, _bm25_text(query), k, self._arms)
+        return [Hit(_record(row), float(row["distance"])) for row in self._sql(sql).to_records()]
+
+
+def arm_query(
+    tables: Sequence[str], where: str, vector: str, text: str, k: int, arms: Sequence[str]
+) -> str:
+    """Return the fused query of `hotmemory.hotdata._fused_query` with only `arms`."""
+    depth = _depth(k)
+    rank = "ROW_NUMBER() OVER (ORDER BY {order}) AS rank"
+    parts = []
+    rankings = []
+    if "bm25" in arms:
+        for table in tables:
+            name = f"text_{table}"
+            parts.append(
+                f"{name} AS (SELECT id, {rank.format(order='score DESC, id ASC')} FROM "
+                f"bm25_search('{_ref(table)}', 'content', {_text(text)}, {depth}) WHERE {where})"
+            )
+            rankings.append(name)
+    if "content" in arms:
+        nearest = " UNION ALL ".join(
+            f"SELECT scanned.id, cosine_distance(scanned.content_embedding, {vector}) AS distance "
+            f"FROM {_ref(table)} scanned WHERE {where}"
+            for table in tables
+        )
+        parts.append(
+            f"content_rank AS (SELECT candidate.id, "
+            f"{rank.format(order='candidate.distance ASC, candidate.id ASC')} FROM ("
+            f"SELECT * FROM ({nearest}) ORDER BY distance ASC, id ASC LIMIT {depth}) candidate)"
+        )
+        rankings.append("content_rank")
+    if "cue" in arms:
+        parts.append(
+            f"cue_rank AS (SELECT candidate.id, "
+            f"{rank.format(order='candidate.distance ASC, candidate.id ASC')} FROM ("
+            f"SELECT scanned.id, cosine_distance(scanned.cues_embedding, {vector}) AS distance "
+            f"FROM {_ref(CUE_TABLE)} scanned WHERE scanned.id IN ({_union(tables, 'id', where)}) "
+            f"ORDER BY distance ASC, id ASC LIMIT {depth}) candidate)"
+        )
+        rankings.append("cue_rank")
+    ranked = " UNION ALL ".join(
+        f"SELECT id, rank FROM {name} WHERE rank <= {depth}" for name in rankings
+    )
+    parts.append(
+        f"fused AS (SELECT id, sum(1.0 / ({RRF_CONSTANT} + rank)) AS score "
+        f"FROM ({ranked}) GROUP BY id)"
+    )
+    columns = ", ".join(READ_COLUMNS)
+    rows = _union(tables, f"{columns}, content_embedding", where)
+    return (
+        f"WITH {', '.join(parts)} "
+        f"SELECT {', '.join(f'r.{column}' for column in READ_COLUMNS)}, "
+        f"cosine_distance(r.content_embedding, {vector}) AS distance FROM fused "
+        f"JOIN ({rows}) r ON r.id = fused.id "
+        f"ORDER BY fused.score DESC, distance ASC, r.created_at DESC, r.id ASC LIMIT {k}"
+    )

bench/diagnose.py

@@ -38,6 +38,7 @@
     existing_databases,
     fact_source,
     local_store,
+    session_key,
 )
 from bench.embed import CACHE_DIR
 from bench.models import JUDGE, CachedModels, Offline, Remote
@@ -132,12 +133,17 @@ def memory_config(result: dict[str, Any]) -> Config:
         fact_sources=found.get("fact_sources", False),
         session_episodes=found.get("session_episodes", False),
         episodes=found.get("episodes", "none"),
+        chunk_chars=found.get("chunk_chars"),
     )
 
 
 def evidence_in_block(question: Question, shown: Sequence[str]) -> bool:
-    """Return True if an episode key in `shown` is that of an evidence session of `question`."""
-    return any(episode_key(session_id) in shown for session_id in question.answer_session_ids)
+    """Return True if an episode key in `shown` is that of an evidence session of `question`.
+
+    A chunk key counts as the key of its session.
+    """
+    keys = {session_key(key) for key in shown}
+    return any(episode_key(session_id) in keys for session_id in question.answer_session_ids)
 
 
 def table(rows: Sequence[dict[str, Any]]) -> str:
@@ -305,10 +311,10 @@ def session_of_fact(fact: Record, sessions: dict[str, str]) -> str | None:
     """Return the session id of `fact` from its sources, or None if no source names one.
 
     `sessions` maps the results of `fact_source` and of `episode_key` to the session id. A
-    source that is an episode id matches by its key.
+    source that is an episode id matches by its key, and a chunk by the key of its session.
     """
     for source in fact.sources:
-        key = source.rpartition("/")[2].rpartition("@")[0]
+        key = session_key(source.rpartition("/")[2].rpartition("@")[0])
         for name in (source, key):
             if name in sessions:
                 return sessions[name]