diff --git a/README.md b/README.md index 99d7712..c6fa595 100644 --- a/README.md +++ b/README.md @@ -138,18 +138,19 @@ as an episode in `HotdataStore` and searched with the question. The 419 scored q had these results with the vector ranking. For 98.3% of them, an evidence session was in the top 5. For 99.3% of them, every evidence session was in the top 10. -The end-to-end score ran on a sample of 54 questions. The reader and the judge were -`gpt-4o-2024-08-06`. - -| What the reader got | Correct | -|---|---| -| Every session of the question, about 114,000 tokens | 39 of 54 (72.2%) | -| The 10 sessions that the search ranked first | 45 of 54 (83.3%) | -| The block of `recall`, after `capture` with a `gpt-4o-mini` extractor | 19 of 54 (35.2%) | - -The memory contract, with this extractor, scores below the full history. In 19 of its 35 -misses, the extractor did not store the answer. With 54 questions, each score can be about -10 points off. In the cloud, a search took 361 ms at p50 with the vector ranking, and +The end-to-end score ran on two samples of 54 questions, sample A on 2026-10-10 and +sample B on 2026-10-14. The reader and the judge were `gpt-4o-2024-08-06`. + +| What the reader got | Sample A | Sample B | +|---|---|---| +| Every session of the question, about 114,000 tokens | 39 of 54 (72.2%) | 43 of 54 (79.6%) | +| The 10 sessions that the search ranked first | 45 of 54 (83.3%) | 48 of 54 (88.9%) | +| The block of `recall`, after `capture` with a `gpt-4o-mini` extractor | 19 of 54 (35.2%) | 27 of 54 (50.0%) | + +The memory contract, with this extractor, scores below the full history. The extractor +did not store the answer in 19 of the 35 misses on sample A, and in 16 of the 27 misses on +sample B. With 54 questions, each score can be about 10 points off, and the memory arm +differs by 14.8 points between the two samples. In the cloud, a search took 361 ms at p50 with the vector ranking, and 998 ms with the fused ranking. ## Documents diff --git a/bench/README.md b/bench/README.md index 5af7d17..84751ce 100644 --- a/bench/README.md +++ b/bench/README.md @@ -48,27 +48,54 @@ sessions of each question. Git ignores `bench/results/`, and ## The end-to-end score -This command runs three arms on the sample of 54 questions in `bench/sample.json`, 9 of -each question type: +This command runs three arms on a sample of 54 questions, 9 of each question type: ```sh uv run --group bench --env-file .env python -m bench.e2e ``` -Without `--run`, it prints the cost estimate of each arm and calls no model. With `--run`, -it stops if the estimate is over 50 dollars, and else it calls the models. Set -`OPENAI_MODEL_KEY` in `.env` first. +Without `--run` or `--replay`, it prints the cost estimate of each arm and calls no model. +With `--run`, it stops if the estimate is over 50 dollars, and else it calls the models. +Set `OPENAI_MODEL_KEY` in `.env` first. With `--replay`, it takes every reply from the +cache, and it stops at the first reply that the cache does not hold. The estimate does not +subtract the cached replies, so the stop at 50 dollars applies to a replay too. | Arm | What the reader gets | |---|---| | `full` | Every session of the question. | -| `retrieval` | The 10 sessions that the fused search ranks first. | +| `retrieval` | The 10 sessions that the search ranks first. | | `memory` | The block of `recall`, after `capture` ran on each session in date order. | The reader and the judge are `gpt-4o-2024-08-06` at temperature 0, with the prompts of -the reference code. The extractor of the memory arm is `gpt-4o-mini-2024-07-18`. Each -reply is kept in `bench/cache/`, so a second run with the same prompts calls no model. -`--arms` runs some of the arms, and `--limit` runs the first N questions of the sample. +the reference code. The baseline extractor of the memory arm is `gpt-4o-mini-2024-07-18`. +Each reply is kept in `bench/cache/`, so a second run with the same prompts calls no model. + +There are two samples. Sample A is `bench/sample.json`, drawn with the seed `20261010`. +Each phase iterates on sample A. Sample B is `bench/sample-b.json`, drawn with the seed +`20261013` from the questions that are not in sample A. A phase runs sample B once, at its +end. + +The options are these: + +| Option | Values | Default | +|---|---|---| +| `--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` | +| `--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` | +| `--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` | +| `--tag` | a suffix for the name of the result file | none | +| `--limit` | the first N questions of the sample | every question | + +Each worker gets the questions in turn, and runs them over a new local database of its +own. 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. The result file is +`bench/results/e2e-sample-[-].json`, and it records each option. ## The checks of the results @@ -79,11 +106,26 @@ check does: uv run --group bench python -m bench.compare bench/results/.json bench/results/.json ``` -This command finds the cause of each miss of the memory arm in `e2e-sample.json`. It runs -`capture` again from the cached replies, so the extractor is not called again. GPT-4o marks -the stored facts that hold each answer. If the replay of `capture` spends more than 1 -dollar, the cache missed, and the command stops. `--result` names another result file: +This command finds the cause of each miss of the memory arm in a result file. The default +file is `e2e-sample-a.json`, and `--result` names another one: ```sh -uv run --group bench --env-file .env python -m bench.diagnose +uv run --group bench --env-file .env python -m bench.diagnose --result bench/results/e2e-sample-b.json ``` + +The command runs `capture` again with the ranking, the `k`, the `budget`, and the extractor +variant that the result file records. The replay calls no model, and it stops at the first +reply that the cache does not hold. GPT-4o marks the stored facts that hold each answer. +Each miss then gets one of four causes: + +| Cause | Meaning | +|---|---| +| not extracted | No stored fact holds the answer. | +| extracted, ranked below k | A fact holds the answer, but the search ranks it below `k`. | +| in the top k, cut by as_of or the budget | A fact is in the top `k`, but `recall` leaves it out. | +| recalled, reader missed | A fact that holds the answer is in the block. | + +For the last cause, GPT-4o tells whether the marked facts in the block hold every value +that the answer needs ("all present") or only some of them ("some present"). The command +prints the groups for each question type, and writes them next to the result file, as +`-diagnosis.json`. diff --git a/bench/diagnose.py b/bench/diagnose.py index 528b51a..b4bef3e 100644 --- a/bench/diagnose.py +++ b/bench/diagnose.py @@ -1,32 +1,41 @@ """The diagnosis of the memory arm: was the answer extracted, and where did recall rank it? For each question that the memory arm missed in a result file, this runs `capture` again -with the cached extractor replies, and keeps every fact. A model marks the facts that hold -the answer, and the search ranks every fact for the question. - - uv run --group bench --env-file .env python -m bench.diagnose +with the cached extractor replies, and keeps every fact. The replay calls no model, and it +stops at the first reply that the cache does not hold. It uses the ranking, `k`, `budget`, +and extractor variant that the result file records. A model marks the facts that hold the +answer, and the search ranks every fact for the question. If a marked fact is in the block, +a second model call tells whether the block holds every needed fact ("all present") or only +some of them ("some present"). + + uv run --group bench --env-file .env python -m bench.diagnose --result """ from __future__ import annotations import argparse import json +import uuid from collections import Counter from collections.abc import Sequence -from datetime import datetime +from datetime import UTC, datetime from pathlib import Path from typing import Any from bench import prompts -from bench.data import DATA_DIR, DATA_FILE, load -from bench.e2e import ACTOR, RECALL_BUDGET, RECALL_K, extractor, local_store +from bench.data import DATA_DIR, DATA_FILE, QUESTION_TYPES, load +from bench.e2e import ACTOR, BASELINE, RECALL_BUDGET, RECALL_K, Config, extractor, local_store from bench.embed import CACHE_DIR -from bench.models import JUDGE, CachedModels, Remote -from bench.retrieval import RESULTS_DIR, _commit -from hotmemory import Filter, Memory, Record +from bench.models import JUDGE, CachedModels, Offline, Remote +from bench.retrieval import RANKINGS, RESULTS_DIR, _commit +from hotmemory import Filter, Memory, Record, Store from hotmemory.contract import RECALLED_KINDS, fit_lines, is_valid_at, render_line -SPEND_LIMIT = 1.0 +NOT_EXTRACTED = "not extracted" +BELOW_K = "extracted, ranked below k" +CUT = "in the top k, cut by as_of or the budget" +RECALLED = "recalled, reader missed" +CAUSES = (NOT_EXTRACTED, BELOW_K, CUT, RECALLED) MARK = """A question about a user, its correct answer, and a numbered list of facts that a \ memory holds about the user follow. Return the numbers of the facts that state the correct \ answer itself, or a value that the answer is computed from. A fact on the same topic that \ @@ -45,6 +54,23 @@ "required": ["facts"], "additionalProperties": False, } +COMPLETE = """A question about a user, its correct answer, and a list of facts follow. Tell \ +whether the facts together state every value that the correct answer needs, so that a \ +reader can give the correct answer from these facts alone. Return true if they do, and \ +false if a needed value is missing. + +Question: {question} + +Correct answer: {answer} + +Facts: +{facts}""" +COMPLETE_SCHEMA = { + "type": "object", + "properties": {"complete": {"type": "boolean"}}, + "required": ["complete"], + "additionalProperties": False, +} def block_ids( @@ -60,37 +86,75 @@ def block_ids( return {record.id for record in valid[: fit_lines(lines, budget)]} -def verdict(holding: set[str], best_rank: int | None, block: set[str]) -> str: +def verdict(holding: set[str], best_rank: int | None, block: set[str], k: int = RECALL_K) -> str: """Return the cause of a miss from the facts that hold the answer and the block. `best_rank` is the best rank of those facts in the full ranking, before any filter. """ if not holding: - return "not extracted" + return NOT_EXTRACTED if holding & block: - return "recalled, reader missed" - if best_rank is not None and best_rank <= RECALL_K: - return "in the top k, cut by as_of or the budget" - return "extracted, ranked below k" + return RECALLED + if best_rank is not None and best_rank <= k: + return CUT + return BELOW_K + + +def memory_config(result: dict[str, Any]) -> Config: + """Return the configuration of the memory arm that a result file records. + + A result file of phase 5 records no `memory` section, and its run used `BASELINE`. + """ + found = result["config"].get("memory") + if found is None: + return BASELINE + return Config( + memory_ranking=found["ranking"], + k=found["k"], + budget=found["budget"], + extractor=found["extractor"], + ) + + +def table(rows: Sequence[dict[str, Any]]) -> str: + """Return the groups of misses for each question type, and in total, as Markdown.""" + columns = ("missed", *CAUSES[:3], "recalled, all present", "recalled, some present") + header = "| Type | " + " | ".join(columns) + " |" + lines = [header, "|" + "---|" * (len(columns) + 1)] + kinds = [kind for kind in QUESTION_TYPES if any(row["type"] == kind for row in rows)] + for kind in [*kinds, "all"]: + chosen = [row for row in rows if kind in ("all", row["type"])] + groups = Counter(group(row) for row in chosen) + counts = [len(chosen), *(groups[name] for name in columns[1:])] + lines.append(f"| {kind} | " + " | ".join(str(count) for count in counts) + " |") + return "\n".join(lines) + + +def group(row: dict[str, Any]) -> str: + """Return the group of one diagnosed miss, with the recalled group split in two.""" + if row["cause"] != RECALLED: + return str(row["cause"]) + return f"recalled, {row['present']} present" def main(argv: Sequence[str] | None = None) -> Path: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) - parser.add_argument("--result", type=Path, default=RESULTS_DIR / "e2e-sample.json") + parser.add_argument("--result", type=Path, default=RESULTS_DIR / "e2e-sample-a.json") parser.add_argument("--url", default="http://localhost:3000") args = parser.parse_args(argv) commit = _commit() result = json.loads(args.result.read_text()) + config = memory_config(result) missed = [row["id"] for row in result["questions"] if not row["arms"]["memory"]["correct"]] questions = {question.id: question for question in load(DATA_DIR / DATA_FILE)} - models = CachedModels(Remote(), CACHE_DIR / "replies.sqlite") + replay = CachedModels(Offline(), CACHE_DIR / "replies.sqlite") + marker = CachedModels(Remote(), CACHE_DIR / "replies.sqlite") rows: list[dict[str, Any]] = [] - run_id = "diagnose-" + result["config"]["run_id"] - marked_spend = 0.0 - with local_store(args.url, run_id) as (store, empty): - memory = Memory(store) - extract = extractor(models) + run_id = f"diagnose-{datetime.now(UTC):%Y%m%d%H%M%S}-{uuid.uuid4().hex[:6]}" + with local_store(args.url, run_id) as (stores, empty): + memory = Memory(stores[config.memory_ranking]) + extract = extractor(replay, config.variant) for count, question_id in enumerate(missed, start=1): question = questions[question_id] namespace = ("bench", "lme_s", run_id, question.id) @@ -102,30 +166,23 @@ def main(argv: Sequence[str] | None = None) -> Path: actor=ACTOR, observed_at=session.date, ) - if models.spent - marked_spend > SPEND_LIMIT: - raise RuntimeError(f"the replay spent ${models.spent:.2f}: the cache missed") + store = stores[config.memory_ranking] facts = store.list(namespace, Filter(kind=RECALLED_KINDS), limit=10_000) - numbered = "\n".join(f"{n}. {fact.content}" for n, fact in enumerate(facts, 1)) - replayed = models.spent - reply = models( - JUDGE, - MARK.format(question=question.question, answer=question.answer, facts=numbered), - 1000, - MARK_SCHEMA, - ) - marked_spend += models.spent - replayed - marked = sorted({n for n in json.loads(reply.text)["facts"] if 1 <= n <= len(facts)}) + marked = _mark(marker, question.question, question.answer, facts) holding = {facts[n - 1].id for n in marked} ranks: dict[str, int | None] = {} - block: set[str] = set() - for name in ("fused", "vector"): - ranked = _search(store, name, question.question, namespace, len(facts)) + for name in RANKINGS: + ranked = _search(stores[name], question.question, namespace, len(facts)) found = [n for n, r in enumerate(ranked, start=1) if r.id in holding] ranks[name] = found[0] if found else None - if name == "fused": - top = _search(store, name, question.question, namespace, RECALL_K) - block = block_ids(top, question.as_of) - cause = verdict(holding, ranks["fused"], block) + top = _search(store, question.question, namespace, config.k) + block = block_ids(top, question.as_of, config.k, config.budget) + cause = verdict(holding, ranks[config.memory_ranking], block, config.k) + present = None + if cause == RECALLED: + shown = [fact.content for fact in facts if fact.id in holding & block] + whole = _complete(marker, question.question, question.answer, shown) + present = "all" if whole else "some" rows.append( { "id": question.id, @@ -134,46 +191,59 @@ def main(argv: Sequence[str] | None = None) -> Path: "holding": [facts[n - 1].content for n in marked], "best_rank": ranks, "cause": cause, + "present": present, } ) empty(namespace) - print(f"{count}/{len(missed)} {question.id} {cause} {ranks}", flush=True) - models.close() + print(f"{count}/{len(missed)} {question.id} {group(rows[-1])} {ranks}", flush=True) + replay.close() + marker.close() causes = Counter(row["cause"] for row in rows) output = { "config": { - "source": str(args.result.relative_to(Path.cwd())), - "recall_k": RECALL_K, - "recall_budget": RECALL_BUDGET, + "source": str(args.result.resolve().relative_to(Path.cwd())), + "memory": { + "ranking": config.memory_ranking, + "k": config.k, + "budget": config.budget, + "extractor": config.extractor, + }, "marker": JUDGE, "hotmemory_commit": commit, - "spent_usd": round(models.spent, 4), - "model_calls": models.calls, + "spent_usd": round(marker.spent, 4), + "model_calls": marker.calls, }, - "causes": dict(causes.most_common()), + "causes": {cause: causes[cause] for cause in CAUSES}, + "groups": dict(Counter(group(row) for row in rows).most_common()), "questions": rows, } - path = RESULTS_DIR / "e2e-sample-diagnosis.json" + path: Path = args.result.with_name(args.result.stem + "-diagnosis.json") path.write_text(json.dumps(output, indent=1) + "\n") - print(json.dumps(output["causes"]), f"spent ${models.spent:.4f}; wrote {path}") + print(table(rows)) + print(f"spent ${marker.spent:.4f}; wrote {path}") return path -def _search(store: Any, ranking: str, query: str, namespace: Sequence[str], k: int) -> list[Any]: - """Return every fact under `namespace`, ranked for `query` by `ranking`.""" - from hotmemory.hotdata import HotdataStore +def _mark(marker: CachedModels, question: str, answer: str, facts: Sequence[Record]) -> list[int]: + """Return the numbers, from 1, of the facts that the marker says hold the answer.""" + numbered = "\n".join(f"{n}. {fact.content}" for n, fact in enumerate(facts, 1)) + prompt = MARK.format(question=question, answer=answer, facts=numbered) + reply = marker(JUDGE, prompt, 1000, MARK_SCHEMA) + return sorted({n for n in json.loads(reply.text)["facts"] if 1 <= n <= len(facts)}) - ranked = HotdataStore( - store.client, - store.database, - embedder=store._embedder, - model=store._model, - dimensions=store._dimensions, - ranking=ranking, # type: ignore[arg-type] - ) - return [ - hit.record for hit in ranked.search(query, [namespace], Filter(kind=RECALLED_KINDS), k=k) - ] + +def _complete(marker: CachedModels, question: str, answer: str, facts: Sequence[str]) -> bool: + """Return whether the marker says that `facts` hold every value that the answer needs.""" + listed = "\n".join(f"- {fact}" for fact in facts) + prompt = COMPLETE.format(question=question, answer=answer, facts=listed) + reply = marker(JUDGE, prompt, 100, COMPLETE_SCHEMA) + return bool(json.loads(reply.text)["complete"]) + + +def _search(store: Store, query: str, namespace: Sequence[str], k: int) -> list[Record]: + """Return the first `k` facts under `namespace`, ranked for `query` by `store`.""" + hits = store.search(query, [namespace], Filter(kind=RECALLED_KINDS), k=k) + return [hit.record for hit in hits] if __name__ == "__main__": diff --git a/bench/e2e.py b/bench/e2e.py index 90345e3..2c9d6e2 100644 --- a/bench/e2e.py +++ b/bench/e2e.py @@ -1,10 +1,12 @@ """The end-to-end score: three arms on the sample, with one reader and one judge. Run it from the root of the repository. Without `--run`, it prints the cost estimate and -stops: +stops. With `--replay`, it calls no model, and it stops at the first reply that the cache +does not hold: uv run --group bench --env-file .env python -m bench.e2e uv run --group bench --env-file .env python -m bench.e2e --run + uv run --group bench --env-file .env python -m bench.e2e --replay --arms memory It writes one JSON file into `bench/results/` with the configuration and the results. """ @@ -13,10 +15,12 @@ import argparse import json +import multiprocessing import random import time import uuid -from collections.abc import Callable, Iterator, Sequence +from collections.abc import Callable, Collection, Iterator, Mapping, Sequence +from concurrent.futures import ProcessPoolExecutor from contextlib import contextmanager from dataclasses import dataclass from datetime import UTC, datetime @@ -36,7 +40,6 @@ from bench.embed import CACHE_DIR, Tokens from bench.episodes import session_of, write_episodes from bench.models import ( - EXTRACTOR, EXTRACTOR_MAX_TOKENS, JUDGE, JUDGE_MAX_TOKENS, @@ -46,12 +49,11 @@ CachedModels, Reply, ) -from bench.retrieval import RESULTS_DIR, _commit -from hotmemory import Fact, Filter, Memory, Record, Store +from bench.retrieval import RANKINGS, RESULTS_DIR, _commit +from hotmemory import Embedder, Fact, Filter, Memory, Record, Store -SAMPLE_FILE = Path(__file__).resolve().parent / "sample.json" +HERE = Path(__file__).resolve().parent PER_TYPE = 9 -SEED = 20261010 ARMS = ("full", "retrieval", "memory") RETRIEVED = 10 RECALL_K = 30 @@ -62,26 +64,99 @@ Models = Callable[..., Reply] +@dataclass(frozen=True) +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 + `budget` characters, and extracts with the variant `extractor` of `prompts.VARIANTS`. + """ + + retrieval_ranking: str = "fused" + memory_ranking: str = "fused" + k: int = RECALL_K + budget: int = RECALL_BUDGET + extractor: str = "baseline" + + def __post_init__(self) -> None: + for ranking in (self.retrieval_ranking, self.memory_ranking): + 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: + raise ValueError(f"no extractor variant {self.extractor!r}") + if self.k < 1 or self.budget < 1: + raise ValueError("k and budget must be at least 1") + + @property + def variant(self) -> prompts.Variant: + """Return the extractor variant.""" + return prompts.VARIANTS[self.extractor] + + +BASELINE = Config() +"""The configuration of the phase 5 run.""" + + +@dataclass(frozen=True) +class Sample: + """A fixed sample of questions: its file, its seed, and the samples it leaves out.""" + + file: str + seed: int + excludes: tuple[str, ...] = () + + +SAMPLES = { + "a": Sample("bench/sample.json", 20261010), + "b": Sample("bench/sample-b.json", 20261013, excludes=("a",)), +} + + def choose_sample( - questions: Sequence[Question], per_type: int = PER_TYPE, seed: int = SEED + questions: Sequence[Question], + per_type: int = PER_TYPE, + seed: int = SAMPLES["a"].seed, + exclude: Collection[str] = (), ) -> list[str]: - """Return `per_type` question ids of each type, drawn with `seed`, sorted by type and id.""" + """Return `per_type` question ids of each type, drawn with `seed`, sorted by type and id. + + The draw leaves out the ids in `exclude`. + """ rng = random.Random(seed) + left_out = set(exclude) chosen = [] for kind in QUESTION_TYPES: - ids = sorted(question.id for question in questions if question.type == kind) + ids = sorted(q.id for q in questions if q.type == kind and q.id not in left_out) chosen.extend(sorted(rng.sample(ids, per_type))) return chosen -def extractor(models: Models) -> Callable[[str, datetime | None, list[Record]], list[Fact]]: - """Return the extractor callable of `capture`: one call to EXTRACTOR for each session.""" +def sample_ids(name: str, questions: Sequence[Question], root: Path = HERE.parent) -> list[str]: + """Return the ids of the sample `name`, and write its file first if it does not exist. + + The draw of a sample leaves out every id of the samples that it excludes. + """ + sample = SAMPLES[name] + path = root / sample.file + if not path.exists(): + exclude = {i for other in sample.excludes for i in sample_ids(other, questions, root)} + chosen = choose_sample(questions, seed=sample.seed, exclude=exclude) + path.write_text(json.dumps(chosen, indent=1) + "\n") + ids: list[str] = json.loads(path.read_text()) + return ids + + +def extractor( + models: Models, variant: prompts.Variant = prompts.VARIANTS["baseline"] +) -> Callable[[str, datetime | None, list[Record]], list[Fact]]: + """Return the extractor callable of `capture`: one call to the variant's model per session.""" 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) - reply = models(EXTRACTOR, prompt, EXTRACTOR_MAX_TOKENS, prompts.EXTRACT_SCHEMA) + prompt = prompts.extract(text, date, known, variant.prompt) + reply = models(variant.model, prompt, EXTRACTOR_MAX_TOKENS, prompts.EXTRACT_SCHEMA) facts = json.loads(reply.text)["facts"] return [ Fact(kind="fact", subject=item["subject"].strip(), content=item["content"].strip()) @@ -118,18 +193,23 @@ def answer( def run_question( question: Question, models: Models, - store: Store, + stores: Mapping[str, Store], empty: Callable[[Sequence[str]], None], run_id: str, arms: Sequence[str] = ARMS, + config: Config = BASELINE, ) -> dict[str, Arm]: - """Run each arm on `question`, and empty its namespaces after.""" + """Run each arm on `question`, and empty its namespaces after. + + `stores` maps each ranking to a store, and every store holds the same records. + """ result: dict[str, Arm] = {} if "full" in arms: prompt = prompts.read_chats(question, question.sessions) result["full"] = answer(models, question, prompt, len(prompt)) if "retrieval" in arms: namespace = ("bench", "lme_s", run_id, "retrieval", question.id) + store = stores[config.retrieval_ranking] write_episodes(store, namespace, question) hits = store.search(question.question, [namespace], Filter(kind="episode"), k=RETRIEVED) empty(namespace) @@ -139,8 +219,8 @@ def run_question( result["retrieval"] = answer(models, question, prompt, len(prompt)) if "memory" in arms: namespace = ("bench", "lme_s", run_id, "memory", question.id) - memory = Memory(store) - extract = extractor(models) + memory = Memory(stores[config.memory_ranking]) + extract = extractor(models, config.variant) written = 0 for session in question.sessions: written += len( @@ -153,7 +233,7 @@ def run_question( ) ) _, block = memory.recall( - question.question, [namespace], as_of=question.as_of, budget=RECALL_BUDGET, k=RECALL_K + question.question, [namespace], as_of=question.as_of, budget=config.budget, k=config.k ) empty(namespace) prompt = prompts.read_facts(question, block) @@ -162,18 +242,21 @@ def run_question( def estimate( - questions: Sequence[Question], tokens: Tokens, arms: Sequence[str] = ARMS + questions: Sequence[Question], + tokens: Tokens, + arms: Sequence[str] = ARMS, + config: Config = BASELINE, ) -> dict[str, float]: """Return an upper estimate of the price of each arm in US dollars, from token counts. The full arm counts its prompts. The retrieval arm counts the RETRIEVED longest sessions of each question. The memory arm counts each session with EXTRACT_OVERHEAD - tokens of prompt and 500 tokens of output, and a reader prompt of RECALL_BUDGET + tokens of prompt and 500 tokens of output, and a reader prompt of `config.budget` characters. Every reader reply counts READER_MAX_TOKENS, and every judge prompt 400 - tokens. + tokens. The estimate does not subtract the replies that the cache holds. """ reader_in, reader_out = PRICES[READER] - extract_in, extract_out = PRICES[EXTRACTOR] + extract_in, extract_out = PRICES[config.variant.model] read_out = READER_MAX_TOKENS * reader_out / 1e6 judge = (400 * reader_in + JUDGE_MAX_TOKENS * reader_out) / 1e6 cost = dict.fromkeys(arms, 0.0) @@ -189,7 +272,7 @@ def estimate( for session in question.sessions: size = tokens.count(prompts.session_text(session)) + EXTRACT_OVERHEAD cost["memory"] += (size * extract_in + 500 * extract_out) / 1e6 - cost["memory"] += (RECALL_BUDGET // 3 + 300) * reader_in / 1e6 + read_out + judge + cost["memory"] += (config.budget // 3 + 300) * reader_in / 1e6 + read_out + judge return cost @@ -211,17 +294,24 @@ def accuracy(found: Sequence[bool]) -> dict[str, Any]: @contextmanager -def local_store(url: str, run_id: str) -> Iterator[tuple[Store, Callable[[Sequence[str]], None]]]: - """Yield a fused HotdataStore over a new local database, and a call that empties a namespace.""" +def local_store( + url: str, run_id: str, embedder: Embedder | None = None +) -> Iterator[tuple[dict[str, Store], Callable[[Sequence[str]], None]]]: + """Yield a HotdataStore for each ranking, and a call that empties a namespace. + + The stores share one new local database, and the database is deleted after. Without + `embedder`, the stores embed with `embed.openai_embedder()`. + """ import pyarrow as pa from hotdata_framework import HotdataClient from hotmemory.hotdata import CUE_TABLE, RECORD_TABLES, HotdataStore client = HotdataClient("local", "local", host=url) + embedder = embedder or embed.openai_embedder() store = HotdataStore.provision( f"hotmemory-bench-{run_id}", - embedder=embed.openai_embedder(), + embedder=embedder, model=embed.MODEL, dimensions=embed.DIMENSIONS, client=client, @@ -236,13 +326,117 @@ def empty(namespace: Sequence[str]) -> None: rows = pa.table({"id": ids}) store._load_all([(table, rows, "delete"), (CUE_TABLE, rows, "delete")]) + stores: dict[str, Store] = { + name: HotdataStore( + client, + store.database, + embedder=embedder, + model=embed.MODEL, + dimensions=embed.DIMENSIONS, + ranking=name, # type: ignore[arg-type] + ) + for name in RANKINGS + } try: - yield store, empty + yield stores, empty finally: client.delete_managed_database(store.database) client.close() +@dataclass(frozen=True) +class Job: + """What every worker of one run shares.""" + + url: str + arms: tuple[str, ...] + config: Config + replay: bool + total: int + + +@dataclass(frozen=True) +class Shard: + """The questions of one worker, and the run id that names its database.""" + + questions: list[Question] + run_id: str + + +@dataclass(frozen=True) +class Done: + """The rows of one shard, and what its model calls cost.""" + + rows: list[dict[str, Any]] + spent: float + calls: int + embedder_misses: int + + +def shards_of(questions: Sequence[Question], workers: int, run_id: str) -> list[Shard]: + """Deal `questions` over at most `workers` shards, in turn, so each shard gets every type. + + With one shard, the shard keeps `run_id`. With more, shard n gets `-w`. + """ + count = max(1, min(workers, len(questions))) + return [ + Shard(list(questions[n::count]), f"{run_id}-w{n}" if count > 1 else run_id) + for n in range(count) + ] + + +def row_of(question: Question, result: dict[str, Arm]) -> dict[str, Any]: + """Return the row of the result file for one question.""" + return { + "id": question.id, + "type": question.type, + "abstention": question.is_abstention, + "arms": { + arm: { + "correct": found.correct, + "context_chars": found.context_chars, + "facts_written": found.extracted, + "reader_input_tokens": found.answer.input_tokens, + "answer": found.answer.text, + "verdict": found.verdict.text, + } + for arm, found in result.items() + }, + } + + +def run_shard(job: Job, shard: Shard) -> Done: + """Run the questions of `shard` in order, over a new local database of its own. + + The worker opens its own model cache, embedder, and store, so it can run in its own + process. It prints one line for each question. + """ + from bench.models import Offline, Remote + + models = CachedModels(Offline() if job.replay else Remote(), CACHE_DIR / "replies.sqlite") + embedder = embed.openai_embedder() + rows: list[dict[str, Any]] = [] + try: + with local_store(job.url, shard.run_id, embedder) as (stores, empty): + for question in shard.questions: + result = run_question( + question, models, stores, empty, shard.run_id, job.arms, job.config + ) + rows.append(row_of(question, result)) + marks = " ".join( + f"{arm}={'yes' if found.correct else 'no'}" for arm, found in result.items() + ) + print( + f"{shard.run_id} {len(rows)}/{len(shard.questions)} of {job.total} " + f"{question.id} {marks} spent ${models.spent:.2f}", + flush=True, + ) + finally: + models.close() + embedder.close() + return Done(rows, models.spent, models.calls, embedder.misses) + + def main(argv: Sequence[str] | None = None) -> Path | None: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument( @@ -252,22 +446,48 @@ def main(argv: Sequence[str] | None = None) -> Path | None: parser.add_argument( "--limit", type=int, default=None, help="the first N questions of the sample" ) + parser.add_argument("--sample", choices=tuple(SAMPLES), default="a", help="the sample to run") + parser.add_argument("--retrieval-ranking", choices=RANKINGS, default="fused") + parser.add_argument("--memory-ranking", choices=RANKINGS, default="fused") + 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" + ) + parser.add_argument("--extractor", choices=tuple(prompts.VARIANTS), default="baseline") + parser.add_argument( + "--replay", action="store_true", help="call no model: use cached replies, stop at a miss" + ) + parser.add_argument( + "--workers", + type=int, + default=1, + help="processes that run the questions, each over a database of its own", + ) + parser.add_argument("--tag", default="", help="a suffix for the name of the result file") parser.add_argument("--url", default="http://localhost:3000") args = parser.parse_args(argv) arms = tuple(arm for arm in args.arms.split(",") if arm) if not set(arms) <= set(ARMS): parser.error(f"arms must be among {', '.join(ARMS)}") + if args.workers < 1: + parser.error("workers must be at least 1") + try: + config = Config( + retrieval_ranking=args.retrieval_ranking, + memory_ranking=args.memory_ranking, + k=args.recall_k, + budget=args.recall_budget, + extractor=args.extractor, + ) + except ValueError as error: + parser.error(str(error)) all_questions = {question.id: question for question in load(DATA_DIR / DATA_FILE)} - if not SAMPLE_FILE.exists(): - SAMPLE_FILE.write_text( - json.dumps(choose_sample(list(all_questions.values())), indent=1) + "\n" - ) - sample = json.loads(SAMPLE_FILE.read_text()) + sample = sample_ids(args.sample, list(all_questions.values())) questions = [all_questions[question_id] for question_id in sample][: args.limit] tokens = embed.TiktokenTokens(encoding="o200k_base") - cost = estimate(questions, tokens, arms) + cost = estimate(questions, tokens, arms, config) total = sum(cost.values()) for arm in arms: print(f"estimate {arm}: ${cost[arm]:.2f}") @@ -278,46 +498,27 @@ def main(argv: Sequence[str] | None = None) -> Path | None: if total > COST_LIMIT: print("the estimate is over the limit: stop") return None - if not args.run: - print("no model was called; pass --run to run") + if not args.run and not args.replay: + print("no model was called; pass --run to run, or --replay to run from the cache") return None - from bench.models import Remote - + commit = _commit() started = datetime.now(UTC) run_id = f"{started:%Y%m%d%H%M%S}-{uuid.uuid4().hex[:6]}" - models = CachedModels(Remote(), CACHE_DIR / "replies.sqlite") - rows: list[dict[str, Any]] = [] clock_start = time.perf_counter() - with local_store(args.url, run_id) as (store, empty): - for count, question in enumerate(questions, start=1): - result = run_question(question, models, store, empty, run_id, arms) - rows.append( - { - "id": question.id, - "type": question.type, - "abstention": question.is_abstention, - "arms": { - arm: { - "correct": found.correct, - "context_chars": found.context_chars, - "facts_written": found.extracted, - "reader_input_tokens": found.answer.input_tokens, - "answer": found.answer.text, - "verdict": found.verdict.text, - } - for arm, found in result.items() - }, - } - ) - marks = " ".join( - f"{arm}={'yes' if found.correct else 'no'}" for arm, found in result.items() - ) - print( - f"{count}/{len(questions)} {question.id} {marks} spent ${models.spent:.2f}", - flush=True, - ) - models.close() + shards = shards_of(questions, args.workers, run_id) + job = Job(args.url, arms, config, args.replay, len(questions)) + if len(shards) == 1: + done = [run_shard(job, shards[0])] + else: + context = multiprocessing.get_context("spawn") + with ProcessPoolExecutor(len(shards), mp_context=context) as pool: + done = list(pool.map(run_shard, [job] * len(shards), shards)) + by_id = {row["id"]: row for found in done for row in found.rows} + rows = [by_id[question.id] for question in questions] + spent = sum(found.spent for found in done) + calls = sum(found.calls for found in done) + misses = sum(found.embedder_misses for found in done) summaries = {arm: summarize(rows, arm) for arm in arms} result_file = { @@ -326,7 +527,13 @@ def main(argv: Sequence[str] | None = None) -> Path | None: "data_file": DATA_FILE, "data_revision": DATA_REVISION, "data_sha256": DATA_SHA256, - "sample": {"file": "bench/sample.json", "per_type": PER_TYPE, "seed": SEED}, + "sample": { + "name": args.sample, + "file": SAMPLES[args.sample].file, + "per_type": PER_TYPE, + "seed": SAMPLES[args.sample].seed, + "excludes": list(SAMPLES[args.sample].excludes), + }, "reader": { "model": READER, "temperature": 0, @@ -339,25 +546,42 @@ def main(argv: Sequence[str] | None = None) -> Path | None: "max_tokens": JUDGE_MAX_TOKENS, "prompt": "reference, by question type", }, - "extractor": {"model": EXTRACTOR, "max_tokens": EXTRACTOR_MAX_TOKENS}, + "extractor": { + "variant": config.extractor, + "model": config.variant.model, + "max_tokens": EXTRACTOR_MAX_TOKENS, + }, "embedder": {"model": embed.MODEL, "dimensions": embed.DIMENSIONS}, "driver": "hotdata", - "ranking": "fused", "location": "local", + "retrieval": {"ranking": config.retrieval_ranking, "sessions": RETRIEVED}, + "memory": { + "ranking": config.memory_ranking, + "k": config.k, + "budget": config.budget, + "extractor": config.extractor, + }, "arms": { "full": "every session of the question", - "retrieval": f"the top {RETRIEVED} sessions of the fused search over episodes", + "retrieval": ( + f"the top {RETRIEVED} sessions of the {config.retrieval_ranking} search " + "over episodes" + ), "memory": ( - "capture on each session in date order, then recall with k " - f"{RECALL_K}, budget {RECALL_BUDGET} characters, as_of the end of " - "the question day" + f"capture on each session in date order with the {config.extractor} " + f"extractor, then recall with the {config.memory_ranking} ranking, k " + f"{config.k}, budget {config.budget} characters, as_of the end of the " + "question day" ), }, "questions": len(questions), "estimate_usd": round(total, 2), - "spent_usd": round(models.spent, 2), - "model_calls": models.calls, - "hotmemory_commit": _commit(), + "replay": args.replay, + "workers": len(shards), + "spent_usd": round(spent, 2), + "model_calls": calls, + "embedder_misses": misses, + "hotmemory_commit": commit, "run_id": run_id, "started_at": started.isoformat(), "duration_s": round(time.perf_counter() - clock_start, 1), @@ -366,13 +590,14 @@ def main(argv: Sequence[str] | None = None) -> Path | None: "questions": rows, } RESULTS_DIR.mkdir(exist_ok=True) - name = "e2e-sample" + (f"-first{args.limit}" if args.limit is not None else "") + name = f"e2e-sample-{args.sample}" + (f"-{args.tag}" if args.tag else "") + name += f"-first{args.limit}" if args.limit is not None else "" path = RESULTS_DIR / f"{name}.json" path.write_text(json.dumps(result_file, indent=1) + "\n") for arm in arms: found = summaries[arm]["all"] print(f"{arm}: {found['correct']}/{found['questions']} = {found['accuracy']:.3f}") - print(f"spent ${models.spent:.2f}; wrote {path}") + print(f"spent ${spent:.2f}; wrote {path}") return path diff --git a/bench/embed.py b/bench/embed.py index aa331b4..d01cf4d 100644 --- a/bench/embed.py +++ b/bench/embed.py @@ -73,7 +73,7 @@ def __init__( self._max_input = max_input self._max_request = max_request path.parent.mkdir(parents=True, exist_ok=True) - self._db = sqlite3.connect(path) + self._db = sqlite3.connect(path, timeout=60) self._db.execute("CREATE TABLE IF NOT EXISTS vectors (key TEXT PRIMARY KEY, vector BLOB)") def __call__(self, texts: Sequence[str]) -> list[list[float]]: diff --git a/bench/models.py b/bench/models.py index dbd046e..197f576 100644 --- a/bench/models.py +++ b/bench/models.py @@ -58,7 +58,7 @@ def __init__(self, complete: Complete, path: Path) -> None: self.calls = 0 self._complete = complete path.parent.mkdir(parents=True, exist_ok=True) - self._db = sqlite3.connect(path) + self._db = sqlite3.connect(path, timeout=60) self._db.execute("CREATE TABLE IF NOT EXISTS replies (key TEXT PRIMARY KEY, reply TEXT)") def __call__( @@ -85,8 +85,25 @@ def close(self) -> None: self._db.close() +class Offline: + """`Complete` that calls no model: it raises `CacheMiss` for each prompt. + + Wrapped in `CachedModels`, it replays the cached replies and stops at the first prompt + that the cache does not hold. + """ + + def __call__( + self, model: str, prompt: str, max_tokens: int, schema: dict[str, Any] | None + ) -> Reply: + raise CacheMiss(f"the cache holds no reply of {model} to this prompt: {prompt[:80]!r}") + + +class CacheMiss(RuntimeError): + """A replay needed a reply that the cache does not hold.""" + + class Remote: - """`Complete` over the OpenAI chat API, at temperature 0. + """`Complete` over the OpenAI chat API, at temperature 0, with 120 seconds for each try. The client reads the key from `OPENAI_MODEL_KEY`. `OPENAI_API_KEY` stays the key of the embedder. With a schema, the reply is JSON that the schema allows. @@ -95,7 +112,9 @@ class Remote: def __init__(self) -> None: import openai - self._openai = openai.OpenAI(api_key=os.environ["OPENAI_MODEL_KEY"], max_retries=6) + self._openai = openai.OpenAI( + api_key=os.environ["OPENAI_MODEL_KEY"], max_retries=6, timeout=120.0 + ) def __call__( self, model: str, prompt: str, max_tokens: int, schema: dict[str, Any] | None diff --git a/bench/prompts.py b/bench/prompts.py index 4b5876f..1cd452b 100644 --- a/bench/prompts.py +++ b/bench/prompts.py @@ -8,8 +8,10 @@ import json from collections.abc import Sequence +from dataclasses import dataclass from bench.data import DATE_FORMAT, Question, Session +from bench.models import EXTRACTOR READ_CHATS = ( "I will give you several history chats between you and a user. Please answer the " @@ -104,6 +106,18 @@ } +@dataclass(frozen=True) +class Variant: + """An extractor variant: the prompt template and the model that reads it.""" + + prompt: str + model: str + + +VARIANTS = {"baseline": Variant(EXTRACT, EXTRACTOR)} +"""The extractor variants by name. `baseline` is the extractor of phase 5.""" + + def history(sessions: Sequence[Session]) -> str: """Return the sessions as the reference reader shows them, sorted by date. @@ -143,7 +157,7 @@ def session_text(session: Session) -> str: return "\n\n".join(f"{turn.role}: {turn.content.strip()}" for turn in session.turns) -def extract(session: str, date: str, known: Sequence[str]) -> str: +def extract(session: str, date: str, known: Sequence[str], template: str = EXTRACT) -> str: """Return the extractor prompt for one session and the facts already in memory.""" listed = "\n".join(known) if known else "(none)" - return EXTRACT.format(known=listed, date=date, session=session) + return template.format(known=listed, date=date, session=session) diff --git a/bench/sample-b.json b/bench/sample-b.json new file mode 100644 index 0000000..96b6cb9 --- /dev/null +++ b/bench/sample-b.json @@ -0,0 +1,56 @@ +[ + "19b5f2b3_abs", + "29f2956b_abs", + "545bd2b5", + "7527f7e2", + "75499fd8", + "8ebdbe50", + "a82c026e", + "c14c00dd", + "ec81a493", + "16c90bf4", + "1903aded", + "1b9b7252", + "51b23612", + "8464fc84", + "d596882b", + "dc439ea3", + "e3fc4d6e", + "fca762bc", + "09d032c9", + "0a34ad58", + "35a27287", + "8a2466db", + "afdc33df", + "b0479f84", + "caf03d32", + "d24813b1", + "d6233ab6", + "09ba9854_abs", + "2318644b", + "2788b940", + "37f165cf", + "4adc0475", + "b3c15d39", + "d3ab962e", + "e56a43b9", + "gpt4_372c3eed_abs", + "59524333", + "5c40ec5b", + "603deb26", + "7a87bd0c", + "830ce83f", + "a2f3aa27", + "ba61f0b9", + "d7c942c3", + "eace081b", + "0bb5a684", + "bbf86515", + "e4e14d04", + "gpt4_1e4a8aec", + "gpt4_8c8961ae", + "gpt4_b4a80587", + "gpt4_d31cdae3", + "gpt4_e414231e", + "gpt4_f420262c" +] diff --git a/bench/tests/test_diagnose.py b/bench/tests/test_diagnose.py index 4b1ac00..2f621b7 100644 --- a/bench/tests/test_diagnose.py +++ b/bench/tests/test_diagnose.py @@ -1,6 +1,12 @@ from datetime import UTC, datetime, timedelta +from pathlib import Path +from typing import Any -from bench.diagnose import block_ids, verdict +import pytest + +from bench.diagnose import block_ids, group, memory_config, table, verdict +from bench.e2e import BASELINE +from bench.models import READER, CachedModels, CacheMiss, Offline, Reply from hotmemory import Record AS_OF = datetime(2023, 5, 30, 23, 59, 59, tzinfo=UTC) @@ -36,3 +42,39 @@ def test_verdict() -> None: assert verdict({"f"}, 3, {"f"}) == "recalled, reader missed" assert verdict({"f"}, 3, set()) == "in the top k, cut by as_of or the budget" assert verdict({"f"}, 31, set()) == "extracted, ranked below k" + assert verdict({"f"}, 31, set(), k=40) == "in the top k, cut by as_of or the budget" + + +def test_memory_config_reads_the_result_file() -> None: + assert memory_config({"config": {}}) == BASELINE + found = memory_config( + { + "config": { + "memory": {"ranking": "vector", "k": 50, "budget": 900, "extractor": "baseline"} + } + } + ) + assert (found.memory_ranking, found.k, found.budget) == ("vector", 50, 900) + + +def test_table_splits_the_recalled_group() -> None: + rows: list[dict[str, Any]] = [ + {"type": "multi-session", "cause": "not extracted", "present": None}, + {"type": "multi-session", "cause": "recalled, reader missed", "present": "some"}, + {"type": "knowledge-update", "cause": "recalled, reader missed", "present": "all"}, + ] + assert group(rows[1]) == "recalled, some present" + lines = table(rows).splitlines() + assert lines[2] == "| multi-session | 2 | 1 | 0 | 0 | 0 | 1 |" + assert lines[3] == "| knowledge-update | 1 | 0 | 0 | 0 | 1 | 0 |" + assert lines[4] == "| all | 3 | 1 | 0 | 0 | 1 | 1 |" + + +def test_offline_replays_the_cache_and_stops_at_a_miss(tmp_path: Path) -> None: + path = tmp_path / "replies.sqlite" + CachedModels(lambda *_: Reply("yes", 1, 1), path)(READER, "p", 10) + replay = CachedModels(Offline(), path) + assert replay(READER, "p", 10).text == "yes" + with pytest.raises(CacheMiss): + replay(READER, "q", 10) + assert (replay.calls, replay.spent) == (0, 0.0) diff --git a/bench/tests/test_e2e.py b/bench/tests/test_e2e.py index 7d765cd..3c143ff 100644 --- a/bench/tests/test_e2e.py +++ b/bench/tests/test_e2e.py @@ -8,7 +8,16 @@ from bench import prompts from bench.data import QUESTION_TYPES, Question, load -from bench.e2e import choose_sample, estimate, run_question, summarize +from bench.e2e import ( + SAMPLES, + Config, + choose_sample, + estimate, + run_question, + sample_ids, + shards_of, + summarize, +) from bench.models import EXTRACTOR, JUDGE, READER, CachedModels, Reply from hotmemory import MemoryStore from tests.conftest import FakeClock, fake_embedder @@ -72,6 +81,23 @@ def test_choose_sample_is_fixed(questions: dict[str, Question]) -> None: assert first == choose_sample(pool, per_type=2, seed=1) assert len(first) == 12 assert first != choose_sample(pool, per_type=2, seed=2) + other = choose_sample(pool, per_type=2, seed=1, exclude=first) + assert len(other) == 12 + assert not set(other) & set(first) + + +def test_sample_b_leaves_out_sample_a(questions: dict[str, Question], tmp_path: Path) -> None: + base = questions["q_one"] + pool = [ + replace(base, id=f"{kind}-{n}", type=kind) for kind in QUESTION_TYPES for n in range(20) + ] + (tmp_path / "bench").mkdir() + a = sample_ids("a", pool, root=tmp_path) + b = sample_ids("b", pool, root=tmp_path) + assert len(a) == len(b) == 54 + assert not set(a) & set(b) + assert json.loads((tmp_path / SAMPLES["b"].file).read_text()) == b + assert sample_ids("b", pool[:60], root=tmp_path) == b def test_run_question_on_every_arm(questions: dict[str, Question]) -> None: @@ -85,7 +111,7 @@ def empty(namespace: Sequence[str]) -> None: for record in store.list(namespace, limit=1000): store.delete(record.namespace, record.key) - result = run_question(question, models, store, empty, "r1") + result = run_question(question, models, {"fused": store, "vector": store}, empty, "r1") assert set(result) == {"full", "retrieval", "memory"} assert all(arm.correct for arm in result.values()) assert result["memory"].extracted == 6 @@ -102,6 +128,49 @@ def empty(namespace: Sequence[str]) -> None: assert JUDGE == READER +def test_config_refuses_unknown_values() -> None: + with pytest.raises(ValueError, match="ranking"): + Config(memory_ranking="bm25") + with pytest.raises(ValueError, match="variant"): + Config(extractor="none") + with pytest.raises(ValueError, match="at least 1"): + Config(k=0) + + +def test_run_question_follows_the_config( + questions: dict[str, Question], monkeypatch: pytest.MonkeyPatch +) -> None: + question = questions["q_one"] + monkeypatch.setitem( + prompts.VARIANTS, + "terse", + prompts.Variant("TERSE {known} {date}\nSession:\n{session}", EXTRACTOR), + ) + fused = MemoryStore(embedder=fake_embedder, clock=FakeClock()) + vector = MemoryStore(embedder=fake_embedder, clock=FakeClock()) + written: list[str] = [] + + def empty(namespace: Sequence[str]) -> None: + for store, name in ((fused, "fused"), (vector, "vector")): + if store.list(namespace, limit=1000): + written.append(f"{namespace[3]}:{name}") + for record in store.list(namespace, limit=1000): + store.delete(record.namespace, record.key) + + models = FakeModels() + config = Config( + retrieval_ranking="vector", memory_ranking="fused", k=2, budget=200, extractor="terse" + ) + result = run_question( + question, models, {"fused": fused, "vector": vector}, empty, "r1", config=config + ) + assert written == ["retrieval:vector", "memory:fused"] + assert all(p.startswith("TERSE ") for m, p in models.prompts if m == EXTRACTOR) + assert 0 < result["memory"].context_chars <= 200 + block = [p for m, p in models.prompts if m == READER and "several facts" in p[:60]][-1] + assert block.count("[sources: none]") <= 2 + + def test_estimate_counts_each_arm(questions: dict[str, Question]) -> None: class Chars: def count(self, text: str) -> int: @@ -113,6 +182,8 @@ def truncate(self, text: str, limit: int) -> str: cost = estimate(list(questions.values()), Chars()) assert set(cost) == {"full", "retrieval", "memory"} assert all(value > 0 for value in cost.values()) + larger = estimate(list(questions.values()), Chars(), ("memory",), Config(budget=80_000)) + assert larger["memory"] > cost["memory"] def test_cached_models_call_once(tmp_path: Path) -> None: @@ -128,3 +199,15 @@ def complete(model: str, prompt: str, max_tokens: int, schema: dict[str, Any] | assert models(READER, "p", 11).text == "yes" assert calls == ["p", "p"] assert models.spent == pytest.approx(5.0) + + +def test_shards_deal_every_question_once(questions: dict[str, Question]) -> None: + base = questions["q_one"] + pool = [replace(base, id=f"q{n}") for n in range(10)] + shards = shards_of(pool, 4, "r1") + assert [shard.run_id for shard in shards] == ["r1-w0", "r1-w1", "r1-w2", "r1-w3"] + dealt = [q.id for shard in shards for q in shard.questions] + assert sorted(dealt) == sorted(q.id for q in pool) + assert [q.id for q in shards[1].questions] == ["q1", "q5", "q9"] + assert [shard.run_id for shard in shards_of(pool, 1, "r1")] == ["r1"] + assert len(shards_of(pool[:2], 4, "r1")) == 2 diff --git a/docs/benchmarks.md b/docs/benchmarks.md index 9274f34..1d93e1f 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -1,8 +1,8 @@ # Benchmarks This page reports the numbers of hotmemory on LongMemEval-S, a public benchmark of 500 -questions over long chat histories. Each number on this page was observed on 2026-10-10. -Each section names the result file that a run wrote into `bench/results/`, with its +questions over long chat histories. Each number on this page was observed on 2026-10-10, +except where a section gives another date. Each section names the result file that a run wrote into `bench/results/`, with its numbers, its configuration, and the commit of hotmemory that ran it. Git ignores that folder, so this page is the record of the numbers. [bench/README.md](../bench/README.md) gives the commands that repeat each run and write the files again. @@ -131,16 +131,20 @@ order as the local run. ## End-to-end accuracy -Three arms answer the same 54 questions, 9 of each question type, with the same reader and -the same judge. A fixed seed chose the questions, and their ids are in -[bench/sample.json](../bench/sample.json). 2 of the 54 are abstention questions. File: -`e2e-sample.json`. +Three arms answer the same questions, with the same reader and the same judge. Each run +uses one of two samples of 54 questions, 9 of each question type. A fixed seed chose each +sample. + +| Sample | File of ids | Seed | Abstention questions | Use | +|---|---|---|---|---| +| A | [bench/sample.json](../bench/sample.json) | `20261010` | 2 | Each phase iterates on it. | +| B | [bench/sample-b.json](../bench/sample-b.json) | `20261013` | 4 | Each phase runs it once, at its end. It holds no question of sample A. | | Item | Value | |---|---| | Reader | `gpt-4o-2024-08-06`, temperature 0, up to 800 output tokens, the chain-of-thought prompt of the reference code | | Judge | `gpt-4o-2024-08-06`, temperature 0, the prompt of the reference code for each question type | -| Extractor of the memory arm | `gpt-4o-mini-2024-07-18`, with the prompt in `bench/prompts.py` | +| Extractor of the memory arm | `gpt-4o-mini-2024-07-18`, with the `baseline` prompt in `bench/prompts.py` | | Driver | `HotdataStore` with the fused ranking, on the local stack | The arms: @@ -152,6 +156,10 @@ The arms: its block, with k of 30, a budget of 8,000 characters, and `as_of` at 23:59:59 on the date of the question. +### Sample A + +File: `e2e-sample.json`. + | Question type | Questions | Full context | Retrieval | Memory | |---|---|---|---|---| | single-session-user | 9 | 7 | 8 | 3 | @@ -166,34 +174,93 @@ At the median, the reader read 113,689 tokens in the full-context arm, 29,878 in retrieval arm, and 1,462 in the memory arm. The run made 2,802 model calls and cost $20.95. +On 2026-10-13, the memory arm ran again on sample A with `--replay`, from the cached +replies, at commit `ce97385` with uncommitted changes to documentation only. It gave the same answer and the same mark on each of the 54 +questions, with no model call and no embedder call. File: `e2e-sample-a.json`. On +2026-10-14, the same replay with `--workers 4` at commit `4d99fbe` gave the same 54 rows +again. It took 519 seconds, and the replay with one process took 2,264 seconds. File: +`e2e-sample-a-workers4.json`. + The LongMemEval paper reports 60.6% for full context with GPT-4o on all 500 questions of the earlier data file. A sample of 54 questions and the cleaned data file do not compare with that number closely. +### Sample B + +Observed on 2026-10-14, at commit `5f2abb9`. File: `e2e-sample-b.json`. + +| Question type | Questions | Full context | Retrieval | Memory | +|---|---|---|---|---| +| single-session-user | 9 | 8 | 9 | 6 | +| single-session-assistant | 9 | 9 | 9 | 4 | +| single-session-preference | 9 | 3 | 6 | 0 | +| multi-session | 9 | 7 | 7 | 5 | +| knowledge-update | 9 | 8 | 8 | 5 | +| temporal-reasoning | 9 | 8 | 9 | 7 | +| All | 54 | 43 (79.6%) | 48 (88.9%) | 27 (50.0%) | + +At the median, the reader read 113,314 tokens in the full-context arm, 30,899 in the +retrieval arm, and 1,466 in the memory arm. The run made 2,869 model calls, cost $21.42, +and took 2.9 hours with one process. + +The memory arm scored 8 questions more on sample B than on sample A, with the same +configuration. That is 14.8 points. A gain on one sample of 54 questions is therefore +weak evidence, and a change counts only if both samples move the same way. + ### Why the memory arm scored low -`bench/diagnose.py` ran `capture` again on each of the 35 questions that the memory arm -missed, from the cached extractor replies. GPT-4o then marked each stored fact that states -the answer or a value that the answer needs, and the search ranked every fact. File: -`e2e-sample-diagnosis.json`. +`bench/diagnose.py` ran `capture` again on each question that the memory arm missed, from +the cached extractor replies. GPT-4o then marked each stored fact that states the answer +or a value that the answer needs, and the search ranked every fact. A second GPT-4o call +then read the marked facts in the block. It told whether they hold every value that the +answer needs. Files: `e2e-sample-a-diagnosis.json`, observed on 2026-10-13, +and `e2e-sample-b-diagnosis.json`, observed on 2026-10-14. -| Cause | Questions | -|---|---| -| No stored fact states the answer or a value that it needs | 19 | -| A stored fact holds it, and the fused search ranked it below 30 | 2 | -| At least one marked fact was in the top 30 | 14 | +| Cause | Sample A | Sample B | +|---|---|---| +| No stored fact states the answer or a value that it needs | 19 | 16 | +| A stored fact holds it, and the fused search ranked it below 30 | 2 | 0 | +| A marked fact was in the top 30, and `recall` left it out | 0 | 0 | +| The block holds every needed fact, and the reader missed | 1 | 1 | +| The block holds some of the needed facts | 13 | 10 | +| Misses | 35 | 27 | + +The three causes of phase 5 (19, 2, and 14 on sample A) are the same in this run, on +each question. The split of the last group is new. + +The extractor is the main loss on both samples. On sample A, phase 5 read the misses. Most +often the extractor stored the interest of the user. It dropped what the assistant said, for example a number, a name, or an item of a +list. In both questions of the second row, the vector ranking ranked the fact higher than +the fused ranking did. When the block held a marked fact, it held every needed fact in +only 2 of 25 misses. A question that needs several facts then lacks at least one of them. + +The misses of each question type, on sample A: -The extractor is the main loss. Most often it stored the interest of the user. It dropped -what the assistant said, for example a number, a name, or an item of a list. In both -questions of the second row, the vector ranking ranked the fact higher than the fused -ranking did. The last row is an upper bound on the misses of the reader. A multi-session -question needs several facts. The marker counts the question when it finds one of them in -the top 30. +| Question type | Misses | Not stored | Ranked below 30 | All present | Some present | +|---|---|---|---|---|---| +| single-session-user | 6 | 4 | 1 | 1 | 0 | +| single-session-assistant | 6 | 5 | 0 | 0 | 1 | +| single-session-preference | 6 | 4 | 1 | 0 | 1 | +| multi-session | 6 | 3 | 0 | 0 | 3 | +| knowledge-update | 5 | 2 | 0 | 0 | 3 | +| temporal-reasoning | 6 | 1 | 0 | 0 | 5 | + +The misses of each question type, on sample B: + +| Question type | Misses | Not stored | Ranked below 30 | All present | Some present | +|---|---|---|---|---|---| +| single-session-user | 3 | 3 | 0 | 0 | 0 | +| single-session-assistant | 5 | 4 | 0 | 0 | 1 | +| single-session-preference | 9 | 6 | 0 | 0 | 3 | +| multi-session | 4 | 1 | 0 | 0 | 3 | +| knowledge-update | 4 | 1 | 0 | 1 | 2 | +| temporal-reasoning | 2 | 1 | 0 | 0 | 1 | ## Known limits -- The end-to-end numbers come from a sample of 54 questions, so one question moves an - arm by 1.9 points. A run on all 500 questions waits for the owner's agreement. +- The end-to-end numbers come from samples of 54 questions, so one question moves an arm + by 1.9 points. The memory arm differs by 14.8 points between sample A and sample B. A + run on all 500 questions waits for the owner's agreement. - The full-context arm gives the reader every session. The reference code keeps the last 50 sessions, so the two differ for histories longer than 50 sessions. - The extractor of the memory arm is the only one tried. A stronger extractor or another diff --git a/docs/internal/plan.md b/docs/internal/plan.md index 11be943..bf3cedb 100644 --- a/docs/internal/plan.md +++ b/docs/internal/plan.md @@ -1,7 +1,7 @@ # Plan: the second workstream, memory accuracy -Status: proposed, 2026-10-13. The owner agrees to each decision before the first issue -opens. This file holds the open workstream, phases 6 to 13. Each phase becomes one GitHub +Status: agreed, 2026-10-13. The owner agreed to each decision and to the target on +2026-10-13. Phase 6 is open, as issue #16. This file holds the open workstream, phases 6 to 13. Each phase becomes one GitHub issue and one pull request, worked in task order. When a phase closes, `roadmap.md` gets its paragraph, and this file keeps only the phases still open. The plan of phase 5 is in git history, at commit `4a6fb36`. @@ -12,7 +12,7 @@ When this workstream closes, the memory arm meets the target below, and hotmemor PyPI. The memory arm scored 19 of 54 on sample A (35.2 percent). The retrieval arm scored 45 of 54. -The target (proposed) is a distance from the lowest published score of a memory layer with +The target (agreed) is a distance from the lowest published score of a memory layer with a GPT-4o reader: Zep at 71.2 percent (read, not reproduced). The published range goes up to 84.2 percent. @@ -62,24 +62,24 @@ shape the phases below: ## Decisions -Each item is proposed. The owner agrees or changes it before phase 6 opens. +The owner agreed to each item on 2026-10-13. -- Deterministic library (proposed). No model call enters `src/hotmemory`. A model-backed +- Deterministic library (agreed). No model call enters `src/hotmemory`. A model-backed extractor in an extra, and decision calls such as a classifier for supersede, wait in Deferred until this workstream closes. -- Two samples (proposed). Sample A is the committed sample of 54, seed `20261010`, in +- Two samples (agreed). Sample A is the committed sample of 54, seed `20261010`, in `bench/sample.json`. Each phase iterates on sample A. Sample B is a second sample of 54, 9 of each type, with a new fixed seed, and with no question of sample A. A phase runs sample B once, at its end, to check that a gain on A is not a fit to A. -- What counts as a gain (proposed). One question moves an arm by 1.9 points on 54 +- What counts as a gain (agreed). One question moves an arm by 1.9 points on 54 questions. Two conditions make a change count. Sample A and sample B move the same way. The diagnosis shows the targeted group of misses getting smaller. A phase reports every variant it tried, not only the best one. -- Fixed models (proposed). The reader and the judge stay `gpt-4o-2024-08-06` at +- Fixed models (agreed). The reader and the judge stay `gpt-4o-2024-08-06` at temperature 0, with the reference prompts. The baseline extractor stays `gpt-4o-mini-2024-07-18`. Phase 7 tries at most one other extractor model, chosen by the owner. -- The full-context arm (proposed). It runs once on sample B, in phase 6, as the baseline +- The full-context arm (agreed). It runs once on sample B, in phase 6, as the baseline of B. No later phase runs it again on 54 questions. It costs the most, at about 114,000 tokens per question. - The cost limit (agreed in phase 5, kept). The runner prints an estimate and stops over @@ -126,6 +126,35 @@ Stop and ask if: - The cached rerun of sample A does not reproduce the published numbers. - The estimate for sample B is over 50 dollars. +Decisions taken during the work. The plan left these details open. The work took the +defaults below, and the owner can change any of them in review. + +- Sample B uses the seed `20261013` and the draw rule of sample A, over the questions that + are not in sample A. It holds 4 abstention questions, and sample A holds 2 (observed). +- One `Config` in `bench/e2e.py` holds the options of tasks 3 and 4, so the two tasks are + one commit. The memory arm uses one ranking for the search inside `capture` and for + `recall`. The options are `--retrieval-ranking`, `--memory-ranking`, `--recall-k`, + `--recall-budget`, `--extractor`, and `--tag`. +- A result file is named `e2e-sample-[-].json`. The phase 5 file + `e2e-sample.json` stays as the record of phase 5. +- `--replay` runs the arms with a model stand-in that raises at the first reply that the + cache does not hold. The diagnosis replays `capture` the same way, in place of the stop + at 1 dollar of spend. +- The diagnosis keeps the marker prompt of phase 5 unchanged, so the cache reproduces the + three groups. A second GPT-4o call splits the group "recalled, reader missed". It gets + the marked facts that are in the block. It tells whether they hold every value that the + answer needs ("all present") or not ("some present"). +- The diagnosis file is named after the result file, as `-diagnosis.json`. +- `--workers N` runs the questions in N processes, and each process has its own local + database. The plan did not ask for it. A run of the three arms took 3.6 hours in + phase 5 with one process, and each later phase runs the samples again. The runner reads + the commit at the start of a run, so a commit during the run does not change the commit + that the result file records. +- The phase spent $21.80 on models, from the result files: $21.42 for the run of sample + B, $0.37 for its diagnosis, and $0.01 for the split of the diagnosis of sample A. The + replays of sample A spent nothing. The first diagnosis run of sample A stalled on one + call and was stopped, and the cost of that call is not observed. + ## Phase 7. The extractor Goal: the extractor stores the facts that the questions ask about. The target is the @@ -395,7 +424,6 @@ These wait, and `roadmap.md` lists them under Deferred: ## Questions for the owner -- Agree or change each decision above. - Phase 7: name a second extractor model to try, or none. - Phase 8: agree the budgets to try, or change them from 8,000, 32,000, and 64,000 characters. diff --git a/docs/internal/roadmap.md b/docs/internal/roadmap.md index 3470288..8afa0a7 100644 --- a/docs/internal/roadmap.md +++ b/docs/internal/roadmap.md @@ -23,13 +23,13 @@ that proves it. ## The second workstream: memory accuracy -Proposed on 2026-10-13. `plan.md` holds the tasks, the acceptance criteria, and the stop +Agreed on 2026-10-13. `plan.md` holds the tasks, the acceptance criteria, and the stop conditions of each phase. The library stays deterministic: no code in `src/hotmemory` calls a model. Latency is not a goal, and its work waits under Deferred. | Phase | Delivers | Status | |---|---|---| -| 6. The accuracy loop | A second sample of 54 questions, and the memory arm configured from the command line. The diagnosis split for each question type, and the baselines of both samples. | proposed | +| 6. The accuracy loop | A second sample of 54 questions, and the memory arm configured from the command line. The diagnosis split for each question type, and the baselines of both samples. | open, issue #16 | | 7. The extractor | A benchmark extractor that keeps what the assistant said, with sources, measured against the 19 misses where no stored fact held the answer. No library change. | proposed | | 8. Episodes behind facts | `capture` stores its text as an episode that its facts cite. `recall` can add the cited or ranked episodes to its block, within one budget. | proposed | | 9. The ranking | An ablation of the fused ranking, and a default ranking chosen by the benchmark. | proposed |