diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..5f88071 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,35 @@ +name: CI + +on: + push: + pull_request: + +jobs: + test-and-typecheck: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - uses: oven-sh/setup-bun@v2 + with: + bun-version: 1.3.6 + + - name: Install backend dependencies + run: bun install --frozen-lockfile + + - name: Install UI dependencies + working-directory: ui + run: bun install --frozen-lockfile + + - name: Run tests + run: bun test + + - name: Typecheck backend + run: bunx tsc --noEmit + + - name: Typecheck UI + working-directory: ui + run: bunx tsc -p tsconfig.json --noEmit + + - name: Check formatting + run: bun run format:check diff --git a/README.md b/README.md index 9b0ac51..d47e2ae 100644 --- a/README.md +++ b/README.md @@ -68,19 +68,28 @@ GOOGLE_API_KEY= | `list-questions` | Browse benchmark questions | | `show-failures` | Debug failed questions | | `serve` | Start web UI | +| `beam prepare` | Download, verify, and convert pinned public BEAM data | | `help` | Show help (`help providers`, `help models`, `help benchmarks`) | ## Options ``` -p, --provider Memory provider (supermemory, mem0, zep) --b, --benchmark Benchmark (locomo, longmemeval, convomem, beam-1m, beam-10m) +-b, --benchmark Benchmark (locomo, longmemeval, convomem, beam-1m, beam-10m, beam-1m-10m) -j, --judge Judge model (gpt-4o, sonnet-4, gemini-2.5-flash, etc.) -r, --run-id Run identifier (auto-generated if omitted) -m, --answering-model Model for answer generation (default: gpt-4o) -l, --limit Limit number of questions -q, --question-id Specific question (for test command) --force Clear checkpoint and restart +--data-path Prepared BEAM snapshot root or snapshot path +--dataset-revision Pin the prepared BEAM dataset fingerprint +--retrieval-top-k BEAM paper cutoff (5, 10, 15, or 20; default: 5) +--answer-cutoff Evidence shown to the answerer in the experimental mem0-nugget profile +--evaluation-profile Experimental BEAM evaluation profile (`mem0-nugget`) +--source-run Reuse validated completed ingest/index builds in a new run +--concurrency Default phase/build concurrency (supported by `run` and `ingest`) +--ingest-batch-size Ordered sessions per provider request before its readiness barrier ``` ## Examples @@ -104,6 +113,15 @@ bun run src/index.ts run -p zep -b longmemeval -j sonnet-4 -m gemini-2.5-flash # Compare multiple providers bun run src/index.ts compare -p supermemory,mem0,zep -b locomo -s 5 +# Prepare and run the supported BEAM 1M tier +bun run src/index.ts beam prepare --tiers 1M +# Copy the dataset fingerprint printed by `beam prepare` into this command. +bun run src/index.ts run -p supermemory -b beam-1m -j gpt-4.1-mini --retrieval-top-k 5 --data-path ./data/benchmarks/beam --dataset-revision DATASET_FINGERPRINT_PRINTED_ABOVE + +# Experimental direct-K50 mem0-style scoring run reusing the completed 1M build. +# This does not alter or claim parity with the BEAM paper protocol. +bun run src/index.ts run -p supermemory -b beam-1m --source-run beam-1m-ingest-c35-b5 -r beam-1m-sm-gpt5-direct-k50-mem0-nugget-v3 --from-phase search --evaluation-profile mem0-nugget --retrieval-top-k 50 --answer-cutoff 50 -m gpt-5 -j gpt-5 --concurrency-search 10 --concurrency-answer 10 --concurrency-evaluate 10 + # Test single question bun run src/index.ts test -r my-test -q question_42 @@ -123,7 +141,35 @@ bun run src/index.ts show-failures -r my-test 6. REPORT Aggregate scores → Output accuracy + latency ``` -Each phase checkpoints independently. Failed runs resume from last successful point. +Build phases checkpoint once per shared haystack; search, answer, and evaluation checkpoint per +question. Failed runs resume only when the pinned dataset, protocol, provider adapter, models, and +retrieval configuration still match. + +BEAM uses a per-session causal build barrier: each conversation stays ordered as +`add -> ready -> checkpoint -> next session`, while separate conversations may run concurrently. + +For BEAM, the report keeps the paper score (macro-average across the ten memory abilities) separate +from pass accuracy (`question score >= 0.5`). Event-ordering questions use the published normalized +Kendall tau-b score rather than a nugget average. Combined 1M/10M runs +report each tier separately and label their cross-tier macro as a MemoryBench aggregate, not as a paper score. +Limited or sampled runs are labeled `beamScorePartial` and cannot enter the ranked leaderboard; only +the exact validated 700-question 1M or 200-question 10M cohort receives the official `beamScore` key. + +The explicit `mem0-nugget` profile is an experimental comparison protocol. It uses the public mem0 +answer/judge prompts, GPT-5 judge identity, mem0's numeric score clamp, and ordinary nugget averages +for all ten abilities—including event ordering. Its primary metric is `mem0NuggetAverage`, never +`beamScore`. GPT-5 answering and judging use OpenAI Chat Completions, omit temperature and reasoning +effort, allow 4,096 completion tokens, and use up to five outer attempts with a 120-second deadline +and 2/4/6/8-second backoff. Each outer attempt retains the pinned client's two inner transport +retries. The direct-K50 command above requests and answers with 50 results. It is not an exact +reproduction of mem0's published Top-50 result: the pinned public mem0 runner retrieves 200 and then +applies an answer cutoff of 50, while this harness currently permits at most 100 results per direct +provider request. If all five answering attempts exhaust without non-empty text—including transport +failures—this profile preserves Mem0's terminal behavior by checkpointing the empty hypothesis and +evaluating it; other protocols remain fail-closed. The profile also fails closed on schema-invalid judge output instead of using +mem0's raw-text `1.0`/`0.5` marker fallback. Source-build reuse validates the dataset, ordered +haystacks, ingestion policy, provider adapter/configuration, and completed indexing before retaining +the existing containers. ## MemScore @@ -135,7 +181,7 @@ accuracy% / latencyMs / contextTokens | Component | What it measures | |-----------|-----------------| -| **Quality** | Answer accuracy — `(correct / total) * 100` from judge evaluations | +| **Quality** | The benchmark protocol's primary quality metric (legacy benchmarks use binary accuracy; BEAM uses its continuous paper score) | | **Latency** | Average search response time in milliseconds | | **Tokens** | Average context tokens sent to the answering model (counted client-side) | diff --git a/bun.lock b/bun.lock index 637c144..e5df2de 100644 --- a/bun.lock +++ b/bun.lock @@ -12,7 +12,10 @@ "@getzep/zep-cloud": "^3.13.0", "ai": "^5.0.115", "drizzle-orm": "^0.45.1", + "hyparquet": "^1.27.1", + "hyparquet-compressors": "^1.1.1", "js-tiktoken": "^1.0.21", + "json5": "2.2.3", "mem0ai": "^2.1.38", "supermemory": "^4.0.0", "zod": "^3.24.4", @@ -324,6 +327,8 @@ "function-bind": ["function-bind@1.1.2", "", {}, "sha512-7XHNxH7qX9xG5mIwxkhumTox/MIRNcOgDrxWsMt2pAr23WHp6MrRlN7FBSFpCpr+oVO0F744iUgR82nJMfG2SA=="], + "fzstd": ["fzstd@0.1.1", "", {}, "sha512-dkuVSOKKwh3eas5VkJy1AW1vFpet8TA/fGmVA5krThl8YcOVE/8ZIoEA1+U1vEn5ckxxhLirSdY837azmbaNHA=="], + "gauge": ["gauge@4.0.4", "", { "dependencies": { "aproba": "^1.0.3 || ^2.0.0", "color-support": "^1.1.3", "console-control-strings": "^1.1.0", "has-unicode": "^2.0.1", "signal-exit": "^3.0.7", "string-width": "^4.2.3", "strip-ansi": "^6.0.1", "wide-align": "^1.1.5" } }, "sha512-f9m+BEN5jkg6a0fZjleidjN51VE1X+mPFQ2DJ0uv1V39oCLCbsGe6yjbBnp7eK7z/+GAon99a3nHuqbuuthyPg=="], "gaxios": ["gaxios@7.1.3", "", { "dependencies": { "extend": "^3.0.2", "https-proxy-agent": "^7.0.1", "node-fetch": "^3.3.2", "rimraf": "^5.0.1" } }, "sha512-YGGyuEdVIjqxkxVH1pUTMY/XtmmsApXrCVv5EU25iX6inEPbV+VakJfLealkBtJN69AQmh1eGOdCl9Sm1UP6XQ=="], @@ -370,6 +375,12 @@ "humanize-ms": ["humanize-ms@1.2.1", "", { "dependencies": { "ms": "^2.0.0" } }, "sha512-Fl70vYtsAFb/C06PTS9dZBo7ihau+Tu/DNCk/OyHhea07S+aeMWpFFkUaXRa8fI+ScZbEI8dfSxwY7gxZ9SAVQ=="], + "hyparquet": ["hyparquet@1.27.1", "", {}, "sha512-kvMGVKlB/xa9UHxWLJkySJYGC436nPeBWC9Fzo/BZtMEzFUjmzD1KuLz67EOpA3Ci6V7yF/yTG7QRSDSx5EDSg=="], + + "hyparquet-compressors": ["hyparquet-compressors@1.1.1", "", { "dependencies": { "fzstd": "0.1.1", "hysnappy": "1.0.0" } }, "sha512-yx7aA3Rhj0YycbdV71+XznQSLAefa4cT0urpgNXy4aM6eSeCknaVDNne8y45Uz74Fb15yyXUzOStlceOJBan7A=="], + + "hysnappy": ["hysnappy@1.0.0", "", {}, "sha512-MNrC4NfwDGPb889O6gIfEtbvEZCSWUsSEhsz4Oq2FRcpGtXHfeVz3KciSPp5Pnnz1NjFMgDQNfxdJozymJEDDA=="], + "iceberg-js": ["iceberg-js@0.8.1", "", {}, "sha512-1dhVQZXhcHje7798IVM+xoo/1ZdVfzOMIc8/rgVSijRK38EDqOJoGula9N/8ZI5RD8QTxNQtK/Gozpr+qUqRRA=="], "iconv-lite": ["iconv-lite@0.6.3", "", { "dependencies": { "safer-buffer": ">= 2.1.2 < 3.0.0" } }, "sha512-4fCk79wshMdzMp2rH06qWrJE4iolqLhCUH+OiuIgU++RB0+94NlDL81atO7GX55uUKueo0txHNtvEyI6D7WdMw=="], @@ -420,6 +431,8 @@ "json-schema": ["json-schema@0.4.0", "", {}, "sha512-es94M3nTIfsEPisRafak+HDLfHXnKBhV3vU5eqPcS3flIWqcxJWgXHXiey3YrpaNsanY5ei1VoYEbOzijuq9BA=="], + "json5": ["json5@2.2.3", "", { "bin": { "json5": "lib/cli.js" } }, "sha512-XmOWe7eyHYH14cLdVPoyg+GOH3rYX++KpzrylJwSW98t3Nk+U8XOl8FWKOgwtzdb8lXGf6zYwDUzeHMWfxasyg=="], + "jwa": ["jwa@2.0.1", "", { "dependencies": { "buffer-equal-constant-time": "^1.0.1", "ecdsa-sig-formatter": "1.0.11", "safe-buffer": "^5.0.1" } }, "sha512-hRF04fqJIP8Abbkq5NKGN0Bbr3JxlQ+qhZufXVr0DvujKy93ZCbXZMHDL4EOtodSbCWxOqR8MS1tXA5hwqCXDg=="], "jws": ["jws@4.0.1", "", { "dependencies": { "jwa": "^2.0.1", "safe-buffer": "^5.0.1" } }, "sha512-EKI/M/yqPncGUUh44xz0PxSidXFr/+r0pA70+gIYhjv+et7yxM+s29Y+VGDkovRofQem0fs7Uvf4+YmAdyRduA=="], diff --git a/package.json b/package.json index 8c1ab9f..1474dcb 100644 --- a/package.json +++ b/package.json @@ -16,7 +16,10 @@ "@getzep/zep-cloud": "^3.13.0", "ai": "^5.0.115", "drizzle-orm": "^0.45.1", + "hyparquet": "^1.27.1", + "hyparquet-compressors": "^1.1.1", "js-tiktoken": "^1.0.21", + "json5": "2.2.3", "mem0ai": "^2.1.38", "supermemory": "^4.0.0", "zod": "^3.24.4" diff --git a/src/benchmarks/README.md b/src/benchmarks/README.md index e6b8dbc..a809cfa 100644 --- a/src/benchmarks/README.md +++ b/src/benchmarks/README.md @@ -7,11 +7,15 @@ Benchmark dataset adapters. Each benchmark implements the `Benchmark` interface. ```typescript interface Benchmark { name: string + scope: BenchmarkScope + protocol: BenchmarkProtocol load(config?: BenchmarkConfig): Promise getQuestions(filter?: QuestionFilter): UnifiedQuestion[] getHaystackSessions(questionId: string): UnifiedSession[] getGroundTruth(questionId: string): string getQuestionTypes(): QuestionTypeRegistry + getIngestionGroupId?(questionId: string): string + getDatasetIdentity?(): DatasetIdentity | undefined } ``` @@ -28,6 +32,9 @@ interface Benchmark { - `getHaystackSessions()` - Return `UnifiedSession[]` for a question - `getGroundTruth()` - Return expected answer string - `getQuestionTypes()` - Return `{ [id]: { id, alias, description } }` +- `protocol` - Own ingestion formatting, retrieval, answer formatting, evaluation, and aggregation +- `getIngestionGroupId()` - Optional shared-build hint; the orchestrator independently fingerprints + every member's ordered provider-visible haystack before accepting the group ## Existing Benchmarks @@ -36,7 +43,7 @@ interface Benchmark { | `locomo` | GitHub snap-research/locomo | Long context memory benchmark | | `longmemeval` | HuggingFace xiaowu0162/longmemeval-cleaned | Long-term memory evaluation | | `convomem` | HuggingFace Salesforce/ConvoMem | Conversational memory benchmark | -| `beam-1m` / `beam-10m` | HuggingFace Mohammadta/BEAM | Beyond a Million Tokens benchmark (1M and 10M token tiers) | +| `beam-1m` / `beam-10m` / `beam-1m-10m` | Pinned Hugging Face BEAM repositories | Supported BEAM public 1M/10M tiers | ## Question Types @@ -69,7 +76,50 @@ interface Benchmark { | `implicit_connection_evidence` | implicit | Implicit reasoning | | `abstention_evidence` | abstain | Unanswerable questions | -### BEAM +### BEAM 1M/10M + +BEAM does not download silently during a run. Prepare an immutable, hashed snapshot first: + +```bash +bun run src/index.ts beam prepare --tiers 1M,10M +``` + +The converter verifies the pinned Parquet hashes, validates 35 chats / 700 questions for 1M and +10 chats / 200 questions for 10M, requires 20 questions and all ten abilities per chat, and fails +closed on missing or malformed data. Use `--dataset-revision ` for a pinned run. + +The pinned 10M rows store each `turns[]` item as a variable-length alternating message block, not +as one message pair. The converter deterministically splits every complete block into strict +user/assistant sessions. The reviewed revision also contains exactly two incomplete blocks: +`10M:1:plan-7:batch-10:source-turn-19` and +`10M:2:plan-7:batch-8:source-turn-51`. Each ends with a literal `followup_question` user message +whose assistant response is absent. Matching the authors' pinned `pair_chunk` behavior in +`src/answer_probing_questions/long_term_memory_methods.py`, only those two stable source identities +receive an `ASSISTANT: N/A` placeholder. Canonical sessions mark the padding and the manifest records +its per-tier count. Any other odd, non-alternating, or identity-mismatched block fails closed. + +The supported scope is explicitly BEAM 1M/10M, not every smaller tier in the paper. The paper score +and the additional `>= 0.5` pass accuracy are both reported; they are not interchangeable. +Only a complete validated tier is labeled with the official `beamScore` and is leaderboard-eligible. +Question-limited or sampled runs remain useful diagnostics but are labeled `beamScorePartial`. + +BEAM ingestion is causal within each conversation: add one ordered user/assistant session, wait for +both document processing and memory dreaming/indexing to complete, durably checkpoint it, and only +then add the next session. Independent conversations still build concurrently. Supermemory requests +top-level `dreaming: "instant"`; readiness requires both `status` and `dreamingStatus` to be `done`. + +For a deliberately non-paper comparison, `--evaluation-profile mem0-nugget` selects a separate +versioned protocol. It supports direct Top-K values through 100, a distinct answer cutoff, GPT-5 as +judge, mem0's nugget clamp, and nugget-average scoring for event-ordering questions. Reports use +`mem0NuggetAverage`, not `beamScore`. Its GPT-5 calls use Chat Completions with 4,096 completion +tokens, no temperature or reasoning-effort override, five outer attempts, and a 120-second per-attempt +deadline. Each outer attempt allows two inner transport retries. A direct K50 run is an ablation; +mem0's published Top-50 setup +retrieves 200 before applying cutoff 50, which is outside the current direct-search limit. +After five answer attempts exhaust without non-empty text, this profile alone checkpoints and +evaluates the empty answer, matching the pinned Mem0 runner; the paper and default protocols continue +to fail closed. + | Type | Alias | Description | |------|-------|-------------| | `abstention` | abstain | Withhold answers when evidence is missing | diff --git a/src/benchmarks/beam/dataset.ts b/src/benchmarks/beam/dataset.ts new file mode 100644 index 0000000..553df12 --- /dev/null +++ b/src/benchmarks/beam/dataset.ts @@ -0,0 +1,1614 @@ +import { createHash } from "node:crypto" +import { createReadStream, existsSync, readFileSync, statSync } from "node:fs" +import { basename, isAbsolute, join, relative, resolve } from "node:path" +import JSON5 from "json5" +import { z } from "zod" +import type { + BeamBatch, + BeamCanonicalChat, + BeamCanonicalFileManifest, + BeamCanonicalMessage, + BeamCanonicalQuestion, + BeamCanonicalSession, + BeamDatasetManifest, + BeamDatasetSource, + BeamMessage, + BeamQuestionType, + BeamScale, + BeamTierCounts, + PreparedBeamDataset, +} from "./types" +import { + BEAM_CANONICAL_SCHEMA_VERSION, + BEAM_CONVERTER_VERSION, + BEAM_MANIFEST_SCHEMA_VERSION, + BEAM_QUESTION_TYPE_IDS, +} from "./types" +import { decodeBeamParquetWithHyparquet, type BeamParquetDecoder } from "./parquet" + +export const BEAM_DATASET_SOURCES: Record = { + "1M": { + repository: "Mohammadta/BEAM", + split: "1M", + revision: "3205395e897e7318c7b094ef4e6047b9b82dbb03", + parquetFiles: [ + { + path: "data/1M-00000-of-00001.parquet", + url: "https://huggingface.co/datasets/Mohammadta/BEAM/resolve/3205395e897e7318c7b094ef4e6047b9b82dbb03/data/1M-00000-of-00001.parquet", + expectedSha256: "41b5acbbb55a586b1305514ef9d9fb03365d9b3331b598a1c2dd7603d93ef533", + }, + ], + }, + "10M": { + repository: "Mohammadta/BEAM-10M", + split: "10M", + revision: "9b2096193fe74e2837e4713e483351e19817773c", + parquetFiles: [ + { + path: "data/10M-00000-of-00002.parquet", + url: "https://huggingface.co/datasets/Mohammadta/BEAM-10M/resolve/9b2096193fe74e2837e4713e483351e19817773c/data/10M-00000-of-00002.parquet", + expectedSha256: "31d96fd47ec56221d202e68792f26c00e49467dd4b36ee105c36ebd19ef78ad5", + }, + { + path: "data/10M-00001-of-00002.parquet", + url: "https://huggingface.co/datasets/Mohammadta/BEAM-10M/resolve/9b2096193fe74e2837e4713e483351e19817773c/data/10M-00001-of-00002.parquet", + expectedSha256: "a4f13fe25af51d57405ae41008689c31d1421377f3efde56a024b441deb2ee65", + }, + ], + }, +} + +export const BEAM_EXPECTED_COUNTS: Record = { + "1M": { chats: 35, questions: 700 }, + "10M": { chats: 10, questions: 200 }, +} + +function hashRunningConverterImplementation(): string { + const hash = createHash("sha256") + for (const source of ["./dataset.ts", "./parquet.ts", "./prepare.ts"]) { + hash.update(source) + hash.update(readFileSync(new URL(source, import.meta.url))) + } + return hash.digest("hex") +} + +/** Hashes the exact converter and publication source files executing this run. */ +export const BEAM_CONVERTER_IMPLEMENTATION_HASH = hashRunningConverterImplementation() + +const QUESTION_TYPE_SET = new Set(BEAM_QUESTION_TYPE_IDS) + +const QUESTION_TYPE_ALIASES: Record = { + abstention: "abstention", + contradiction_resolution: "contradiction_resolution", + event_ordering: "event_ordering", + information_extraction: "information_extraction", + instruction_following: "instruction_following", + knowledge_update: "knowledge_update", + multi_hop_reasoning: "multi_session_reasoning", + multi_session_reasoning: "multi_session_reasoning", + preference_following: "preference_following", + summarization: "summarization", + temporal_reasoning: "temporal_reasoning", +} + +const MONTHS: Record = { + january: 1, + february: 2, + march: 3, + april: 4, + may: 5, + june: 6, + july: 7, + august: 8, + september: 9, + october: 10, + november: 11, + december: 12, +} + +export interface CanonicalBeamTier { + scale: BeamScale + chats: BeamCanonicalChat[] + questions: BeamCanonicalQuestion[] + counts: BeamTierCounts +} + +export interface LoadPreparedBeamDatasetOptions { + snapshotPath: string + tiers: BeamScale[] + expectedDatasetFingerprint?: string +} + +interface ValidatePreparedBeamSnapshotContentsOptions extends LoadPreparedBeamDatasetOptions { + allowInjectedTestSourceIdentity: boolean +} + +const canonicalMessageSchema = z + .object({ + role: z.enum(["user", "assistant"]), + content: z.string(), + timeAnchor: z.string().optional(), + }) + .strict() + +const canonicalSessionSchema = z + .object({ + sessionId: z.string().min(1), + planNumber: z.number().int().positive().optional(), + batchNumber: z.number().int().nonnegative(), + turnIndex: z.number().int().positive(), + documentDate: z.string().optional(), + hadInvalidTimeAnchor: z.boolean().optional(), + hasPaddedAssistant: z.literal(true).optional(), + messages: z.array(canonicalMessageSchema).length(2), + }) + .strict() + +const canonicalChatSchema = z + .object({ + schemaVersion: z.literal(BEAM_CANONICAL_SCHEMA_VERSION), + scale: z.enum(["1M", "10M"]), + chatId: z.string().min(1), + sessions: z.array(canonicalSessionSchema), + }) + .strict() + +const canonicalQuestionSchema = z + .object({ + schemaVersion: z.literal(BEAM_CANONICAL_SCHEMA_VERSION), + scale: z.enum(["1M", "10M"]), + chatId: z.string(), + questionId: z.string(), + questionType: z.enum(BEAM_QUESTION_TYPE_IDS), + question: z.string(), + rubric: z.array(z.string()), + difficulty: z.string().optional(), + referenceAnswer: z.string().optional(), + }) + .strict() + +const sourceFileSchema = z + .object({ + path: z.string(), + snapshotPath: z.string(), + url: z.string(), + byteSize: z.number().int().nonnegative(), + sha256: z.string().regex(/^[a-f0-9]{64}$/), + }) + .strict() + +const canonicalFileSchema = z + .object({ + path: z.string(), + byteSize: z.number().int().nonnegative(), + sha256: z.string().regex(/^[a-f0-9]{64}$/), + rowCount: z.number().int().nonnegative(), + }) + .strict() + +const tierCountsSchema = z + .object({ + chats: z.number().int().nonnegative(), + questions: z.number().int().nonnegative(), + sessions: z.number().int().nonnegative(), + sessionsWithDocumentDate: z.number().int().nonnegative(), + sessionsWithoutDocumentDate: z.number().int().nonnegative(), + sessionsWithInvalidTimeAnchor: z.number().int().nonnegative(), + sessionsWithPaddedAssistant: z.number().int().nonnegative(), + byQuestionType: z.record(z.string(), z.number().int().nonnegative()), + byChat: z.record( + z.string(), + z + .object({ + sessions: z.number().int().nonnegative(), + questions: z.number().int().nonnegative(), + byQuestionType: z.record(z.string(), z.number().int().nonnegative()), + }) + .strict() + ), + }) + .strict() + +const manifestSchema = z + .object({ + manifestSchemaVersion: z.literal(BEAM_MANIFEST_SCHEMA_VERSION), + canonicalSchemaVersion: z.literal(BEAM_CANONICAL_SCHEMA_VERSION), + converter: z + .object({ + name: z.literal("memorybench-beam"), + version: z.string(), + implementationHash: z.string().regex(/^[a-f0-9]{64}$/), + }) + .strict(), + includedTiers: z.array(z.enum(["1M", "10M"])), + sources: z.array( + z + .object({ + tier: z.enum(["1M", "10M"]), + sourceIdentity: z.enum(["reviewed-published", "injected-test-fixture"]), + repository: z.string(), + split: z.enum(["1M", "10M"]), + revision: z.string().regex(/^[a-f0-9]{40}$/), + files: z.array(sourceFileSchema), + }) + .strict() + ), + canonicalFiles: z.array(canonicalFileSchema), + counts: z.record(z.string(), tierCountsSchema), + orderedChatIds: z.record(z.string(), z.array(z.string())), + orderedChatIdsDigest: z.record(z.string(), z.string().regex(/^[a-f0-9]{64}$/)), + orderedQuestionIds: z.record(z.string(), z.array(z.string())), + orderedQuestionIdsDigest: z.record(z.string(), z.string().regex(/^[a-f0-9]{64}$/)), + datasetFingerprint: z.string().regex(/^[a-f0-9]{64}$/), + manifestHash: z.string().regex(/^[a-f0-9]{64}$/), + }) + .strict() + +function compareStrings(left: string, right: string): number { + return left < right ? -1 : left > right ? 1 : 0 +} + +function canonicalize(value: unknown): unknown { + if (Array.isArray(value)) return value.map(canonicalize) + if (value && typeof value === "object") { + return Object.fromEntries( + Object.entries(value as Record) + .filter(([, item]) => item !== undefined) + .sort(([left], [right]) => compareStrings(left, right)) + .map(([key, item]) => [key, canonicalize(item)]) + ) + } + if (typeof value === "number" && !Number.isFinite(value)) { + throw new Error("BEAM canonical JSON cannot contain a non-finite number") + } + return value +} + +export function stableBeamStringify(value: unknown): string { + return JSON.stringify(canonicalize(value)) +} + +export function sha256Text(value: string): string { + return createHash("sha256").update(value, "utf8").digest("hex") +} + +export function sha256Bytes(value: Uint8Array): string { + return createHash("sha256").update(value).digest("hex") +} + +function asRecord(value: unknown, context: string): Record { + const parsed = parseJsonValue(value, context) + if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) { + throw new Error(`${context} must be an object`) + } + return parsed as Record +} + +function normalizePythonLiteralKeywords(value: string): string { + let result = "" + let quote: "'" | '"' | undefined + let escaped = false + + for (let index = 0; index < value.length; ) { + const character = value[index]! + if (quote) { + result += character + if (escaped) escaped = false + else if (character === "\\") escaped = true + else if (character === quote) quote = undefined + index += 1 + continue + } + if (character === "'" || character === '"') { + quote = character + result += character + index += 1 + continue + } + if (/[A-Za-z_]/.test(character)) { + let end = index + 1 + while (end < value.length && /[A-Za-z0-9_]/.test(value[end]!)) end += 1 + const token = value.slice(index, end) + result += + token === "None" ? "null" : token === "True" ? "true" : token === "False" ? "false" : token + index = end + continue + } + result += character + index += 1 + } + return result +} + +function parseJsonValue(value: unknown, context: string): unknown { + if (typeof value !== "string") return value + const trimmed = value.trim() + if (!trimmed.startsWith("{") && !trimmed.startsWith("[")) return value + try { + return JSON.parse(trimmed) + } catch (jsonError) { + try { + return JSON5.parse(normalizePythonLiteralKeywords(trimmed)) + } catch (literalError) { + throw new Error( + `${context} contains invalid JSON/Python-style data: ${String(jsonError)}; ${String(literalError)}` + ) + } + } +} + +function pickAlias( + record: Record, + aliases: string[], + context: string, + required = true +): unknown { + const matches = aliases.filter((key) => record[key] !== undefined && record[key] !== null) + if (matches.length === 0) { + if (!required) return undefined + throw new Error(`${context} is missing; expected one of: ${aliases.join(", ")}`) + } + const first = record[matches[0]] + for (const key of matches.slice(1)) { + if (stableBeamStringify(record[key]) !== stableBeamStringify(first)) { + throw new Error(`${context} is ambiguous; conflicting fields: ${matches.join(", ")}`) + } + } + return first +} + +function requireNonEmptyString(value: unknown, context: string): string { + if (typeof value !== "string" || value.trim().length === 0) { + throw new Error(`${context} must be a non-empty string`) + } + return value +} + +function normalizeChatId(value: unknown, context: string): string { + if (typeof value !== "string" && typeof value !== "number") { + throw new Error(`${context} must be a string or number`) + } + const chatId = String(value).trim() + if (!chatId) throw new Error(`${context} must not be empty`) + if (!/^[a-zA-Z0-9_-]+$/.test(chatId)) { + throw new Error(`${context} contains unsupported characters: ${chatId}`) + } + return chatId +} + +function normalizeQuestionType(value: unknown, context: string): BeamQuestionType { + const source = requireNonEmptyString(value, context) + const normalized = source + .trim() + .toLowerCase() + .replace(/[^a-z0-9]+/g, "_") + .replace(/^_+|_+$/g, "") + const questionType = QUESTION_TYPE_ALIASES[normalized] + if (!questionType || !QUESTION_TYPE_SET.has(questionType)) { + throw new Error(`${context} has unknown BEAM question type: ${source}`) + } + return questionType +} + +function validDateParts(year: number, month: number, day: number): string | undefined { + if (!Number.isInteger(year) || !Number.isInteger(month) || !Number.isInteger(day)) + return undefined + if (month < 1 || month > 12 || day < 1 || day > 31) return undefined + const date = new Date(Date.UTC(year, month - 1, day, 12)) + if ( + date.getUTCFullYear() !== year || + date.getUTCMonth() + 1 !== month || + date.getUTCDate() !== day + ) { + return undefined + } + return `${String(year).padStart(4, "0")}-${String(month).padStart(2, "0")}-${String(day).padStart(2, "0")}` +} + +export function parseBeamTimeAnchorStrict(value: unknown): string | undefined { + if (typeof value !== "string") return undefined + const anchor = value.trim() + const iso = anchor.match(/^(\d{4})-(\d{2})-(\d{2})$/) + if (iso) return validDateParts(Number(iso[1]), Number(iso[2]), Number(iso[3])) + + const named = anchor.match(/^([A-Za-z]+)-(\d{1,2})-(\d{4})$/) + if (!named) return undefined + const month = MONTHS[named[1].toLowerCase()] + if (!month) return undefined + return validDateParts(Number(named[3]), month, Number(named[2])) +} + +function isMessageRecord(value: unknown): value is Record { + if (!value || typeof value !== "object" || Array.isArray(value)) return false + const record = value as Record + return typeof record.role === "string" && typeof record.content === "string" +} + +function flatMessagesToBatch( + rawMessages: unknown[], + batchNumber: number, + context: string +): Record { + if (rawMessages.length === 0 || rawMessages.length % 2 !== 0) { + throw new Error(`${context} must contain complete user/assistant pairs`) + } + const turns: unknown[][] = [] + for (let index = 0; index < rawMessages.length; index += 2) { + const user = asRecord(rawMessages[index], `${context}[${index}]`) + const assistant = asRecord(rawMessages[index + 1], `${context}[${index + 1}]`) + if (user.role !== "user" || assistant.role !== "assistant") { + throw new Error( + `${context} must alternate user then assistant; found ${String(user.role)}/${String(assistant.role)} at messages ${index}/${index + 1}` + ) + } + turns.push([user, assistant]) + } + return { batch_number: batchNumber, turns } +} + +function flattenBatches(value: unknown, context: string): Record[] { + const parsed = parseJsonValue(value, context) + if (Array.isArray(parsed)) { + if (parsed.length === 0) throw new Error(`${context} must not be empty`) + if (parsed.every(isMessageRecord)) { + return [flatMessagesToBatch(parsed, 1, context)] + } + if ( + parsed.every((item) => Array.isArray(item) && item.length > 0 && item.every(isMessageRecord)) + ) { + return parsed.map((messages, index) => + flatMessagesToBatch(messages as unknown[], index + 1, `${context}[${index}]`) + ) + } + return parsed.flatMap((item, index) => flattenBatches(item, `${context}[${index}]`)) + } + const record = asRecord(parsed, context) + if (record.batch_number !== undefined || record.batchNumber !== undefined) return [record] + return Object.keys(record) + .sort(compareStrings) + .flatMap((key) => { + const nested = record[key] + if (nested === undefined || nested === null) return [] + const batches = flattenBatches(nested, `${context}.${key}`) + const planMatch = key.match(/^plan-(\d+)$/) + if (!planMatch) return batches + const planNumber = Number(planMatch[1]) + if (!Number.isInteger(planNumber) || planNumber < 1) { + throw new Error(`${context}.${key} has an invalid plan number`) + } + return batches.map((batch) => { + const existing = batch.plan_number ?? batch.planNumber + if (existing !== undefined && Number(existing) !== planNumber) { + throw new Error(`${context}.${key} conflicts with nested plan number ${String(existing)}`) + } + return { ...batch, plan_number: planNumber } + }) + }) +} + +const BEAM_MISSING_ASSISTANT_PLACEHOLDER = "N/A" + +export const BEAM_10M_PINNED_PADDED_ASSISTANT_SOURCE_IDENTITIES = [ + "10M:1:plan-7:batch-10:source-turn-19", + "10M:2:plan-7:batch-8:source-turn-51", +] as const + +const BEAM_10M_PINNED_PADDED_ASSISTANT_SOURCE_IDENTITY_SET = new Set( + BEAM_10M_PINNED_PADDED_ASSISTANT_SOURCE_IDENTITIES +) + +function structuredTurnSourceIdentity(input: { + scale: BeamScale + chatId: string + planNumber?: number + batchNumber: number + sourceTurnNumber: number +}): string { + return `${input.scale}:${input.chatId}:plan-${input.planNumber ?? 0}:batch-${input.batchNumber}:source-turn-${input.sourceTurnNumber}` +} + +function normalizeBatch( + value: Record, + context: string, + scale: BeamScale, + chatId: string +): BeamBatch { + const rawPlanNumber = pickAlias( + value, + ["plan_number", "planNumber"], + `${context}.plan_number`, + false + ) + const planNumber = + rawPlanNumber === undefined + ? undefined + : typeof rawPlanNumber === "number" + ? rawPlanNumber + : Number(rawPlanNumber) + if (planNumber !== undefined && (!Number.isInteger(planNumber) || planNumber < 1)) { + throw new Error(`${context}.plan_number must be a positive integer`) + } + const rawBatchNumber = pickAlias( + value, + ["batch_number", "batchNumber"], + `${context}.batch_number` + ) + const batchNumber = typeof rawBatchNumber === "number" ? rawBatchNumber : Number(rawBatchNumber) + if (!Number.isInteger(batchNumber) || batchNumber < 0) { + throw new Error(`${context}.batch_number must be a non-negative integer`) + } + + const rawTurns = parseJsonValue( + pickAlias(value, ["turns"], `${context}.turns`), + `${context}.turns` + ) + if (!Array.isArray(rawTurns) || rawTurns.length === 0) { + throw new Error(`${context}.turns must be a non-empty array`) + } + + const turns = rawTurns.flatMap((rawTurn, turnIndex) => { + const parsedTurn = parseJsonValue(rawTurn, `${context}.turns[${turnIndex}]`) + if (!Array.isArray(parsedTurn) || parsedTurn.length < 2) { + throw new Error( + `${context}.turns[${turnIndex}] must contain at least one user message followed by one assistant message` + ) + } + const messages = parsedTurn.map((rawMessage, messageIndex): BeamMessage => { + const message = asRecord(rawMessage, `${context}.turns[${turnIndex}][${messageIndex}]`) + const role = requireNonEmptyString( + pickAlias(message, ["role"], `${context}.turns[${turnIndex}][${messageIndex}].role`), + `${context}.turns[${turnIndex}][${messageIndex}].role` + ) + if (role !== "user" && role !== "assistant") { + throw new Error( + `${context}.turns[${turnIndex}][${messageIndex}].role must be user or assistant` + ) + } + const content = requireNonEmptyString( + pickAlias(message, ["content"], `${context}.turns[${turnIndex}][${messageIndex}].content`), + `${context}.turns[${turnIndex}][${messageIndex}].content` + ) + const timeAnchor = pickAlias( + message, + ["time_anchor", "timeAnchor"], + `${context}.turns[${turnIndex}][${messageIndex}].time_anchor`, + false + ) + if (timeAnchor !== undefined && typeof timeAnchor !== "string") { + throw new Error( + `${context}.turns[${turnIndex}][${messageIndex}].time_anchor must be a string` + ) + } + return { + role, + content, + ...(timeAnchor ? { time_anchor: timeAnchor } : {}), + } + }) + for (let messageIndex = 0; messageIndex < messages.length; messageIndex++) { + const expectedRole = messageIndex % 2 === 0 ? "user" : "assistant" + if (messages[messageIndex]?.role !== expectedRole) { + throw new Error( + `${context}.turns[${turnIndex}] must alternate user then assistant; expected ${expectedRole} at message ${messageIndex}` + ) + } + } + + if (messages.length % 2 !== 0) { + const trailingSourceMessage = asRecord( + parsedTurn[parsedTurn.length - 1], + `${context}.turns[${turnIndex}][${parsedTurn.length - 1}]` + ) + const trailingQuestionType = pickAlias( + trailingSourceMessage, + ["question_type", "questionType"], + `${context}.turns[${turnIndex}][${parsedTurn.length - 1}].question_type`, + false + ) + const sourceIdentity = structuredTurnSourceIdentity({ + scale, + chatId, + planNumber, + batchNumber, + sourceTurnNumber: turnIndex + 1, + }) + + // The pinned 10M source has exactly the two stable source identities + // listed above whose three-message blocks end with a follow-up user + // message persisted without its assistant response. + // The authors' pair-chunk implementation represents this case as + // `ASSISTANT: N/A`; reproduce that explicit source policy while rejecting + // every other incomplete or odd structured block. + if ( + scale !== "10M" || + messages.length !== 3 || + trailingQuestionType !== "followup_question" || + !BEAM_10M_PINNED_PADDED_ASSISTANT_SOURCE_IDENTITY_SET.has(sourceIdentity) + ) { + throw new Error( + `${context}.turns[${turnIndex}] (${sourceIdentity}) must contain complete user/assistant pairs` + ) + } + messages.push({ + role: "assistant", + content: BEAM_MISSING_ASSISTANT_PLACEHOLDER, + isPaddedAssistant: true, + }) + } + + const pairs: BeamMessage[][] = [] + for (let messageIndex = 0; messageIndex < messages.length; messageIndex += 2) { + pairs.push([messages[messageIndex]!, messages[messageIndex + 1]!]) + } + return pairs + }) + + const timeAnchor = pickAlias( + value, + ["time_anchor", "timeAnchor"], + `${context}.time_anchor`, + false + ) + if (timeAnchor !== undefined && typeof timeAnchor !== "string") { + throw new Error(`${context}.time_anchor must be a string when present`) + } + + return { + ...(planNumber !== undefined ? { plan_number: planNumber } : {}), + batch_number: batchNumber, + turns, + ...(timeAnchor !== undefined ? { time_anchor: timeAnchor } : {}), + } +} + +function resolveBatchDate(batch: BeamBatch): { date?: string; hadInvalidTimeAnchor: boolean } { + let hadInvalidTimeAnchor = false + if (batch.time_anchor) { + const parsed = parseBeamTimeAnchorStrict(batch.time_anchor) + if (parsed) return { date: parsed, hadInvalidTimeAnchor } + hadInvalidTimeAnchor = true + } + + for (const turn of batch.turns) { + for (const message of turn) { + if (!message.time_anchor) continue + const parsed = parseBeamTimeAnchorStrict(message.time_anchor) + if (parsed) return { date: parsed, hadInvalidTimeAnchor } + hadInvalidTimeAnchor = true + } + } + return { hadInvalidTimeAnchor } +} + +function canonicalizeChat(scale: BeamScale, chatId: string, rawChat: unknown): BeamCanonicalChat { + const batches = flattenBatches(rawChat, `BEAM ${scale}/${chatId} transcript`) + .map((batch, index) => + normalizeBatch(batch, `BEAM ${scale}/${chatId} batch[${index}]`, scale, chatId) + ) + .sort( + (left, right) => + (left.plan_number ?? 0) - (right.plan_number ?? 0) || left.batch_number - right.batch_number + ) + + const batchIdentities = new Set() + const sessions: BeamCanonicalSession[] = [] + for (const batch of batches) { + const batchIdentity = `${batch.plan_number ?? 0}:${batch.batch_number}` + if (batchIdentities.has(batchIdentity)) { + throw new Error( + `BEAM ${scale}/${chatId} has duplicate ${batch.plan_number ? `plan ${batch.plan_number} ` : ""}batch ${batch.batch_number}` + ) + } + batchIdentities.add(batchIdentity) + const date = resolveBatchDate(batch) + for (let turnIndex = 0; turnIndex < batch.turns.length; turnIndex++) { + const hasPaddedAssistant = batch.turns[turnIndex][1]?.isPaddedAssistant === true + const messages: BeamCanonicalMessage[] = batch.turns[turnIndex].map((message) => ({ + role: message.role, + content: message.content, + ...(message.time_anchor ? { timeAnchor: message.time_anchor } : {}), + })) + const sessionPrefix = batch.plan_number + ? `beam-${scale}-${chatId}-plan-${batch.plan_number}-batch-${batch.batch_number}` + : `beam-${scale}-${chatId}-batch-${batch.batch_number}` + sessions.push({ + sessionId: `${sessionPrefix}-turn-${turnIndex + 1}`, + ...(batch.plan_number ? { planNumber: batch.plan_number } : {}), + batchNumber: batch.batch_number, + turnIndex: turnIndex + 1, + ...(date.date ? { documentDate: date.date } : {}), + ...(date.hadInvalidTimeAnchor ? { hadInvalidTimeAnchor: true } : {}), + ...(hasPaddedAssistant ? { hasPaddedAssistant: true as const } : {}), + messages, + }) + } + } + + return { + schemaVersion: BEAM_CANONICAL_SCHEMA_VERSION, + scale, + chatId, + sessions, + } +} + +function normalizeQuestionCollections( + value: unknown, + context: string +): Map { + const parsed = parseJsonValue(value, context) + const grouped = new Map() + + if (Array.isArray(parsed)) { + for (let index = 0; index < parsed.length; index++) { + const question = asRecord(parsed[index], `${context}[${index}]`) + const type = normalizeQuestionType( + pickAlias( + question, + ["question_type", "questionType", "type", "ability"], + `${context}[${index}].question_type` + ), + `${context}[${index}].question_type` + ) + const entries = grouped.get(type) ?? [] + entries.push(question) + grouped.set(type, entries) + } + return grouped + } + + const record = asRecord(parsed, context) + for (const [rawType, rawQuestions] of Object.entries(record)) { + const type = normalizeQuestionType(rawType, `${context} key`) + const questions = parseJsonValue(rawQuestions, `${context}.${rawType}`) + if (!Array.isArray(questions)) { + throw new Error(`${context}.${rawType} must be an array`) + } + if (grouped.has(type)) { + throw new Error(`${context} contains duplicate aliases for ${type}`) + } + grouped.set(type, questions) + } + return grouped +} + +function getReferenceAnswer(record: Record, context: string): string | undefined { + const aliases = ["answer", "ideal_answer", "ideal_response", "ideal_summary"] + const raw = pickAlias(record, aliases, `${context}.referenceAnswer`, false) + if (raw === undefined) return undefined + return requireNonEmptyString(raw, `${context}.referenceAnswer`) +} + +function canonicalizeQuestions( + scale: BeamScale, + chatId: string, + rawQuestions: unknown +): BeamCanonicalQuestion[] { + const collections = normalizeQuestionCollections( + rawQuestions, + `BEAM ${scale}/${chatId} questions` + ) + const questions: BeamCanonicalQuestion[] = [] + + for (const questionType of BEAM_QUESTION_TYPE_IDS) { + const entries = collections.get(questionType) + if (!entries) continue + for (let index = 0; index < entries.length; index++) { + const context = `BEAM ${scale}/${chatId}/${questionType}[${index}]` + const record = asRecord(entries[index], context) + const question = requireNonEmptyString( + pickAlias(record, ["question"], `${context}.question`), + `${context}.question` + ) + const rawRubric = parseJsonValue( + pickAlias(record, ["rubric"], `${context}.rubric`), + `${context}.rubric` + ) + if (!Array.isArray(rawRubric) || rawRubric.length === 0) { + throw new Error(`${context}.rubric must be a non-empty array`) + } + const rubric = rawRubric.map((item, rubricIndex) => + requireNonEmptyString(item, `${context}.rubric[${rubricIndex}]`) + ) + const difficultyRaw = pickAlias(record, ["difficulty"], `${context}.difficulty`, false) + const difficulty = + difficultyRaw === undefined + ? undefined + : requireNonEmptyString(difficultyRaw, `${context}.difficulty`) + const referenceAnswer = getReferenceAnswer(record, context) + const contentHash = sha256Text( + stableBeamStringify({ question, rubric, referenceAnswer: referenceAnswer ?? null }) + ) + questions.push({ + schemaVersion: BEAM_CANONICAL_SCHEMA_VERSION, + scale, + chatId, + questionId: `beam:${scale}:${chatId}:${questionType}:${contentHash}`, + questionType, + question, + rubric, + ...(difficulty ? { difficulty } : {}), + ...(referenceAnswer ? { referenceAnswer } : {}), + }) + } + } + + for (const type of collections.keys()) { + if (!QUESTION_TYPE_SET.has(type)) { + throw new Error(`BEAM ${scale}/${chatId} contains unsupported question type ${type}`) + } + } + return questions.sort((left, right) => compareStrings(left.questionId, right.questionId)) +} + +function extractSourceRow( + scale: BeamScale, + rawRow: unknown, + rowIndex: number +): { chatId: string; chat: unknown; questions: unknown } { + let row = asRecord(rawRow, `BEAM ${scale} source row ${rowIndex}`) + if (row.row !== undefined && Object.keys(row).length === 1) { + row = asRecord(row.row, `BEAM ${scale} source row ${rowIndex}.row`) + } + const chatId = normalizeChatId( + pickAlias(row, ["conversation_id"], `BEAM ${scale} source row ${rowIndex} chat id`), + `BEAM ${scale} source row ${rowIndex} chat id` + ) + const chat = pickAlias(row, ["chat"], `BEAM ${scale}/${chatId} published transcript`) + + const questions = pickAlias( + row, + ["probing_questions"], + `BEAM ${scale}/${chatId} probing questions` + ) + return { chatId, chat, questions } +} + +export function canonicalizeBeamRows(scale: BeamScale, rows: unknown[]): CanonicalBeamTier { + if (!Array.isArray(rows)) throw new Error(`BEAM ${scale} source rows must be an array`) + const chats: BeamCanonicalChat[] = [] + const questions: BeamCanonicalQuestion[] = [] + + for (let rowIndex = 0; rowIndex < rows.length; rowIndex++) { + const source = extractSourceRow(scale, rows[rowIndex], rowIndex) + chats.push(canonicalizeChat(scale, source.chatId, source.chat)) + questions.push(...canonicalizeQuestions(scale, source.chatId, source.questions)) + } + + chats.sort((left, right) => compareStrings(left.chatId, right.chatId)) + questions.sort((left, right) => compareStrings(left.questionId, right.questionId)) + return validateCanonicalBeamTier(scale, chats, questions) +} + +function emptyQuestionTypeCounts(): Record { + return Object.fromEntries(BEAM_QUESTION_TYPE_IDS.map((type) => [type, 0])) as Record< + BeamQuestionType, + number + > +} + +function assertCanonicalQuestionId(question: BeamCanonicalQuestion): void { + const contentHash = sha256Text( + stableBeamStringify({ + question: question.question, + rubric: question.rubric, + referenceAnswer: question.referenceAnswer ?? null, + }) + ) + const expected = `beam:${question.scale}:${question.chatId}:${question.questionType}:${contentHash}` + if (question.questionId !== expected) { + throw new Error(`BEAM question has unstable or tampered ID: ${question.questionId}`) + } +} + +function expectedCanonicalSessionId( + scale: BeamScale, + chatId: string, + session: BeamCanonicalSession +): string { + const prefix = session.planNumber + ? `beam-${scale}-${chatId}-plan-${session.planNumber}-batch-${session.batchNumber}` + : `beam-${scale}-${chatId}-batch-${session.batchNumber}` + return `${prefix}-turn-${session.turnIndex}` +} + +export function validateCanonicalBeamTier( + scale: BeamScale, + chatsInput: BeamCanonicalChat[], + questionsInput: BeamCanonicalQuestion[] +): CanonicalBeamTier { + const chats = chatsInput.map((chat, index) => { + const parsed = canonicalChatSchema.parse(chat) + if (parsed.scale !== scale) + throw new Error(`BEAM ${scale} chat ${index} has tier ${parsed.scale}`) + return parsed + }) + const questions = questionsInput.map((question, index) => { + const parsed = canonicalQuestionSchema.parse(question) + if (parsed.scale !== scale) { + throw new Error(`BEAM ${scale} question ${index} has tier ${parsed.scale}`) + } + return parsed + }) + + const expected = BEAM_EXPECTED_COUNTS[scale] + if (chats.length !== expected.chats) { + throw new Error(`BEAM ${scale} expected ${expected.chats} chats, found ${chats.length}`) + } + if (questions.length !== expected.questions) { + throw new Error( + `BEAM ${scale} expected ${expected.questions} questions, found ${questions.length}` + ) + } + + const chatIds = new Set() + const sessionIds = new Set() + let sessionCount = 0 + let sessionsWithDocumentDate = 0 + let sessionsWithInvalidTimeAnchor = 0 + let sessionsWithPaddedAssistant = 0 + for (const chat of chats) { + if (normalizeChatId(chat.chatId, `BEAM ${scale} chat ID`) !== chat.chatId) { + throw new Error(`BEAM ${scale} chat ID is not canonical: ${chat.chatId}`) + } + if (chatIds.has(chat.chatId)) throw new Error(`BEAM ${scale} has duplicate chat ${chat.chatId}`) + chatIds.add(chat.chatId) + if (chat.sessions.length === 0) throw new Error(`BEAM ${scale}/${chat.chatId} has no sessions`) + + let previousPlan = -1 + let previousBatch = -1 + let previousTurn = 0 + for (const session of chat.sessions) { + const expectedSessionId = expectedCanonicalSessionId(scale, chat.chatId, session) + if (session.sessionId !== expectedSessionId) { + throw new Error( + `BEAM ${scale}/${chat.chatId} has non-canonical session ID ${JSON.stringify(session.sessionId)}; expected ${expectedSessionId}` + ) + } + if (sessionIds.has(session.sessionId)) { + throw new Error(`BEAM ${scale} has duplicate session ${session.sessionId}`) + } + sessionIds.add(session.sessionId) + sessionCount++ + if ( + session.messages.length !== 2 || + session.messages[0]?.role !== "user" || + session.messages[1]?.role !== "assistant" + ) { + throw new Error( + `BEAM ${scale}/${chat.chatId}/${session.sessionId} must contain exactly one user message followed by one assistant message` + ) + } + for (const message of session.messages) { + if (!message.content.trim()) { + throw new Error( + `BEAM ${scale}/${chat.chatId}/${session.sessionId} has empty message content` + ) + } + } + if (session.hasPaddedAssistant) { + if ( + scale !== "10M" || + session.messages[1]?.role !== "assistant" || + session.messages[1].content !== BEAM_MISSING_ASSISTANT_PLACEHOLDER + ) { + throw new Error( + `BEAM ${scale}/${chat.chatId}/${session.sessionId} has invalid padded-assistant metadata` + ) + } + sessionsWithPaddedAssistant++ + } + const planNumber = session.planNumber ?? 0 + if ( + planNumber < previousPlan || + (planNumber === previousPlan && session.batchNumber < previousBatch) || + (planNumber === previousPlan && + session.batchNumber === previousBatch && + session.turnIndex <= previousTurn) + ) { + throw new Error(`BEAM ${scale}/${chat.chatId} sessions are not chronologically ordered`) + } + previousPlan = planNumber + previousBatch = session.batchNumber + previousTurn = session.turnIndex + if (session.documentDate) { + if (parseBeamTimeAnchorStrict(session.documentDate) !== session.documentDate) { + throw new Error( + `BEAM ${scale}/${chat.chatId}/${session.sessionId} has invalid documentDate ${session.documentDate}` + ) + } + sessionsWithDocumentDate++ + } + if (session.hadInvalidTimeAnchor) sessionsWithInvalidTimeAnchor++ + } + } + + const questionIds = new Set() + const questionsByChat = new Map() + const byQuestionType = emptyQuestionTypeCounts() + for (const question of questions) { + if (!chatIds.has(question.chatId)) { + throw new Error( + `BEAM ${scale} question ${question.questionId} references unknown chat ${question.chatId}` + ) + } + if (questionIds.has(question.questionId)) { + throw new Error(`BEAM ${scale} has duplicate question ${question.questionId}`) + } + questionIds.add(question.questionId) + assertCanonicalQuestionId(question) + if (!question.question.trim()) { + throw new Error(`BEAM ${scale} question ${question.questionId} is empty`) + } + if (question.rubric.length === 0 || question.rubric.some((nugget) => !nugget.trim())) { + throw new Error(`BEAM ${scale} question ${question.questionId} has an empty rubric`) + } + byQuestionType[question.questionType]++ + const grouped = questionsByChat.get(question.chatId) ?? [] + grouped.push(question) + questionsByChat.set(question.chatId, grouped) + } + + for (const chat of chats) { + const chatQuestions = questionsByChat.get(chat.chatId) ?? [] + if (chatQuestions.length !== 20) { + throw new Error( + `BEAM ${scale}/${chat.chatId} expected 20 questions, found ${chatQuestions.length}` + ) + } + for (const questionType of BEAM_QUESTION_TYPE_IDS) { + const count = chatQuestions.filter( + (question) => question.questionType === questionType + ).length + if (count !== 2) { + throw new Error( + `BEAM ${scale}/${chat.chatId} expected 2 ${questionType} questions, found ${count}` + ) + } + } + } + + const byChat = Object.fromEntries( + chats.map((chat) => { + const chatQuestions = questionsByChat.get(chat.chatId) ?? [] + const chatQuestionTypes = emptyQuestionTypeCounts() + for (const question of chatQuestions) chatQuestionTypes[question.questionType]++ + return [ + chat.chatId, + { + sessions: chat.sessions.length, + questions: chatQuestions.length, + byQuestionType: chatQuestionTypes, + }, + ] + }) + ) + + const counts: BeamTierCounts = { + chats: chats.length, + questions: questions.length, + sessions: sessionCount, + sessionsWithDocumentDate, + sessionsWithoutDocumentDate: sessionCount - sessionsWithDocumentDate, + sessionsWithInvalidTimeAnchor, + sessionsWithPaddedAssistant, + byQuestionType, + byChat, + } + return { scale, chats, questions, counts } +} + +export function assertReviewedBeamPaddedAssistantCounts( + scale: BeamScale, + counts: BeamTierCounts +): void { + const expected = scale === "10M" ? BEAM_10M_PINNED_PADDED_ASSISTANT_SOURCE_IDENTITIES.length : 0 + if (counts.sessionsWithPaddedAssistant !== expected) { + throw new Error( + `Reviewed BEAM ${scale} source expected ${expected} padded missing-assistant sessions, found ${counts.sessionsWithPaddedAssistant}` + ) + } +} + +export function serializeBeamJsonl(rows: unknown[]): string { + return rows.map((row) => stableBeamStringify(row)).join("\n") + "\n" +} + +export function parseBeamJsonl(content: string, schema: z.ZodType, context: string): T[] { + const rows: T[] = [] + const lines = content.split("\n") + for (let index = 0; index < lines.length; index++) { + const line = lines[index] + if (!line) continue + try { + rows.push(schema.parse(JSON.parse(line))) + } catch (error) { + throw new Error(`${context} line ${index + 1} is malformed: ${String(error)}`) + } + } + return rows +} + +export function computeCanonicalFileManifest( + relativePath: string, + bytes: Uint8Array, + rowCount: number +): BeamCanonicalFileManifest { + return { + path: relativePath, + byteSize: bytes.byteLength, + sha256: sha256Bytes(bytes), + rowCount, + } +} + +export function datasetFingerprintPayload( + manifest: Omit +): unknown { + return { + manifestSchemaVersion: manifest.manifestSchemaVersion, + canonicalSchemaVersion: manifest.canonicalSchemaVersion, + converter: manifest.converter, + includedTiers: manifest.includedTiers, + sources: manifest.sources.map((source) => ({ + tier: source.tier, + sourceIdentity: source.sourceIdentity, + repository: source.repository, + split: source.split, + revision: source.revision, + files: source.files + .map((file) => ({ + path: file.path, + byteSize: file.byteSize, + sha256: file.sha256, + })) + .sort((left, right) => compareStrings(left.path, right.path)), + })), + canonicalFiles: manifest.canonicalFiles, + counts: manifest.counts, + orderedChatIds: manifest.orderedChatIds, + orderedChatIdsDigest: manifest.orderedChatIdsDigest, + orderedQuestionIds: manifest.orderedQuestionIds, + orderedQuestionIdsDigest: manifest.orderedQuestionIdsDigest, + } +} + +export function computeDatasetFingerprint( + manifest: Omit +): string { + return sha256Text(stableBeamStringify(datasetFingerprintPayload(manifest))) +} + +export function computeManifestHash(manifest: Omit): string { + return sha256Text(stableBeamStringify(manifest)) +} + +function assertManifestSourcePins( + manifest: BeamDatasetManifest, + tiers: BeamScale[], + allowTestSourceIdentity: boolean +): void { + for (const tier of tiers) { + const expected = BEAM_DATASET_SOURCES[tier] + const matchingSources = manifest.sources.filter((entry) => entry.tier === tier) + if (matchingSources.length !== 1) { + throw new Error( + `BEAM snapshot must include exactly one source identity for ${tier}; found ${matchingSources.length}` + ) + } + const source = matchingSources[0]! + if ( + source.repository !== expected.repository || + source.split !== expected.split || + source.revision !== expected.revision + ) { + throw new Error(`BEAM ${tier} source pin does not match the reviewed official revision`) + } + const expectedPaths = expected.parquetFiles.map((file) => file.path).sort(compareStrings) + const actualPaths = source.files.map((file) => file.path).sort(compareStrings) + if (stableBeamStringify(expectedPaths) !== stableBeamStringify(actualPaths)) { + throw new Error(`BEAM ${tier} source file set does not match the reviewed source descriptor`) + } + + if (source.sourceIdentity === "injected-test-fixture") { + if (!allowTestSourceIdentity) { + throw new Error( + `BEAM ${tier} snapshot uses an injected test-source identity and cannot be used for a scored run` + ) + } + continue + } + + for (const expectedFile of expected.parquetFiles) { + if (!expectedFile.expectedSha256) { + throw new Error( + `BEAM ${tier} reviewed source ${expectedFile.path} is missing a SHA-256 pin` + ) + } + const actualFile = source.files.find((file) => file.path === expectedFile.path) + if (!actualFile || actualFile.sha256 !== expectedFile.expectedSha256) { + throw new Error( + `BEAM ${tier} source ${expectedFile.path} SHA-256 does not match the reviewed published pin` + ) + } + if (actualFile.url !== expectedFile.url) { + throw new Error( + `BEAM ${tier} source ${expectedFile.path} URL does not match the reviewed pin` + ) + } + } + } +} + +function assertRunningConverterIdentity(manifest: BeamDatasetManifest): void { + if (stableBeamStringify(manifest.converter) !== stableBeamStringify(BEAM_CONVERTER_IDENTITY)) { + throw new Error( + `BEAM snapshot converter identity does not match this runtime; prepare a new snapshot with converter ${BEAM_CONVERTER_IDENTITY.version}` + ) + } +} + +function resolveSnapshotFile(snapshotPath: string, filePath: string): string { + const absolutePath = resolve(snapshotPath, filePath) + const snapshotRoot = resolve(snapshotPath) + const relativePath = relative(snapshotRoot, absolutePath) + if (relativePath.startsWith("..") || isAbsolute(relativePath)) { + throw new Error(`BEAM manifest contains an unsafe snapshot path: ${filePath}`) + } + return absolutePath +} + +function verifyFile(snapshotPath: string, file: BeamCanonicalFileManifest): Uint8Array { + const absolutePath = resolveSnapshotFile(snapshotPath, file.path) + if (!existsSync(absolutePath)) throw new Error(`BEAM canonical file is missing: ${file.path}`) + const stat = statSync(absolutePath) + if (!stat.isFile() || stat.size !== file.byteSize) { + throw new Error(`BEAM canonical file size mismatch: ${file.path}`) + } + const bytes = readFileSync(absolutePath) + if (sha256Bytes(bytes) !== file.sha256) { + throw new Error(`BEAM canonical file hash mismatch: ${file.path}`) + } + return bytes +} + +async function hashFile(filePath: string): Promise { + const hash = createHash("sha256") + for await (const chunk of createReadStream(filePath)) hash.update(chunk) + return hash.digest("hex") +} + +async function verifySourceFiles( + snapshotPath: string, + manifest: BeamDatasetManifest +): Promise { + for (const source of manifest.sources) { + for (const file of source.files) { + const absolutePath = resolveSnapshotFile(snapshotPath, file.snapshotPath) + if (!existsSync(absolutePath)) { + throw new Error(`BEAM source file is missing: ${file.snapshotPath}`) + } + const fileStat = statSync(absolutePath) + if (!fileStat.isFile() || fileStat.size !== file.byteSize) { + throw new Error(`BEAM source file size mismatch: ${file.snapshotPath}`) + } + if ((await hashFile(absolutePath)) !== file.sha256) { + throw new Error(`BEAM source file hash mismatch: ${file.snapshotPath}`) + } + } + } +} + +function assertCompleteMarker(snapshotPath: string, manifest: BeamDatasetManifest): void { + const completePath = join(snapshotPath, ".complete") + if (!existsSync(completePath)) { + throw new Error(`BEAM snapshot is incomplete at ${snapshotPath}; run the BEAM prepare command`) + } + let marker: unknown + try { + marker = JSON.parse(readFileSync(completePath, "utf8")) + } catch (error) { + throw new Error(`BEAM .complete marker is malformed: ${String(error)}`) + } + const record = asRecord(marker, "BEAM .complete marker") + if ( + record.datasetFingerprint !== manifest.datasetFingerprint || + record.manifestHash !== manifest.manifestHash + ) { + throw new Error("BEAM .complete marker does not match manifest identity") + } +} + +/** + * Validate every source, canonical file, manifest identity, and dataset invariant. + * Preparation uses this before it writes `.complete`; scored runs must call + * `loadPreparedBeamDataset`, which additionally requires the completion marker. + */ +export async function validatePreparedBeamSnapshotContents( + options: ValidatePreparedBeamSnapshotContentsOptions +): Promise { + const snapshotPath = resolve(options.snapshotPath) + const manifestPath = join(snapshotPath, "manifest.json") + if (!existsSync(manifestPath)) { + throw new Error(`BEAM manifest not found at ${manifestPath}; run the BEAM prepare command`) + } + + let rawManifest: unknown + try { + rawManifest = JSON.parse(readFileSync(manifestPath, "utf8")) + } catch (error) { + throw new Error(`BEAM manifest is malformed: ${String(error)}`) + } + const manifest = manifestSchema.parse(rawManifest) as BeamDatasetManifest + const withoutHash = { ...manifest } as Partial + delete withoutHash.manifestHash + const actualManifestHash = computeManifestHash( + withoutHash as Omit + ) + if (actualManifestHash !== manifest.manifestHash) { + throw new Error("BEAM manifest hash mismatch") + } + + const withoutIdentity = { ...withoutHash } as Partial + delete withoutIdentity.datasetFingerprint + const actualFingerprint = computeDatasetFingerprint( + withoutIdentity as Omit + ) + if (actualFingerprint !== manifest.datasetFingerprint) { + throw new Error("BEAM dataset fingerprint mismatch") + } + if ( + options.expectedDatasetFingerprint && + options.expectedDatasetFingerprint !== manifest.datasetFingerprint + ) { + throw new Error( + `BEAM dataset revision mismatch: expected ${options.expectedDatasetFingerprint}, found ${manifest.datasetFingerprint}` + ) + } + const manifestTiers = [...new Set(manifest.includedTiers)] + if ( + manifestTiers.length !== manifest.includedTiers.length || + manifest.sources.length !== manifestTiers.length + ) { + throw new Error("BEAM manifest tier/source identities must be unique and one-to-one") + } + assertRunningConverterIdentity(manifest) + assertManifestSourcePins(manifest, manifestTiers, options.allowInjectedTestSourceIdentity) + + const expectedTierKeys = [...manifestTiers].sort(compareStrings) + for (const [name, record] of Object.entries({ + counts: manifest.counts, + orderedChatIds: manifest.orderedChatIds, + orderedChatIdsDigest: manifest.orderedChatIdsDigest, + orderedQuestionIds: manifest.orderedQuestionIds, + orderedQuestionIdsDigest: manifest.orderedQuestionIdsDigest, + })) { + if ( + stableBeamStringify(Object.keys(record).sort(compareStrings)) !== + stableBeamStringify(expectedTierKeys) + ) { + throw new Error(`BEAM manifest ${name} tier keys do not match includedTiers`) + } + } + + const expectedCanonicalPaths = manifestTiers + .flatMap((tier) => [`canonical/${tier}/chats.jsonl`, `canonical/${tier}/questions.jsonl`]) + .sort(compareStrings) + const actualCanonicalPaths = manifest.canonicalFiles.map((file) => file.path).sort(compareStrings) + if ( + new Set(actualCanonicalPaths).size !== actualCanonicalPaths.length || + stableBeamStringify(actualCanonicalPaths) !== stableBeamStringify(expectedCanonicalPaths) + ) { + throw new Error( + "BEAM manifest canonical file set must contain exactly one chats and questions file per tier" + ) + } + + const tiers = [...new Set(options.tiers)] + for (const tier of tiers) { + if (!manifest.includedTiers.includes(tier)) { + throw new Error(`BEAM snapshot ${manifest.datasetFingerprint} does not include tier ${tier}`) + } + } + await verifySourceFiles(snapshotPath, manifest) + + const chatsByTier: Partial> = {} + const questionsByTier: Partial> = {} + for (const tier of tiers) { + const chatsPath = `canonical/${tier}/chats.jsonl` + const questionsPath = `canonical/${tier}/questions.jsonl` + const chatsFile = manifest.canonicalFiles.find((file) => file.path === chatsPath) + const questionsFile = manifest.canonicalFiles.find((file) => file.path === questionsPath) + if (!chatsFile || !questionsFile) { + throw new Error(`BEAM ${tier} canonical file entries are missing from manifest`) + } + const chatsBytes = verifyFile(snapshotPath, chatsFile) + const questionsBytes = verifyFile(snapshotPath, questionsFile) + const chats = parseBeamJsonl( + Buffer.from(chatsBytes).toString("utf8"), + canonicalChatSchema, + `BEAM ${tier} chats` + ) + const questions = parseBeamJsonl( + Buffer.from(questionsBytes).toString("utf8"), + canonicalQuestionSchema, + `BEAM ${tier} questions` + ) + if (chats.length !== chatsFile.rowCount || questions.length !== questionsFile.rowCount) { + throw new Error(`BEAM ${tier} canonical row count does not match manifest`) + } + const validated = validateCanonicalBeamTier(tier, chats, questions) + if (stableBeamStringify(validated.counts) !== stableBeamStringify(manifest.counts[tier])) { + throw new Error(`BEAM ${tier} validated counts do not match manifest`) + } + const source = manifest.sources.find((entry) => entry.tier === tier) + if (source?.sourceIdentity === "reviewed-published") { + assertReviewedBeamPaddedAssistantCounts(tier, validated.counts) + } + const orderedChatIds = validated.chats.map((chat) => chat.chatId) + if ( + stableBeamStringify(orderedChatIds) !== stableBeamStringify(manifest.orderedChatIds[tier]) + ) { + throw new Error(`BEAM ${tier} ordered chat identity does not match manifest`) + } + const chatIdDigest = sha256Text(orderedChatIds.join("\n")) + if (chatIdDigest !== manifest.orderedChatIdsDigest[tier]) { + throw new Error(`BEAM ${tier} chat identity digest does not match manifest`) + } + const orderedQuestionIds = validated.questions.map((question) => question.questionId) + if ( + stableBeamStringify(orderedQuestionIds) !== + stableBeamStringify(manifest.orderedQuestionIds[tier]) + ) { + throw new Error(`BEAM ${tier} ordered question identity does not match manifest`) + } + const questionIdDigest = sha256Text(orderedQuestionIds.join("\n")) + if (questionIdDigest !== manifest.orderedQuestionIdsDigest[tier]) { + throw new Error(`BEAM ${tier} question identity digest does not match manifest`) + } + chatsByTier[tier] = validated.chats + questionsByTier[tier] = validated.questions + } + + return { snapshotPath, manifest, chatsByTier, questionsByTier } +} + +/** + * Re-run the canonical converter over authenticated source Parquet rows and + * compare exact deterministic JSONL bytes. This couples source and canonical + * files instead of trusting a self-asserted manifest relationship. + */ +export async function verifyPreparedBeamSourceDerivation( + prepared: PreparedBeamDataset, + tiers: BeamScale[], + decodeParquet: BeamParquetDecoder = decodeBeamParquetWithHyparquet +): Promise { + for (const tier of [...new Set(tiers)]) { + const source = prepared.manifest.sources.find((entry) => entry.tier === tier) + if (!source) throw new Error(`BEAM ${tier} source identity is missing during derivation check`) + const descriptor = BEAM_DATASET_SOURCES[tier] + const sourceRows: unknown[] = [] + for (const expectedFile of descriptor.parquetFiles) { + const sourceFile = source.files.find((file) => file.path === expectedFile.path) + if (!sourceFile) { + throw new Error( + `BEAM ${tier} source ${expectedFile.path} is missing during derivation check` + ) + } + const rows = await decodeParquet( + resolveSnapshotFile(prepared.snapshotPath, sourceFile.snapshotPath), + tier + ) + if (!Array.isArray(rows)) { + throw new Error(`BEAM ${tier} derivation decoder did not return an array of rows`) + } + sourceRows.push(...rows) + } + + const canonical = canonicalizeBeamRows(tier, sourceRows) + if (source.sourceIdentity === "reviewed-published") { + assertReviewedBeamPaddedAssistantCounts(tier, canonical.counts) + } + for (const [kind, expectedContent] of [ + ["chats", serializeBeamJsonl(canonical.chats)], + ["questions", serializeBeamJsonl(canonical.questions)], + ] as const) { + const relativePath = `canonical/${tier}/${kind}.jsonl` + const actualContent = readFileSync( + resolveSnapshotFile(prepared.snapshotPath, relativePath), + "utf8" + ) + if (actualContent !== expectedContent) { + throw new Error( + `BEAM ${tier} ${kind} source-to-canonical derivation mismatch; prepare a new snapshot from the pinned source` + ) + } + } + } +} + +const verifiedPublishedDerivations = new Set() + +export function getUnverifiedBeamDerivationTiers( + datasetFingerprint: string, + tiers: BeamScale[], + verifiedDerivations: ReadonlySet = verifiedPublishedDerivations +): BeamScale[] { + return [...new Set(tiers)].filter( + (tier) => !verifiedDerivations.has(`${datasetFingerprint}:${tier}`) + ) +} + +export async function loadPreparedBeamDataset( + options: LoadPreparedBeamDatasetOptions +): Promise { + const prepared = await validatePreparedBeamSnapshotContents({ + ...options, + allowInjectedTestSourceIdentity: false, + }) + assertCompleteMarker(prepared.snapshotPath, prepared.manifest) + const unverifiedTiers = getUnverifiedBeamDerivationTiers( + prepared.manifest.datasetFingerprint, + options.tiers + ) + for (const tier of unverifiedTiers) { + await verifyPreparedBeamSourceDerivation(prepared, [tier]) + verifiedPublishedDerivations.add(`${prepared.manifest.datasetFingerprint}:${tier}`) + } + return prepared +} + +/** Explicit fixture-only loader; scored benchmark code must never call this. */ +export async function loadPreparedBeamTestFixture( + options: LoadPreparedBeamDatasetOptions +): Promise { + const prepared = await validatePreparedBeamSnapshotContents({ + ...options, + allowInjectedTestSourceIdentity: true, + }) + assertCompleteMarker(prepared.snapshotPath, prepared.manifest) + return prepared +} + +export function resolvePreparedSnapshotPath(dataPath: string, datasetRevision?: string): string { + const fullPath = resolve(dataPath) + if (existsSync(join(fullPath, "manifest.json"))) return fullPath + if (datasetRevision) return join(fullPath, datasetRevision) + + throw new Error( + `BEAM data path ${fullPath} is not a prepared snapshot. Pass --dataset-revision or run the BEAM prepare command.` + ) +} + +export function describeBeamSnapshot(manifest: BeamDatasetManifest): string { + return `${manifest.includedTiers.join("+")} @ ${manifest.datasetFingerprint.slice(0, 12)}` +} + +export function describeBeamTemporalCoverage( + counts: BeamDatasetManifest["counts"], + tiers: BeamScale[] +): string { + const tierSummaries = [...new Set(tiers)].map((tier) => { + const tierCounts = counts[tier] + if (!tierCounts) throw new Error(`BEAM temporal counts are missing tier ${tier}`) + return `${tier}: ${tierCounts.sessionsWithoutDocumentDate}/${tierCounts.sessions} sessions without a valid date; ${tierCounts.sessionsWithInvalidTimeAnchor}/${tierCounts.sessions} encountered an invalid source time anchor; ${tierCounts.sessionsWithPaddedAssistant}/${tierCounts.sessions} use an audited N/A assistant padding` + }) + return `BEAM temporal/source coverage — ${tierSummaries.join(" | ")}. Invalid-anchor counts are separate, not additive: a session may have a valid fallback date after an invalid anchor. Padded-assistant counts identify the two pinned 10M source follow-ups whose responses are absent.` +} + +export function isPreparedBeamSnapshot(path: string): boolean { + return existsSync(join(path, "manifest.json")) && existsSync(join(path, ".complete")) +} + +export function snapshotDirectoryName(manifest: BeamDatasetManifest): string { + return basename(manifest.datasetFingerprint) +} + +export const BEAM_DATASET_MANIFEST_SCHEMA = manifestSchema +export const BEAM_CANONICAL_CHAT_SCHEMA = canonicalChatSchema +export const BEAM_CANONICAL_QUESTION_SCHEMA = canonicalQuestionSchema +export const BEAM_CONVERTER_IDENTITY = { + name: "memorybench-beam" as const, + version: BEAM_CONVERTER_VERSION, + implementationHash: BEAM_CONVERTER_IMPLEMENTATION_HASH, +} diff --git a/src/benchmarks/beam/index.ts b/src/benchmarks/beam/index.ts index 6fe6854..0f112ae 100644 --- a/src/benchmarks/beam/index.ts +++ b/src/benchmarks/beam/index.ts @@ -1,17 +1,36 @@ -import { existsSync, readFileSync, readdirSync } from "fs" -import { join } from "path" -import type { Benchmark, BenchmarkConfig, QuestionFilter } from "../../types/benchmark" +import type { + Benchmark, + BenchmarkConfig, + BenchmarkScope, + DatasetIdentity, + QuestionFilter, +} from "../../types/benchmark" +import type { BenchmarkProtocol } from "../../types/protocol" import type { QuestionTypeRegistry, UnifiedMessage, UnifiedQuestion, UnifiedSession, } from "../../types/unified" +import { BeamPaperProtocol } from "../../protocols/beam-paper" +import { BEAM_MEM0_NUGGET_PROFILE, BeamMem0NuggetProtocol } from "../../protocols/beam-mem0" import { logger } from "../../utils/logger" -import { formatBeamDate, parseBeamTimeAnchor } from "../../prompts/beam" -import type { BeamBatch, BeamChatFile, BeamProbingQuestionsFile, BeamScale } from "./types" +import { + computeDatasetFingerprint, + computeManifestHash, + describeBeamSnapshot, + describeBeamTemporalCoverage, + loadPreparedBeamDataset, + resolvePreparedSnapshotPath, +} from "./dataset" +import type { + BeamCanonicalChat, + BeamCanonicalQuestion, + BeamDatasetManifest, + BeamScale, +} from "./types" -const DEFAULT_DATA_PATH = "./data/benchmarks/beam/chats" +const DEFAULT_DATA_PATH = "./data/benchmarks/beam" export const BEAM_QUESTION_TYPES: QuestionTypeRegistry = { abstention: { @@ -66,237 +85,224 @@ export const BEAM_QUESTION_TYPES: QuestionTypeRegistry = { }, } -function flattenChatFile(chatFile: BeamChatFile): BeamBatch[] { - if (Array.isArray(chatFile)) { - return chatFile.flatMap((entry) => { - if (isBeamBatch(entry)) return [entry] - return flattenChatFile(entry) +function selectTierRecord( + record: Partial>, + scales: readonly BeamScale[] +): Partial> { + return Object.fromEntries( + scales.map((scale) => { + const value = record[scale] + if (value === undefined) throw new Error(`BEAM manifest identity is missing tier ${scale}`) + return [scale, value] }) - } - - return Object.keys(chatFile) - .sort((a, b) => a.localeCompare(b, undefined, { numeric: true })) - .flatMap((key) => chatFile[key] || []) -} - -function isBeamBatch(value: unknown): value is BeamBatch { - return ( - typeof value === "object" && - value !== null && - "batch_number" in value && - "turns" in value && - Array.isArray((value as BeamBatch).turns) ) } -function createGroundTruth(question: unknown): string { - if (typeof question === "object" && question !== null) { - const record = question as Record - const answer = getQuestionAnswer(record) - if (answer) return answer - - // Fall back to the rubric so retrieval-eval gets a useful expected-answer - // signal for types like instruction_following/preference_following that - // describe expected behavior via rubric items instead of a single answer. - const rubric = record.rubric - if (Array.isArray(rubric) && rubric.every((item) => typeof item === "string")) { - return rubric.join("\n") - } - - return JSON.stringify(question) +export function createBeamDatasetIdentity( + manifest: BeamDatasetManifest, + scales: readonly BeamScale[] +): DatasetIdentity { + const selectedScales = [...scales] + const selectedScaleSet = new Set(selectedScales) + const sources = manifest.sources.filter((source) => selectedScaleSet.has(source.tier)) + const canonicalFiles = manifest.canonicalFiles.filter((file) => + selectedScales.some((scale) => file.path.startsWith(`canonical/${scale}/`)) + ) + const effectiveManifestCore: Omit = { + manifestSchemaVersion: manifest.manifestSchemaVersion, + canonicalSchemaVersion: manifest.canonicalSchemaVersion, + converter: manifest.converter, + includedTiers: selectedScales, + sources, + canonicalFiles, + counts: selectTierRecord(manifest.counts, selectedScales), + orderedChatIds: selectTierRecord(manifest.orderedChatIds, selectedScales), + orderedChatIdsDigest: selectTierRecord(manifest.orderedChatIdsDigest, selectedScales), + orderedQuestionIds: selectTierRecord(manifest.orderedQuestionIds, selectedScales), + orderedQuestionIdsDigest: selectTierRecord(manifest.orderedQuestionIdsDigest, selectedScales), + } + const datasetFingerprint = computeDatasetFingerprint(effectiveManifestCore) + const manifestHash = computeManifestHash({ ...effectiveManifestCore, datasetFingerprint }) + return { + datasetFingerprint, + manifestHash, + snapshotFingerprint: manifest.datasetFingerprint, + snapshotManifestHash: manifest.manifestHash, + manifestSchemaVersion: manifest.manifestSchemaVersion, + canonicalSchemaVersion: manifest.canonicalSchemaVersion, + converterVersion: manifest.converter.version, + converterImplementationHash: manifest.converter.implementationHash, + includedTiers: selectedScales, + counts: effectiveManifestCore.counts as Record, + orderedQuestionIdsDigest: effectiveManifestCore.orderedQuestionIdsDigest as Record< + string, + string + >, + sourceFiles: sources.flatMap((source) => + source.files.map((file) => ({ + path: `${source.tier}/${file.path}`, + byteSize: file.byteSize, + sha256: file.sha256, + })) + ), + canonicalFiles: canonicalFiles.map((file) => ({ + path: file.path, + byteSize: file.byteSize, + sha256: file.sha256, + })), + sources: sources.map((source) => ({ + repository: source.repository, + split: source.split, + revision: source.revision, + sourceIdentity: source.sourceIdentity, + })), } +} - return JSON.stringify(question) +function createSessions(chat: BeamCanonicalChat): UnifiedSession[] { + return chat.sessions.map((session) => ({ + sessionId: session.sessionId, + messages: session.messages.map( + (message): UnifiedMessage => ({ + role: message.role, + content: message.content, + speaker: message.role, + timestamp: message.timeAnchor, + }) + ), + metadata: { + scale: chat.scale, + chatId: chat.chatId, + ...(session.planNumber ? { planNumber: session.planNumber } : {}), + batchNumber: session.batchNumber, + turnIndex: session.turnIndex, + ...(session.documentDate + ? { date: session.documentDate, documentDate: session.documentDate } + : {}), + ...(session.hadInvalidTimeAnchor ? { hadInvalidTimeAnchor: true } : {}), + ...(session.hasPaddedAssistant ? { hasPaddedAssistant: true } : {}), + }, + })) } -function getQuestionAnswer(question: Record): string | undefined { - const answer = - question.answer || question.ideal_answer || question.ideal_response || question.ideal_summary - return typeof answer === "string" ? answer : undefined +function groundTruth(question: BeamCanonicalQuestion): string { + return question.referenceAnswer || question.rubric.join("\n") } export class BeamBenchmark implements Benchmark { - name: string - private scales: BeamScale[] + readonly name: string + readonly scope: BenchmarkScope + protocol: BenchmarkProtocol + private readonly scales: BeamScale[] private questions: UnifiedQuestion[] = [] - private sessionsMap: Map = new Map() - private ingestionGroupMap: Map = new Map() - private dataPath: string = "" + private sessionsByQuestion = new Map() + private ingestionGroupByQuestion = new Map() + private datasetIdentity?: DatasetIdentity - constructor(scales: BeamScale[] = ["1M", "10M"], name = "beam") { + constructor(scales: BeamScale[], name: string) { this.scales = scales this.name = name - } - - async load(config?: BenchmarkConfig): Promise { - this.dataPath = config?.dataPath || DEFAULT_DATA_PATH - const fullPath = join(process.cwd(), this.dataPath) - - if (!existsSync(fullPath)) { - throw new Error( - `BEAM dataset not found at ${fullPath}. Expected chats under ${DEFAULT_DATA_PATH}/{1M,10M}.` - ) + this.scope = { + displayName: scales.length === 1 ? `BEAM ${scales[0]}` : `BEAM ${scales.join("/")}`, + includedTiers: [...scales], + coverage: "subset", } - - for (const scale of this.scales) { - this.loadScale(fullPath, scale) - } - - logger.info( - `Loaded ${this.questions.length} questions from BEAM (${this.scales.join(", ")})` - ) + this.protocol = new BeamPaperProtocol() } - private loadScale(basePath: string, scale: BeamScale): void { - const scalePath = join(basePath, scale) - if (!existsSync(scalePath)) { - throw new Error(`BEAM ${scale} dataset not found at ${scalePath}`) - } - - const chatDirs = readdirSync(scalePath, { withFileTypes: true }) - .filter((entry) => entry.isDirectory() && /^\d+$/.test(entry.name)) - .map((entry) => entry.name) - .sort((a, b) => Number(a) - Number(b)) - - for (const chatId of chatDirs) { - this.loadChat(scalePath, scale, chatId) - } - } - - private loadChat(scalePath: string, scale: BeamScale, chatId: string): void { - const chatDir = join(scalePath, chatId) - const truncatedPath = join(chatDir, "chat_trunecated.json") - const fullChatPath = join(chatDir, "chat.json") - const chatPath = existsSync(truncatedPath) ? truncatedPath : fullChatPath - const probingPath = join(chatDir, "probing_questions", "probing_questions.json") - - if (!existsSync(chatPath) || !existsSync(probingPath)) { - logger.warn(`Skipping BEAM ${scale}/${chatId}: missing chat or probing questions`) - return + async load(config: BenchmarkConfig = {}): Promise { + if (config.evaluationProfile === BEAM_MEM0_NUGGET_PROFILE) { + this.protocol = new BeamMem0NuggetProtocol({ + retrievalTopK: config.retrievalTopK, + answerCutoff: config.answerCutoff, + }) + } else { + if (config.evaluationProfile) { + throw new Error(`Unsupported BEAM evaluation profile: ${config.evaluationProfile}`) + } + if (config.answerCutoff !== undefined) { + throw new Error("--answer-cutoff is only valid with --evaluation-profile mem0-nugget") + } + this.protocol = new BeamPaperProtocol({ retrievalTopK: config.retrievalTopK }) } + this.questions = [] + this.sessionsByQuestion.clear() + this.ingestionGroupByQuestion.clear() + const dataPath = config.dataPath || DEFAULT_DATA_PATH + const snapshotPath = resolvePreparedSnapshotPath(dataPath, config.datasetRevision) + const prepared = await loadPreparedBeamDataset({ + snapshotPath, + tiers: this.scales, + expectedDatasetFingerprint: config.datasetRevision, + }) + this.datasetIdentity = createBeamDatasetIdentity(prepared.manifest, this.scales) + logger.info(describeBeamTemporalCoverage(prepared.manifest.counts, this.scales)) - const batches = flattenChatFile(JSON.parse(readFileSync(chatPath, "utf8")) as BeamChatFile) - const sessions = this.extractSessions(scale, chatId, batches) - const probingQuestions = JSON.parse( - readFileSync(probingPath, "utf8") - ) as BeamProbingQuestionsFile - const sessionIds = sessions.map((session) => session.sessionId) - const ingestionGroupId = `beam-${scale}-${chatId}` - - for (const questionType of Object.keys(probingQuestions)) { - const questionsForType = probingQuestions[questionType] || [] - - for (let i = 0; i < questionsForType.length; i++) { - const probingQuestion = questionsForType[i] - const questionId = `${ingestionGroupId}-${questionType}-${i}` - const answer = getQuestionAnswer(probingQuestion) - - this.questions.push({ - questionId, - question: probingQuestion.question, - questionType, - groundTruth: createGroundTruth(probingQuestion), + for (const scale of this.scales) { + const chats = prepared.chatsByTier[scale] + const questions = prepared.questionsByTier[scale] + if (!chats || !questions) throw new Error(`Prepared BEAM snapshot is missing ${scale}`) + const sessionsByChat = new Map( + chats.map((chat) => [chat.chatId, createSessions(chat)] as const) + ) + for (const sourceQuestion of questions) { + const sessions = sessionsByChat.get(sourceQuestion.chatId) + if (!sessions) { + throw new Error( + `BEAM ${scale} question ${sourceQuestion.questionId} references missing chat ${sourceQuestion.chatId}` + ) + } + const sessionIds = sessions.map((session) => session.sessionId) + const ingestionGroupId = `beam-${scale}-${sourceQuestion.chatId}` + const question: UnifiedQuestion = { + questionId: sourceQuestion.questionId, + question: sourceQuestion.question, + questionType: sourceQuestion.questionType, + groundTruth: groundTruth(sourceQuestion), haystackSessionIds: sessionIds, metadata: { scale, - chatId, + chatId: sourceQuestion.chatId, ingestionGroupId, - rubric: probingQuestion.rubric, - difficulty: probingQuestion.difficulty, - answer, + rubric: sourceQuestion.rubric, + difficulty: sourceQuestion.difficulty, + referenceAnswer: sourceQuestion.referenceAnswer, }, - }) - - this.sessionsMap.set(questionId, sessions) - this.ingestionGroupMap.set(questionId, ingestionGroupId) - } - } - } - - private extractSessions(scale: BeamScale, chatId: string, batches: BeamBatch[]): UnifiedSession[] { - const sessions: UnifiedSession[] = [] - - for (const batch of batches) { - // mem0's `get_time_anchor_epoch` finds the earliest non-null `time_anchor` - // across all messages in a batch and tags every memory derived from that - // batch with it. Most turns in BEAM don't carry their own anchor, so - // hoisting the batch-level anchor here gives dates to every session in - // the batch (matching mem0's per-memory dating). - let batchTimeAnchor: string | null | undefined = batch.time_anchor ?? null - if (!batchTimeAnchor) { - for (const turn of batch.turns) { - const msgWithAnchor = turn.find((m) => m.time_anchor) - if (msgWithAnchor?.time_anchor) { - batchTimeAnchor = msgWithAnchor.time_anchor - break - } } - } - const batchDateIso = parseBeamTimeAnchor(batchTimeAnchor) - const batchDateFormatted = batchDateIso ? formatBeamDate(batchDateIso) : undefined - - for (let turnIndex = 0; turnIndex < batch.turns.length; turnIndex++) { - const turn = batch.turns[turnIndex] - const messages = turn - .filter((message) => message.content) - .map( - (message): UnifiedMessage => ({ - role: message.role, - content: message.content, - speaker: message.role, - timestamp: message.time_anchor, - }) - ) - - if (messages.length === 0) continue - - sessions.push({ - sessionId: `beam-${scale}-${chatId}-batch-${batch.batch_number}-turn-${turnIndex + 1}`, - messages, - metadata: { - scale, - chatId, - batchNumber: batch.batch_number, - turnIndex: turnIndex + 1, - // Match LocoMo / LongMemEval: `date` (ISO) + `formattedDate` - // (readable). The Supermemory provider reads these fields to - // (a) attach `metadata.date` to the document and (b) prefix the - // ingested content with a natural-language date sentence. - date: batchDateIso, - formattedDate: batchDateFormatted, - }, - }) + this.protocol.validateQuestion(question) + this.questions.push(question) + this.sessionsByQuestion.set(question.questionId, sessions) + this.ingestionGroupByQuestion.set(question.questionId, ingestionGroupId) } } - return sessions + logger.info( + `Loaded ${this.questions.length} validated questions from ${describeBeamSnapshot(prepared.manifest)} (${this.scales.join(", ")})` + ) } getQuestions(filter?: QuestionFilter): UnifiedQuestion[] { - let result = [...this.questions] - + let questions = [...this.questions] if (filter?.questionTypes?.length) { - result = result.filter((q) => filter.questionTypes!.includes(q.questionType)) - } - - if (filter?.offset) { - result = result.slice(filter.offset) - } - - if (filter?.limit) { - result = result.slice(0, filter.limit) + questions = questions.filter((question) => + filter.questionTypes!.includes(question.questionType) + ) } - - return result + if (filter?.offset != null) questions = questions.slice(filter.offset) + if (filter?.limit != null) questions = questions.slice(0, filter.limit) + return questions } getHaystackSessions(questionId: string): UnifiedSession[] { - return this.sessionsMap.get(questionId) || [] + const sessions = this.sessionsByQuestion.get(questionId) + if (!sessions) throw new Error(`Unknown BEAM question: ${questionId}`) + return sessions } getGroundTruth(questionId: string): string { - const question = this.questions.find((q) => q.questionId === questionId) - return question?.groundTruth || "" + const question = this.questions.find((candidate) => candidate.questionId === questionId) + if (!question) throw new Error(`Unknown BEAM question: ${questionId}`) + return question.groundTruth } getQuestionTypes(): QuestionTypeRegistry { @@ -304,7 +310,13 @@ export class BeamBenchmark implements Benchmark { } getIngestionGroupId(questionId: string): string { - return this.ingestionGroupMap.get(questionId) || questionId + const groupId = this.ingestionGroupByQuestion.get(questionId) + if (!groupId) throw new Error(`Unknown BEAM question: ${questionId}`) + return groupId + } + + getDatasetIdentity(): DatasetIdentity | undefined { + return this.datasetIdentity } } @@ -320,4 +332,10 @@ export class Beam10MBenchmark extends BeamBenchmark { } } +export class Beam1M10MBenchmark extends BeamBenchmark { + constructor() { + super(["1M", "10M"], "beam-1m-10m") + } +} + export default BeamBenchmark diff --git a/src/benchmarks/beam/parquet.ts b/src/benchmarks/beam/parquet.ts new file mode 100644 index 0000000..34cbf30 --- /dev/null +++ b/src/benchmarks/beam/parquet.ts @@ -0,0 +1,64 @@ +import { open, stat } from "node:fs/promises" +import type { BeamScale } from "./types" + +export type BeamParquetDecoder = (filePath: string, tier: BeamScale) => Promise + +/** + * Decode an on-disk Parquet file through range reads so preparation and + * scored-run provenance verification execute the same decoder. + */ +export async function decodeBeamParquetWithHyparquet( + filePath: string, + _tier: BeamScale +): Promise { + const hyparquetModuleName = "hyparquet" + const compressorsModuleName = "hyparquet-compressors" + let hyparquet: Record + let compressorModule: Record + try { + hyparquet = (await import(hyparquetModuleName)) as Record + compressorModule = (await import(compressorsModuleName)) as Record + } catch (error) { + throw new Error( + `BEAM preparation and provenance verification require hyparquet and hyparquet-compressors: ${String(error)}` + ) + } + + const parquetReadObjects = hyparquet.parquetReadObjects + if (typeof parquetReadObjects !== "function") { + throw new Error("Installed hyparquet package does not export parquetReadObjects") + } + const compressors = compressorModule.compressors ?? compressorModule.default + if (!compressors || typeof compressors !== "object") { + throw new Error("Installed hyparquet-compressors package does not export compressors") + } + + const fileStat = await stat(filePath) + const handle = await open(filePath, "r") + const asyncBuffer = { + byteLength: fileStat.size, + slice: async (start: number, end: number): Promise => { + if (!Number.isInteger(start) || !Number.isInteger(end) || start < 0 || end < start) { + throw new Error(`Invalid Parquet byte range: ${start}-${end}`) + } + const length = Math.min(end, fileStat.size) - start + const target = Buffer.alloc(Math.max(0, length)) + const { bytesRead } = await handle.read(target, 0, target.byteLength, start) + if (bytesRead !== target.byteLength) { + throw new Error(`Short Parquet read at ${start}-${end}: got ${bytesRead} bytes`) + } + return target.buffer.slice(target.byteOffset, target.byteOffset + target.byteLength) + }, + } + + try { + return (await ( + parquetReadObjects as (options: { + file: typeof asyncBuffer + compressors: unknown + }) => Promise + )({ file: asyncBuffer, compressors })) as unknown[] + } finally { + await handle.close() + } +} diff --git a/src/benchmarks/beam/prepare.ts b/src/benchmarks/beam/prepare.ts new file mode 100644 index 0000000..4260dc6 --- /dev/null +++ b/src/benchmarks/beam/prepare.ts @@ -0,0 +1,371 @@ +import { createHash, randomUUID } from "node:crypto" +import { mkdir, open, readFile, rename, rm, writeFile } from "node:fs/promises" +import { dirname, join, resolve } from "node:path" +import { + BEAM_CONVERTER_IDENTITY, + BEAM_DATASET_SOURCES, + assertReviewedBeamPaddedAssistantCounts, + canonicalizeBeamRows, + computeCanonicalFileManifest, + computeDatasetFingerprint, + computeManifestHash, + loadPreparedBeamDataset, + loadPreparedBeamTestFixture, + serializeBeamJsonl, + sha256Text, + stableBeamStringify, + validatePreparedBeamSnapshotContents, +} from "./dataset" +import type { + BeamCanonicalFileManifest, + BeamDatasetManifest, + BeamScale, + BeamSourceFileManifest, +} from "./types" +import { BEAM_CANONICAL_SCHEMA_VERSION, BEAM_MANIFEST_SCHEMA_VERSION } from "./types" +import { decodeBeamParquetWithHyparquet, type BeamParquetDecoder } from "./parquet" + +export { decodeBeamParquetWithHyparquet } from "./parquet" + +export interface PrepareBeamDatasetOptions { + tiers: BeamScale[] + outputRoot: string + fetchImpl?: typeof fetch + parquetDecoder?: BeamParquetDecoder + /** + * Only for deterministic fixtures that inject both transport and decoding. + * The resulting manifest is permanently marked as an injected test fixture + * and ordinary scored-run loading rejects it. + */ + unsafeSkipPublishedHashCheckForTests?: boolean +} + +export interface PrepareBeamDatasetResult { + snapshotPath: string + manifest: BeamDatasetManifest + reused: boolean +} + +interface DownloadResult extends BeamSourceFileManifest { + localPath: string +} + +function compareStrings(left: string, right: string): number { + return left < right ? -1 : left > right ? 1 : 0 +} + +function normalizeTiers(tiers: BeamScale[]): BeamScale[] { + const unique = [...new Set(tiers)] + if (unique.length === 0) throw new Error("BEAM prepare requires at least one tier") + for (const tier of unique) { + if (tier !== "1M" && tier !== "10M") throw new Error(`Unsupported BEAM tier: ${tier}`) + } + return unique.sort((left, right) => (left === right ? 0 : left === "1M" ? -1 : 1)) +} + +function sanitizeSourceName(sourcePath: string): string { + const name = sourcePath.split("/").pop() || "source.parquet" + if (!/^[a-zA-Z0-9_.-]+$/.test(name)) { + throw new Error(`BEAM source path has unsafe filename: ${sourcePath}`) + } + return name +} + +async function downloadPinnedFile( + fetchImpl: typeof fetch, + url: string, + sourcePath: string, + snapshotPath: string, + destination: string, + expectedSha256?: string +): Promise { + const response = await fetchImpl(url, { redirect: "follow" }) + if (!response.ok || !response.body) { + throw new Error(`Failed to download pinned BEAM source ${sourcePath}: HTTP ${response.status}`) + } + + await mkdir(dirname(destination), { recursive: true }) + const handle = await open(destination, "wx") + const hash = createHash("sha256") + let byteSize = 0 + try { + const reader = response.body.getReader() + while (true) { + const { done, value } = await reader.read() + if (done) break + if (!value || value.byteLength === 0) continue + await handle.write(value) + hash.update(value) + byteSize += value.byteLength + } + } finally { + await handle.close() + } + + if (byteSize < 8) throw new Error(`Downloaded BEAM source ${sourcePath} is empty or truncated`) + const contentLength = response.headers.get("content-length") + if (contentLength && Number(contentLength) !== byteSize) { + throw new Error( + `Downloaded BEAM source ${sourcePath} size mismatch: expected ${contentLength}, got ${byteSize}` + ) + } + const sha256 = hash.digest("hex") + if (expectedSha256 && sha256 !== expectedSha256) { + throw new Error(`Downloaded BEAM source ${sourcePath} SHA-256 mismatch`) + } + + const file = await open(destination, "r") + try { + const first = Buffer.alloc(4) + const last = Buffer.alloc(4) + await file.read(first, 0, 4, 0) + await file.read(last, 0, 4, byteSize - 4) + if (first.toString("ascii") !== "PAR1" || last.toString("ascii") !== "PAR1") { + throw new Error(`Downloaded BEAM source ${sourcePath} is not a complete Parquet file`) + } + } finally { + await file.close() + } + + return { path: sourcePath, snapshotPath, url, byteSize, sha256, localPath: destination } +} + +async function writeCanonicalFile( + stagingPath: string, + relativePath: string, + content: string, + rowCount: number +): Promise { + const absolutePath = join(stagingPath, relativePath) + await mkdir(dirname(absolutePath), { recursive: true }) + const bytes = Buffer.from(content, "utf8") + await writeFile(absolutePath, bytes, { flag: "wx" }) + return computeCanonicalFileManifest(relativePath, bytes, rowCount) +} + +function manifestWithoutIdentity(input: { + tiers: BeamScale[] + sourceIdentity: "reviewed-published" | "injected-test-fixture" + sourceFiles: Partial> + canonicalFiles: BeamCanonicalFileManifest[] + counts: BeamDatasetManifest["counts"] + orderedChatIds: BeamDatasetManifest["orderedChatIds"] + orderedChatIdsDigest: BeamDatasetManifest["orderedChatIdsDigest"] + orderedQuestionIds: BeamDatasetManifest["orderedQuestionIds"] + orderedQuestionIdsDigest: BeamDatasetManifest["orderedQuestionIdsDigest"] +}): Omit { + return { + manifestSchemaVersion: BEAM_MANIFEST_SCHEMA_VERSION, + canonicalSchemaVersion: BEAM_CANONICAL_SCHEMA_VERSION, + converter: BEAM_CONVERTER_IDENTITY, + includedTiers: input.tiers, + sources: input.tiers.map((tier) => { + const descriptor = BEAM_DATASET_SOURCES[tier] + return { + tier, + sourceIdentity: input.sourceIdentity, + repository: descriptor.repository, + split: descriptor.split, + revision: descriptor.revision, + files: (input.sourceFiles[tier] ?? []) + .map(({ localPath: _localPath, ...file }) => file) + .sort((left, right) => compareStrings(left.path, right.path)), + } + }), + canonicalFiles: [...input.canonicalFiles].sort((left, right) => + compareStrings(left.path, right.path) + ), + counts: input.counts, + orderedChatIds: input.orderedChatIds, + orderedChatIdsDigest: input.orderedChatIdsDigest, + orderedQuestionIds: input.orderedQuestionIds, + orderedQuestionIdsDigest: input.orderedQuestionIdsDigest, + } +} + +async function publishSnapshot( + stagingPath: string, + outputRoot: string, + manifest: BeamDatasetManifest, + allowTestSourceIdentity: boolean +): Promise<{ snapshotPath: string; reused: boolean }> { + const snapshotPath = join(outputRoot, manifest.datasetFingerprint) + try { + await rename(stagingPath, snapshotPath) + return { snapshotPath, reused: false } + } catch (error) { + const code = (error as NodeJS.ErrnoException).code + if (code !== "EEXIST" && code !== "ENOTEMPTY") throw error + + const existing = await (allowTestSourceIdentity + ? loadPreparedBeamTestFixture({ + snapshotPath, + tiers: manifest.includedTiers, + expectedDatasetFingerprint: manifest.datasetFingerprint, + }) + : loadPreparedBeamDataset({ + snapshotPath, + tiers: manifest.includedTiers, + expectedDatasetFingerprint: manifest.datasetFingerprint, + })) + if (existing.manifest.manifestHash !== manifest.manifestHash) { + throw new Error( + `BEAM snapshot ${manifest.datasetFingerprint} already exists with a different manifest` + ) + } + await rm(stagingPath, { recursive: true }) + return { snapshotPath, reused: true } + } +} + +export async function prepareBeamDataset( + options: PrepareBeamDatasetOptions +): Promise { + const tiers = normalizeTiers(options.tiers) + if ( + options.unsafeSkipPublishedHashCheckForTests && + (!options.fetchImpl || !options.parquetDecoder) + ) { + throw new Error( + "Skipping BEAM published source hashes requires injected test transport and decoder" + ) + } + const outputRoot = resolve(options.outputRoot) + await mkdir(outputRoot, { recursive: true }) + const stagingPath = join(outputRoot, `.staging-${randomUUID()}`) + await mkdir(stagingPath, { recursive: false }) + + const fetchImpl = options.fetchImpl ?? fetch + const decodeParquet = options.parquetDecoder ?? decodeBeamParquetWithHyparquet + const sourceIdentity = options.unsafeSkipPublishedHashCheckForTests + ? "injected-test-fixture" + : "reviewed-published" + try { + const sourceFiles: Partial> = {} + const canonicalFiles: BeamCanonicalFileManifest[] = [] + const counts: BeamDatasetManifest["counts"] = {} + const orderedChatIds: BeamDatasetManifest["orderedChatIds"] = {} + const orderedChatIdsDigest: BeamDatasetManifest["orderedChatIdsDigest"] = {} + const orderedQuestionIds: BeamDatasetManifest["orderedQuestionIds"] = {} + const orderedQuestionIdsDigest: BeamDatasetManifest["orderedQuestionIdsDigest"] = {} + + for (const tier of tiers) { + const descriptor = BEAM_DATASET_SOURCES[tier] + const tierSourceFiles: DownloadResult[] = [] + const sourceRows: unknown[] = [] + for (let index = 0; index < descriptor.parquetFiles.length; index++) { + const source = descriptor.parquetFiles[index] + const sourceName = sanitizeSourceName(source.path) + const sourceSnapshotPath = `source/${tier}/${index}-${sourceName}` + const localPath = join(stagingPath, sourceSnapshotPath) + const downloaded = await downloadPinnedFile( + fetchImpl, + source.url, + source.path, + sourceSnapshotPath, + localPath, + options.unsafeSkipPublishedHashCheckForTests ? undefined : source.expectedSha256 + ) + tierSourceFiles.push(downloaded) + const decoded = await decodeParquet(localPath, tier) + if (!Array.isArray(decoded)) { + throw new Error(`BEAM ${tier} Parquet decoder did not return an array of rows`) + } + sourceRows.push(...decoded) + } + sourceFiles[tier] = tierSourceFiles + + const canonical = canonicalizeBeamRows(tier, sourceRows) + if (sourceIdentity === "reviewed-published") { + assertReviewedBeamPaddedAssistantCounts(tier, canonical.counts) + } + const chatsContent = serializeBeamJsonl(canonical.chats) + const questionsContent = serializeBeamJsonl(canonical.questions) + canonicalFiles.push( + await writeCanonicalFile( + stagingPath, + `canonical/${tier}/chats.jsonl`, + chatsContent, + canonical.chats.length + ), + await writeCanonicalFile( + stagingPath, + `canonical/${tier}/questions.jsonl`, + questionsContent, + canonical.questions.length + ) + ) + counts[tier] = canonical.counts + orderedChatIds[tier] = canonical.chats.map((chat) => chat.chatId) + orderedChatIdsDigest[tier] = sha256Text(orderedChatIds[tier]!.join("\n")) + orderedQuestionIds[tier] = canonical.questions.map((question) => question.questionId) + orderedQuestionIdsDigest[tier] = sha256Text(orderedQuestionIds[tier]!.join("\n")) + } + + const manifestCore = manifestWithoutIdentity({ + tiers, + sourceIdentity, + sourceFiles, + canonicalFiles, + counts, + orderedChatIds, + orderedChatIdsDigest, + orderedQuestionIds, + orderedQuestionIdsDigest, + }) + const datasetFingerprint = computeDatasetFingerprint(manifestCore) + const manifestWithoutHash: Omit = { + ...manifestCore, + datasetFingerprint, + } + const manifest: BeamDatasetManifest = { + ...manifestWithoutHash, + manifestHash: computeManifestHash(manifestWithoutHash), + } + + await writeFile(join(stagingPath, "manifest.json"), stableBeamStringify(manifest) + "\n", { + flag: "wx", + }) + + // Validate every source byte, canonical byte, identity, schema, and count + // before a completion marker can make this staging directory publishable. + await validatePreparedBeamSnapshotContents({ + snapshotPath: stagingPath, + tiers, + expectedDatasetFingerprint: datasetFingerprint, + allowInjectedTestSourceIdentity: sourceIdentity === "injected-test-fixture", + }) + + await writeFile( + join(stagingPath, ".complete"), + stableBeamStringify({ + datasetFingerprint: manifest.datasetFingerprint, + manifestHash: manifest.manifestHash, + }) + "\n", + { flag: "wx" } + ) + + const published = await publishSnapshot( + stagingPath, + outputRoot, + manifest, + sourceIdentity === "injected-test-fixture" + ) + return { snapshotPath: published.snapshotPath, manifest, reused: published.reused } + } catch (error) { + await rm(stagingPath, { recursive: true }).catch(() => undefined) + throw error + } +} + +export async function verifyPreparedBeamSnapshot( + snapshotPath: string, + tiers: BeamScale[] +): Promise { + return (await loadPreparedBeamDataset({ snapshotPath, tiers })).manifest +} + +export async function sourceFileSha256(filePath: string): Promise { + const bytes = await readFile(filePath) + return createHash("sha256").update(bytes).digest("hex") +} diff --git a/src/benchmarks/beam/types.ts b/src/benchmarks/beam/types.ts index 522cbb0..4b9ce4d 100644 --- a/src/benchmarks/beam/types.ts +++ b/src/benchmarks/beam/types.ts @@ -1,5 +1,143 @@ export type BeamScale = "1M" | "10M" +export const BEAM_SCALES = ["1M", "10M"] as const + +export const BEAM_CANONICAL_SCHEMA_VERSION = 3 +export const BEAM_MANIFEST_SCHEMA_VERSION = 4 +export const BEAM_CONVERTER_VERSION = "1.5.0" + +export type BeamSourceIdentity = "reviewed-published" | "injected-test-fixture" + +export const BEAM_QUESTION_TYPE_IDS = [ + "abstention", + "contradiction_resolution", + "event_ordering", + "information_extraction", + "instruction_following", + "knowledge_update", + "multi_session_reasoning", + "preference_following", + "summarization", + "temporal_reasoning", +] as const + +export type BeamQuestionType = (typeof BEAM_QUESTION_TYPE_IDS)[number] + +export interface BeamDatasetSource { + repository: string + split: BeamScale + revision: string + parquetFiles: Array<{ + path: string + url: string + expectedSha256?: string + }> +} + +export interface BeamCanonicalMessage { + role: "user" | "assistant" + content: string + timeAnchor?: string +} + +export interface BeamCanonicalSession { + sessionId: string + planNumber?: number + batchNumber: number + turnIndex: number + documentDate?: string + hadInvalidTimeAnchor?: boolean + hasPaddedAssistant?: true + messages: BeamCanonicalMessage[] +} + +export interface BeamCanonicalChat { + schemaVersion: number + scale: BeamScale + chatId: string + sessions: BeamCanonicalSession[] +} + +export interface BeamCanonicalQuestion { + schemaVersion: number + scale: BeamScale + chatId: string + questionId: string + questionType: BeamQuestionType + question: string + rubric: string[] + difficulty?: string + referenceAnswer?: string +} + +export interface BeamSourceFileManifest { + path: string + snapshotPath: string + url: string + byteSize: number + sha256: string +} + +export interface BeamCanonicalFileManifest { + path: string + byteSize: number + sha256: string + rowCount: number +} + +export interface BeamTierCounts { + chats: number + questions: number + sessions: number + sessionsWithDocumentDate: number + sessionsWithoutDocumentDate: number + sessionsWithInvalidTimeAnchor: number + sessionsWithPaddedAssistant: number + byQuestionType: Record + byChat: Record< + string, + { + sessions: number + questions: number + byQuestionType: Record + } + > +} + +export interface BeamDatasetManifest { + manifestSchemaVersion: number + canonicalSchemaVersion: number + converter: { + name: "memorybench-beam" + version: string + implementationHash: string + } + includedTiers: BeamScale[] + sources: Array<{ + tier: BeamScale + sourceIdentity: BeamSourceIdentity + repository: string + split: BeamScale + revision: string + files: BeamSourceFileManifest[] + }> + canonicalFiles: BeamCanonicalFileManifest[] + counts: Partial> + orderedChatIds: Partial> + orderedChatIdsDigest: Partial> + orderedQuestionIds: Partial> + orderedQuestionIdsDigest: Partial> + datasetFingerprint: string + manifestHash: string +} + +export interface PreparedBeamDataset { + snapshotPath: string + manifest: BeamDatasetManifest + chatsByTier: Partial> + questionsByTier: Partial> +} + export interface BeamMessage { role: "user" | "assistant" id?: number @@ -7,9 +145,11 @@ export interface BeamMessage { time_anchor?: string index?: string question_type?: string + isPaddedAssistant?: true } export interface BeamBatch { + plan_number?: number batch_number: number time_anchor?: string | null turns: BeamMessage[][] diff --git a/src/benchmarks/convomem/index.ts b/src/benchmarks/convomem/index.ts index 9c253fa..ffc9253 100644 --- a/src/benchmarks/convomem/index.ts +++ b/src/benchmarks/convomem/index.ts @@ -8,6 +8,7 @@ import type { QuestionTypeRegistry, } from "../../types/unified" import { logger } from "../../utils/logger" +import { legacyBenchmarkProtocol } from "../../protocols/legacy" const DEFAULT_DATA_PATH = "./data/benchmarks/convomem" const HF_BASE_URL = @@ -68,6 +69,8 @@ export const CONVOMEM_QUESTION_TYPES: QuestionTypeRegistry = { export class ConvoMemBenchmark implements Benchmark { name = "convomem" + scope = { displayName: "ConvoMem", includedTiers: [], coverage: "full" as const } + protocol = legacyBenchmarkProtocol private questions: UnifiedQuestion[] = [] private sessionsMap: Map = new Map() private dataPath: string = "" diff --git a/src/benchmarks/index.ts b/src/benchmarks/index.ts index 56cd0c4..ebcb52f 100644 --- a/src/benchmarks/index.ts +++ b/src/benchmarks/index.ts @@ -2,7 +2,7 @@ import type { Benchmark, BenchmarkName } from "../types/benchmark" import { LoCoMoBenchmark } from "./locomo" import { LongMemEvalBenchmark } from "./longmemeval" import { ConvoMemBenchmark } from "./convomem" -import { Beam1MBenchmark, Beam10MBenchmark, BeamBenchmark } from "./beam" +import { Beam1MBenchmark, Beam10MBenchmark, Beam1M10MBenchmark } from "./beam" const benchmarks: Record Benchmark> = { locomo: LoCoMoBenchmark, @@ -10,7 +10,7 @@ const benchmarks: Record Benchmark> = { convomem: ConvoMemBenchmark, "beam-1m": Beam1MBenchmark, "beam-10m": Beam10MBenchmark, - beam: BeamBenchmark, + "beam-1m-10m": Beam1M10MBenchmark, } export function createBenchmark(name: BenchmarkName): Benchmark { @@ -25,4 +25,11 @@ export function getAvailableBenchmarks(): BenchmarkName[] { return Object.keys(benchmarks) as BenchmarkName[] } -export { LoCoMoBenchmark, LongMemEvalBenchmark, ConvoMemBenchmark, BeamBenchmark } +export { + LoCoMoBenchmark, + LongMemEvalBenchmark, + ConvoMemBenchmark, + Beam1MBenchmark, + Beam10MBenchmark, + Beam1M10MBenchmark, +} diff --git a/src/benchmarks/locomo/index.ts b/src/benchmarks/locomo/index.ts index 7b8c86c..7a6deea 100644 --- a/src/benchmarks/locomo/index.ts +++ b/src/benchmarks/locomo/index.ts @@ -9,6 +9,7 @@ import type { } from "../../types/unified" import type { LoCoMoItem, LoCoMoMessage } from "./types" import { logger } from "../../utils/logger" +import { legacyBenchmarkProtocol } from "../../protocols/legacy" const DEFAULT_DATA_PATH = "./data/benchmarks/locomo/locomo10.json" const GITHUB_DATASET_URL = @@ -76,6 +77,8 @@ const CATEGORY_TO_TYPE: Record = { export class LoCoMoBenchmark implements Benchmark { name = "locomo" + scope = { displayName: "LoCoMo", includedTiers: [], coverage: "full" as const } + protocol = legacyBenchmarkProtocol private data: LoCoMoItem[] = [] private questions: UnifiedQuestion[] = [] private sessionsMap: Map = new Map() diff --git a/src/benchmarks/longmemeval/index.ts b/src/benchmarks/longmemeval/index.ts index 158d3f9..7414864 100644 --- a/src/benchmarks/longmemeval/index.ts +++ b/src/benchmarks/longmemeval/index.ts @@ -9,6 +9,7 @@ import type { } from "../../types/unified" import type { LongMemEvalItem } from "./types" import { logger } from "../../utils/logger" +import { legacyBenchmarkProtocol } from "../../protocols/legacy" const DEFAULT_DATA_PATH = "./data/benchmarks/longmemeval/datasets" const HF_DATASET_URL = @@ -81,6 +82,8 @@ export const LONGMEMEVAL_QUESTION_TYPES: QuestionTypeRegistry = { export class LongMemEvalBenchmark implements Benchmark { name = "longmemeval" + scope = { displayName: "LongMemEval", includedTiers: [], coverage: "full" as const } + protocol = legacyBenchmarkProtocol private data: LongMemEvalItem[] = [] private questions: UnifiedQuestion[] = [] private sessionsMap: Map = new Map() diff --git a/src/cli/commands/beam.ts b/src/cli/commands/beam.ts new file mode 100644 index 0000000..7bfd7de --- /dev/null +++ b/src/cli/commands/beam.ts @@ -0,0 +1,44 @@ +import { prepareBeamDataset } from "../../benchmarks/beam/prepare" +import type { BeamScale } from "../../benchmarks/beam/types" + +const DEFAULT_OUTPUT_ROOT = "./data/benchmarks/beam" + +function parseTiers(value: string): BeamScale[] { + const tiers = value.split(",").map((tier) => tier.trim()) + if (tiers.length === 0 || tiers.some((tier) => tier !== "1M" && tier !== "10M")) { + throw new Error(`Invalid BEAM tiers ${value}; use 1M, 10M, or 1M,10M`) + } + return tiers as BeamScale[] +} + +export async function beamCommand(args: string[]): Promise { + const subcommand = args[0] + if (subcommand !== "prepare") { + console.log("Usage: bun run src/index.ts beam prepare [--tiers 1M,10M] [--output ]") + return + } + + let tiers: BeamScale[] = ["1M", "10M"] + let outputRoot = DEFAULT_OUTPUT_ROOT + for (let index = 1; index < args.length; index++) { + const argument = args[index] + if (argument === "--tiers") { + const value = args[++index] + if (!value) throw new Error("--tiers requires a value") + tiers = parseTiers(value) + } else if (argument === "--output" || argument === "--data-path") { + const value = args[++index] + if (!value) throw new Error(`${argument} requires a value`) + outputRoot = value + } else { + throw new Error(`Unknown beam prepare option: ${argument}`) + } + } + + const prepared = await prepareBeamDataset({ tiers, outputRoot }) + console.log(`Prepared BEAM ${tiers.join("/")} snapshot: ${prepared.snapshotPath}`) + console.log(`Dataset fingerprint: ${prepared.manifest.datasetFingerprint}`) + console.log( + `Run with --data-path ${outputRoot} --dataset-revision ${prepared.manifest.datasetFingerprint}` + ) +} diff --git a/src/cli/commands/compare.ts b/src/cli/commands/compare.ts index fefae08..d307731 100644 --- a/src/cli/commands/compare.ts +++ b/src/cli/commands/compare.ts @@ -3,9 +3,8 @@ import type { BenchmarkName } from "../../types/benchmark" import type { SamplingConfig, SampleType } from "../../types/checkpoint" import { batchManager } from "../../orchestrator/batch" import { getAvailableProviders } from "../../providers" -import { getAvailableBenchmarks } from "../../benchmarks" +import { createBenchmark, getAvailableBenchmarks } from "../../benchmarks" import { DEFAULT_ANSWERING_MODEL } from "../../utils/models" -import { logger } from "../../utils/logger" const DEFAULT_JUDGE_MODEL = "gpt-4o" @@ -19,13 +18,13 @@ interface CompareArgs { sampleType?: SampleType limit?: number force?: boolean + dataPath?: string + datasetRevision?: string + retrievalTopK?: number } export function parseCompareArgs(args: string[]): CompareArgs | null { - const parsed: Partial = { - judgeModel: DEFAULT_JUDGE_MODEL, - answeringModel: DEFAULT_ANSWERING_MODEL, - } + const parsed: Partial = { answeringModel: DEFAULT_ANSWERING_MODEL } for (let i = 0; i < args.length; i++) { const arg = args[i] @@ -47,13 +46,18 @@ export function parseCompareArgs(args: string[]): CompareArgs | null { if (type === "consecutive" || type === "random") { parsed.sampleType = type } else { - logger.error(`Invalid sample type: ${type}. Valid types: consecutive, random`) - return null + throw new Error(`Invalid sample type: ${type}. Valid types: consecutive, random`) } } else if (arg === "-l" || arg === "--limit") { parsed.limit = parseInt(args[++i], 10) } else if (arg === "--force") { parsed.force = true + } else if (arg === "--data-path") { + parsed.dataPath = args[++i] + } else if (arg === "--dataset-revision") { + parsed.datasetRevision = args[++i] + } else if (arg === "--top-k" || arg === "--retrieval-top-k") { + parsed.retrievalTopK = parseInt(args[++i], 10) } } @@ -61,6 +65,13 @@ export function parseCompareArgs(args: string[]): CompareArgs | null { return null } + parsed.judgeModel = + parsed.judgeModel || + (parsed.benchmark + ? createBenchmark(parsed.benchmark as BenchmarkName).protocol.requiredJudge?.modelAlias + : undefined) || + DEFAULT_JUDGE_MODEL + return parsed as CompareArgs } @@ -86,11 +97,18 @@ export async function compareCommand(args: string[]): Promise { console.log(" -l, --limit Limit total number of questions") console.log(" --compare-id Compare ID (for resuming)") console.log(" --force Clear existing comparison and start fresh") + console.log(" --data-path PATH Prepared dataset snapshot root") + console.log(" --dataset-revision ID Expected dataset fingerprint") + console.log(" --retrieval-top-k K Retrieval cutoff (BEAM: 5, 10, 15, or 20)") + console.log(" --top-k N Benchmark retrieval Top-K") console.log("") console.log("Examples:") console.log(" bun run src/index.ts compare -p supermemory,mem0,zep -b locomo -s 5") console.log(" bun run src/index.ts compare -p supermemory,filesystem,rag -b locomo -s 5") console.log(" bun run src/index.ts compare --compare-id compare-20251222-103045") + if (!args.includes("--help") && !args.includes("-h")) { + throw new Error("Invalid or incomplete compare arguments") + } return } @@ -102,18 +120,16 @@ export async function compareCommand(args: string[]): Promise { } else { for (const provider of parsed.providers!) { if (!getAvailableProviders().includes(provider as ProviderName)) { - logger.error( + throw new Error( `Invalid provider: ${provider}. Available: ${getAvailableProviders().join(", ")}` ) - return } } if (!getAvailableBenchmarks().includes(parsed.benchmark as BenchmarkName)) { - logger.error( + throw new Error( `Invalid benchmark: ${parsed.benchmark}. Available: ${getAvailableBenchmarks().join(", ")}` ) - return } let sampling: SamplingConfig | undefined @@ -137,13 +153,19 @@ export async function compareCommand(args: string[]): Promise { answeringModel: parsed.answeringModel, sampling, force: parsed.force, + dataPath: parsed.dataPath, + datasetRevision: parsed.datasetRevision, + retrievalTopK: parsed.retrievalTopK, }) } if (result.successes > 0) { batchManager.printComparisonReport(result.manifest) } + if (result.failures > 0) { + throw new Error(`${result.failures} provider comparison run(s) failed`) + } } catch (error) { - logger.error(`${error}`) + throw error } } diff --git a/src/cli/commands/ingest.ts b/src/cli/commands/ingest.ts index 3d4d1a9..428a897 100644 --- a/src/cli/commands/ingest.ts +++ b/src/cli/commands/ingest.ts @@ -1,5 +1,6 @@ import type { ProviderName } from "../../types/provider" import type { BenchmarkName } from "../../types/benchmark" +import type { ConcurrencyConfig } from "../../types/concurrency" import { orchestrator, CheckpointManager } from "../../orchestrator" import { getAvailableProviders } from "../../providers" import { getAvailableBenchmarks } from "../../benchmarks" @@ -10,6 +11,14 @@ interface IngestArgs { benchmark?: string runId: string force?: boolean + dataPath?: string + datasetRevision?: string + retrievalTopK?: number + judgeModel?: string + answeringModel?: string + concurrency?: ConcurrencyConfig + ingestBatchSize?: number + ingestReadinessTimeoutMs?: number } function generateRunId(): string { @@ -21,6 +30,7 @@ function generateRunId(): string { export function parseIngestArgs(args: string[]): IngestArgs | null { const parsed: Partial = {} + const concurrency: Partial = {} for (let i = 0; i < args.length; i++) { const arg = args[i] @@ -32,6 +42,30 @@ export function parseIngestArgs(args: string[]): IngestArgs | null { parsed.runId = args[++i] } else if (arg === "--force") { parsed.force = true + } else if (arg === "--data-path") { + parsed.dataPath = args[++i] + } else if (arg === "--dataset-revision") { + parsed.datasetRevision = args[++i] + } else if (arg === "--top-k" || arg === "--retrieval-top-k") { + parsed.retrievalTopK = parseInt(args[++i], 10) + } else if (arg === "--concurrency") { + concurrency.default = parseInt(args[++i], 10) + } else if (arg === "--concurrency-ingest") { + concurrency.ingest = parseInt(args[++i], 10) + } else if (arg === "--concurrency-indexing") { + concurrency.indexing = parseInt(args[++i], 10) + } else if (arg === "--ingest-batch-size") { + const value = Number(args[++i]) + if (!Number.isInteger(value) || value < 1 || value > 600) { + throw new Error("--ingest-batch-size must be an integer between 1 and 600") + } + parsed.ingestBatchSize = value + } else if (arg === "--ingest-timeout-seconds") { + const value = Number(args[++i]) + if (!Number.isInteger(value) || value < 1) { + throw new Error("--ingest-timeout-seconds must be a positive integer") + } + parsed.ingestReadinessTimeoutMs = value * 1000 } } @@ -44,6 +78,10 @@ export function parseIngestArgs(args: string[]): IngestArgs | null { parsed.runId = generateRunId() } + if (Object.keys(concurrency).length > 0) { + parsed.concurrency = concurrency as ConcurrencyConfig + } + return parsed as IngestArgs } @@ -62,6 +100,17 @@ export async function ingestCommand(args: string[]): Promise { console.log(` -b, --benchmark Benchmark: ${getAvailableBenchmarks().join(", ")}`) console.log(" -r, --run-id Run identifier") console.log(" --force Clear existing checkpoint and start fresh") + console.log(" --data-path PATH Prepared dataset snapshot root") + console.log(" --dataset-revision ID Expected dataset fingerprint") + console.log(" --concurrency N Concurrency for conversation builds") + console.log(" --concurrency-ingest N Concurrency for document submission") + console.log(" --concurrency-indexing N Concurrency for readiness polling") + console.log(" --ingest-batch-size N Ordered sessions per provider batch (1-600)") + console.log(" --ingest-timeout-seconds N Per-readiness-call timeout (default: 300)") + console.log( + " --retrieval-top-k K Retrieval configuration recorded in the run protocol identity" + ) + console.log(" --top-k N Benchmark retrieval Top-K") return } @@ -85,6 +134,11 @@ export async function ingestCommand(args: string[]): Promise { parsed.provider = checkpoint.provider parsed.benchmark = checkpoint.benchmark + parsed.dataPath = parsed.dataPath || checkpoint.dataPath + parsed.datasetRevision = parsed.datasetRevision || checkpoint.datasetRevision + parsed.retrievalTopK = parsed.retrievalTopK ?? checkpoint.retrievalTopK + parsed.judgeModel = checkpoint.judge + parsed.answeringModel = checkpoint.answeringModel logger.info( `Continuing ingest for ${parsed.runId} (${checkpoint.provider}/${checkpoint.benchmark})` ) @@ -110,5 +164,13 @@ export async function ingestCommand(args: string[]): Promise { benchmark: parsed.benchmark as BenchmarkName, runId: parsed.runId, force: parsed.force, + dataPath: parsed.dataPath, + datasetRevision: parsed.datasetRevision, + retrievalTopK: parsed.retrievalTopK, + judgeModel: parsed.judgeModel, + answeringModel: parsed.answeringModel, + concurrency: parsed.concurrency, + ingestBatchSize: parsed.ingestBatchSize, + ingestReadinessTimeoutMs: parsed.ingestReadinessTimeoutMs, }) } diff --git a/src/cli/commands/list-questions.ts b/src/cli/commands/list-questions.ts index 6da7e69..8196750 100644 --- a/src/cli/commands/list-questions.ts +++ b/src/cli/commands/list-questions.ts @@ -6,6 +6,8 @@ interface ListQuestionsArgs { offset: number limit: number type?: string + dataPath?: string + datasetRevision?: string } export function parseListQuestionsArgs(args: string[]): ListQuestionsArgs | null { @@ -24,6 +26,10 @@ export function parseListQuestionsArgs(args: string[]): ListQuestionsArgs | null parsed.limit = parseInt(args[++i], 10) } else if (arg === "-t" || arg === "--type") { parsed.type = args[++i] + } else if (arg === "--data-path") { + parsed.dataPath = args[++i] + } else if (arg === "--dataset-revision") { + parsed.datasetRevision = args[++i] } } @@ -52,6 +58,8 @@ export async function listQuestionsCommand(args: string[]): Promise { console.log(" -o, --offset Start from question number (default: 0)") console.log(" -l, --limit Number of questions to show (default: 50)") console.log(" -t, --type Filter by question type") + console.log(" --data-path PATH Prepared dataset snapshot root") + console.log(" --dataset-revision ID Expected dataset fingerprint") console.log("") console.log("Examples:") console.log(" bun run src/index.ts list-questions -b locomo") @@ -61,18 +69,20 @@ export async function listQuestionsCommand(args: string[]): Promise { } if (!getAvailableBenchmarks().includes(parsed.benchmark as BenchmarkName)) { - console.error(`Invalid benchmark: ${parsed.benchmark}`) - console.error(`Available: ${getAvailableBenchmarks().join(", ")}`) - return + throw new Error( + `Invalid benchmark: ${parsed.benchmark}. Available: ${getAvailableBenchmarks().join(", ")}` + ) } const benchmark = createBenchmark(parsed.benchmark as BenchmarkName) try { - await benchmark.load() + await benchmark.load({ + dataPath: parsed.dataPath, + datasetRevision: parsed.datasetRevision, + }) } catch (e: any) { - console.error(`Failed to load benchmark: ${e.message}`) - return + throw new Error(`Failed to load benchmark: ${e.message}`, { cause: e }) } let questions = benchmark.getQuestions() @@ -88,7 +98,7 @@ export async function listQuestionsCommand(args: string[]): Promise { const pageQuestions = questions.slice(start, end) console.log("") - console.log(`Benchmark: ${parsed.benchmark}`) + console.log(`Benchmark: ${benchmark.scope.displayName} (${parsed.benchmark})`) console.log( `Total questions: ${totalQuestions}${parsed.type ? ` (${filteredTotal} matching type "${parsed.type}")` : ""}` ) diff --git a/src/cli/commands/run.ts b/src/cli/commands/run.ts index 46d1f14..c8fecb1 100644 --- a/src/cli/commands/run.ts +++ b/src/cli/commands/run.ts @@ -5,9 +5,10 @@ import type { ConcurrencyConfig } from "../../types/concurrency" import { PHASE_ORDER, getPhasesFromPhase } from "../../types/checkpoint" import { orchestrator, CheckpointManager } from "../../orchestrator" import { getAvailableProviders } from "../../providers" -import { getAvailableBenchmarks } from "../../benchmarks" +import { createBenchmark, getAvailableBenchmarks } from "../../benchmarks" import { listAvailableModels, DEFAULT_ANSWERING_MODEL } from "../../utils/models" import { logger } from "../../utils/logger" +import { BEAM_MEM0_NUGGET_PROFILE } from "../../protocols/beam-mem0" const DEFAULT_JUDGE_MODEL = "gpt-4o" @@ -23,6 +24,12 @@ interface RunArgs { force?: boolean fromPhase?: PhaseId concurrency?: ConcurrencyConfig + dataPath?: string + datasetRevision?: string + retrievalTopK?: number + answerCutoff?: number + evaluationProfile?: string + sourceRunId?: string } function generateRunId(): string { @@ -80,6 +87,24 @@ export function parseRunArgs(args: string[]): RunArgs | null { concurrency.answer = parseInt(args[++i], 10) } else if (arg === "--concurrency-evaluate") { concurrency.evaluate = parseInt(args[++i], 10) + } else if (arg === "--data-path") { + parsed.dataPath = args[++i] + } else if (arg === "--dataset-revision") { + parsed.datasetRevision = args[++i] + } else if (arg === "--top-k" || arg === "--retrieval-top-k") { + parsed.retrievalTopK = parseInt(args[++i], 10) + } else if (arg === "--answer-cutoff") { + parsed.answerCutoff = parseInt(args[++i], 10) + } else if (arg === "--evaluation-profile") { + parsed.evaluationProfile = args[++i] + if (parsed.evaluationProfile !== BEAM_MEM0_NUGGET_PROFILE) { + logger.error( + `Invalid evaluation profile: ${parsed.evaluationProfile}. Valid profile: ${BEAM_MEM0_NUGGET_PROFILE}` + ) + return null + } + } else if (arg === "--source-run") { + parsed.sourceRunId = args[++i] } else if (arg === "--force") { parsed.force = true } @@ -110,6 +135,9 @@ export async function runCommand(args: string[]): Promise { ) console.log(" Continue run: bun run src/index.ts run -r [-j ] [-m ]") console.log(" From phase: bun run src/index.ts run -r -f ") + console.log( + " Reuse builds: bun run src/index.ts run --source-run -r -f search ..." + ) console.log("") console.log("Options:") console.log(` -p, --provider Memory provider: ${getAvailableProviders().join(", ")}`) @@ -130,6 +158,14 @@ export async function runCommand(args: string[]): Promise { console.log(" --concurrency-answer N Concurrency for answer phase") console.log(" --concurrency-evaluate N Concurrency for evaluate phase") console.log(" --force Clear existing checkpoint and start fresh") + console.log(" --data-path PATH Prepared benchmark snapshot root or directory") + console.log(" --dataset-revision ID Expected prepared dataset fingerprint") + console.log(" --retrieval-top-k K Paper: 5/10/15/20; mem0-nugget direct ablation: 1-100") + console.log(" --top-k N Alias for --retrieval-top-k") + console.log(" --answer-cutoff N Evidence cutoff for the mem0-nugget profile") + console.log(" --evaluation-profile P Experimental BEAM profile: mem0-nugget") + console.log(" --source-run ID Reuse validated completed builds; requires a new run ID") + console.log(" BEAM paper judge: gpt-4.1-mini (selected by default for BEAM)") console.log("") console.log(`Available models: ${listAvailableModels().join(", ")}`) return @@ -137,8 +173,42 @@ export async function runCommand(args: string[]): Promise { const checkpointManager = new CheckpointManager() - // Check if run exists - if (checkpointManager.exists(parsed.runId)) { + const targetExists = checkpointManager.exists(parsed.runId) + if (parsed.sourceRunId) { + if (parsed.sourceRunId === parsed.runId) { + logger.error("Source and target run IDs must be different") + return + } + if (targetExists) { + logger.error(`Target run ${parsed.runId} already exists`) + return + } + if (parsed.force) { + logger.error("--source-run cannot be combined with --force") + return + } + if (parsed.fromPhase !== "search") { + logger.error("--source-run requires --from-phase search") + return + } + const source = checkpointManager.load(parsed.sourceRunId) + if (!source) { + logger.error(`Source run not found: ${parsed.sourceRunId}`) + return + } + if (parsed.provider && parsed.provider !== source.provider) { + logger.error(`Source run uses provider ${source.provider}, not ${parsed.provider}`) + return + } + if (parsed.benchmark && parsed.benchmark !== source.benchmark) { + logger.error(`Source run uses benchmark ${source.benchmark}, not ${parsed.benchmark}`) + return + } + parsed.provider = source.provider + parsed.benchmark = source.benchmark + parsed.dataPath = parsed.dataPath || source.dataPath + parsed.datasetRevision = parsed.datasetRevision || source.datasetRevision + } else if (targetExists) { const checkpoint = checkpointManager.load(parsed.runId)! // If provider/benchmark provided, validate they match @@ -160,6 +230,11 @@ export async function runCommand(args: string[]): Promise { parsed.benchmark = checkpoint.benchmark parsed.judgeModel = parsed.judgeModel || checkpoint.judge parsed.answeringModel = parsed.answeringModel || checkpoint.answeringModel + parsed.dataPath = parsed.dataPath || checkpoint.dataPath + parsed.datasetRevision = parsed.datasetRevision || checkpoint.datasetRevision + parsed.retrievalTopK = parsed.retrievalTopK ?? checkpoint.retrievalTopK + parsed.answerCutoff = parsed.answerCutoff ?? checkpoint.answerCutoff + parsed.evaluationProfile = parsed.evaluationProfile ?? checkpoint.evaluationProfile logger.info(`Continuing run ${parsed.runId} (${checkpoint.provider}/${checkpoint.benchmark})`) } else { @@ -178,9 +253,30 @@ export async function runCommand(args: string[]): Promise { console.error(`Invalid benchmark: ${parsed.benchmark}`) return } + } - // Apply defaults for new run - parsed.judgeModel = parsed.judgeModel || DEFAULT_JUDGE_MODEL + if ( + parsed.evaluationProfile && + parsed.benchmark !== "beam-1m" && + parsed.benchmark !== "beam-10m" && + parsed.benchmark !== "beam-1m-10m" + ) { + logger.error("--evaluation-profile mem0-nugget is only valid for BEAM benchmarks") + return + } + if (parsed.answerCutoff !== undefined && !parsed.evaluationProfile) { + logger.error("--answer-cutoff requires --evaluation-profile mem0-nugget") + return + } + if (!parsed.judgeModel) { + parsed.judgeModel = + parsed.evaluationProfile === BEAM_MEM0_NUGGET_PROFILE + ? "gpt-5" + : createBenchmark(parsed.benchmark as BenchmarkName).protocol.requiredJudge?.modelAlias || + DEFAULT_JUDGE_MODEL + } + if (!parsed.answeringModel && parsed.evaluationProfile === BEAM_MEM0_NUGGET_PROFILE) { + parsed.answeringModel = "gpt-5" } const phases = parsed.fromPhase ? getPhasesFromPhase(parsed.fromPhase) : undefined @@ -209,5 +305,11 @@ export async function runCommand(args: string[]): Promise { concurrency: parsed.concurrency, force: parsed.force, phases, + dataPath: parsed.dataPath, + datasetRevision: parsed.datasetRevision, + retrievalTopK: parsed.retrievalTopK, + answerCutoff: parsed.answerCutoff, + evaluationProfile: parsed.evaluationProfile, + sourceRunId: parsed.sourceRunId, }) } diff --git a/src/cli/commands/search.ts b/src/cli/commands/search.ts index 3fd0c0f..fb2cc25 100644 --- a/src/cli/commands/search.ts +++ b/src/cli/commands/search.ts @@ -70,5 +70,10 @@ export async function searchCommand(args: string[]): Promise { provider: parsed.provider as ProviderName, benchmark: parsed.benchmark as BenchmarkName, runId: parsed.runId, + judgeModel: checkpoint.judge, + answeringModel: checkpoint.answeringModel, + dataPath: checkpoint.dataPath, + datasetRevision: checkpoint.datasetRevision, + retrievalTopK: checkpoint.retrievalTopK, }) } diff --git a/src/cli/commands/show-failures.ts b/src/cli/commands/show-failures.ts index 758b360..4461fff 100644 --- a/src/cli/commands/show-failures.ts +++ b/src/cli/commands/show-failures.ts @@ -88,7 +88,7 @@ export async function showFailuresCommand(args: string[]): Promise { if (searchResults.length > 0) { console.log(`Search Results (${searchResults.length}):`) for (let j = 0; j < Math.min(searchResults.length, 3); j++) { - const r = searchResults[j] as Record + const r = searchResults[j] as unknown as Record const content = String(r.content || r.text || r.memory || JSON.stringify(r)).slice(0, 100) const score = typeof r.score === "number" ? r.score.toFixed(2) : "N/A" console.log(` ${j + 1}. [${score}] "${content}${content.length >= 100 ? "..." : ""}"`) diff --git a/src/cli/commands/test-question.ts b/src/cli/commands/test-question.ts index e188542..73d20c1 100644 --- a/src/cli/commands/test-question.ts +++ b/src/cli/commands/test-question.ts @@ -94,5 +94,8 @@ export async function testQuestionCommand(args: string[]): Promise { runId: parsed.runId, questionId: parsed.questionId, answeringModel: parsed.answeringModel, + dataPath: checkpoint.dataPath, + datasetRevision: checkpoint.datasetRevision, + retrievalTopK: checkpoint.retrievalTopK, }) } diff --git a/src/cli/index.ts b/src/cli/index.ts index e6ae3db..c17ba54 100644 --- a/src/cli/index.ts +++ b/src/cli/index.ts @@ -7,6 +7,7 @@ import { statusCommand } from "./commands/status" import { listQuestionsCommand } from "./commands/list-questions" import { showFailuresCommand } from "./commands/show-failures" import { serveCommand } from "./commands/serve" +import { beamCommand } from "./commands/beam" import { getAvailableProviders } from "../providers" import { getAvailableBenchmarks } from "../benchmarks" import { listModelsByProvider, MODEL_ALIASES, DEFAULT_ANSWERING_MODEL } from "../utils/models" @@ -27,6 +28,7 @@ Commands: show-failures Show failed questions from a run with full debugging data status Check run status serve Start the web UI server + beam prepare Download, convert, hash, and validate pinned BEAM 1M/10M data help Show help (use 'help providers', 'help models', 'help benchmarks' for details) Examples: @@ -44,6 +46,9 @@ Options: -r, --run-id Run identifier -m, --answering-model Answering model (default: ${DEFAULT_ANSWERING_MODEL}) -q, --question-id Question ID (for test command) + --evaluation-profile Experimental BEAM profile (mem0-nugget) + --answer-cutoff Retrieved evidence exposed to the answer model + --source-run Reuse validated completed ingestion/indexing builds --force Clear checkpoint and start fresh Run 'bun run src/index.ts help ' for more details: @@ -151,9 +156,11 @@ Available benchmark datasets for evaluation: Tests: user facts, assistant facts, preferences, implicit connections Source: HuggingFace Salesforce/ConvoMem (downloaded on first use) - beam BEAM - Beyond a Million Tokens benchmark + beam-1m / beam-10m / beam-1m-10m + BEAM 1M/10M - intentional public-tier subset Tests: abstention, contradiction, event ordering, extraction, instructions, knowledge update, multi-session, preferences, summarization, temporal - Source: HuggingFace Mohammadta/BEAM (downloaded on first use) + Source: pinned Hugging Face revisions; prepare explicitly with: + bun run src/index.ts beam prepare --tiers 1M,10M Scales: beam-1m (700 q / 35 chats), beam-10m (200 q / 10 chats) Usage: @@ -162,6 +169,7 @@ Usage: -b convomem Run ConvoMem benchmark -b beam-1m Run BEAM 1M-token tier -b beam-10m Run BEAM 10M-token tier + -b beam-1m-10m Run both supported tiers (reported separately) `) } @@ -197,6 +205,9 @@ export async function cli(args: string[]): Promise { case "serve": await serveCommand(commandArgs) break + case "beam": + await beamCommand(commandArgs) + break case "help": case "--help": case "-h": diff --git a/src/index.ts b/src/index.ts index df0930c..b9c2648 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,4 +1,7 @@ import { cli } from "./cli" const args = process.argv.slice(2) -cli(args).catch(console.error) +cli(args).catch((error) => { + console.error(error) + process.exitCode = 1 +}) diff --git a/src/judges/anthropic.ts b/src/judges/anthropic.ts index e8cad9e..9f994eb 100644 --- a/src/judges/anthropic.ts +++ b/src/judges/anthropic.ts @@ -1,6 +1,7 @@ import { createAnthropic } from "@ai-sdk/anthropic" import { generateText } from "ai" import type { Judge, JudgeConfig, JudgeInput, JudgeResult } from "../types/judge" +import type { ModelTransport } from "../types/protocol" import type { ProviderPrompts } from "../types/prompts" import { buildJudgePrompt, parseJudgeResponse, getJudgePrompt } from "./base" import { logger } from "../utils/logger" @@ -46,8 +47,11 @@ export class AnthropicJudge implements Judge { return getJudgePrompt(questionType, providerPrompts) } - getModel() { + getModel(transport: ModelTransport = "provider-default") { if (!this.client || !this.modelConfig) throw new Error("Judge not initialized") + if (transport !== "provider-default") { + throw new Error(`Anthropic judge does not support model transport ${transport}`) + } return this.client(this.modelConfig.id) } } diff --git a/src/judges/google.ts b/src/judges/google.ts index 868dbb7..9bbc2d2 100644 --- a/src/judges/google.ts +++ b/src/judges/google.ts @@ -1,6 +1,7 @@ import { createGoogleGenerativeAI } from "@ai-sdk/google" import { generateText } from "ai" import type { Judge, JudgeConfig, JudgeInput, JudgeResult } from "../types/judge" +import type { ModelTransport } from "../types/protocol" import type { ProviderPrompts } from "../types/prompts" import { buildJudgePrompt, parseJudgeResponse, getJudgePrompt } from "./base" import { logger } from "../utils/logger" @@ -46,8 +47,11 @@ export class GoogleJudge implements Judge { return getJudgePrompt(questionType, providerPrompts) } - getModel() { + getModel(transport: ModelTransport = "provider-default") { if (!this.client || !this.modelConfig) throw new Error("Judge not initialized") + if (transport !== "provider-default") { + throw new Error(`Google judge does not support model transport ${transport}`) + } return this.client(this.modelConfig.id) } } diff --git a/src/judges/openai.ts b/src/judges/openai.ts index 0d6a0e0..2873898 100644 --- a/src/judges/openai.ts +++ b/src/judges/openai.ts @@ -1,6 +1,7 @@ import { createOpenAI } from "@ai-sdk/openai" import { generateText } from "ai" import type { Judge, JudgeConfig, JudgeInput, JudgeResult } from "../types/judge" +import type { ModelTransport } from "../types/protocol" import type { ProviderPrompts } from "../types/prompts" import { buildJudgePrompt, parseJudgeResponse, getJudgePrompt } from "./base" import { logger } from "../utils/logger" @@ -47,8 +48,11 @@ export class OpenAIJudge implements Judge { return getJudgePrompt(questionType, providerPrompts) } - getModel() { + getModel(transport: ModelTransport = "provider-default") { if (!this.client || !this.modelConfig) throw new Error("Judge not initialized") + if (transport === "openai-chat-completions") { + return this.client.chat(this.modelConfig.id) + } return this.client(this.modelConfig.id) } } diff --git a/src/orchestrator/batch.ts b/src/orchestrator/batch.ts index 239c6ae..10f94e7 100644 --- a/src/orchestrator/batch.ts +++ b/src/orchestrator/batch.ts @@ -1,13 +1,28 @@ import type { ProviderName } from "../types/provider" -import type { BenchmarkName } from "../types/benchmark" +import type { BenchmarkName, BenchmarkScope, DatasetIdentity } from "../types/benchmark" import type { SamplingConfig } from "../types/checkpoint" +import type { ProtocolIdentity } from "../types/protocol" +import type { AnsweringRuntimeIdentity } from "../types/model" import type { BenchmarkResult } from "../types/unified" -import { orchestrator, CheckpointManager } from "./index" +import { + orchestrator, + CheckpointManager, + resolveEffectiveRetrievalTopK, + type OrchestratorOptions, +} from "./index" import { createBenchmark } from "../benchmarks" +import { createProvider } from "../providers" +import { fingerprintProviderPrompts } from "../providers/prompt-identity" import { logger } from "../utils/logger" +import { resolveAnsweringRuntimeIdentity, resolveModel } from "../utils/models" +import { stableSha256 } from "../utils/stable" import { existsSync, mkdirSync, readFileSync, writeFileSync, rmSync } from "fs" import { join } from "path" import { startRun, endRun } from "../server/runState" +import { + canonicalizeSelectedQuestionIds, + fingerprintSelectedBenchmarkInput, +} from "./input-identity" const checkpointManager = new CheckpointManager() @@ -21,11 +36,21 @@ export interface CompareManifest { benchmark: string judge: string answeringModel: string + answeringRuntimeIdentity: AnsweringRuntimeIdentity sampling?: SamplingConfig targetQuestionIds: string[] + dataPath?: string + datasetRevision?: string + retrievalTopK: number + datasetIdentity?: DatasetIdentity + benchmarkInputFingerprint: string + benchmarkScope: BenchmarkScope + protocolIdentity: ProtocolIdentity + selectedQuestionIdsDigest: string runs: Array<{ provider: string runId: string + providerPromptFingerprint: string }> } @@ -36,6 +61,9 @@ export interface CompareOptions { answeringModel: string sampling?: SamplingConfig force?: boolean + dataPath?: string + datasetRevision?: string + retrievalTopK?: number } export interface CompareResult { @@ -45,6 +73,221 @@ export interface CompareResult { failures: number } +type ComparisonReport = { provider: string; report: BenchmarkResult } + +type ScalarPrimaryMetric = NonNullable + +export interface ComparisonInputIdentity { + benchmark: string + benchmarkScopeFingerprint: string + datasetFingerprint: string + benchmarkInputFingerprint: string + questionSetFingerprint: string + protocolFingerprint: string + retrievalTopK: number + judge: string + answeringModel: string + answeringRuntimeFingerprint: string + primaryMetricKey: string + primaryMetricHigherIsBetter: boolean + cohortKey: string +} + +export interface PrimaryMetricComparisonRow extends ComparisonReport { + primaryMetric?: ScalarPrimaryMetric + passAccuracy: number + deltaFromBest?: number +} + +export interface PrimaryMetricComparison { + comparable: boolean + identity?: Pick + identities: string[] + mismatchReasons: string[] + rows: PrimaryMetricComparisonRow[] + bestValue?: number + winners: string[] +} + +function resolveReportRetrievalTopK(report: BenchmarkResult): number { + const recorded = [ + ...new Set( + report.questionMetrics + .map((metric) => metric.configuredTopK) + .filter((value) => Number.isInteger(value) && value > 0) + ), + ] + if (recorded.length !== 1) { + throw new Error( + recorded.length === 0 + ? "report has no recorded retrieval Top-K" + : `report mixes retrieval Top-K values (${recorded.join(", ")})` + ) + } + if (report.retrievalTopK != null && report.retrievalTopK !== recorded[0]) { + throw new Error( + `report retrieval Top-K ${report.retrievalTopK} disagrees with question metrics ${recorded[0]}` + ) + } + return recorded[0] +} + +export function getComparisonInputIdentity(report: BenchmarkResult): ComparisonInputIdentity { + const primaryMetric = report.quality.primaryMetric + if (!primaryMetric) throw new Error("report has no scalar primary metric") + if (!Number.isFinite(primaryMetric.value)) { + throw new Error("report primary metric is not finite") + } + if (!report.selectedQuestionIdsDigest?.trim()) { + throw new Error("report has no selected-question fingerprint") + } + if (!report.benchmarkInputFingerprint?.trim()) { + throw new Error("report has no derived benchmark-input identity") + } + if (!report.judge?.trim()) throw new Error("report has no judge model") + if (!report.answeringModel?.trim()) throw new Error("report has no answering model") + if (!report.answeringRuntimeIdentity) { + throw new Error("report has no resolved answering runtime identity") + } + if ( + report.protocolIdentity.id === "memorybench.legacy" && + !report.providerPromptFingerprint?.trim() + ) { + throw new Error("legacy report has no provider-prompt fingerprint") + } + + const datasetFingerprint = report.datasetIdentity + ? typeof report.datasetIdentity.datasetFingerprint === "string" && + report.datasetIdentity.datasetFingerprint.trim() + ? report.datasetIdentity.datasetFingerprint + : stableSha256(report.datasetIdentity) + : `derived:${report.benchmarkInputFingerprint}` + const identityWithoutKey = { + benchmark: report.benchmark, + benchmarkScopeFingerprint: stableSha256(report.benchmarkScope), + datasetFingerprint, + benchmarkInputFingerprint: report.benchmarkInputFingerprint, + questionSetFingerprint: report.selectedQuestionIdsDigest, + protocolFingerprint: stableSha256(report.protocolIdentity), + retrievalTopK: resolveReportRetrievalTopK(report), + judge: report.judge, + answeringModel: report.answeringModel, + answeringRuntimeFingerprint: stableSha256(report.answeringRuntimeIdentity), + providerPromptFingerprint: + report.protocolIdentity.id === "memorybench.legacy" ? report.providerPromptFingerprint : null, + primaryMetricKey: primaryMetric.key, + primaryMetricHigherIsBetter: primaryMetric.higherIsBetter, + } + return { + ...identityWithoutKey, + cohortKey: stableSha256(identityWithoutKey), + } +} + +export function comparePrimaryMetrics( + reports: ComparisonReport[], + expectedReportCount = reports.length +): PrimaryMetricComparison { + const rows = reports.map(({ provider, report }) => ({ + provider, + report, + primaryMetric: report.quality.primaryMetric, + passAccuracy: report.summary.accuracy, + })) + const mismatchReasons: string[] = [] + const inputIdentities = rows.flatMap(({ provider, report }) => { + try { + return [getComparisonInputIdentity(report)] + } catch (error) { + mismatchReasons.push( + `${provider}: ${error instanceof Error ? error.message : "invalid comparison identity"}` + ) + return [] + } + }) + const identities = [...new Set(inputIdentities.map((identity) => identity.cohortKey))] + if (rows.length === 0) mismatchReasons.push("no reports") + if (rows.length !== expectedReportCount) { + mismatchReasons.push( + `only ${rows.length} of ${expectedReportCount} provider reports are complete` + ) + } + if (expectedReportCount < 2) mismatchReasons.push("comparison requires at least two providers") + if (inputIdentities.length === rows.length && identities.length > 1) { + mismatchReasons.push( + "dataset, transformed benchmark input, question set, protocol, retrieval Top-K, judge, answering model, provider prompt, or primary metric semantics differ" + ) + } + if ( + rows.length === 0 || + rows.length !== expectedReportCount || + expectedReportCount < 2 || + inputIdentities.length !== rows.length || + identities.length !== 1 + ) { + return { comparable: false, identities, mismatchReasons, rows, winners: [] } + } + + const firstPrimaryMetric = rows[0].primaryMetric! + const identity = { + key: firstPrimaryMetric.key, + higherIsBetter: firstPrimaryMetric.higherIsBetter, + } + const ranked = rows + .map((row, originalIndex) => ({ row, originalIndex })) + .sort((left, right) => { + const delta = identity.higherIsBetter + ? right.row.primaryMetric!.value - left.row.primaryMetric!.value + : left.row.primaryMetric!.value - right.row.primaryMetric!.value + return delta || left.originalIndex - right.originalIndex + }) + .map(({ row }) => row) + const bestValue = ranked[0].primaryMetric!.value + const rankedWithDeltas = ranked.map((row) => ({ + ...row, + deltaFromBest: row.primaryMetric!.value - bestValue, + })) + + return { + comparable: true, + identity, + identities, + mismatchReasons, + rows: rankedWithDeltas, + bestValue, + winners: rankedWithDeltas + .filter(({ primaryMetric }) => primaryMetric!.value === bestValue) + .map(({ provider }) => provider), + } +} + +function getQuestionTypeQuality( + report: BenchmarkResult, + questionType: string +): { + value?: number + key: string + passAccuracy?: number +} { + const qualitySlice = report.quality.bySlice?.[questionType] + if (typeof qualitySlice?.averageScore === "number") { + return { + value: qualitySlice.averageScore, + key: "averageScore", + passAccuracy: + typeof qualitySlice.passAccuracy === "number" + ? qualitySlice.passAccuracy + : report.byQuestionType[questionType]?.accuracy, + } + } + const legacyAccuracy = report.byQuestionType[questionType]?.accuracy + return { + value: legacyAccuracy, + key: "accuracy", + passAccuracy: legacyAccuracy, + } +} + function generateCompareId(): string { const now = new Date() const date = now.toISOString().slice(0, 10).replace(/-/g, "") @@ -82,7 +325,81 @@ function selectQuestionsBySampling( return allQuestions.map((q) => q.questionId) } +export function assertComparisonReportMatchesManifest( + manifest: CompareManifest, + run: CompareManifest["runs"][number], + report: BenchmarkResult +): void { + const mismatches: string[] = [] + if (report.runId !== run.runId) mismatches.push("run ID") + if (report.provider !== run.provider) mismatches.push("provider") + if (report.benchmark !== manifest.benchmark) mismatches.push("benchmark") + if (report.judge !== manifest.judge) mismatches.push("judge model") + if (report.answeringModel !== manifest.answeringModel) mismatches.push("answering model") + if ( + stableSha256(report.answeringRuntimeIdentity ?? null) !== + stableSha256(manifest.answeringRuntimeIdentity) + ) { + mismatches.push("answering runtime") + } + if (report.providerPromptFingerprint !== run.providerPromptFingerprint) { + mismatches.push("provider-prompt fingerprint") + } + if (stableSha256(report.benchmarkScope) !== stableSha256(manifest.benchmarkScope)) { + mismatches.push("benchmark scope") + } + if ( + stableSha256(report.datasetIdentity ?? null) !== stableSha256(manifest.datasetIdentity ?? null) + ) { + mismatches.push("dataset identity") + } + if (report.benchmarkInputFingerprint !== manifest.benchmarkInputFingerprint) { + mismatches.push("benchmark input") + } + if (stableSha256(report.protocolIdentity) !== stableSha256(manifest.protocolIdentity)) { + mismatches.push("protocol identity") + } + if (report.selectedQuestionIdsDigest !== manifest.selectedQuestionIdsDigest) { + mismatches.push("selected-question fingerprint") + } + let reportedTopK: number | undefined + try { + reportedTopK = resolveReportRetrievalTopK(report) + } catch (error) { + mismatches.push(error instanceof Error ? error.message : "retrieval Top-K") + } + if (reportedTopK !== manifest.retrievalTopK) mismatches.push("retrieval Top-K") + if (report.summary.totalQuestions !== manifest.targetQuestionIds.length) { + mismatches.push("question count") + } + const evaluationQuestionIds = report.evaluations.map((evaluation) => evaluation.questionId) + if ( + evaluationQuestionIds.length !== manifest.targetQuestionIds.length || + stableSha256(evaluationQuestionIds) !== manifest.selectedQuestionIdsDigest + ) { + mismatches.push("evaluation question set/order") + } + const metricQuestionIds = report.questionMetrics.map((metric) => metric.questionId) + if ( + metricQuestionIds.length !== manifest.targetQuestionIds.length || + stableSha256(metricQuestionIds) !== manifest.selectedQuestionIdsDigest + ) { + mismatches.push("question-metric set/order") + } + if (mismatches.length > 0) { + throw new Error( + `Report ${report.runId} does not match comparison ${manifest.compareId}: ${[ + ...new Set(mismatches), + ].join(", ")}` + ) + } +} + export class BatchManager { + constructor( + private readonly runner: { run(options: OrchestratorOptions): Promise } = orchestrator + ) {} + private getComparePath(compareId: string): string { return join(COMPARE_DIR, compareId) } @@ -115,11 +432,12 @@ export class BatchManager { } delete(compareId: string): void { + // Read the run ids before removing the manifest that owns them. + const manifest = this.loadManifest(compareId) const comparePath = this.getComparePath(compareId) if (existsSync(comparePath)) { rmSync(comparePath, { recursive: true }) } - const manifest = this.loadManifest(compareId) if (manifest) { for (const run of manifest.runs) { const runPath = join(RUNS_DIR, run.runId) @@ -142,24 +460,74 @@ export class BatchManager { async compare(options: CompareOptions): Promise { const manifest = await this.createManifest(options) - return this.executeRuns(manifest) + try { + await this.preflightRuns(manifest) + return this.executeRuns(manifest) + } catch (error) { + this.delete(manifest.compareId) + throw error + } } async createManifest(options: CompareOptions): Promise { - const { providers, benchmark, judgeModel, answeringModel, sampling } = options + const { + providers, + benchmark, + judgeModel, + answeringModel, + sampling, + dataPath, + datasetRevision, + retrievalTopK, + } = options const compareId = generateCompareId() logger.info(`Loading benchmark: ${benchmark}`) const benchmarkInstance = createBenchmark(benchmark) - await benchmarkInstance.load() + await benchmarkInstance.load({ dataPath, datasetRevision, retrievalTopK }) + const requiredJudge = benchmarkInstance.protocol.requiredJudge + if (requiredJudge) { + const resolvedJudge = resolveModel(judgeModel) + if ( + resolvedJudge.provider !== requiredJudge.provider || + resolvedJudge.id !== requiredJudge.modelId + ) { + throw new Error( + `Protocol ${benchmarkInstance.protocol.identity.id} requires judge ${requiredJudge.provider}/${requiredJudge.modelId}; received ${resolvedJudge.provider}/${resolvedJudge.id}` + ) + } + } const allQuestions = benchmarkInstance.getQuestions() - let targetQuestionIds: string[] + let requestedQuestionIds: string[] if (sampling) { - targetQuestionIds = selectQuestionsBySampling(allQuestions, sampling) + requestedQuestionIds = selectQuestionsBySampling(allQuestions, sampling) } else { - targetQuestionIds = allQuestions.map((q) => q.questionId) + requestedQuestionIds = allQuestions.map((q) => q.questionId) + } + const targetQuestionIds = canonicalizeSelectedQuestionIds(allQuestions, requestedQuestionIds) + if (targetQuestionIds.length === 0) { + throw new Error("Comparison question selection is empty") + } + if (new Set(providers).size !== providers.length) { + throw new Error("Comparison providers must be unique") } + const questionById = new Map(allQuestions.map((question) => [question.questionId, question])) + const selectedQuestions = targetQuestionIds.map((questionId) => { + const question = questionById.get(questionId) + if (!question) throw new Error(`Unknown comparison question ID: ${questionId}`) + benchmarkInstance.protocol.validateQuestion(question) + return question + }) + const benchmarkInputFingerprint = fingerprintSelectedBenchmarkInput( + benchmarkInstance, + selectedQuestions + ) + const resolvedRetrievalTopK = resolveEffectiveRetrievalTopK( + benchmarkInstance.protocol, + selectedQuestions, + retrievalTopK + ) const manifest: CompareManifest = { compareId, @@ -168,11 +536,21 @@ export class BatchManager { benchmark, judge: judgeModel, answeringModel, + answeringRuntimeIdentity: resolveAnsweringRuntimeIdentity(answeringModel), sampling, targetQuestionIds, + dataPath, + datasetRevision, + retrievalTopK: resolvedRetrievalTopK, + datasetIdentity: benchmarkInstance.getDatasetIdentity?.(), + benchmarkInputFingerprint, + benchmarkScope: benchmarkInstance.scope, + protocolIdentity: benchmarkInstance.protocol.identity, + selectedQuestionIdsDigest: stableSha256(targetQuestionIds), runs: providers.map((provider) => ({ provider, runId: `${compareId}-${provider}`, + providerPromptFingerprint: fingerprintProviderPrompts(createProvider(provider).prompts), })), } @@ -184,6 +562,44 @@ export class BatchManager { return manifest } + private runOptions( + manifest: CompareManifest, + run: CompareManifest["runs"][number], + preflightOnly = false + ): OrchestratorOptions { + return { + provider: run.provider as ProviderName, + benchmark: manifest.benchmark as BenchmarkName, + judgeModel: manifest.judge, + runId: run.runId, + answeringModel: manifest.answeringModel, + questionIds: manifest.targetQuestionIds, + dataPath: manifest.dataPath, + datasetRevision: manifest.datasetRevision, + retrievalTopK: manifest.retrievalTopK, + preflightOnly, + } + } + + /** Durably validate every run without initializing providers or starting ingestion. */ + async preflightRuns(manifest: CompareManifest): Promise { + const results = await Promise.allSettled( + manifest.runs.map((run) => this.runner.run(this.runOptions(manifest, run, true))) + ) + const failures = results.flatMap((result, index) => + result.status === "rejected" + ? [ + `${manifest.runs[index].provider}: ${ + result.reason instanceof Error ? result.reason.message : String(result.reason) + }`, + ] + : [] + ) + if (failures.length > 0) { + throw new Error(`Comparison preflight failed (${failures.join("; ")})`) + } + } + async resume(compareId: string, force?: boolean): Promise { if (force) { this.delete(compareId) @@ -210,19 +626,13 @@ export class BatchManager { const results = await Promise.allSettled( manifest.runs.map(async (run) => { try { - return await orchestrator.run({ - provider: run.provider as ProviderName, - benchmark: manifest.benchmark as BenchmarkName, - judgeModel: manifest.judge, - runId: run.runId, - answeringModel: manifest.answeringModel, - questionIds: manifest.targetQuestionIds, - }) + return await this.runner.run(this.runOptions(manifest, run)) } catch (error) { // Update checkpoint status to persist the failure state const checkpoint = checkpointManager.load(run.runId) if (checkpoint) { checkpointManager.updateStatus(checkpoint, "failed") + await checkpointManager.flush(run.runId) } throw error } finally { @@ -264,6 +674,7 @@ export class BatchManager { for (const run of manifest.runs) { const report = this.loadReport(run.runId) if (report) { + assertComparisonReportMatchesManifest(manifest, run, report) reports.push({ provider: run.provider, report }) } } @@ -289,47 +700,68 @@ export class BatchManager { ) console.log("═".repeat(80)) - const sortedByAccuracy = [...reports].sort( - (a, b) => b.report.summary.accuracy - a.report.summary.accuracy + const primaryComparison = comparePrimaryMetrics(reports, manifest.runs.length) + const metricWidth = Math.max( + "Metric".length, + ...primaryComparison.rows.map(({ primaryMetric }) => (primaryMetric?.key ?? "none").length) ) - const bestAccuracy = sortedByAccuracy[0]?.provider - - console.log("\nOVERALL ACCURACY") - console.log( - "┌" + "─".repeat(17) + "┬" + "─".repeat(10) + "┬" + "─".repeat(9) + "┬" + "─".repeat(10) + "┐" + const passValues = primaryComparison.rows.map( + ({ report, passAccuracy }) => + `${(passAccuracy * 100).toFixed(1)}% (${report.summary.correctCount}/${report.summary.totalQuestions})` ) + const passWidth = Math.max("Pass accuracy".length, ...passValues.map((value) => value.length)) + const qualityBorder = (left: string, middle: string, right: string) => + left + + [17, 12, metricWidth + 2, 12, passWidth + 2].map((width) => "─".repeat(width)).join(middle) + + right + const metricDirection = primaryComparison.identity + ? `${primaryComparison.identity.key} (${primaryComparison.identity.higherIsBetter ? "higher" : "lower"} is better)` + : "incomparable primary metrics" + + console.log(`\nQUALITY — ${metricDirection}`) + if (!primaryComparison.comparable) { + console.log( + `Reports are not like-for-like (${primaryComparison.mismatchReasons.join("; ")}); preserving provider order and suppressing ranking, deltas, and winner.` + ) + } + console.log(qualityBorder("┌", "┬", "┐")) console.log( "│ " + pad("Provider", 15) + " │ " + - pad("Correct", 8) + + pad("Primary", 10) + " │ " + - pad("Total", 7) + + pad("Metric", metricWidth) + + " │ " + + pad("Δ best", 10) + " │ " + - pad("Accuracy", 8) + + pad("Pass accuracy", passWidth) + " │" ) - console.log( - "├" + "─".repeat(17) + "┼" + "─".repeat(10) + "┼" + "─".repeat(9) + "┼" + "─".repeat(10) + "┤" - ) - for (const { provider, report } of sortedByAccuracy) { - const best = provider === bestAccuracy ? " ←" : "" + console.log(qualityBorder("├", "┼", "┤")) + for (const [index, row] of primaryComparison.rows.entries()) { + const passAccuracy = passValues[index] + const delta = + row.deltaFromBest === undefined + ? "—" + : `${row.deltaFromBest > 0 ? "+" : ""}${row.deltaFromBest.toFixed(4)}` + const best = primaryComparison.winners.includes(row.provider) ? " ←" : "" + const primaryValue = row.primaryMetric ? `${row.primaryMetric.value.toFixed(4)}${best}` : "—" console.log( "│ " + - pad(provider, 15) + + pad(row.provider, 15) + + " │ " + + primaryValue.padStart(10) + " │ " + - padNum(report.summary.correctCount, 8) + + pad(row.primaryMetric?.key ?? "none", metricWidth) + " │ " + - padNum(report.summary.totalQuestions, 7) + + delta.padStart(10) + " │ " + - padPct(report.summary.accuracy, 7) + - best.padEnd(2) + + passAccuracy.padStart(passWidth) + " │" ) } - console.log( - "└" + "─".repeat(17) + "┴" + "─".repeat(10) + "┴" + "─".repeat(9) + "┴" + "─".repeat(10) + "┘" - ) + console.log(qualityBorder("└", "┴", "┘")) console.log("\nLATENCY (avg ms)") console.log( @@ -552,9 +984,10 @@ export class BatchManager { } if (allTypes.size > 0) { - console.log("\nBY QUESTION TYPE") - const providerWidth = 13 - const headerRow = ["│ " + pad("Type", 17)] + console.log("\nBY QUESTION TYPE (primary slice score / pass accuracy secondary)") + const typeWidth = Math.max("Type".length, ...[...allTypes].map((type) => type.length)) + const providerWidth = Math.max(40, ...reports.map(({ provider }) => provider.length)) + const headerRow = ["│ " + pad("Type", typeWidth)] for (const { provider } of reports) { headerRow.push(pad(provider, providerWidth)) } @@ -562,21 +995,21 @@ export class BatchManager { const borderTop = "┌" + - "─".repeat(19) + + "─".repeat(typeWidth + 2) + reports.map(() => "┬" + "─".repeat(providerWidth + 2)).join("") + "┬" + "─".repeat(15) + "┐" const borderMid = "├" + - "─".repeat(19) + + "─".repeat(typeWidth + 2) + reports.map(() => "┼" + "─".repeat(providerWidth + 2)).join("") + "┼" + "─".repeat(15) + "┤" const borderBot = "└" + - "─".repeat(19) + + "─".repeat(typeWidth + 2) + reports.map(() => "┴" + "─".repeat(providerWidth + 2)).join("") + "┴" + "─".repeat(15) + @@ -587,21 +1020,39 @@ export class BatchManager { console.log(borderMid) for (const type of [...allTypes].sort()) { - const row = ["│ " + pad(type, 17)] - let bestProvider = "" - let bestAccuracyForType = -1 - - for (const { provider, report } of reports) { - const stats = report.byQuestionType[type] - if (stats) { - row.push(padPct(stats.accuracy, providerWidth)) - if (stats.accuracy > bestAccuracyForType) { - bestAccuracyForType = stats.accuracy - bestProvider = provider - } - } else { + const row = ["│ " + pad(type, typeWidth)] + const values = reports.map(({ provider, report }) => ({ + provider, + ...getQuestionTypeQuality(report, type), + })) + const valueKeys = new Set(values.flatMap(({ value, key }) => (value == null ? [] : [key]))) + const allValuesPresent = values.every( + ({ value }) => value != null && Number.isFinite(value) + ) + const comparableSlice = + primaryComparison.comparable && allValuesPresent && valueKeys.size === 1 + const sliceValues = values.map(({ value }) => value as number) + const bestValue = comparableSlice + ? primaryComparison.identity!.higherIsBetter + ? Math.max(...sliceValues) + : Math.min(...sliceValues) + : undefined + const bestProvider = + bestValue === undefined + ? "" + : (values.find(({ value }) => value === bestValue)?.provider ?? "") + + for (const value of values) { + if (value.value == null) { row.push(pad("N/A", providerWidth)) + continue } + const primary = `${value.value.toFixed(3)} ${value.key}` + const pass = + value.passAccuracy == null + ? "pass N/A" + : `${(value.passAccuracy * 100).toFixed(1)}% pass` + row.push(pad(`${primary} / ${pass}`, providerWidth)) } row.push(pad(bestProvider, 13) + " │") console.log(row.join(" │ ")) @@ -610,10 +1061,21 @@ export class BatchManager { } console.log("\n" + "═".repeat(80)) - if (bestAccuracy) { - const bestReport = reports.find((r) => r.provider === bestAccuracy)?.report + if (primaryComparison.comparable && primaryComparison.winners.length > 0) { + const winnerRows = primaryComparison.rows.filter(({ provider }) => + primaryComparison.winners.includes(provider) + ) + const winnerNames = winnerRows.map(({ provider }) => provider).join(", ") + const winnerLabel = winnerRows.length > 1 ? "TIE" : "WINNER" + const passSummary = winnerRows + .map(({ provider, passAccuracy }) => `${provider} ${(passAccuracy * 100).toFixed(1)}%`) + .join(", ") + console.log( + `${winnerLabel}: ${winnerNames} (${primaryComparison.identity!.key}=${primaryComparison.bestValue!.toFixed(4)}; pass accuracy secondary: ${passSummary})` + ) + } else { console.log( - `WINNER: ${bestAccuracy} (${(bestReport!.summary.accuracy * 100).toFixed(1)}% overall accuracy)` + `NO WINNER: reports are not like-for-like (${primaryComparison.mismatchReasons.join("; ")})` ) } console.log("═".repeat(80) + "\n") diff --git a/src/orchestrator/builds.ts b/src/orchestrator/builds.ts new file mode 100644 index 0000000..e194a3e --- /dev/null +++ b/src/orchestrator/builds.ts @@ -0,0 +1,726 @@ +import type { Benchmark } from "../types/benchmark" +import type { BuildCheckpoint, HaystackIdentity, RunCheckpoint } from "../types/checkpoint" +import type { IngestionExecutionPolicy } from "../types/protocol" +import type { CanonicalIngestionDocument, UnifiedQuestion } from "../types/unified" +import { stableSha256 } from "../utils/stable" + +const HAYSTACK_SCHEMA_VERSION = 2 as const +const BUILD_SCHEMA_VERSION = 4 +const CONTAINER_TAG_PATTERN = /^[a-zA-Z0-9_:-]+$/ + +export interface ValidatedBuildPlan { + buildId: string + ingestionGroupId: string + memberQuestionIds: string[] + containerTag: string + haystack: HaystackIdentity + buildFingerprint: string + providerIngestionConfigFingerprint: string + ingestionExecutionPolicy: IngestionExecutionPolicy + ingestBatchSize: number + documents: CanonicalIngestionDocument[] + questions: Array<{ + questionId: string + question: string + groundTruth: string + questionType: string + questionDate?: string + }> +} + +interface QuestionPlan { + question: UnifiedQuestion + ingestionGroupId: string + documents: CanonicalIngestionDocument[] + haystack: HaystackIdentity +} + +function calculateHaystackFingerprint( + orderedSessionIds: readonly string[], + sessionFingerprints: readonly string[] +): string { + return stableSha256({ + schemaVersion: HAYSTACK_SCHEMA_VERSION, + orderedSessions: orderedSessionIds.map((sessionId, index) => ({ + index, + sessionId, + sessionFingerprint: sessionFingerprints[index], + })), + }) +} + +function createHaystackIdentity(documents: CanonicalIngestionDocument[]): HaystackIdentity { + const sessionFingerprints = documents.map((document) => + stableSha256({ + schemaVersion: HAYSTACK_SCHEMA_VERSION, + customId: document.customId, + content: document.content, + metadata: document.metadata, + // Hash the complete canonical message objects. Chat-oriented adapters + // consume role/content, while extraction-based adapters also render + // speaker/timestamp. Omitting either would let provider-visible ingestion + // requests drift inside one supposedly shared build. + messages: document.messages, + }) + ) + const orderedSessionIds = documents.map((document) => document.metadata.sessionId) + + return { + schemaVersion: HAYSTACK_SCHEMA_VERSION, + algorithm: "sha256", + fingerprint: calculateHaystackFingerprint(orderedSessionIds, sessionFingerprints), + orderedSessionIds, + sessionFingerprints, + } +} + +function assertOrderedSessionIds( + question: UnifiedQuestion, + sessionIds: readonly string[], + source: string +): void { + if (new Set(sessionIds).size !== sessionIds.length) { + throw new Error(`Question ${question.questionId} has duplicate session IDs in ${source}`) + } + if (new Set(question.haystackSessionIds).size !== question.haystackSessionIds.length) { + throw new Error(`Question ${question.questionId} has duplicate haystackSessionIds`) + } + if ( + sessionIds.length !== question.haystackSessionIds.length || + sessionIds.some((sessionId, index) => sessionId !== question.haystackSessionIds[index]) + ) { + throw new Error( + `Question ${question.questionId} haystackSessionIds do not exactly match ${source}` + ) + } +} + +function assertQuestionPlan( + question: UnifiedQuestion, + sourceSessions: readonly string[], + documents: CanonicalIngestionDocument[] +): void { + const plannedIds = documents.map((document) => document.metadata.sessionId) + assertOrderedSessionIds(question, plannedIds, "the ordered ingestion plan") + if ( + plannedIds.length !== sourceSessions.length || + plannedIds.some((sessionId, index) => sessionId !== sourceSessions[index]) + ) { + throw new Error( + `Question ${question.questionId} ingestion plan does not preserve every ordered session returned by getHaystackSessions()` + ) + } + + for (const [index, document] of documents.entries()) { + if (!document.customId || !document.content || !document.metadata.sessionId) { + throw new Error( + `Question ${question.questionId} has malformed ingestion document at index ${index}` + ) + } + if (document.customId !== document.metadata.sessionId) { + throw new Error( + `Question ${question.questionId} document ${index} must use customId=sessionId` + ) + } + } +} + +function firstHaystackDifference(left: QuestionPlan, right: QuestionPlan): string { + const max = Math.max( + left.haystack.orderedSessionIds.length, + right.haystack.orderedSessionIds.length + ) + for (let index = 0; index < max; index++) { + const leftId = left.haystack.orderedSessionIds[index] + const rightId = right.haystack.orderedSessionIds[index] + const leftFingerprint = left.haystack.sessionFingerprints[index] + const rightFingerprint = right.haystack.sessionFingerprints[index] + if (leftId !== rightId || leftFingerprint !== rightFingerprint) { + return `index ${index}: ${leftId ?? ""}/${leftFingerprint ?? ""} != ${rightId ?? ""}/${rightFingerprint ?? ""}` + } + } + return "final fingerprint differs" +} + +function createContainerTag(buildId: string): string { + const tag = `mb:${stableSha256({ schemaVersion: 1, buildId }).slice(0, 48)}` + if (tag.length > 100 || !CONTAINER_TAG_PATTERN.test(tag)) { + throw new Error(`Generated invalid containerTag: ${tag}`) + } + return tag +} + +export function prepareValidatedBuildPlans(input: { + benchmark: Benchmark + questions: UnifiedQuestion[] + provider: string + providerAdapterVersion: string + providerPromptFingerprint: string + providerIngestionConfigFingerprint: string + dataSourceRunId: string + ingestBatchSize?: number +}): ValidatedBuildPlan[] { + const { + benchmark, + questions, + provider, + providerAdapterVersion, + providerIngestionConfigFingerprint, + dataSourceRunId, + } = input + const ingestBatchSize = input.ingestBatchSize ?? 1 + if (!Number.isInteger(ingestBatchSize) || ingestBatchSize < 1 || ingestBatchSize > 600) { + throw new Error( + `Ingest batch size must be an integer between 1 and 600; received ${ingestBatchSize}` + ) + } + const independentlyPlanned = questions.map((question): QuestionPlan => { + benchmark.protocol.validateQuestion(question) + const sessions = benchmark.getHaystackSessions(question.questionId) + const sourceSessionIds = sessions.map((session) => session.sessionId) + assertOrderedSessionIds( + question, + sourceSessionIds, + "the ordered sessions returned by getHaystackSessions()" + ) + const documents = benchmark.protocol.createIngestionPlan({ question, sessions }) + assertQuestionPlan(question, sourceSessionIds, documents) + return { + question, + ingestionGroupId: benchmark.getIngestionGroupId?.(question.questionId) || question.questionId, + documents, + haystack: createHaystackIdentity(documents), + } + }) + + const grouped = new Map() + for (const plan of independentlyPlanned) { + const members = grouped.get(plan.ingestionGroupId) || [] + members.push(plan) + grouped.set(plan.ingestionGroupId, members) + } + + return [...grouped.entries()].map(([ingestionGroupId, members]) => { + const reference = members[0] + for (const member of members.slice(1)) { + if (member.haystack.fingerprint !== reference.haystack.fingerprint) { + throw new Error( + `Ingestion group ${ingestionGroupId} has different haystacks for ${reference.question.questionId} (${reference.haystack.fingerprint}) and ${member.question.questionId} (${member.haystack.fingerprint}); ${firstHaystackDifference(reference, member)}` + ) + } + } + + const buildFingerprint = stableSha256({ + schemaVersion: BUILD_SCHEMA_VERSION, + haystackFingerprint: reference.haystack.fingerprint, + datasetFingerprint: benchmark.getDatasetIdentity?.()?.datasetFingerprint ?? null, + provider, + providerAdapterVersion, + providerIngestionConfigFingerprint, + protocolIngestionPolicyHash: benchmark.protocol.identity.ingestionPolicyHash, + ingestionExecutionPolicy: benchmark.protocol.ingestionExecutionPolicy, + ...(ingestBatchSize === 1 ? {} : { ingestBatchSize }), + }) + const buildId = `build:${stableSha256({ ingestionGroupId, buildFingerprint, dataSourceRunId }).slice(0, 48)}` + + return { + buildId, + ingestionGroupId, + memberQuestionIds: members.map((member) => member.question.questionId).sort(), + containerTag: createContainerTag(buildId), + haystack: reference.haystack, + buildFingerprint, + providerIngestionConfigFingerprint, + ingestionExecutionPolicy: benchmark.protocol.ingestionExecutionPolicy, + ingestBatchSize, + documents: reference.documents, + questions: members + .map(({ question }) => ({ + questionId: question.questionId, + question: question.question, + groundTruth: question.groundTruth, + questionType: question.questionType, + ...(typeof question.metadata?.questionDate === "string" + ? { questionDate: question.metadata.questionDate } + : {}), + })) + .sort((left, right) => left.questionId.localeCompare(right.questionId)), + } + }) +} + +export function createBuildCheckpoint(plan: ValidatedBuildPlan): BuildCheckpoint { + return { + buildId: plan.buildId, + ingestionGroupId: plan.ingestionGroupId, + memberQuestionIds: [...plan.memberQuestionIds], + containerTag: plan.containerTag, + haystack: plan.haystack, + buildFingerprint: plan.buildFingerprint, + providerIngestionConfigFingerprint: plan.providerIngestionConfigFingerprint, + ingestionExecutionPolicy: plan.ingestionExecutionPolicy, + ingestBatchSize: plan.ingestBatchSize, + sessions: plan.documents.map((document) => ({ + sessionId: document.metadata.sessionId, + documentDate: document.metadata.documentDate, + messageCount: document.messages?.length ?? 0, + })), + missingDocumentDateCount: plan.documents.filter((document) => !document.metadata.documentDate) + .length, + reused: false, + ingest: { + status: "pending", + completedSessionIds: [], + documentIds: [], + taskIds: [], + deferredSessions: [], + attempts: [], + }, + indexing: { + status: "pending", + completedIds: [], + failedIds: [], + attempts: [], + }, + } +} + +export function assertCompletedSessionsAreOrderedPrefix(build: BuildCheckpoint): void { + const completed = build.ingest.completedSessionIds + if (new Set(completed).size !== completed.length) { + throw new Error(`Build ${build.buildId} has duplicate completedSessionIds`) + } + if (completed.some((sessionId, index) => sessionId !== build.haystack.orderedSessionIds[index])) { + throw new Error( + `Build ${build.buildId} completedSessionIds must be an ordered prefix of its haystack` + ) + } +} + +function assertUniqueIds(buildId: string, name: string, ids: readonly string[]): void { + if (ids.some((id) => typeof id !== "string" || id.length === 0)) { + throw new Error(`Build ${buildId} has an empty or invalid ${name} entry`) + } + if (new Set(ids).size !== ids.length) { + throw new Error(`Build ${buildId} has duplicate ${name}`) + } +} + +export function assertBuildCheckpointConsistency(build: BuildCheckpoint): void { + if (!build.providerIngestionConfigFingerprint?.trim()) { + throw new Error(`Build ${build.buildId} has no provider ingestion configuration fingerprint`) + } + if ( + !["after-build", "after-each-document"].includes( + build.ingestionExecutionPolicy?.readinessBarrier + ) || + !["provider-default", "instant"].includes(build.ingestionExecutionPolicy?.processingMode) + ) { + throw new Error(`Build ${build.buildId} has an invalid ingestion execution policy`) + } + const ingestBatchSize = build.ingestBatchSize ?? 1 + if (!Number.isInteger(ingestBatchSize) || ingestBatchSize < 1 || ingestBatchSize > 600) { + throw new Error(`Build ${build.buildId} has an invalid ingest batch size`) + } + if (build.haystack.schemaVersion !== HAYSTACK_SCHEMA_VERSION) { + throw new Error( + `Build ${build.buildId} uses unsupported haystack schema ${String(build.haystack.schemaVersion)}` + ) + } + if (build.haystack.algorithm !== "sha256") { + throw new Error( + `Build ${build.buildId} uses unsupported haystack algorithm ${String(build.haystack.algorithm)}` + ) + } + assertCompletedSessionsAreOrderedPrefix(build) + const deferredSessions = build.ingest.deferredSessions ?? [] + const deferredSequences = new Set() + for (const deferred of deferredSessions) { + if ( + !Number.isInteger(deferred.sequence) || + deferred.sequence < 0 || + deferred.sequence >= build.haystack.orderedSessionIds.length || + deferredSequences.has(deferred.sequence) || + build.haystack.orderedSessionIds[deferred.sequence] !== deferred.sessionId || + !deferred.customId?.trim() || + (deferred.stage !== "submission" && deferred.stage !== "readiness") || + !Number.isInteger(deferred.attempts) || + deferred.attempts < 1 || + !deferred.firstFailedAt?.trim() || + !deferred.lastFailedAt?.trim() || + !deferred.lastError?.trim() + ) { + throw new Error(`Build ${build.buildId} has invalid deferred ingest state`) + } + deferredSequences.add(deferred.sequence) + assertUniqueIds(build.buildId, "deferred document IDs", deferred.documentIds) + assertUniqueIds(build.buildId, "deferred task IDs", deferred.taskIds) + } + assertUniqueIds(build.buildId, "haystack session IDs", build.haystack.orderedSessionIds) + if (build.haystack.sessionFingerprints.length !== build.haystack.orderedSessionIds.length) { + throw new Error(`Build ${build.buildId} has mismatched haystack fingerprint arrays`) + } + if ( + calculateHaystackFingerprint( + build.haystack.orderedSessionIds, + build.haystack.sessionFingerprints + ) !== build.haystack.fingerprint + ) { + throw new Error(`Build ${build.buildId} has a tampered haystack fingerprint`) + } + if ( + build.ingest.status === "completed" && + build.ingest.completedSessionIds.length !== build.haystack.orderedSessionIds.length + ) { + throw new Error(`Build ${build.buildId} is marked ingested before every session completed`) + } + if (build.ingest.status === "completed" && deferredSessions.length > 0) { + throw new Error(`Build ${build.buildId} is marked ingested with deferred sessions remaining`) + } + if ( + build.ingestionExecutionPolicy.readinessBarrier === "after-build" && + build.ingest.status !== "completed" && + build.indexing.status !== "pending" + ) { + throw new Error(`Build ${build.buildId} started indexing before ingestion completed`) + } + if (build.indexing.status === "completed" && build.ingest.status !== "completed") { + throw new Error(`Build ${build.buildId} completed indexing before ingestion completed`) + } + if ( + build.ingestionExecutionPolicy.readinessBarrier === "after-each-document" && + build.ingest.status === "completed" && + build.indexing.status !== "completed" + ) { + throw new Error( + `Build ${build.buildId} completed causal ingestion without completing its per-document indexing barriers` + ) + } + + assertUniqueIds(build.buildId, "document IDs", build.ingest.documentIds) + assertUniqueIds(build.buildId, "task IDs", build.ingest.taskIds) + const duplicatePhysicalId = build.ingest.documentIds.find((id) => + build.ingest.taskIds.includes(id) + ) + if (duplicatePhysicalId) { + throw new Error( + `Build ${build.buildId} uses ${duplicatePhysicalId} as both a document ID and task ID` + ) + } + assertUniqueIds(build.buildId, "completed indexing IDs", build.indexing.completedIds) + assertUniqueIds(build.buildId, "failed indexing IDs", build.indexing.failedIds) + const expectedIds = new Set([...build.ingest.documentIds, ...build.ingest.taskIds]) + for (const deferred of deferredSessions) { + for (const id of [...deferred.documentIds, ...deferred.taskIds]) { + if (!expectedIds.has(id)) { + throw new Error(`Build ${build.buildId} deferred unknown physical ID ${id}`) + } + } + } + const completedIds = new Set(build.indexing.completedIds) + const failedIds = new Set(build.indexing.failedIds) + for (const id of [...completedIds, ...failedIds]) { + if (!expectedIds.has(id)) { + throw new Error(`Build ${build.buildId} has indexing progress for unknown ID ${id}`) + } + } + for (const id of completedIds) { + if (failedIds.has(id)) { + throw new Error(`Build ${build.buildId} indexed ID ${id} is both completed and failed`) + } + } + if (build.indexing.status === "completed") { + if (failedIds.size > 0) { + throw new Error(`Build ${build.buildId} is marked indexed with failed IDs`) + } + if ( + completedIds.size !== expectedIds.size || + [...expectedIds].some((id) => !completedIds.has(id)) + ) { + throw new Error(`Build ${build.buildId} is marked indexed before every ID completed`) + } + } +} + +function assertCompletedQuestionPayload(question: RunCheckpoint["questions"][string]): void { + const { search, answer, evaluate } = question.phases + if (search.status === "completed") { + if (!search.retrievalPlan || !Array.isArray(search.results)) { + throw new Error(`Question ${question.questionId} has incomplete completed-search state`) + } + const requestedTopK = search.retrievalPlan.requestedTopK + const counts = [ + search.requestedCount, + search.rawReturnedCount, + search.returnedCount, + search.normalizedCount, + search.droppedCount, + search.answerCutoff, + ] + if (counts.some((value) => !Number.isInteger(value) || value! < 0)) { + throw new Error(`Question ${question.questionId} has invalid completed-search counts`) + } + if ( + search.requestedCount !== requestedTopK || + search.answerCutoff !== search.retrievalPlan.answerCutoff || + search.returnedCount !== search.results.length || + search.normalizedCount !== search.results.length || + search.rawReturnedCount! - search.normalizedCount! !== search.droppedCount || + search.rawReturnedCount! > requestedTopK || + !Array.isArray(search.providerRequests) || + search.providerRequests.reduce((sum, request) => sum + request.limit, 0) !== requestedTopK || + !search.resultFile?.trim() + ) { + throw new Error(`Question ${question.questionId} has inconsistent completed-search state`) + } + const ranks = new Set() + for (const result of search.results) { + if ( + !result.text?.trim() || + !Number.isInteger(result.rank) || + result.rank < 1 || + ranks.has(result.rank) + ) { + throw new Error(`Question ${question.questionId} has malformed persisted search evidence`) + } + ranks.add(result.rank) + } + } + + if (search.answerEvidenceCount !== undefined) { + if ( + !Number.isInteger(search.answerEvidenceCount) || + search.answerEvidenceCount < 0 || + search.answerEvidenceCount > (search.answerCutoff ?? -1) + ) { + throw new Error(`Question ${question.questionId} has invalid answer evidence count`) + } + } + + if (answer.status === "completed") { + const hasOrdinaryHypothesis = + typeof answer.hypothesis === "string" && answer.hypothesis.trim().length > 0 + const hasAcceptedTerminalEmpty = + answer.hypothesis === "" && answer.terminalEmptyAccepted === true + if (!hasOrdinaryHypothesis && !hasAcceptedTerminalEmpty) { + throw new Error(`Question ${question.questionId} has no completed answer hypothesis`) + } + if (hasOrdinaryHypothesis && answer.terminalEmptyAccepted === true) { + throw new Error( + `Question ${question.questionId} marks a non-empty hypothesis as terminal-empty` + ) + } + for (const [name, value] of Object.entries({ + promptTokens: answer.promptTokens, + basePromptTokens: answer.basePromptTokens, + contextTokens: answer.contextTokens, + evidenceCount: answer.evidenceCount, + })) { + if (!Number.isInteger(value) || value! < 0) { + throw new Error(`Question ${question.questionId} has invalid completed-answer ${name}`) + } + } + if ( + answer.evidenceCount !== search.answerEvidenceCount || + answer.evidenceCount! > (search.answerCutoff ?? -1) + ) { + throw new Error(`Question ${question.questionId} has inconsistent completed-answer evidence`) + } + } + + if (evaluate.status === "completed") { + const evaluation = evaluate.evaluation + if ( + !evaluation || + evaluation.questionId !== question.questionId || + evaluation.questionType !== question.questionType || + !Number.isFinite(evaluation.primaryScore) || + evaluation.primaryScore < 0 || + evaluation.primaryScore > 1 || + typeof evaluation.passed !== "boolean" || + evaluate.score !== evaluation.primaryScore + ) { + throw new Error(`Question ${question.questionId} has incomplete completed-evaluation state`) + } + if (evaluate.label !== (evaluation.passed ? "correct" : "incorrect")) { + throw new Error(`Question ${question.questionId} has inconsistent completed-evaluation label`) + } + } +} + +export function assertCheckpointReferences(checkpoint: RunCheckpoint): void { + const questionIds = Object.keys(checkpoint.questions) + const knownQuestionIds = new Set(questionIds) + if (checkpoint.targetQuestionIds) { + if ( + new Set(checkpoint.targetQuestionIds).size !== checkpoint.targetQuestionIds.length || + stableSha256([...checkpoint.targetQuestionIds].sort()) !== + stableSha256([...questionIds].sort()) + ) { + throw new Error("Checkpoint targetQuestionIds do not match its question records") + } + } + for (const [buildKey, build] of Object.entries(checkpoint.builds)) { + if (buildKey !== build.buildId) { + throw new Error(`Checkpoint build key ${buildKey} does not match ${build.buildId}`) + } + assertUniqueIds(build.buildId, "member question IDs", build.memberQuestionIds) + for (const memberQuestionId of build.memberQuestionIds) { + const member = checkpoint.questions[memberQuestionId] + if (!knownQuestionIds.has(memberQuestionId) || member?.buildId !== build.buildId) { + throw new Error( + `Build ${build.buildId} has invalid member question reference ${memberQuestionId}` + ) + } + } + assertBuildCheckpointConsistency(build) + } + + for (const [questionKey, question] of Object.entries(checkpoint.questions)) { + if (questionKey !== question.questionId) { + throw new Error( + `Checkpoint question key ${questionKey} does not match ${question.questionId}` + ) + } + const build = checkpoint.builds[question.buildId] + if (!build || !build.memberQuestionIds.includes(question.questionId)) { + throw new Error( + `Checkpoint question ${question.questionId} has an invalid build reference ${question.buildId}` + ) + } + assertCompletedQuestionPayload(question) + if (question.phases.search.status !== "pending" && build.indexing.status !== "completed") { + throw new Error(`Question ${question.questionId} searched before its build was fully indexed`) + } + if ( + question.phases.answer.status !== "pending" && + question.phases.search.status !== "completed" + ) { + throw new Error(`Question ${question.questionId} answered before search completed`) + } + if ( + question.phases.evaluate.status !== "pending" && + question.phases.answer.status !== "completed" + ) { + throw new Error(`Question ${question.questionId} evaluated before answering completed`) + } + } +} + +export function assertResumeBuilds(checkpoint: RunCheckpoint, plans: ValidatedBuildPlan[]): void { + assertCheckpointReferences(checkpoint) + const plannedById = new Map(plans.map((plan) => [plan.buildId, plan])) + const checkpointIds = Object.keys(checkpoint.builds).sort() + const planIds = [...plannedById.keys()].sort() + if (stableSha256(checkpointIds) !== stableSha256(planIds)) { + throw new Error( + "Checkpoint build identities do not match the selected dataset and configuration" + ) + } + + const expectedQuestionToBuild = new Map() + for (const plan of plans) { + for (const questionId of plan.memberQuestionIds) { + if (expectedQuestionToBuild.has(questionId)) { + throw new Error(`Question ${questionId} belongs to more than one validated build`) + } + expectedQuestionToBuild.set(questionId, plan.buildId) + } + } + const checkpointQuestionIds = Object.keys(checkpoint.questions).sort() + const expectedQuestionIds = [...expectedQuestionToBuild.keys()].sort() + if (stableSha256(checkpointQuestionIds) !== stableSha256(expectedQuestionIds)) { + throw new Error("Checkpoint questions do not match the selected dataset and configuration") + } + if (checkpoint.targetQuestionIds) { + if (new Set(checkpoint.targetQuestionIds).size !== checkpoint.targetQuestionIds.length) { + throw new Error("Checkpoint targetQuestionIds contains duplicates") + } + if ( + stableSha256([...checkpoint.targetQuestionIds].sort()) !== stableSha256(expectedQuestionIds) + ) { + throw new Error("Checkpoint targetQuestionIds do not match its question records") + } + } + + for (const [buildKey, build] of Object.entries(checkpoint.builds)) { + const plan = plannedById.get(build.buildId) + if ( + buildKey !== build.buildId || + !plan || + plan.ingestionGroupId !== build.ingestionGroupId || + plan.buildFingerprint !== build.buildFingerprint || + plan.providerIngestionConfigFingerprint !== build.providerIngestionConfigFingerprint || + stableSha256(plan.ingestionExecutionPolicy) !== + stableSha256(build.ingestionExecutionPolicy) || + plan.ingestBatchSize !== (build.ingestBatchSize ?? 1) || + plan.containerTag !== build.containerTag || + plan.haystack.fingerprint !== build.haystack.fingerprint || + stableSha256(plan.haystack.orderedSessionIds) !== + stableSha256(build.haystack.orderedSessionIds) || + stableSha256(plan.haystack.sessionFingerprints) !== + stableSha256(build.haystack.sessionFingerprints) || + stableSha256(plan.memberQuestionIds) !== stableSha256(build.memberQuestionIds) + ) { + throw new Error( + `Checkpoint build ${build.buildId} no longer matches the validated build plan` + ) + } + const expectedSessions = plan.documents.map((document) => ({ + sessionId: document.metadata.sessionId, + documentDate: document.metadata.documentDate, + messageCount: document.messages?.length ?? 0, + })) + if ( + stableSha256(expectedSessions) !== stableSha256(build.sessions) || + build.missingDocumentDateCount !== + plan.documents.filter((document) => !document.metadata.documentDate).length + ) { + throw new Error(`Checkpoint build ${build.buildId} has inconsistent session metadata`) + } + assertBuildCheckpointConsistency(build) + + for (const expectedQuestion of plan.questions) { + const question = checkpoint.questions[expectedQuestion.questionId] + if ( + !question || + question.buildId !== plan.buildId || + question.question !== expectedQuestion.question || + question.groundTruth !== expectedQuestion.groundTruth || + question.questionType !== expectedQuestion.questionType || + question.questionDate !== expectedQuestion.questionDate + ) { + throw new Error( + `Checkpoint question ${expectedQuestion.questionId} no longer matches its validated build plan` + ) + } + } + } +} + +/** + * Reuse only completed provider builds. Query-time state is intentionally not + * copied, allowing a new retrieval/answer/evaluation protocol to start cleanly + * while retaining the exact validated ingestion containers. + */ +export function cloneCompletedBuildsForReuse( + source: RunCheckpoint, + plans: ValidatedBuildPlan[] +): BuildCheckpoint[] { + assertResumeBuilds(source, plans) + return plans.map((plan) => { + const sourceBuild = source.builds[plan.buildId] + if (!sourceBuild) throw new Error(`Source run is missing build ${plan.buildId}`) + if (sourceBuild.ingest.status !== "completed") { + throw new Error(`Cannot reuse build ${plan.buildId}; ingestion is incomplete`) + } + if (sourceBuild.indexing.status !== "completed" || sourceBuild.indexing.failedIds.length > 0) { + throw new Error(`Cannot reuse build ${plan.buildId}; indexing is incomplete or failed`) + } + const build = structuredClone(sourceBuild) + build.sourceRunId = source.runId + build.reused = true + build.reusedPhases = { ingest: true, indexing: true } + return build + }) +} diff --git a/src/orchestrator/checkpoint.ts b/src/orchestrator/checkpoint.ts index aa00835..1aac0cc 100644 --- a/src/orchestrator/checkpoint.ts +++ b/src/orchestrator/checkpoint.ts @@ -1,28 +1,162 @@ import { + closeSync, + cpSync, existsSync, - readFileSync, - writeFileSync, + fsyncSync, mkdirSync, - rmSync, + openSync, + readFileSync, readdirSync, - cpSync, renameSync, + rmSync, + truncateSync, unlinkSync, -} from "fs" -import { join } from "path" -import type { - RunCheckpoint, - QuestionCheckpoint, - PhaseStatus, - PhaseId, - RunStatus, - SamplingConfig, -} from "../types/checkpoint" + writeFileSync, +} from "node:fs" +import { join } from "node:path" +import type { BenchmarkScope, DatasetIdentity } from "../types/benchmark" import type { ConcurrencyConfig } from "../types/concurrency" -import { PHASE_ORDER } from "../types/checkpoint" +import { + CHECKPOINT_SCHEMA_VERSION, + PHASE_ORDER, + type BuildCheckpoint, + type PhaseId, + type PhaseStatus, + type QuestionCheckpoint, + type RunCheckpoint, + type RunStatus, + type SamplingConfig, +} from "../types/checkpoint" +import type { ProtocolIdentity } from "../types/protocol" import { logger } from "../utils/logger" +import { resolveAnsweringRuntimeIdentity } from "../utils/models" +import { sha256Text } from "../utils/stable" +import { assertBuildCheckpointConsistency, assertCheckpointReferences } from "./builds" const RUNS_DIR = "./data/runs" +const INGEST_PROGRESS_JOURNAL_SCHEMA_VERSION = 3 + +interface IngestProgressJournalDeferredFailure { + customId: string + stage: "submission" | "readiness" + attempts: number + firstFailedAt: string + lastFailedAt: string + lastError: string +} + +interface IngestProgressJournalPayload { + schemaVersion: 2 | typeof INGEST_PROGRESS_JOURNAL_SCHEMA_VERSION + buildId: string + buildFingerprint: string + sequence: number + sessionId: string + documentIds: string[] + taskIds: string[] + /** True only after the provider confirmed every physical ID is query-ready. */ + readyForNextSession: boolean + /** Present when the ordered step advanced into the durable end-of-build retry queue. */ + deferredFailure?: IngestProgressJournalDeferredFailure +} + +interface IngestProgressJournalRecord extends IngestProgressJournalPayload { + checksum: string +} + +function createIngestProgressJournalRecord( + payload: IngestProgressJournalPayload +): IngestProgressJournalRecord { + return { ...payload, checksum: sha256Text(JSON.stringify(payload)) } +} + +function assertStringArray(value: unknown, field: string): asserts value is string[] { + if (!Array.isArray(value) || value.some((item) => typeof item !== "string" || !item.trim())) { + throw new Error(`Ingest progress journal has invalid ${field}`) + } +} + +function parseIngestProgressJournalRecord( + line: string, + journalPath: string, + lineNumber: number +): IngestProgressJournalRecord { + let candidate: unknown + try { + candidate = JSON.parse(line) + } catch (error) { + throw new Error( + `Ingest progress journal ${journalPath} line ${lineNumber} is malformed: ${String(error)}` + ) + } + if (!candidate || typeof candidate !== "object") { + throw new Error(`Ingest progress journal ${journalPath} line ${lineNumber} is not an object`) + } + const record = candidate as Partial + if ( + (record.schemaVersion !== 2 && + record.schemaVersion !== INGEST_PROGRESS_JOURNAL_SCHEMA_VERSION) || + typeof record.buildId !== "string" || + !record.buildId || + typeof record.buildFingerprint !== "string" || + !record.buildFingerprint || + typeof record.sequence !== "number" || + !Number.isInteger(record.sequence) || + (record.sequence ?? -1) < 0 || + typeof record.sessionId !== "string" || + !record.sessionId || + typeof record.readyForNextSession !== "boolean" || + typeof record.checksum !== "string" + ) { + throw new Error(`Ingest progress journal ${journalPath} line ${lineNumber} is invalid`) + } + assertStringArray(record.documentIds, "documentIds") + assertStringArray(record.taskIds, "taskIds") + let deferredFailure: IngestProgressJournalDeferredFailure | undefined + if (record.schemaVersion === INGEST_PROGRESS_JOURNAL_SCHEMA_VERSION) { + const candidateFailure = record.deferredFailure as + | Partial + | undefined + if (candidateFailure !== undefined) { + if ( + !candidateFailure || + typeof candidateFailure.customId !== "string" || + !candidateFailure.customId || + (candidateFailure.stage !== "submission" && candidateFailure.stage !== "readiness") || + !Number.isInteger(candidateFailure.attempts) || + (candidateFailure.attempts ?? 0) < 1 || + typeof candidateFailure.firstFailedAt !== "string" || + !candidateFailure.firstFailedAt || + typeof candidateFailure.lastFailedAt !== "string" || + !candidateFailure.lastFailedAt || + typeof candidateFailure.lastError !== "string" || + !candidateFailure.lastError + ) { + throw new Error( + `Ingest progress journal ${journalPath} line ${lineNumber} has invalid deferred failure` + ) + } + deferredFailure = candidateFailure as IngestProgressJournalDeferredFailure + } + } + const payload: IngestProgressJournalPayload = { + schemaVersion: record.schemaVersion, + buildId: record.buildId, + buildFingerprint: record.buildFingerprint, + sequence: record.sequence, + sessionId: record.sessionId, + documentIds: record.documentIds, + taskIds: record.taskIds, + readyForNextSession: record.readyForNextSession, + ...(deferredFailure ? { deferredFailure } : {}), + } + const expectedChecksum = sha256Text(JSON.stringify(payload)) + if (record.checksum !== expectedChecksum) { + throw new Error( + `Ingest progress journal ${journalPath} line ${lineNumber} checksum does not match` + ) + } + return { ...payload, checksum: record.checksum } +} export class CheckpointManager { private basePath: string @@ -44,6 +178,17 @@ export class CheckpointManager { return join(this.getRunPath(runId), "results") } + getIngestProgressJournalPath(runId: string, buildId: string): string { + return join(this.getRunPath(runId), "progress", "ingest", `${sha256Text(buildId)}.jsonl`) + } + + getQuestionResultsPath(runId: string, questionId: string): string { + // Canonical BEAM IDs contain colons, which are invalid in Windows file + // names. Keep the original ID inside the artifact and use a portable, + // collision-resistant filename on disk. + return join(this.getResultsDir(runId), `${sha256Text(questionId)}.json`) + } + exists(runId: string): boolean { return existsSync(this.getCheckpointPath(runId)) } @@ -52,68 +197,111 @@ export class CheckpointManager { const path = this.getCheckpointPath(runId) if (!existsSync(path)) return null + let parsed: unknown try { - const data = readFileSync(path, "utf8") - return JSON.parse(data) as RunCheckpoint - } catch (e) { - logger.warn(`Failed to load checkpoint: ${e}`) - return null + parsed = JSON.parse(readFileSync(path, "utf8")) + } catch (error) { + throw new Error(`Checkpoint ${runId} is unreadable: ${String(error)}`) } + + const schemaVersion = (parsed as { schemaVersion?: unknown })?.schemaVersion + if (schemaVersion !== CHECKPOINT_SCHEMA_VERSION) { + throw new Error( + `Checkpoint ${runId} uses unsupported schema ${String(schemaVersion ?? "legacy")}; PR #44 shared-build runs require a new schema-${CHECKPOINT_SCHEMA_VERSION} run` + ) + } + + const checkpoint = parsed as Partial + this.replayIngestProgressJournals(checkpoint as RunCheckpoint) + const missingIdentityFields = [ + typeof checkpoint.benchmarkInputFingerprint === "string" && + checkpoint.benchmarkInputFingerprint.trim() + ? null + : "benchmarkInputFingerprint", + checkpoint.answeringRuntimeIdentity && typeof checkpoint.answeringRuntimeIdentity === "object" + ? null + : "answeringRuntimeIdentity", + Number.isInteger(checkpoint.retrievalTopK) && (checkpoint.retrievalTopK ?? 0) > 0 + ? null + : "retrievalTopK", + typeof checkpoint.protocolIdentity?.ingestionPolicyHash === "string" && + checkpoint.protocolIdentity.ingestionPolicyHash.trim() + ? null + : "protocolIdentity.ingestionPolicyHash", + ].filter((value): value is string => value !== null) + if (missingIdentityFields.length > 0) { + throw new Error( + `Checkpoint ${runId} has incomplete schema-${CHECKPOINT_SCHEMA_VERSION} identity: ${missingIdentityFields.join(", ")}` + ) + } + + return checkpoint as RunCheckpoint } save(checkpoint: RunCheckpoint): void { const currentQueue = this.saveLock.get(checkpoint.runId) || Promise.resolve() - const nextQueue = currentQueue.then(() => this._performSave(checkpoint)) + // A later full-checkpoint save can safely recover from an earlier failed + // write because it contains the complete current state. Keep the rejected + // tail registered until flush() observes it, and suppress only the runtime's + // unhandled-rejection warning—not the error returned to flush(). + const nextQueue = currentQueue.catch(() => undefined).then(() => this.performSave(checkpoint)) this.saveLock.set(checkpoint.runId, nextQueue) - - nextQueue.finally(() => { - if (this.saveLock.get(checkpoint.runId) === nextQueue) { - this.saveLock.delete(checkpoint.runId) - } - }) + void nextQueue.catch(() => undefined) } - private async _performSave(checkpoint: RunCheckpoint): Promise { + private async performSave(checkpoint: RunCheckpoint): Promise { const runPath = this.getRunPath(checkpoint.runId) const path = this.getCheckpointPath(checkpoint.runId) - const tempPath = path + ".tmp" - - if (!existsSync(runPath)) { - mkdirSync(runPath, { recursive: true }) - } - + const tempPath = `${path}.tmp` + mkdirSync(runPath, { recursive: true }) checkpoint.updatedAt = new Date().toISOString() - let lastError: any - - // Windows often locks files briefly (EPERM/EBUSY), so we retry a few times + let lastError: unknown for (let attempt = 0; attempt < 5; attempt++) { try { writeFileSync(tempPath, JSON.stringify(checkpoint, null, 2)) renameSync(tempPath, path) - return // Success - } catch (e: any) { - lastError = e - if (e.code !== "EPERM" && e.code !== "EBUSY") { - break // Don't retry other errors - } - // Wait with exponential backoff: 50, 100, 200, 400, 800ms - await new Promise((resolve) => setTimeout(resolve, 50 * Math.pow(2, attempt))) + return + } catch (error) { + lastError = error + const code = (error as NodeJS.ErrnoException).code + if (code !== "EPERM" && code !== "EBUSY") break + await new Promise((resolve) => setTimeout(resolve, 50 * 2 ** attempt)) } } - // If we get here, all retries failed or it was a non-retriable error try { unlinkSync(tempPath) - } catch { } + } catch { + // Best-effort cleanup only. + } throw lastError } async flush(runId?: string): Promise { if (runId) { - await this.saveLock.get(runId) - } else { - await Promise.all(Array.from(this.saveLock.values())) + while (true) { + const pending = this.saveLock.get(runId) + if (!pending) return + try { + await pending + } catch (error) { + // If another save was queued after this failure, it is a full-state + // retry. Await that newer tail before deciding persistence failed. + if (this.saveLock.get(runId) !== pending) continue + this.saveLock.delete(runId) + throw error + } + if (this.saveLock.get(runId) === pending) { + this.saveLock.delete(runId) + return + } + } + } + + while (this.saveLock.size > 0) { + const runIds = [...this.saveLock.keys()] + await Promise.all(runIds.map((pendingRunId) => this.flush(pendingRunId))) } } @@ -123,42 +311,66 @@ export class CheckpointManager { benchmark: string, judge: string, answeringModel: string, - options?: { + options: { + providerAdapterVersion: string + providerPromptFingerprint: string + benchmarkScope: BenchmarkScope + protocolIdentity: ProtocolIdentity + selectedQuestionIdsDigest: string + datasetIdentity?: DatasetIdentity + benchmarkInputFingerprint: string + dataPath?: string + datasetRevision?: string + retrievalTopK: number + evaluationProfile?: string + answerCutoff?: number limit?: number sampling?: SamplingConfig targetQuestionIds?: string[] dataSourceRunId?: string status?: RunStatus concurrency?: ConcurrencyConfig + ingestBatchSize?: number + ingestReadinessTimeoutMs?: number } ): RunCheckpoint { + const now = new Date().toISOString() const checkpoint: RunCheckpoint = { + schemaVersion: CHECKPOINT_SCHEMA_VERSION, runId, - dataSourceRunId: options?.dataSourceRunId || runId, - status: options?.status || "initializing", + dataSourceRunId: options.dataSourceRunId || runId, + status: options.status || "initializing", provider, + providerAdapterVersion: options.providerAdapterVersion, + providerPromptFingerprint: options.providerPromptFingerprint, benchmark, + benchmarkScope: options.benchmarkScope, + datasetIdentity: options.datasetIdentity, + benchmarkInputFingerprint: options.benchmarkInputFingerprint, + selectedQuestionIdsDigest: options.selectedQuestionIdsDigest, + protocolIdentity: options.protocolIdentity, judge, answeringModel, - createdAt: new Date().toISOString(), - updatedAt: new Date().toISOString(), - limit: options?.limit, - sampling: options?.sampling, - targetQuestionIds: options?.targetQuestionIds, - concurrency: options?.concurrency, + answeringRuntimeIdentity: resolveAnsweringRuntimeIdentity(answeringModel), + createdAt: now, + updatedAt: now, + dataPath: options.dataPath, + datasetRevision: options.datasetRevision, + retrievalTopK: options.retrievalTopK, + evaluationProfile: options.evaluationProfile, + answerCutoff: options.answerCutoff, + limit: options.limit, + sampling: options.sampling, + targetQuestionIds: options.targetQuestionIds, + concurrency: options.concurrency, + ingestBatchSize: options.ingestBatchSize, + ingestReadinessTimeoutMs: options.ingestReadinessTimeoutMs, + buildPhaseAttempts: [], + builds: {}, questions: {}, } - const runPath = this.getRunPath(runId) - const resultsDir = this.getResultsDir(runId) - - if (!existsSync(runPath)) { - mkdirSync(runPath, { recursive: true }) - } - if (!existsSync(resultsDir)) { - mkdirSync(resultsDir, { recursive: true }) - } - + mkdirSync(this.getResultsDir(runId), { recursive: true }) this.save(checkpoint) return checkpoint } @@ -179,15 +391,19 @@ export class CheckpointManager { listRuns(): string[] { if (!existsSync(this.basePath)) return [] return readdirSync(this.basePath, { withFileTypes: true }) - .filter((d) => d.isDirectory()) - .map((d) => d.name) + .filter((entry) => entry.isDirectory()) + .map((entry) => entry.name) .sort() } + initBuild(checkpoint: RunCheckpoint, build: BuildCheckpoint): void { + if (!checkpoint.builds[build.buildId]) checkpoint.builds[build.buildId] = build + } + initQuestion( checkpoint: RunCheckpoint, questionId: string, - containerTag: string, + buildId: string, metadata: { question: string groundTruth: string @@ -195,46 +411,272 @@ export class CheckpointManager { questionDate?: string } ): void { - if (!checkpoint.questions[questionId]) { - checkpoint.questions[questionId] = { - questionId, - containerTag, - question: metadata.question, - groundTruth: metadata.groundTruth, - questionType: metadata.questionType, - questionDate: metadata.questionDate, - phases: { - ingest: { status: "pending", completedSessions: [] }, - indexing: { status: "pending" }, - search: { status: "pending" }, - answer: { status: "pending" }, - evaluate: { status: "pending" }, - }, - } + if (checkpoint.questions[questionId]) return + checkpoint.questions[questionId] = { + questionId, + buildId, + question: metadata.question, + groundTruth: metadata.groundTruth, + questionType: metadata.questionType, + questionDate: metadata.questionDate, + phases: { + search: { status: "pending" }, + answer: { status: "pending" }, + evaluate: { status: "pending" }, + }, } } - updateSessions( + updateBuild( checkpoint: RunCheckpoint, - questionId: string, - sessions: Array<{ sessionId: string; date?: string; messageCount: number }> + buildId: string, + update: (build: BuildCheckpoint) => void ): void { - const q = checkpoint.questions[questionId] - if (!q) return - q.sessions = sessions + const build = checkpoint.builds[buildId] + if (!build) throw new Error(`Unknown build ${buildId}`) + update(build) this.save(checkpoint) } + /** + * Persist one attempted document step without rewriting the full run checkpoint. + * A causal step advances after readiness or after its failure has been durably + * placed in the end-of-build retry queue. + * The append-only record is fsynced before in-memory progress advances, so load() + * can replay the ordered prefix after a process crash. + */ + recordIngestProgress( + checkpoint: RunCheckpoint, + buildId: string, + input: { + sequence: number + sessionId: string + documentIds: string[] + taskIds?: string[] + readyForNextSession: boolean + deferredFailure?: { + customId: string + stage: "submission" | "readiness" + attempts?: number + firstFailedAt?: string + lastFailedAt?: string + error: string + } + } + ): void { + const build = checkpoint.builds[buildId] + if (!build) throw new Error(`Unknown build ${buildId}`) + const expectedSequence = build.ingest.completedSessionIds.length + const expectedSessionId = build.haystack.orderedSessionIds[expectedSequence] + if (input.sequence !== expectedSequence || input.sessionId !== expectedSessionId) { + throw new Error( + `Build ${buildId} ingest progress is out of order at ${input.sequence}: ${input.sessionId} != ${expectedSessionId ?? ""}` + ) + } + assertStringArray(input.documentIds, "documentIds") + const taskIds = input.taskIds ?? [] + assertStringArray(taskIds, "taskIds") + const requiresSessionBarrier = + build.ingestionExecutionPolicy.readinessBarrier === "after-each-document" + if ( + input.deferredFailure + ? input.readyForNextSession + : input.readyForNextSession !== requiresSessionBarrier + ) { + throw new Error( + `Build ${buildId} progress readiness does not match its ingestion execution policy` + ) + } + if ( + input.deferredFailure && + (!input.deferredFailure.customId?.trim() || + !input.deferredFailure.error?.trim() || + (input.deferredFailure.attempts !== undefined && + (!Number.isInteger(input.deferredFailure.attempts) || + input.deferredFailure.attempts < 1))) + ) { + throw new Error(`Build ${buildId} has invalid deferred ingest progress`) + } + const failedAt = input.deferredFailure?.lastFailedAt ?? new Date().toISOString() + const record = createIngestProgressJournalRecord({ + schemaVersion: INGEST_PROGRESS_JOURNAL_SCHEMA_VERSION, + buildId, + buildFingerprint: build.buildFingerprint, + sequence: input.sequence, + sessionId: input.sessionId, + documentIds: [...input.documentIds], + taskIds: [...taskIds], + readyForNextSession: input.readyForNextSession, + ...(input.deferredFailure + ? { + deferredFailure: { + customId: input.deferredFailure.customId, + stage: input.deferredFailure.stage, + attempts: input.deferredFailure.attempts ?? 1, + firstFailedAt: input.deferredFailure.firstFailedAt ?? failedAt, + lastFailedAt: failedAt, + lastError: input.deferredFailure.error, + }, + } + : {}), + }) + const journalPath = this.getIngestProgressJournalPath(checkpoint.runId, buildId) + mkdirSync(join(this.getRunPath(checkpoint.runId), "progress", "ingest"), { recursive: true }) + const fileDescriptor = openSync(journalPath, "a") + try { + writeFileSync(fileDescriptor, `${JSON.stringify(record)}\n`) + fsyncSync(fileDescriptor) + } finally { + closeSync(fileDescriptor) + } + this.applyIngestProgressRecord(checkpoint, record, journalPath) + } + + clearIngestProgressJournal(runId: string, buildId: string): void { + const journalPath = this.getIngestProgressJournalPath(runId, buildId) + if (existsSync(journalPath)) unlinkSync(journalPath) + } + + private replayIngestProgressJournals(checkpoint: RunCheckpoint): void { + if (!checkpoint?.runId || !checkpoint.builds || typeof checkpoint.builds !== "object") return + for (const build of Object.values(checkpoint.builds)) { + if (!build?.buildId) continue + const journalPath = this.getIngestProgressJournalPath(checkpoint.runId, build.buildId) + if (!existsSync(journalPath)) continue + let contents = readFileSync(journalPath, "utf8") + if (contents && !contents.endsWith("\n")) { + const lastCommittedOffset = contents.lastIndexOf("\n") + 1 + logger.warn( + `Discarding an incomplete trailing ingest-progress record for build ${build.buildId}` + ) + truncateSync(journalPath, lastCommittedOffset) + contents = contents.slice(0, lastCommittedOffset) + } + const seen = new Map() + const lines = contents.split("\n") + for (let index = 0; index < lines.length; index++) { + const line = lines[index] + if (!line) continue + const record = parseIngestProgressJournalRecord(line, journalPath, index + 1) + const previousChecksum = seen.get(record.sequence) + if (previousChecksum) { + if (previousChecksum !== record.checksum) { + throw new Error( + `Ingest progress journal ${journalPath} has conflicting sequence ${record.sequence}` + ) + } + continue + } + seen.set(record.sequence, record.checksum) + this.applyIngestProgressRecord(checkpoint, record, journalPath) + } + } + } + + private applyIngestProgressRecord( + checkpoint: RunCheckpoint, + record: IngestProgressJournalRecord, + journalPath: string + ): void { + const build = checkpoint.builds[record.buildId] + if (!build) throw new Error(`Ingest progress journal ${journalPath} references unknown build`) + if (record.buildFingerprint !== build.buildFingerprint) { + throw new Error(`Ingest progress journal ${journalPath} build fingerprint does not match`) + } + const requiresSessionBarrier = + build.ingestionExecutionPolicy.readinessBarrier === "after-each-document" + if ( + record.deferredFailure + ? record.readyForNextSession + : record.readyForNextSession !== requiresSessionBarrier + ) { + throw new Error( + `Ingest progress journal ${journalPath} readiness does not match the build execution policy` + ) + } + const expectedSessionId = build.haystack.orderedSessionIds[record.sequence] + if (expectedSessionId !== record.sessionId) { + throw new Error( + `Ingest progress journal ${journalPath} session ${record.sessionId} is not ordered session ${record.sequence}` + ) + } + const completedCount = build.ingest.completedSessionIds.length + if (record.sequence < completedCount) { + if (build.ingest.completedSessionIds[record.sequence] !== record.sessionId) { + throw new Error(`Ingest progress journal ${journalPath} conflicts with checkpoint progress`) + } + for (const documentId of record.documentIds) { + if (!build.ingest.documentIds.includes(documentId)) { + throw new Error( + `Ingest progress journal ${journalPath} document ID conflicts with compacted checkpoint` + ) + } + } + for (const taskId of record.taskIds) { + if (!build.ingest.taskIds.includes(taskId)) { + throw new Error( + `Ingest progress journal ${journalPath} task ID conflicts with compacted checkpoint` + ) + } + } + if (record.readyForNextSession) { + for (const indexedId of [...record.documentIds, ...record.taskIds]) { + if (!build.indexing.completedIds.includes(indexedId)) { + throw new Error( + `Ingest progress journal ${journalPath} indexing state conflicts with compacted checkpoint` + ) + } + } + } + return + } + if (record.sequence !== completedCount) { + throw new Error( + `Ingest progress journal ${journalPath} has a gap before sequence ${record.sequence}` + ) + } + build.ingest.completedSessionIds.push(record.sessionId) + build.ingest.documentIds = [...new Set([...build.ingest.documentIds, ...record.documentIds])] + build.ingest.taskIds = [...new Set([...build.ingest.taskIds, ...record.taskIds])] + if (record.deferredFailure) { + const deferredSessions = (build.ingest.deferredSessions ??= []) + if (deferredSessions.some((deferred) => deferred.sequence === record.sequence)) { + throw new Error( + `Ingest progress journal ${journalPath} repeats deferred sequence ${record.sequence}` + ) + } + deferredSessions.push({ + sequence: record.sequence, + sessionId: record.sessionId, + customId: record.deferredFailure.customId, + documentIds: [...record.documentIds], + taskIds: [...record.taskIds], + stage: record.deferredFailure.stage, + attempts: record.deferredFailure.attempts, + firstFailedAt: record.deferredFailure.firstFailedAt, + lastFailedAt: record.deferredFailure.lastFailedAt, + lastError: record.deferredFailure.lastError, + }) + } + if (record.readyForNextSession) { + build.indexing.completedIds = [ + ...new Set([...build.indexing.completedIds, ...record.documentIds, ...record.taskIds]), + ] + build.indexing.failedIds = build.indexing.failedIds.filter( + (id) => !record.documentIds.includes(id) && !record.taskIds.includes(id) + ) + } + } + updatePhase

( checkpoint: RunCheckpoint, questionId: string, phase: P, updates: Partial ): void { - const q = checkpoint.questions[questionId] - if (!q) return - - Object.assign(q.phases[phase], updates) + const question = checkpoint.questions[questionId] + if (!question) throw new Error(`Unknown question ${questionId}`) + Object.assign(question.phases[phase], updates) this.save(checkpoint) } @@ -248,56 +690,50 @@ export class CheckpointManager { getSummary(checkpoint: RunCheckpoint): { total: number + builds: number ingested: number indexed: number searched: number answered: number evaluated: number - indexingEpisodes?: { - total: number - completed: number - failed: number - } + indexingEpisodes?: { total: number; completed: number; failed: number } } { const questions = Object.values(checkpoint.questions) - - let episodesTotal = 0 - let episodesCompleted = 0 - let episodesFailed = 0 - - for (const q of questions) { - const ingestResult = q.phases.ingest.ingestResult - const total = (ingestResult?.documentIds?.length || 0) + (ingestResult?.taskIds?.length || 0) - episodesTotal += total - - const indexing = q.phases.indexing - episodesCompleted += indexing?.completedIds?.length || 0 - episodesFailed += indexing?.failedIds?.length || 0 - } + const builds = Object.values(checkpoint.builds) + const episodeTotal = builds.reduce( + (sum, build) => sum + new Set([...build.ingest.documentIds, ...build.ingest.taskIds]).size, + 0 + ) return { total: questions.length, - ingested: questions.filter((q) => q.phases.ingest.status === "completed").length, - indexed: questions.filter((q) => q.phases.indexing?.status === "completed").length, - searched: questions.filter((q) => q.phases.search.status === "completed").length, - answered: questions.filter((q) => q.phases.answer.status === "completed").length, - evaluated: questions.filter((q) => q.phases.evaluate.status === "completed").length, - ...(episodesTotal > 0 + builds: builds.length, + ingested: builds.filter((build) => build.ingest.status === "completed").length, + indexed: builds.filter((build) => build.indexing.status === "completed").length, + searched: questions.filter((question) => question.phases.search.status === "completed") + .length, + answered: questions.filter((question) => question.phases.answer.status === "completed") + .length, + evaluated: questions.filter((question) => question.phases.evaluate.status === "completed") + .length, + ...(episodeTotal > 0 ? { - indexingEpisodes: { - total: episodesTotal, - completed: episodesCompleted, - failed: episodesFailed, - }, - } + indexingEpisodes: { + total: episodeTotal, + completed: builds.reduce( + (sum, build) => sum + new Set(build.indexing.completedIds).size, + 0 + ), + failed: builds.reduce( + (sum, build) => sum + new Set(build.indexing.failedIds).size, + 0 + ), + }, + } : {}), } } - /** - * Copy a checkpoint from sourceRunId to newRunId, resetting phases from fromPhase onwards. - * This allows creating a new run that reuses ingest/indexing data from an existing run. - */ copyCheckpoint( sourceRunId: string, newRunId: string, @@ -305,94 +741,135 @@ export class CheckpointManager { overrides?: { judge?: string; answeringModel?: string } ): RunCheckpoint { const source = this.load(sourceRunId) - if (!source) { - throw new Error(`Source checkpoint not found: ${sourceRunId}`) + if (!source) throw new Error(`Source checkpoint not found: ${sourceRunId}`) + assertCopyPhaseOverrides(source, fromPhase, overrides) + if (fromPhase === "ingest") { + throw new Error("Copying from ingest requires a new validated build; start a new run instead") } - - // Get the index of the phase to start from + assertCheckpointReferences(source) const fromIndex = PHASE_ORDER.indexOf(fromPhase) - const phasesToReset = PHASE_ORDER.slice(fromIndex) - - // Map phase IDs to question phase keys (excluding "report" which isn't a question phase) - const questionPhaseKeys: (keyof QuestionCheckpoint["phases"])[] = [ - "ingest", - "indexing", - "search", - "answer", - "evaluate", - ] - - // Deep copy questions and reset phases from fromPhase onwards - const newQuestions: Record = {} - for (const [qId, q] of Object.entries(source.questions)) { - const newQ: QuestionCheckpoint = JSON.parse(JSON.stringify(q)) - - // Reset phases that are at or after fromPhase - for (const phaseKey of questionPhaseKeys) { - if (phasesToReset.includes(phaseKey as PhaseId)) { - if (phaseKey === "ingest") { - newQ.phases.ingest = { status: "pending", completedSessions: [] } - } else if (phaseKey === "indexing") { - newQ.phases.indexing = { status: "pending" } - } else if (phaseKey === "search") { - newQ.phases.search = { status: "pending" } - } else if (phaseKey === "answer") { - newQ.phases.answer = { status: "pending" } - } else if (phaseKey === "evaluate") { - newQ.phases.evaluate = { status: "pending" } - } + for (const build of Object.values(source.builds)) { + assertBuildCheckpointConsistency(build) + if (build.ingest.status !== "completed") { + throw new Error( + `Cannot copy ${sourceRunId} from ${fromPhase}; build ${build.buildId} has incomplete ingestion` + ) + } + if (fromIndex > PHASE_ORDER.indexOf("indexing") && build.indexing.status !== "completed") { + throw new Error( + `Cannot copy ${sourceRunId} from ${fromPhase}; build ${build.buildId} is not fully indexed` + ) + } + } + + const copy = structuredClone(source) + copy.runId = newRunId + copy.status = "running" + copy.judge = overrides?.judge || source.judge + const rerunsAnswer = fromIndex <= PHASE_ORDER.indexOf("answer") + if (rerunsAnswer) { + copy.answeringModel = overrides?.answeringModel || source.answeringModel + copy.answeringRuntimeIdentity = resolveAnsweringRuntimeIdentity(copy.answeringModel) + } else { + // The answer artifact is being reused, so retain the exact runtime identity + // that produced it. Re-resolving the alias here could silently relabel old + // output after a model alias/default changes. + if (!source.answeringRuntimeIdentity) { + throw new Error( + `Cannot copy ${sourceRunId} from ${fromPhase}; source answer runtime identity is missing` + ) + } + copy.answeringModel = source.answeringModel + copy.answeringRuntimeIdentity = structuredClone(source.answeringRuntimeIdentity) + } + copy.createdAt = new Date().toISOString() + copy.updatedAt = copy.createdAt + copy.buildPhaseAttempts = [] + + for (const build of Object.values(copy.builds)) { + // Keep the immediate reuse provenance. dataSourceRunId separately retains + // the original container namespace across a chain of copied runs. + build.sourceRunId = sourceRunId + if (fromPhase === "indexing") { + build.reused = false + build.reusedPhases = { ingest: true, indexing: false } + build.indexing = { + status: "pending", + completedIds: [], + failedIds: [], + attempts: [], } + } else { + build.reused = true + build.reusedPhases = { ingest: true, indexing: true } } + } - newQuestions[qId] = newQ - } - - // Create new checkpoint - use source's dataSourceRunId (or sourceRunId if source is also a copy) - const newCheckpoint: RunCheckpoint = { - runId: newRunId, - dataSourceRunId: source.dataSourceRunId || sourceRunId, // Keep original data source - status: "running", - provider: source.provider, - benchmark: source.benchmark, - judge: overrides?.judge || source.judge, - answeringModel: overrides?.answeringModel || source.answeringModel, - createdAt: new Date().toISOString(), - updatedAt: new Date().toISOString(), - limit: source.limit, - sampling: source.sampling, - targetQuestionIds: source.targetQuestionIds, - concurrency: source.concurrency, - questions: newQuestions, - } - - // Create directories - const newRunPath = this.getRunPath(newRunId) - const newResultsDir = this.getResultsDir(newRunId) - if (!existsSync(newRunPath)) { - mkdirSync(newRunPath, { recursive: true }) - } - if (!existsSync(newResultsDir)) { - mkdirSync(newResultsDir, { recursive: true }) - } - - // Copy results directory if we're keeping search results (fromPhase is after search) - const sourceResultsDir = this.getResultsDir(sourceRunId) - if (existsSync(sourceResultsDir) && fromIndex > PHASE_ORDER.indexOf("search")) { - // Copy search results files - try { - cpSync(sourceResultsDir, newResultsDir, { recursive: true }) - logger.info(`Copied results from ${sourceRunId} to ${newRunId}`) - } catch (e) { - logger.warn(`Failed to copy results: ${e}`) + for (const question of Object.values(copy.questions)) { + if (fromIndex <= PHASE_ORDER.indexOf("search")) question.phases.search = { status: "pending" } + if (fromIndex <= PHASE_ORDER.indexOf("answer")) question.phases.answer = { status: "pending" } + if (fromIndex <= PHASE_ORDER.indexOf("evaluate")) { + question.phases.evaluate = { status: "pending" } } } - this.save(newCheckpoint) - logger.info( - `Created new checkpoint ${newRunId} from ${sourceRunId}, starting from ${fromPhase}` - ) + mkdirSync(this.getResultsDir(newRunId), { recursive: true }) + const sourceResults = this.getResultsDir(sourceRunId) + if (fromIndex > PHASE_ORDER.indexOf("search")) { + const reusedSearches = Object.values(copy.questions).filter( + (question) => question.phases.search.status === "completed" + ) + for (const question of reusedSearches) { + const sourceResultPath = this.getQuestionResultsPath(sourceRunId, question.questionId) + if (!existsSync(sourceResultPath)) { + throw new Error( + `Cannot copy ${sourceRunId} from ${fromPhase}; search artifact is missing for ${question.questionId}` + ) + } + } + if (reusedSearches.length > 0) { + cpSync(sourceResults, this.getResultsDir(newRunId), { recursive: true }) + } + for (const question of reusedSearches) { + const copiedResultPath = this.getQuestionResultsPath(newRunId, question.questionId) + if (!existsSync(copiedResultPath)) { + throw new Error( + `Cannot copy ${sourceRunId} from ${fromPhase}; copied search artifact is missing for ${question.questionId}` + ) + } + question.phases.search.resultFile = copiedResultPath + } + } + this.save(copy) + return copy + } +} - return newCheckpoint +export function assertCopyPhaseOverrides( + source: Pick, + fromPhase: PhaseId, + overrides?: { judge?: string; answeringModel?: string } +): void { + const fromIndex = PHASE_ORDER.indexOf(fromPhase) + if (fromIndex < 0) throw new Error(`Invalid copy phase: ${String(fromPhase)}`) + + if ( + overrides?.answeringModel && + overrides.answeringModel !== source.answeringModel && + fromIndex > PHASE_ORDER.indexOf("answer") + ) { + throw new Error( + `Cannot change answering model when copying ${source.runId} from ${fromPhase}; rerun answer or an earlier phase` + ) + } + if ( + overrides?.judge && + overrides.judge !== source.judge && + fromIndex > PHASE_ORDER.indexOf("evaluate") + ) { + throw new Error( + `Cannot change judge when copying ${source.runId} from ${fromPhase}; rerun evaluate or an earlier phase` + ) } } diff --git a/src/orchestrator/evaluation-runtime.ts b/src/orchestrator/evaluation-runtime.ts new file mode 100644 index 0000000..a6ae5d4 --- /dev/null +++ b/src/orchestrator/evaluation-runtime.ts @@ -0,0 +1,200 @@ +import { generateObject } from "ai" +import type { Judge } from "../types/judge" +import type { EvaluationRuntime, ModelUsage, StructuredModelRequest } from "../types/protocol" + +export const STRUCTURED_RUNTIME_EXECUTION_VERSION = "chat-transport-outer-retry-v1" + +export async function executeStructuredWithRetries( + request: StructuredModelRequest, + execute: () => Promise +): Promise { + const maxAttempts = request.maxAttempts ?? 3 + const retryBackoffMs = request.retryBackoffMs ?? 0 + if (!Number.isInteger(maxAttempts) || maxAttempts < 1) { + throw new Error(`Structured maxAttempts must be a positive integer; received ${maxAttempts}`) + } + if (!Number.isInteger(retryBackoffMs) || retryBackoffMs < 0) { + throw new Error( + `Structured retryBackoffMs must be a non-negative integer; received ${retryBackoffMs}` + ) + } + let lastError: unknown + + for (let attempt = 1; attempt <= maxAttempts; attempt++) { + try { + return request.schema.parse(await execute()) + } catch (error) { + lastError = error + if (attempt < maxAttempts && retryBackoffMs > 0) { + await new Promise((resolve) => setTimeout(resolve, retryBackoffMs * attempt)) + } + } + } + + const message = lastError instanceof Error ? lastError.message : String(lastError) + throw new Error( + `Structured judge request ${request.schemaName} failed after ${maxAttempts} attempts: ${message}` + ) +} + +export interface StructuredGenerationResult { + object: unknown + usage?: unknown +} + +export type StructuredGenerationExecutor = ( + judge: Judge, + request: StructuredModelRequest +) => Promise + +interface NormalizedTokenUsage { + inputTokens?: number + outputTokens?: number + reasoningTokens?: number + totalTokens?: number +} + +function finiteNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) && value >= 0 ? value : undefined +} + +function normalizeTokenUsage(value: unknown): NormalizedTokenUsage | undefined { + if (!value || typeof value !== "object" || Array.isArray(value)) return undefined + const usage = value as Record + const inputTokens = finiteNumber(usage.inputTokens ?? usage.promptTokens) + const outputTokens = finiteNumber(usage.outputTokens ?? usage.completionTokens) + const reasoningTokens = finiteNumber(usage.reasoningTokens) + const totalTokens = finiteNumber(usage.totalTokens) + if ( + inputTokens === undefined && + outputTokens === undefined && + reasoningTokens === undefined && + totalTokens === undefined + ) { + return undefined + } + return { + ...(inputTokens !== undefined ? { inputTokens } : {}), + ...(outputTokens !== undefined ? { outputTokens } : {}), + ...(reasoningTokens !== undefined ? { reasoningTokens } : {}), + ...(totalTokens !== undefined ? { totalTokens } : {}), + } +} + +/** Extract token usage exposed by AI SDK generation errors without assuming it exists. */ +export function extractTokenUsageFromError(error: unknown): NormalizedTokenUsage | undefined { + const seen = new Set() + const visit = (value: unknown): NormalizedTokenUsage | undefined => { + if (!value || typeof value !== "object" || Array.isArray(value) || seen.has(value)) { + return undefined + } + seen.add(value) + const record = value as Record + const direct = normalizeTokenUsage(record.usage) + if (direct) return direct + const response = record.response + if (response && typeof response === "object" && !Array.isArray(response)) { + const responseUsage = normalizeTokenUsage((response as Record).usage) + if (responseUsage) return responseUsage + } + return visit(record.cause) + } + return visit(error) +} + +export async function generateStructuredObject( + judge: Judge, + request: StructuredModelRequest +): Promise { + const innerMaxRetries = request.innerMaxRetries ?? 0 + if (!Number.isInteger(innerMaxRetries) || innerMaxRetries < 0) { + throw new Error( + `Structured innerMaxRetries must be a non-negative integer; received ${innerMaxRetries}` + ) + } + const result = await generateObject({ + model: judge.getModel(request.transport), + schema: request.schema, + schemaName: request.schemaName, + prompt: request.prompt, + ...(request.system ? { system: request.system } : {}), + ...(request.temperature !== undefined ? { temperature: request.temperature } : {}), + maxOutputTokens: request.maxOutputTokens ?? 512, + maxRetries: innerMaxRetries, + abortSignal: AbortSignal.timeout(request.timeoutMs ?? 120_000), + }) + return { object: result.object, usage: result.usage } +} + +export class JudgeEvaluationRuntime implements EvaluationRuntime { + private usage: ModelUsage = { requestCount: 0 } + + constructor( + private readonly judge: Judge, + private readonly structuredGeneration: StructuredGenerationExecutor = generateStructuredObject + ) {} + + async evaluateLegacy(input: Parameters[0]) { + this.beginRequest() + try { + const result = await this.judge.evaluate(input) + // The legacy JudgeResult contract does not expose token usage. + this.recordTokenUsage(undefined) + return result + } catch (error) { + this.recordTokenUsage(extractTokenUsageFromError(error)) + throw error + } + } + + async generateStructured(request: StructuredModelRequest): Promise { + return executeStructuredWithRetries(request, async () => { + this.beginRequest() + try { + const result = await this.structuredGeneration( + this.judge, + request as StructuredModelRequest + ) + this.recordTokenUsage(normalizeTokenUsage(result.usage)) + return result.object + } catch (error) { + this.recordTokenUsage(extractTokenUsageFromError(error)) + throw error + } + }) + } + + getUsage(): ModelUsage | undefined { + return this.usage.requestCount ? { ...this.usage } : undefined + } + + private beginRequest(): void { + this.usage.requestCount = (this.usage.requestCount ?? 0) + 1 + } + + private recordTokenUsage(usage: NormalizedTokenUsage | undefined): void { + if (!usage) { + this.usage.tokenUsageUnknownRequestCount = (this.usage.tokenUsageUnknownRequestCount ?? 0) + 1 + return + } + + const fields = [usage.inputTokens, usage.outputTokens, usage.totalTokens] + const complete = fields.every((value) => value !== undefined) + const coverageField = complete + ? "tokenUsageCompleteRequestCount" + : "tokenUsagePartialRequestCount" + this.usage[coverageField] = (this.usage[coverageField] ?? 0) + 1 + if (usage.inputTokens !== undefined) { + this.usage.inputTokens = (this.usage.inputTokens ?? 0) + usage.inputTokens + } + if (usage.outputTokens !== undefined) { + this.usage.outputTokens = (this.usage.outputTokens ?? 0) + usage.outputTokens + } + if (usage.reasoningTokens !== undefined) { + this.usage.reasoningTokens = (this.usage.reasoningTokens ?? 0) + usage.reasoningTokens + } + if (usage.totalTokens !== undefined) { + this.usage.totalTokens = (this.usage.totalTokens ?? 0) + usage.totalTokens + } + } +} diff --git a/src/orchestrator/index.ts b/src/orchestrator/index.ts index 4d2b2b0..04b36df 100644 --- a/src/orchestrator/index.ts +++ b/src/orchestrator/index.ts @@ -1,21 +1,37 @@ -import type { ProviderName } from "../types/provider" import type { BenchmarkName } from "../types/benchmark" -import type { JudgeName } from "../types/judge" -import type { RunCheckpoint, SamplingConfig } from "../types/checkpoint" import type { ConcurrencyConfig } from "../types/concurrency" -import { createProvider } from "../providers" +import type { BuildPhaseAttempt, PhaseId, RunCheckpoint, SamplingConfig } from "../types/checkpoint" +import type { JudgeName } from "../types/judge" +import type { AnsweringRuntimeIdentity } from "../types/model" +import type { ProviderName } from "../types/provider" +import type { BenchmarkProtocol } from "../types/protocol" +import type { UnifiedQuestion } from "../types/unified" import { createBenchmark } from "../benchmarks" import { createJudge } from "../judges" -import { CheckpointManager } from "./checkpoint" -import { getProviderConfig, getJudgeConfig } from "../utils/config" -import { resolveModel } from "../utils/models" +import { createProvider } from "../providers" +import { fingerprintProviderPrompts } from "../providers/prompt-identity" +import { getJudgeConfig, getProviderConfig } from "../utils/config" import { logger } from "../utils/logger" -import { runIngestPhase } from "./phases/ingest" -import { runIndexingPhase } from "./phases/indexing" -import { runSearchPhase } from "./phases/search" +import { resolveAnsweringRuntimeIdentity, resolveModel } from "../utils/models" +import { stableSha256 } from "../utils/stable" +import { + assertResumeBuilds, + cloneCompletedBuildsForReuse, + createBuildCheckpoint, + prepareValidatedBuildPlans, +} from "./builds" +import { CheckpointManager } from "./checkpoint" import { runAnswerPhase } from "./phases/answer" import { runEvaluatePhase } from "./phases/evaluate" -import { generateReport, saveReport, printReport } from "./phases/report" +import { runIndexingPhase } from "./phases/indexing" +import { DEFAULT_INGEST_READINESS_TIMEOUT_MS, runIngestPhase } from "./phases/ingest" +import { generateReport, printReport, saveReport } from "./phases/report" +import { runSearchPhase } from "./phases/search" +import { + canonicalizeSelectedQuestionIds, + fingerprintSelectedBenchmarkInput, + resolveEffectiveDatasetRevision, +} from "./input-identity" export interface OrchestratorOptions { provider: ProviderName @@ -23,53 +39,170 @@ export interface OrchestratorOptions { judgeModel: string runId: string answeringModel?: string + dataPath?: string + datasetRevision?: string + retrievalTopK?: number + answerCutoff?: number + evaluationProfile?: string + /** Reuse validated completed ingest/index builds and start a new run at search. */ + sourceRunId?: string limit?: number sampling?: SamplingConfig concurrency?: ConcurrencyConfig + ingestBatchSize?: number + ingestReadinessTimeoutMs?: number force?: boolean questionIds?: string[] - phases?: ("ingest" | "indexing" | "search" | "answer" | "evaluate" | "report")[] + phases?: PhaseId[] + /** Resolves only after dataset/protocol/build/resume validation and checkpoint durability. */ + onPreflightComplete?: () => void + /** Internal comparison barrier: stop after durable preflight, before provider initialization. */ + preflightOnly?: boolean +} + +export function mergeResumeConcurrency( + persisted: ConcurrencyConfig | undefined, + override: ConcurrencyConfig | undefined +): ConcurrencyConfig | undefined { + if (!override) return persisted + return { ...(persisted ?? {}), ...override } as ConcurrencyConfig } function selectQuestionsBySampling( allQuestions: { questionId: string; questionType: string }[], - sampling: SamplingConfig + sampling?: SamplingConfig, + limit?: number, + explicitQuestionIds?: string[] ): string[] { - if (sampling.mode === "full") { - return allQuestions.map((q) => q.questionId) - } - - if (sampling.mode === "limit" && sampling.limit) { - return allQuestions.slice(0, sampling.limit).map((q) => q.questionId) + if (explicitQuestionIds?.length) return [...explicitQuestionIds] + if (sampling?.mode === "limit" && sampling.limit != null) { + return allQuestions.slice(0, sampling.limit).map((question) => question.questionId) } - - if (sampling.mode === "sample" && sampling.perCategory) { - const byType: Record = {} - for (const q of allQuestions) { - if (!byType[q.questionType]) byType[q.questionType] = [] - byType[q.questionType].push(q) + if (sampling?.mode === "sample" && sampling.perCategory != null) { + const grouped = new Map() + for (const question of allQuestions) { + const members = grouped.get(question.questionType) || [] + members.push(question) + grouped.set(question.questionType, members) } + return [...grouped.values()].flatMap((questions) => { + const candidates = + sampling.sampleType === "random" + ? [...questions].sort(() => Math.random() - 0.5) + : questions + return candidates.slice(0, sampling.perCategory).map((question) => question.questionId) + }) + } + if (limit != null) return allQuestions.slice(0, limit).map((question) => question.questionId) + return allQuestions.map((question) => question.questionId) +} - const selected: string[] = [] - for (const questions of Object.values(byType)) { - if (sampling.sampleType === "random") { - const shuffled = [...questions].sort(() => Math.random() - 0.5) - selected.push(...shuffled.slice(0, sampling.perCategory).map((q) => q.questionId)) - } else { - selected.push(...questions.slice(0, sampling.perCategory).map((q) => q.questionId)) - } +/** + * Resolve the actual benchmark-owned retrieval budget from the selected + * protocol plans. A run checkpoint has one Top-K identity, so mixed per-question + * budgets are rejected instead of being hidden behind an optional CLI value. + */ +export function resolveEffectiveRetrievalTopK( + protocol: BenchmarkProtocol, + questions: UnifiedQuestion[], + configuredTopK?: number +): number { + if (questions.length === 0) throw new Error("Question selection is empty") + const requestedTopKs = [ + ...new Set( + questions.map((question) => protocol.createRetrievalPlan({ question }).requestedTopK) + ), + ] + for (const topK of requestedTopKs) { + if (!Number.isInteger(topK) || topK <= 0) { + throw new Error(`Protocol returned invalid retrieval Top-K: ${String(topK)}`) } - return selected } + if (requestedTopKs.length !== 1) { + throw new Error( + `Selected questions require mixed retrieval Top-K values: ${requestedTopKs.join(", ")}` + ) + } + const effectiveTopK = requestedTopKs[0] + if (configuredTopK != null && configuredTopK !== effectiveTopK) { + throw new Error( + `Configured retrieval Top-K ${configuredTopK} differs from protocol plan ${effectiveTopK}` + ) + } + return effectiveTopK +} - return allQuestions.map((q) => q.questionId) +export function assertResumeIdentity( + checkpoint: RunCheckpoint, + input: { + provider: string + providerAdapterVersion: string + providerPromptFingerprint: string + benchmark: string + benchmarkScope: unknown + datasetIdentity: unknown + benchmarkInputFingerprint: string + selectedQuestionIdsDigest: string + protocolIdentity: unknown + retrievalTopK?: number + judge: string + answeringModel: string + answeringRuntimeIdentity: AnsweringRuntimeIdentity + ingestBatchSize?: number + } +): void { + const mismatches: string[] = [] + if (checkpoint.provider !== input.provider) mismatches.push("provider") + if (checkpoint.providerAdapterVersion !== input.providerAdapterVersion) { + mismatches.push("provider adapter version") + } + if ( + checkpoint.protocolIdentity.id === "memorybench.legacy" && + checkpoint.providerPromptFingerprint !== input.providerPromptFingerprint + ) { + mismatches.push("provider prompt") + } + if (checkpoint.benchmark !== input.benchmark) mismatches.push("benchmark") + if (stableSha256(checkpoint.benchmarkScope) !== stableSha256(input.benchmarkScope)) { + mismatches.push("benchmark scope") + } + if ( + stableSha256(checkpoint.datasetIdentity ?? null) !== stableSha256(input.datasetIdentity ?? null) + ) { + mismatches.push("dataset identity") + } + if (checkpoint.benchmarkInputFingerprint !== input.benchmarkInputFingerprint) { + mismatches.push("benchmark input") + } + if (checkpoint.selectedQuestionIdsDigest !== input.selectedQuestionIdsDigest) { + mismatches.push("selected question IDs") + } + if (stableSha256(checkpoint.protocolIdentity) !== stableSha256(input.protocolIdentity)) { + mismatches.push("benchmark protocol") + } + if (checkpoint.retrievalTopK !== input.retrievalTopK) mismatches.push("retrieval Top-K") + if (checkpoint.judge !== input.judge) mismatches.push("judge model") + if (checkpoint.answeringModel !== input.answeringModel) mismatches.push("answering model") + if ((checkpoint.ingestBatchSize ?? 1) !== (input.ingestBatchSize ?? 1)) { + mismatches.push("ingest batch size") + } + if ( + !checkpoint.answeringRuntimeIdentity || + stableSha256(checkpoint.answeringRuntimeIdentity) !== + stableSha256(input.answeringRuntimeIdentity) + ) { + mismatches.push("answering runtime") + } + if (mismatches.length > 0) { + throw new Error(`Cannot resume ${checkpoint.runId}; changed ${mismatches.join(", ")}`) + } } export class Orchestrator { private checkpointManager: CheckpointManager - constructor() { - this.checkpointManager = new CheckpointManager() + constructor(checkpointManager = new CheckpointManager()) { + this.checkpointManager = checkpointManager } async run(options: OrchestratorOptions): Promise { @@ -82,235 +215,318 @@ export class Orchestrator { limit, sampling, concurrency, + ingestBatchSize, + ingestReadinessTimeoutMs, + sourceRunId, force = false, questionIds, phases = ["ingest", "indexing", "search", "answer", "evaluate", "report"], } = options - - const judgeModelInfo = resolveModel(judgeModel) - const judgeName = judgeModelInfo.provider as JudgeName - - logger.info(`Starting MemoryBench run: ${providerName} + ${benchmarkName}`) - logger.info(`Run ID: ${runId}`) - logger.info( - `Judge: ${judgeModelInfo.displayName} (${judgeModelInfo.id}), Answering Model: ${answeringModel}` - ) - logger.info(`Force: ${force}, Phases: ${phases?.join(", ") || "all"}`) - if (sampling) { - logger.info(`Sampling config received: ${JSON.stringify(sampling)}`) - if (sampling.mode === "sample") { - logger.info( - `Sampling: ${sampling.perCategory} per category (${sampling.sampleType || "consecutive"})` - ) - } else if (sampling.mode === "limit") { - logger.info(`Limit: ${sampling.limit} questions`) - } else { - logger.info(`Selection: full (all questions)`) - } - } else if (limit) { - logger.info(`Limit: ${limit} questions`) - } else { - logger.info(`No sampling or limit provided`) + if (sourceRunId && force) { + throw new Error("--source-run cannot be combined with --force") + } + if (sourceRunId === runId) { + throw new Error("Source and target run IDs must be different") } - if (force && this.checkpointManager.exists(runId)) { this.checkpointManager.delete(runId) logger.info("Cleared existing checkpoint (--force)") } - let checkpoint!: RunCheckpoint - let effectiveLimit: number | undefined - let targetQuestionIds: string[] | undefined - let isNewRun = false - - if (!this.checkpointManager.exists(runId)) { - isNewRun = true - checkpoint = this.checkpointManager.create( - runId, - providerName, - benchmarkName, - judgeModel, - answeringModel, - { limit, sampling, concurrency, status: "initializing" } + const existing = this.checkpointManager.exists(runId) + ? this.checkpointManager.load(runId) + : null + const source = sourceRunId ? this.checkpointManager.load(sourceRunId) : null + if (sourceRunId && !source) throw new Error(`Source checkpoint not found: ${sourceRunId}`) + if (source && existing) throw new Error(`Target run ${runId} already exists`) + if (source && phases[0] !== "search") { + throw new Error("Source-build reuse must start from the search phase") + } + if (source && !source.targetQuestionIds?.length) { + throw new Error(`Source run ${source.runId} does not record its selected question IDs`) + } + if (source && (limit !== undefined || sampling !== undefined || questionIds?.length)) { + throw new Error("Source-build reuse must retain the source run's exact question set") + } + if (source && (source.provider !== providerName || source.benchmark !== benchmarkName)) { + throw new Error( + `Source run ${source.runId} is ${source.provider}/${source.benchmark}, not ${providerName}/${benchmarkName}` + ) + } + const dataPath = options.dataPath ?? existing?.dataPath ?? source?.dataPath + const configuredDatasetRevision = + options.datasetRevision ?? existing?.datasetRevision ?? source?.datasetRevision + const configuredRetrievalTopK = options.retrievalTopK ?? existing?.retrievalTopK + const configuredAnswerCutoff = options.answerCutoff ?? existing?.answerCutoff + const configuredEvaluationProfile = options.evaluationProfile ?? existing?.evaluationProfile + const configuredIngestBatchSize = + ingestBatchSize ?? existing?.ingestBatchSize ?? source?.ingestBatchSize ?? 1 + const configuredIngestReadinessTimeoutMs = + ingestReadinessTimeoutMs ?? + existing?.ingestReadinessTimeoutMs ?? + source?.ingestReadinessTimeoutMs ?? + DEFAULT_INGEST_READINESS_TIMEOUT_MS + if ( + !Number.isInteger(configuredIngestBatchSize) || + configuredIngestBatchSize < 1 || + configuredIngestBatchSize > 600 + ) { + throw new Error( + `Ingest batch size must be an integer between 1 and 600; received ${configuredIngestBatchSize}` + ) + } + if ( + !Number.isInteger(configuredIngestReadinessTimeoutMs) || + configuredIngestReadinessTimeoutMs < 1 + ) { + throw new Error( + `Ingest readiness timeout must be a positive integer; received ${configuredIngestReadinessTimeoutMs}` ) - logger.info("Created checkpoint (initializing)") } - const benchmark = createBenchmark(benchmarkName) - await benchmark.load() - const allQuestions = benchmark.getQuestions() - - if (this.checkpointManager.exists(runId) && !isNewRun) { - checkpoint = this.checkpointManager.load(runId)! - - effectiveLimit = checkpoint.limit - targetQuestionIds = checkpoint.targetQuestionIds - - if (!targetQuestionIds) { - const startedQuestions = Object.values(checkpoint.questions) - .filter((q) => Object.values(q.phases).some((p) => p.status !== "pending")) - .map((q) => q.questionId) - - if (startedQuestions.length > 0) { - const pendingQuestions = Object.values(checkpoint.questions) - .filter((q) => Object.values(q.phases).every((p) => p.status === "pending")) - .map((q) => q.questionId) - - if (limit) { - const remainingSlots = limit - startedQuestions.length - targetQuestionIds = [ - ...startedQuestions, - ...pendingQuestions.slice(0, Math.max(0, remainingSlots)), - ] - effectiveLimit = limit - logger.warn( - `Old checkpoint detected. Using CLI limit (${limit}) to determine target questions.` - ) - } else { - targetQuestionIds = startedQuestions - logger.warn( - `Old checkpoint without stored limit. Only processing ${startedQuestions.length} already-started questions.` - ) - } - - checkpoint.limit = effectiveLimit - checkpoint.targetQuestionIds = targetQuestionIds - this.checkpointManager.save(checkpoint) - } else { - if (limit) { - const limitedQuestions = allQuestions.slice(0, limit).map((q) => q.questionId) - targetQuestionIds = limitedQuestions - effectiveLimit = limit - checkpoint.limit = limit - checkpoint.targetQuestionIds = targetQuestionIds - this.checkpointManager.save(checkpoint) - logger.warn( - `Old checkpoint with no progress. Applying limit (${limit}) to first ${limit} questions.` - ) - } - } - } - - const summary = this.checkpointManager.getSummary(checkpoint) - const targetCount = targetQuestionIds?.length || summary.total - - const inProgressQuestions = Object.values(checkpoint.questions) - .filter((q) => Object.values(q.phases).some((p) => p.status === "in_progress")) - .map((q) => q.questionId) - - logger.info( - `Resuming from checkpoint: ${summary.ingested}/${targetCount} ingested, ${summary.evaluated}/${targetCount} evaluated` + const answeringRuntimeIdentity = resolveAnsweringRuntimeIdentity(answeringModel) + + // Dataset and protocol preflight deliberately happens before checkpoint creation + // and before provider initialization. + await benchmark.load({ + dataPath, + datasetRevision: configuredDatasetRevision, + retrievalTopK: configuredRetrievalTopK, + answerCutoff: configuredAnswerCutoff, + evaluationProfile: configuredEvaluationProfile, + }) + const datasetIdentity = benchmark.getDatasetIdentity?.() + const datasetRevision = resolveEffectiveDatasetRevision( + configuredDatasetRevision, + datasetIdentity + ) + const judgeModelInfo = resolveModel(judgeModel) + const requiredJudge = benchmark.protocol.requiredJudge + if ( + requiredJudge && + (judgeModelInfo.provider !== requiredJudge.provider || + judgeModelInfo.id !== requiredJudge.modelId) + ) { + throw new Error( + `Protocol ${benchmark.protocol.identity.id} requires judge ${requiredJudge.provider}/${requiredJudge.modelId}; received ${judgeModelInfo.provider}/${judgeModelInfo.id}` ) - if (inProgressQuestions.length > 0) { - logger.info(`In-progress questions: ${inProgressQuestions.join(", ")}`) - } - + } + const allQuestions = benchmark.getQuestions() + for (const question of allQuestions) benchmark.protocol.validateQuestion(question) + + const requestedQuestionIds = existing?.targetQuestionIds + ? [...existing.targetQuestionIds] + : source?.targetQuestionIds + ? [...source.targetQuestionIds] + : selectQuestionsBySampling(allQuestions, sampling, limit, questionIds) + const targetQuestionIds = canonicalizeSelectedQuestionIds(allQuestions, requestedQuestionIds) + const selectedQuestions = allQuestions.filter((question) => + targetQuestionIds.includes(question.questionId) + ) + const retrievalTopK = resolveEffectiveRetrievalTopK( + benchmark.protocol, + selectedQuestions, + configuredRetrievalTopK + ) + const selectedQuestionIdsDigest = stableSha256(targetQuestionIds) + const benchmarkInputFingerprint = fingerprintSelectedBenchmarkInput( + benchmark, + selectedQuestions + ) + const provider = createProvider(providerName) + const providerConfig = getProviderConfig(providerName) + const providerPromptFingerprint = fingerprintProviderPrompts(provider.prompts) + const providerIngestionConfigFingerprint = + provider.getIngestionConfigFingerprint(providerConfig) + const dataSourceRunId = existing?.dataSourceRunId || source?.dataSourceRunId || runId + const buildPlans = prepareValidatedBuildPlans({ + benchmark, + questions: selectedQuestions, + provider: provider.name, + providerAdapterVersion: provider.adapterVersion, + providerPromptFingerprint, + providerIngestionConfigFingerprint, + dataSourceRunId, + ingestBatchSize: configuredIngestBatchSize, + }) + const reusedBuilds = source ? cloneCompletedBuildsForReuse(source, buildPlans) : undefined + + let checkpoint: RunCheckpoint + if (existing) { + checkpoint = existing + assertResumeIdentity(checkpoint, { + provider: provider.name, + providerAdapterVersion: provider.adapterVersion, + providerPromptFingerprint, + benchmark: benchmark.name, + benchmarkScope: benchmark.scope, + datasetIdentity, + benchmarkInputFingerprint, + selectedQuestionIdsDigest, + protocolIdentity: benchmark.protocol.identity, + retrievalTopK, + judge: judgeModel, + answeringModel, + answeringRuntimeIdentity, + ingestBatchSize: configuredIngestBatchSize, + }) + assertResumeBuilds(checkpoint, buildPlans) + checkpoint.dataPath = dataPath + checkpoint.datasetRevision = datasetRevision + checkpoint.evaluationProfile = configuredEvaluationProfile + checkpoint.answerCutoff = configuredAnswerCutoff + checkpoint.concurrency = mergeResumeConcurrency(checkpoint.concurrency, concurrency) + checkpoint.ingestBatchSize ??= configuredIngestBatchSize + checkpoint.ingestReadinessTimeoutMs = configuredIngestReadinessTimeoutMs this.checkpointManager.updateStatus(checkpoint, "running") } else { - logger.info( - `New run path: isNewRun=${isNewRun}, sampling=${JSON.stringify(sampling)}, limit=${limit}` + checkpoint = this.checkpointManager.create( + runId, + provider.name, + benchmark.name, + judgeModel, + answeringModel, + { + providerAdapterVersion: provider.adapterVersion, + providerPromptFingerprint, + benchmarkScope: benchmark.scope, + protocolIdentity: benchmark.protocol.identity, + selectedQuestionIdsDigest, + datasetIdentity, + benchmarkInputFingerprint, + dataPath, + datasetRevision, + retrievalTopK, + evaluationProfile: configuredEvaluationProfile, + answerCutoff: configuredAnswerCutoff, + dataSourceRunId, + limit, + sampling, + targetQuestionIds, + concurrency, + ingestBatchSize: configuredIngestBatchSize, + ingestReadinessTimeoutMs: configuredIngestReadinessTimeoutMs, + status: "initializing", + } ) - effectiveLimit = limit - - if (questionIds && questionIds.length > 0) { - logger.info(`Using explicit questionIds: ${questionIds.length} questions`) - targetQuestionIds = questionIds - } else if (sampling) { - logger.info(`Using sampling mode: ${sampling.mode}`) - targetQuestionIds = selectQuestionsBySampling(allQuestions, sampling) - checkpoint.sampling = sampling - logger.info( - `Sampling selected ${targetQuestionIds.length} questions from ${allQuestions.length} total` + for (const [index, plan] of buildPlans.entries()) { + this.checkpointManager.initBuild( + checkpoint, + reusedBuilds?.[index] ?? createBuildCheckpoint(plan) ) - } else if (effectiveLimit) { - logger.info(`Using limit: ${effectiveLimit}`) - targetQuestionIds = allQuestions.slice(0, effectiveLimit).map((q) => q.questionId) - } else { - logger.info(`No sampling/limit specified, using all ${allQuestions.length} questions`) } - - checkpoint.targetQuestionIds = targetQuestionIds - checkpoint.limit = effectiveLimit - - const questionsToInit = targetQuestionIds - ? allQuestions.filter((q) => targetQuestionIds!.includes(q.questionId)) - : allQuestions - - for (const q of questionsToInit) { - const ingestionGroupId = benchmark.getIngestionGroupId?.(q.questionId) || q.questionId - const containerTag = `${ingestionGroupId}-${checkpoint.dataSourceRunId}` - this.checkpointManager.initQuestion(checkpoint, q.questionId, containerTag, { - question: q.question, - groundTruth: q.groundTruth, - questionType: q.questionType, + for (const question of selectedQuestions) { + const plan = buildPlans.find((candidate) => + candidate.memberQuestionIds.includes(question.questionId) + ) + if (!plan) throw new Error(`No validated build found for ${question.questionId}`) + this.checkpointManager.initQuestion(checkpoint, question.questionId, plan.buildId, { + question: question.question, + groundTruth: question.groundTruth, + questionType: question.questionType, + questionDate: + typeof question.metadata?.questionDate === "string" + ? question.metadata.questionDate + : undefined, }) } - this.checkpointManager.updateStatus(checkpoint, "running") + await this.checkpointManager.flush(runId) } - const provider = createProvider(providerName) - await provider.initialize(getProviderConfig(providerName)) - - if (phases.includes("ingest")) { - await runIngestPhase( - provider, - benchmark, - checkpoint, - this.checkpointManager, - targetQuestionIds - ) - } - - if (phases.includes("indexing")) { - await runIndexingPhase(provider, checkpoint, this.checkpointManager, targetQuestionIds) - } - - if (phases.includes("search")) { - await runSearchPhase( - provider, - benchmark, - checkpoint, - this.checkpointManager, - targetQuestionIds - ) - } - - if (phases.includes("answer")) { - await runAnswerPhase( - benchmark, - checkpoint, - this.checkpointManager, - targetQuestionIds, - provider - ) - } + const judgeName = judgeModelInfo.provider as JudgeName + await this.checkpointManager.flush(runId) + options.onPreflightComplete?.() + if (options.preflightOnly) return + logger.info( + `Starting ${benchmark.scope.displayName}: ${providerName}, protocol ${benchmark.protocol.identity.id}@${benchmark.protocol.identity.version}, ${selectedQuestions.length} questions across ${buildPlans.length} builds` + ) - if (phases.includes("evaluate")) { - const judge = createJudge(judgeName) - const judgeConfig = getJudgeConfig(judgeName) - judgeConfig.model = judgeModel - await judge.initialize(judgeConfig) - await runEvaluatePhase( - judge, - benchmark, - checkpoint, - this.checkpointManager, - targetQuestionIds, - provider - ) - } + try { + // All dataset, protocol, haystack and resume checks have passed at this point. + // Initialization still belongs to the durable run lifecycle: a provider + // failure must mark and flush the checkpoint as failed. + await provider.initialize(providerConfig) + const runsBuildPhase = phases.includes("ingest") || phases.includes("indexing") + let buildPhaseAttempt: BuildPhaseAttempt | undefined + let buildPhaseStartedMs = 0 + if (runsBuildPhase) { + buildPhaseStartedMs = Date.now() + buildPhaseAttempt = { + startedAt: new Date().toISOString(), + status: "in_progress", + } + checkpoint.buildPhaseAttempts.push(buildPhaseAttempt) + this.checkpointManager.save(checkpoint) + } - if (phases.includes("report")) { - const report = generateReport(benchmark, checkpoint) - saveReport(report) - printReport(report) + if (phases.includes("ingest")) { + await runIngestPhase(provider, checkpoint, this.checkpointManager, buildPlans) + } + if (phases.includes("indexing")) { + await runIndexingPhase(provider, checkpoint, this.checkpointManager) + } + if (buildPhaseAttempt) { + buildPhaseAttempt.status = "completed" + buildPhaseAttempt.completedAt = new Date().toISOString() + buildPhaseAttempt.durationMs = Date.now() - buildPhaseStartedMs + this.checkpointManager.save(checkpoint) + } + const phaseQuestionIds = questionIds?.length ? questionIds : targetQuestionIds + if (phases.includes("search")) { + await runSearchPhase( + provider, + benchmark, + checkpoint, + this.checkpointManager, + phaseQuestionIds + ) + } + if (phases.includes("answer")) { + await runAnswerPhase( + benchmark, + checkpoint, + this.checkpointManager, + phaseQuestionIds, + provider + ) + } + if (phases.includes("evaluate")) { + const judge = createJudge(judgeName) + const judgeConfig = getJudgeConfig(judgeName) + judgeConfig.model = judgeModel + await judge.initialize(judgeConfig) + await runEvaluatePhase( + judge, + benchmark, + checkpoint, + this.checkpointManager, + phaseQuestionIds, + provider + ) + } + if (phases.includes("report")) { + const report = generateReport(benchmark, checkpoint) + saveReport(report) + printReport(report) + } + this.checkpointManager.updateStatus(checkpoint, "completed") + await this.checkpointManager.flush(runId) + logger.success("Run complete!") + } catch (error) { + const activeBuildAttempt = checkpoint.buildPhaseAttempts.at(-1) + if (activeBuildAttempt?.status === "in_progress") { + activeBuildAttempt.status = "failed" + activeBuildAttempt.completedAt = new Date().toISOString() + activeBuildAttempt.durationMs = Math.max( + 0, + Date.parse(activeBuildAttempt.completedAt) - Date.parse(activeBuildAttempt.startedAt) + ) + } + this.checkpointManager.updateStatus(checkpoint, "failed") + await this.checkpointManager.flush(runId) + throw error } - - // Flush all pending checkpoint saves before marking as complete - await this.checkpointManager.flush(checkpoint.runId) - this.checkpointManager.updateStatus(checkpoint, "completed") - logger.success("Run complete!") } async ingest( @@ -318,7 +534,10 @@ export class Orchestrator { ): Promise { await this.run({ ...options, - judgeModel: options.judgeModel || "gpt-4o", + judgeModel: + options.judgeModel || + createBenchmark(options.benchmark).protocol.requiredJudge?.modelAlias || + "gpt-4o", phases: ["ingest", "indexing"], }) } @@ -326,7 +545,14 @@ export class Orchestrator { async search( options: Omit & { judgeModel?: string } ): Promise { - await this.run({ ...options, judgeModel: options.judgeModel || "gpt-4o", phases: ["search"] }) + await this.run({ + ...options, + judgeModel: + options.judgeModel || + createBenchmark(options.benchmark).protocol.requiredJudge?.modelAlias || + "gpt-4o", + phases: ["search"], + }) } async evaluate(options: OrchestratorOptions): Promise { @@ -347,20 +573,22 @@ export class Orchestrator { logger.error(`No run found: ${runId}`) return } - const summary = this.checkpointManager.getSummary(checkpoint) - console.log("\n" + "=".repeat(50)) console.log(`Run: ${runId}`) console.log(`Provider: ${checkpoint.provider}`) - console.log(`Benchmark: ${checkpoint.benchmark}`) - console.log("=".repeat(50)) - console.log(`Total Questions: ${summary.total}`) - console.log(`Ingested: ${summary.ingested}`) - console.log(`Indexed: ${summary.indexed}`) + console.log(`Benchmark: ${checkpoint.benchmarkScope.displayName}`) + console.log( + `Builds: ${summary.builds} (${summary.ingested} ingested, ${summary.indexed} indexed)` + ) + const deferredSessions = Object.values(checkpoint.builds).reduce( + (sum, build) => sum + (build.ingest.deferredSessions?.length ?? 0), + 0 + ) + if (deferredSessions > 0) console.log(`Deferred sessions awaiting retry: ${deferredSessions}`) + console.log(`Questions: ${summary.total}`) console.log(`Searched: ${summary.searched}`) console.log(`Answered: ${summary.answered}`) console.log(`Evaluated: ${summary.evaluated}`) - console.log("=".repeat(50) + "\n") } } diff --git a/src/orchestrator/input-identity.ts b/src/orchestrator/input-identity.ts new file mode 100644 index 0000000..096d898 --- /dev/null +++ b/src/orchestrator/input-identity.ts @@ -0,0 +1,121 @@ +import type { Benchmark, DatasetIdentity } from "../types/benchmark" +import type { UnifiedQuestion } from "../types/unified" +import { stableSha256 } from "../utils/stable" + +/** Persist the enclosing immutable snapshot revision used for path lookup and resume. */ +export function resolveEffectiveDatasetRevision( + configuredRevision: string | undefined, + datasetIdentity: DatasetIdentity | undefined +): string | undefined { + return ( + datasetIdentity?.snapshotFingerprint ?? + datasetIdentity?.datasetFingerprint ?? + configuredRevision + ) +} + +/** Resolve a selected ID set into the benchmark's canonical question order. */ +export function canonicalizeSelectedQuestionIds( + allQuestions: readonly Pick[], + requestedQuestionIds: readonly string[] +): string[] { + const allIds = allQuestions.map((question) => question.questionId) + if (new Set(allIds).size !== allIds.length) { + throw new Error("Benchmark contains duplicate question IDs") + } + if (new Set(requestedQuestionIds).size !== requestedQuestionIds.length) { + throw new Error("Selected question IDs contain duplicates") + } + + const knownIds = new Set(allIds) + const unknownIds = requestedQuestionIds.filter((questionId) => !knownIds.has(questionId)) + if (unknownIds.length > 0) { + throw new Error(`Unknown question IDs: ${unknownIds.join(", ")}`) + } + + const selectedIds = new Set(requestedQuestionIds) + return allIds.filter((questionId) => selectedIds.has(questionId)) +} + +/** + * Provider-independent identity for the exact selected benchmark questions and + * their raw ordered haystacks. Official dataset identity remains authoritative + * when a benchmark (such as BEAM) supplies one. + */ +export function fingerprintSelectedBenchmarkInput( + benchmark: Benchmark, + selectedQuestions: readonly UnifiedQuestion[] +): string { + const allQuestions = benchmark.getQuestions() + const canonicalQuestionIds = canonicalizeSelectedQuestionIds( + allQuestions, + selectedQuestions.map((question) => question.questionId) + ) + const questionById = new Map(allQuestions.map((question) => [question.questionId, question])) + const canonicalQuestions = canonicalQuestionIds.map((questionId) => questionById.get(questionId)!) + const digestBySessionArray = new WeakMap() + const haystacksByDigest = new Map< + string, + { orderedSessionIds: string[]; sessions: ReturnType } + >() + const digestByDeclaredGroup = new Map() + + const questions = canonicalQuestions.map((question) => { + const sessions = benchmark.getHaystackSessions(question.questionId) + const orderedSessionIds = sessions.map((session) => session.sessionId) + if ( + orderedSessionIds.length !== question.haystackSessionIds.length || + orderedSessionIds.some((sessionId, index) => sessionId !== question.haystackSessionIds[index]) + ) { + throw new Error( + `Question ${question.questionId} haystack IDs do not match its ordered benchmark sessions` + ) + } + + let haystackFingerprint = digestBySessionArray.get(sessions) + if (!haystackFingerprint) { + haystackFingerprint = stableSha256({ + schemaVersion: 1, + orderedSessionIds, + sessions, + }) + digestBySessionArray.set(sessions, haystackFingerprint) + } + if (!haystacksByDigest.has(haystackFingerprint)) { + haystacksByDigest.set(haystackFingerprint, { orderedSessionIds, sessions }) + } + + const declaredGroupId = benchmark.getIngestionGroupId?.(question.questionId) + if (declaredGroupId) { + const existingGroupDigest = digestByDeclaredGroup.get(declaredGroupId) + if (existingGroupDigest && existingGroupDigest !== haystackFingerprint) { + throw new Error( + `Ingestion group ${declaredGroupId} resolves to different raw benchmark haystacks` + ) + } + digestByDeclaredGroup.set(declaredGroupId, haystackFingerprint) + } + + return { + questionId: question.questionId, + question: question.question, + questionType: question.questionType, + groundTruth: question.groundTruth, + haystackSessionIds: question.haystackSessionIds, + metadata: question.metadata, + ingestionGroupId: declaredGroupId ?? question.questionId, + haystackFingerprint, + } + }) + + return stableSha256({ + schemaVersion: 1, + kind: "selected-benchmark-questions-and-haystacks", + benchmark: benchmark.name, + benchmarkScope: benchmark.scope, + questions, + haystacks: [...haystacksByDigest.entries()] + .map(([haystackFingerprint, haystack]) => ({ haystackFingerprint, ...haystack })) + .sort((left, right) => left.haystackFingerprint.localeCompare(right.haystackFingerprint)), + }) +} diff --git a/src/orchestrator/phases/answer.ts b/src/orchestrator/phases/answer.ts index 4cfc92a..27163f8 100644 --- a/src/orchestrator/phases/answer.ts +++ b/src/orchestrator/phases/answer.ts @@ -1,71 +1,257 @@ -import { readFileSync, existsSync } from "fs" -import { createOpenAI } from "@ai-sdk/openai" import { createAnthropic } from "@ai-sdk/anthropic" import { createGoogleGenerativeAI } from "@ai-sdk/google" +import { createOpenAI } from "@ai-sdk/openai" import { generateText } from "ai" import type { Benchmark } from "../../types/benchmark" -import type { RunCheckpoint } from "../../types/checkpoint" +import type { + AnswerAttemptMetrics, + ProviderUsage, + QuestionCheckpoint, + RunCheckpoint, +} from "../../types/checkpoint" +import type { ModelRequest, ModelTransport, TerminalEmptyOutputPolicy } from "../../types/protocol" import type { Provider } from "../../types/provider" -import { CheckpointManager } from "../checkpoint" +import { resolveConcurrency } from "../../types/concurrency" import { config } from "../../utils/config" import { logger } from "../../utils/logger" -import { getModelConfig, ModelConfig, DEFAULT_ANSWERING_MODEL } from "../../utils/models" -import { buildDefaultAnswerPrompt } from "../../prompts/defaults" -import { buildBeamAnswerPrompt } from "../../prompts/beam" -import { buildContextString } from "../../types/prompts" -import { ConcurrentExecutor } from "../concurrent" -import { resolveConcurrency } from "../../types/concurrency" +import { + DEFAULT_ANSWERING_MODEL, + getModelConfig, + resolveAnsweringRuntimeIdentity, + type ModelConfig, +} from "../../utils/models" +import { stableSha256 } from "../../utils/stable" import { countTokens } from "../../utils/tokens" +import { CheckpointManager } from "../checkpoint" +import { ConcurrentExecutor } from "../concurrent" +import { extractTokenUsageFromError } from "../evaluation-runtime" -type LanguageModel = +export type LanguageModelFactory = | ReturnType | ReturnType | ReturnType +export const ANSWER_RUNTIME_EXECUTION_VERSION = "chat-transport-durable-outer-retry-v1" + function getAnsweringModel(modelAlias: string): { - client: LanguageModel + client: LanguageModelFactory modelConfig: ModelConfig } { const modelConfig = getModelConfig(modelAlias || DEFAULT_ANSWERING_MODEL) - switch (modelConfig.provider) { case "openai": - return { - client: createOpenAI({ apiKey: config.openaiApiKey }), - modelConfig, - } + return { client: createOpenAI({ apiKey: config.openaiApiKey }), modelConfig } case "anthropic": - return { - client: createAnthropic({ apiKey: config.anthropicApiKey }), - modelConfig, - } + return { client: createAnthropic({ apiKey: config.anthropicApiKey }), modelConfig } case "google": - return { - client: createGoogleGenerativeAI({ apiKey: config.googleApiKey }), - modelConfig, - } + return { client: createGoogleGenerativeAI({ apiKey: config.googleApiKey }), modelConfig } } } -function buildAnswerPrompt( - question: string, - context: unknown[], - questionDate?: string, - provider?: Provider -): string { - if (provider?.prompts?.answerPrompt) { - const customPrompt = provider.prompts.answerPrompt - if (typeof customPrompt === "function") { - return customPrompt(question, context, questionDate) +function requestText(request: ModelRequest): string { + return request.system ? `${request.system}\n\n${request.prompt}` : request.prompt +} + +export function getLanguageModel( + client: LanguageModelFactory, + modelConfig: ModelConfig, + transport: ModelTransport = "provider-default" +) { + if (transport === "openai-chat-completions") { + if (modelConfig.provider !== "openai") { + throw new Error( + `Model transport ${transport} requires an OpenAI answering model; received ${modelConfig.provider}` + ) + } + return (client as ReturnType).chat(modelConfig.id) + } + return client(modelConfig.id) +} + +export function normalizeAnsweringUsage(value: unknown): ProviderUsage { + const usage = + typeof value === "object" && value !== null ? (value as Record) : {} + const token = (name: string): number | undefined => { + const candidate = usage[name] + return typeof candidate === "number" && Number.isFinite(candidate) && candidate >= 0 + ? candidate + : undefined + } + const inputTokens = token("inputTokens") + const outputTokens = token("outputTokens") + const reasoningTokens = token("reasoningTokens") + const reportedTotal = token("totalTokens") + const totalTokens = + reportedTotal ?? + (inputTokens !== undefined && outputTokens !== undefined + ? inputTokens + outputTokens + : undefined) + return { + requestCount: 1, + ...(inputTokens !== undefined ? { inputTokens } : {}), + ...(outputTokens !== undefined ? { outputTokens } : {}), + ...(reasoningTokens !== undefined ? { reasoningTokens } : {}), + ...(totalTokens !== undefined ? { totalTokens } : {}), + } +} + +function withUsageCoverage(value: unknown): ProviderUsage { + const usage = normalizeAnsweringUsage(value) + const knownTokenFields = [usage.inputTokens, usage.outputTokens, usage.totalTokens].filter( + (token) => token !== undefined + ).length + return { + ...usage, + ...(knownTokenFields === 3 + ? { tokenUsageCompleteRequestCount: 1 } + : knownTokenFields > 0 + ? { tokenUsagePartialRequestCount: 1 } + : { tokenUsageUnknownRequestCount: 1 }), + } +} + +export function aggregateAnswerAttemptUsage( + attempts: readonly AnswerAttemptMetrics[] +): ProviderUsage | undefined { + const usages = attempts.flatMap((attempt) => (attempt.usage ? [attempt.usage] : [])) + if (usages.length === 0) return undefined + const aggregate: ProviderUsage = {} + const additiveKeys = [ + "requestCount", + "tokenUsageCompleteRequestCount", + "tokenUsagePartialRequestCount", + "tokenUsageUnknownRequestCount", + "inputTokens", + "outputTokens", + "reasoningTokens", + "totalTokens", + ] as const + for (const key of additiveKeys) { + const values = usages.flatMap((usage) => (usage[key] === undefined ? [] : [usage[key]!])) + if (values.length > 0) aggregate[key] = values.reduce((sum, value) => sum + value, 0) + } + return aggregate +} + +export interface AnswerGenerationResult { + text: string + usage?: unknown + finishReason?: string +} + +export interface AnswerRetryOptions { + maxAttempts?: number + timeoutMs?: number + retryBackoffMs?: number + terminalEmptyOutputPolicy?: TerminalEmptyOutputPolicy + attemptOffset?: number + execute(attempt: number, abortSignal?: AbortSignal): Promise + onAttempt(attempt: AnswerAttemptMetrics): void | Promise + sleep?(delayMs: number): Promise +} + +export interface GeneratedAnswerOutcome { + hypothesis: string + terminalEmptyAccepted: boolean +} + +export function shouldRunAnswerPhase(phases: QuestionCheckpoint["phases"] | undefined): boolean { + return phases?.answer.status !== "completed" && phases?.search.status === "completed" +} + +/** Protocol-owned outer retry loop. Empty model text is retried before terminal policy applies. */ +export async function generateAnswerWithRetries( + options: AnswerRetryOptions +): Promise { + const maxAttempts = options.maxAttempts ?? 1 + const timeoutMs = options.timeoutMs + const retryBackoffMs = options.retryBackoffMs ?? 0 + if (!Number.isInteger(maxAttempts) || maxAttempts < 1) { + throw new Error(`Answer maxAttempts must be a positive integer; received ${maxAttempts}`) + } + if (timeoutMs !== undefined && (!Number.isInteger(timeoutMs) || timeoutMs < 1)) { + throw new Error(`Answer timeoutMs must be a positive integer; received ${timeoutMs}`) + } + if (!Number.isInteger(retryBackoffMs) || retryBackoffMs < 0) { + throw new Error( + `Answer retryBackoffMs must be a non-negative integer; received ${retryBackoffMs}` + ) + } + + const sleep = + options.sleep ?? ((delayMs) => new Promise((resolve) => setTimeout(resolve, delayMs))) + let lastError: unknown + for (let localAttempt = 1; localAttempt <= maxAttempts; localAttempt++) { + const attempt = (options.attemptOffset ?? 0) + localAttempt + const startedAt = new Date().toISOString() + const startedMs = Date.now() + await options.onAttempt({ attempt, startedAt, status: "in_progress" }) + let result: AnswerGenerationResult + try { + result = await options.execute( + localAttempt, + timeoutMs === undefined ? undefined : AbortSignal.timeout(timeoutMs) + ) + } catch (error) { + const completedAt = new Date().toISOString() + const message = error instanceof Error ? error.message : String(error) + await options.onAttempt({ + attempt, + startedAt, + completedAt, + durationMs: Date.now() - startedMs, + status: "failed", + usage: withUsageCoverage(extractTokenUsageFromError(error)), + error: message, + }) + lastError = error + if (localAttempt < maxAttempts && retryBackoffMs > 0) { + await sleep(retryBackoffMs * localAttempt) + } + continue + } + + const completedAt = new Date().toISOString() + const usage = withUsageCoverage(result.usage) + const hypothesis = result.text.trim() + if (!hypothesis) { + const error = "Answering model returned an empty hypothesis" + await options.onAttempt({ + attempt, + startedAt, + completedAt, + durationMs: Date.now() - startedMs, + status: "failed", + finishReason: result.finishReason, + ...(usage.reasoningTokens !== undefined ? { reasoningTokens: usage.reasoningTokens } : {}), + usage, + error, + }) + lastError = new Error(error) + } else { + await options.onAttempt({ + attempt, + startedAt, + completedAt, + durationMs: Date.now() - startedMs, + status: "completed", + finishReason: result.finishReason, + ...(usage.reasoningTokens !== undefined ? { reasoningTokens: usage.reasoningTokens } : {}), + usage, + }) + return { hypothesis, terminalEmptyAccepted: false } } - const contextStr = buildContextString(context) - return customPrompt - .replace("{{question}}", question) - .replace("{{questionDate}}", questionDate || "Not specified") - .replace("{{context}}", contextStr) + if (localAttempt < maxAttempts && retryBackoffMs > 0) { + await sleep(retryBackoffMs * localAttempt) + } + } + + if (options.terminalEmptyOutputPolicy === "accept-and-evaluate") { + return { hypothesis: "", terminalEmptyAccepted: true } } - return buildDefaultAnswerPrompt(question, context, questionDate) + const message = lastError instanceof Error ? lastError.message : String(lastError) + throw new Error(`Answer generation failed after ${maxAttempts} attempts: ${message}`) } export async function runAnswerPhase( @@ -77,16 +263,11 @@ export async function runAnswerPhase( ): Promise { const questions = benchmark.getQuestions() const targetQuestions = questionIds - ? questions.filter((q) => questionIds.includes(q.questionId)) + ? questions.filter((question) => questionIds.includes(question.questionId)) : questions - - const pendingQuestions = targetQuestions.filter((q) => { - const status = checkpointManager.getPhaseStatus(checkpoint, q.questionId, "answer") - const searchStatus = checkpointManager.getPhaseStatus(checkpoint, q.questionId, "search") - const resultFile = checkpoint.questions[q.questionId]?.phases.search.resultFile - return ( - status !== "completed" && searchStatus === "completed" && resultFile && existsSync(resultFile) - ) + const pendingQuestions = targetQuestions.filter((question) => { + const phases = checkpoint.questions[question.questionId]?.phases + return shouldRunAnswerPhase(phases) }) if (pendingQuestions.length === 0) { @@ -94,9 +275,15 @@ export async function runAnswerPhase( return } + const resolvedRuntimeIdentity = resolveAnsweringRuntimeIdentity(checkpoint.answeringModel) + if ( + !checkpoint.answeringRuntimeIdentity || + stableSha256(checkpoint.answeringRuntimeIdentity) !== stableSha256(resolvedRuntimeIdentity) + ) { + throw new Error("Answering runtime identity differs from the checkpoint") + } const { client, modelConfig } = getAnsweringModel(checkpoint.answeringModel) const concurrency = resolveConcurrency("answer", checkpoint.concurrency, provider?.concurrency) - logger.info( `Generating answers for ${pendingQuestions.length} questions using ${modelConfig.displayName} (concurrency: ${concurrency})...` ) @@ -107,83 +294,138 @@ export async function runAnswerPhase( checkpoint.runId, "answer", async ({ item: question, index, total }) => { - const resultFile = checkpoint.questions[question.questionId].phases.search.resultFile! + const questionCheckpoint = checkpoint.questions[question.questionId] + const search = questionCheckpoint.phases.search + if (!search.retrievalPlan) + throw new Error(`Missing retrieval plan for ${question.questionId}`) + const results = search.results || [] + const sessions = benchmark.getHaystackSessions(question.questionId) + const answerPlan = benchmark.protocol.createAnswerPlan({ + question, + sessions, + results, + retrieval: search.retrievalPlan, + questionDate: questionCheckpoint.questionDate, + providerPrompts: provider?.prompts, + }) + if (answerPlan.answerEvidenceCount > search.retrievalPlan.answerCutoff) { + throw new Error( + `Protocol exposed ${answerPlan.answerEvidenceCount} evidence items above answer cutoff ${search.retrievalPlan.answerCutoff}` + ) + } - const startTime = Date.now() + const basePromptTokens = countTokens(requestText(answerPlan.baseRequest), modelConfig) + const promptTokens = countTokens(requestText(answerPlan.request), modelConfig) + const contextTokens = Math.max(0, promptTokens - basePromptTokens) + const startedAt = new Date().toISOString() + const startedMs = Date.now() + const priorAttempts = (questionCheckpoint.phases.answer.attempts ?? []).map((attempt) => + attempt.status === "in_progress" + ? { + ...attempt, + status: "failed" as const, + completedAt: new Date().toISOString(), + usage: attempt.usage ?? { + requestCount: 1, + tokenUsageUnknownRequestCount: 1, + }, + error: "Interrupted before the answering attempt completed", + } + : attempt + ) checkpointManager.updatePhase(checkpoint, question.questionId, "answer", { status: "in_progress", - startedAt: new Date().toISOString(), + startedAt, + evidenceCount: answerPlan.answerEvidenceCount, + attempts: priorAttempts, + terminalEmptyAccepted: undefined, + costUsd: null, + error: undefined, + }) + checkpointManager.updatePhase(checkpoint, question.questionId, "search", { + answerEvidenceCount: answerPlan.answerEvidenceCount, }) try { - const searchData = JSON.parse(readFileSync(resultFile, "utf8")) - const context: unknown[] = searchData.results || [] - const questionDate = checkpoint.questions[question.questionId]?.questionDate - - // BEAM uses mem0's answer-generation prompt verbatim so retrieval + - // judging numbers are apples-to-apples with their published BEAM - // results. The provider-specific prompt is intentionally bypassed. - let basePrompt: string - let prompt: string - if (benchmark.name.startsWith("beam")) { - const sessions = benchmark.getHaystackSessions(question.questionId) - const sessionDateMap = new Map() - for (const s of sessions) { - const date = s.metadata?.date - if (typeof date === "string") sessionDateMap.set(s.sessionId, date) - } - basePrompt = buildBeamAnswerPrompt(question.question, [], sessionDateMap) - prompt = buildBeamAnswerPrompt(question.question, context, sessionDateMap) - } else { - basePrompt = buildAnswerPrompt(question.question, [], questionDate, provider) - prompt = buildAnswerPrompt(question.question, context, questionDate, provider) - } - - const basePromptTokens = countTokens(basePrompt, modelConfig) - const promptTokens = countTokens(prompt, modelConfig) - // Derive contextTokens from the difference so it reflects the actual formatted - // context in the prompt (not the raw JSON), which matters for providers with - // custom prompt functions that transform context (e.g. Zep's XML-like tags). - const contextTokens = Math.max(0, promptTokens - basePromptTokens) - - const params: Record = { - model: client(modelConfig.id), - prompt, - maxTokens: modelConfig.defaultMaxTokens, + const request = answerPlan.request + if ( + request.innerMaxRetries !== undefined && + (!Number.isInteger(request.innerMaxRetries) || request.innerMaxRetries < 0) + ) { + throw new Error( + `Answer innerMaxRetries must be a non-negative integer; received ${request.innerMaxRetries}` + ) } - - if (modelConfig.supportsTemperature) { - params.temperature = modelConfig.defaultTemperature + const params: Parameters[0] = { + model: getLanguageModel(client, modelConfig, request.transport), + prompt: request.prompt, + maxOutputTokens: request.maxOutputTokens ?? modelConfig.defaultMaxTokens, + ...(request.system ? { system: request.system } : {}), + ...(modelConfig.supportsTemperature + ? { temperature: request.temperature ?? modelConfig.defaultTemperature } + : {}), + ...(request.innerMaxRetries !== undefined ? { maxRetries: request.innerMaxRetries } : {}), } - - const { text } = await generateText(params as Parameters[0]) - - const durationMs = Date.now() - startTime + const outcome = await generateAnswerWithRetries({ + maxAttempts: request.maxAttempts, + timeoutMs: request.timeoutMs, + retryBackoffMs: request.retryBackoffMs, + terminalEmptyOutputPolicy: request.terminalEmptyOutputPolicy, + attemptOffset: priorAttempts.length, + execute: async (_attempt, abortSignal) => { + const result = await generateText({ + ...params, + ...(abortSignal ? { abortSignal } : {}), + }) + return { + text: result.text, + usage: result.usage, + finishReason: result.finishReason, + } + }, + onAttempt: (attempt) => { + const attempts = questionCheckpoint.phases.answer.attempts ?? [] + const existingAttempt = attempts.findIndex( + (candidate) => candidate.attempt === attempt.attempt + ) + const updatedAttempts = [...attempts] + if (existingAttempt >= 0) updatedAttempts[existingAttempt] = attempt + else updatedAttempts.push(attempt) + checkpointManager.updatePhase(checkpoint, question.questionId, "answer", { + attempts: updatedAttempts, + }) + }, + }) + const completedAt = new Date().toISOString() + const durationMs = Date.now() - startedMs checkpointManager.updatePhase(checkpoint, question.questionId, "answer", { status: "completed", - hypothesis: text.trim(), + hypothesis: outcome.hypothesis, + terminalEmptyAccepted: outcome.terminalEmptyAccepted, promptTokens, basePromptTokens, contextTokens, - completedAt: new Date().toISOString(), + evidenceCount: answerPlan.answerEvidenceCount, + usage: aggregateAnswerAttemptUsage(questionCheckpoint.phases.answer.attempts ?? []), + completedAt, durationMs, + error: undefined, }) - logger.progress( index + 1, total, - `Answered ${question.questionId} (${durationMs}ms, ${promptTokens} tokens: ${basePromptTokens} base + ${contextTokens} context)` + `Answered ${question.questionId}${outcome.terminalEmptyAccepted ? " (terminal empty accepted)" : ""} (${durationMs}ms, ${promptTokens} tokens: ${basePromptTokens} base + ${contextTokens} context)` ) return { questionId: question.questionId, durationMs } - } catch (e) { - const error = e instanceof Error ? e.message : String(e) + } catch (error) { + const message = error instanceof Error ? error.message : String(error) checkpointManager.updatePhase(checkpoint, question.questionId, "answer", { status: "failed", - error, + usage: aggregateAnswerAttemptUsage(questionCheckpoint.phases.answer.attempts ?? []), + error: message, }) - logger.error(`Failed to answer ${question.questionId}: ${error}`) throw new Error( - `Answer failed at ${question.questionId}: ${error}. Fix the issue and resume with the same run ID.` + `Answer failed at ${question.questionId}: ${message}. Fix the issue and resume with the same run ID.` ) } } diff --git a/src/orchestrator/phases/evaluate.ts b/src/orchestrator/phases/evaluate.ts index 26c990d..8e0e3f8 100644 --- a/src/orchestrator/phases/evaluate.ts +++ b/src/orchestrator/phases/evaluate.ts @@ -1,76 +1,22 @@ -import type { Judge } from "../../types/judge" import type { Benchmark } from "../../types/benchmark" -import type { RunCheckpoint } from "../../types/checkpoint" +import type { AnswerPhaseCheckpoint, RunCheckpoint } from "../../types/checkpoint" +import type { Judge } from "../../types/judge" import type { Provider } from "../../types/provider" -import { generateText } from "ai" -import { CheckpointManager } from "../checkpoint" +import { resolveConcurrency } from "../../types/concurrency" import { logger } from "../../utils/logger" +import { CheckpointManager } from "../checkpoint" import { ConcurrentExecutor } from "../concurrent" -import { resolveConcurrency } from "../../types/concurrency" -import { calculateRetrievalMetrics } from "./retrieval-eval" -import { buildBeamRubricJudgePrompt, parseBeamRubricJudgeResponse } from "../../prompts/beam" - -interface BeamRubricItemResult { - rubricItem: string - score: number - reason: string -} - -function getBeamRubric(question: { metadata?: Record }): string[] | null { - const rubric = question.metadata?.rubric - if (!Array.isArray(rubric) || rubric.some((item) => typeof item !== "string")) { - return null - } - - return rubric -} - -async function evaluateBeamRubricQuestion( - judge: Judge, - question: { question: string; metadata?: Record }, - hypothesis: string -): Promise<{ score: number; label: "correct" | "incorrect"; explanation: string; details: Record }> { - const rubric = getBeamRubric(question) - if (!rubric) { - return { - score: 0, - label: "incorrect", - explanation: "Missing BEAM rubric metadata", - details: {}, - } - } - - const model = judge.getModel() - const results: BeamRubricItemResult[] = [] - - for (const rubricItem of rubric) { - const prompt = buildBeamRubricJudgePrompt(question.question, rubricItem, hypothesis) - const { text } = await generateText({ - model, - prompt, - maxOutputTokens: 512, - temperature: 0, - }) - const parsed = parseBeamRubricJudgeResponse(text) - results.push({ - rubricItem, - score: parsed.score, - reason: parsed.reason, - }) - } +import { JudgeEvaluationRuntime } from "../evaluation-runtime" +import { calculateProtocolRetrievalMetrics } from "./retrieval-eval" - const averageScore = - results.length > 0 ? results.reduce((sum, item) => sum + item.score, 0) / results.length : 0 - - return { - score: averageScore, - label: averageScore >= 1 ? "correct" : "incorrect", - explanation: `BEAM rubric average score: ${averageScore.toFixed(2)}`, - details: { - rubricResults: results, - rubricAverageScore: averageScore, - }, +export function hasEvaluableAnswer(answer: AnswerPhaseCheckpoint | undefined): boolean { + if (!answer || answer.status !== "completed" || typeof answer.hypothesis !== "string") { + return false } + return ( + answer.hypothesis.trim().length > 0 || + (answer.hypothesis === "" && answer.terminalEmptyAccepted === true) + ) } export async function runEvaluatePhase( @@ -83,14 +29,11 @@ export async function runEvaluatePhase( ): Promise { const questions = benchmark.getQuestions() const targetQuestions = questionIds - ? questions.filter((q) => questionIds.includes(q.questionId)) + ? questions.filter((question) => questionIds.includes(question.questionId)) : questions - - const pendingQuestions = targetQuestions.filter((q) => { - const status = checkpointManager.getPhaseStatus(checkpoint, q.questionId, "evaluate") - const answerStatus = checkpointManager.getPhaseStatus(checkpoint, q.questionId, "answer") - const hypothesis = checkpoint.questions[q.questionId]?.phases.answer.hypothesis - return status !== "completed" && answerStatus === "completed" && hypothesis + const pendingQuestions = targetQuestions.filter((question) => { + const phases = checkpoint.questions[question.questionId]?.phases + return phases?.evaluate.status !== "completed" && hasEvaluableAnswer(phases?.answer) }) if (pendingQuestions.length === 0) { @@ -99,7 +42,6 @@ export async function runEvaluatePhase( } const concurrency = resolveConcurrency("evaluate", checkpoint.concurrency, provider?.concurrency) - logger.info( `Evaluating ${pendingQuestions.length} questions with ${judge.name} (concurrency: ${concurrency})...` ) @@ -110,46 +52,78 @@ export async function runEvaluatePhase( checkpoint.runId, "evaluate", async ({ item: question, index, total }) => { - const hypothesis = checkpoint.questions[question.questionId].phases.answer.hypothesis! - - const startTime = Date.now() + const questionCheckpoint = checkpoint.questions[question.questionId] + const hypothesis = questionCheckpoint.phases.answer.hypothesis + if (typeof hypothesis !== "string") { + throw new Error(`Missing completed hypothesis for ${question.questionId}`) + } + const search = questionCheckpoint.phases.search + if (!search.retrievalPlan) + throw new Error(`Missing retrieval plan for ${question.questionId}`) + const results = search.results || [] + const startedAt = new Date().toISOString() + const startedMs = Date.now() checkpointManager.updatePhase(checkpoint, question.questionId, "evaluate", { status: "in_progress", - startedAt: new Date().toISOString(), + startedAt, + costUsd: null, + error: undefined, }) + await checkpointManager.flush(checkpoint.runId) + const runtime = new JudgeEvaluationRuntime(judge) try { - const searchResults = checkpoint.questions[question.questionId].phases.search.results || [] - const rubric = getBeamRubric(question) - - const [result, retrievalMetrics] = await Promise.all([ - rubric - ? evaluateBeamRubricQuestion(judge, question, hypothesis) - : judge.evaluate({ - question: question.question, - questionType: question.questionType, - groundTruth: question.groundTruth, - hypothesis, - providerPrompts: provider?.prompts, - }), - calculateRetrievalMetrics( - judge.getModel(), - question.question, - question.groundTruth, - searchResults - ), - ]) - - const durationMs = Date.now() - startTime + const evaluation = await benchmark.protocol.evaluateQuestion( + { + question, + hypothesis, + results, + retrieval: search.retrievalPlan, + providerPrompts: provider?.prompts, + protocolProgress: questionCheckpoint.phases.evaluate.protocolProgress, + onProtocolProgress: async (protocolProgress) => { + checkpointManager.updatePhase(checkpoint, question.questionId, "evaluate", { + protocolProgress, + }) + await checkpointManager.flush(checkpoint.runId) + }, + }, + runtime + ) + const retrievalMetrics = await calculateProtocolRetrievalMetrics( + benchmark.protocol.auxiliaryRetrievalEvaluation, + runtime, + question.question, + question.groundTruth, + results, + search.retrievalPlan.answerCutoff + ) + if ( + !Number.isFinite(evaluation.primaryScore) || + evaluation.primaryScore < 0 || + evaluation.primaryScore > 1 + ) { + throw new Error(`Protocol returned invalid primary score ${evaluation.primaryScore}`) + } + + const completedAt = new Date().toISOString() + const durationMs = Date.now() - startedMs + const passed = evaluation.passed checkpointManager.updatePhase(checkpoint, question.questionId, "evaluate", { status: "completed", - score: result.score, - label: result.label, - explanation: result.explanation, - details: result.details, - retrievalMetrics, - completedAt: new Date().toISOString(), + evaluation, + score: evaluation.primaryScore, + label: passed ? "correct" : "incorrect", + explanation: evaluation.explanation, + details: { + ...(evaluation.details || {}), + ...(evaluation.metrics ? { protocolMetrics: evaluation.metrics } : {}), + }, + ...(retrievalMetrics ? { retrievalMetrics } : {}), + ...(runtime.getUsage?.() ? { usage: runtime.getUsage!() } : {}), + completedAt, durationMs, + error: undefined, }) const retrievalInfo = retrievalMetrics @@ -158,19 +132,18 @@ export async function runEvaluatePhase( logger.progress( index + 1, total, - `Evaluated ${question.questionId}: ${result.label}${retrievalInfo} (${durationMs}ms)` + `Evaluated ${question.questionId}: ${passed ? "pass" : "fail"} (${evaluation.primaryScore.toFixed(3)})${retrievalInfo} (${durationMs}ms)` ) - - return { questionId: question.questionId, durationMs, label: result.label } - } catch (e) { - const error = e instanceof Error ? e.message : String(e) + return { questionId: question.questionId, durationMs, passed } + } catch (error) { + const message = error instanceof Error ? error.message : String(error) checkpointManager.updatePhase(checkpoint, question.questionId, "evaluate", { status: "failed", - error, + ...(runtime.getUsage?.() ? { usage: runtime.getUsage!() } : {}), + error: message, }) - logger.error(`Failed to evaluate ${question.questionId}: ${error}`) throw new Error( - `Evaluate failed at ${question.questionId}: ${error}. Fix the issue and resume with the same run ID.` + `Evaluate failed at ${question.questionId}: ${message}. Fix the issue and resume with the same run ID.` ) } } diff --git a/src/orchestrator/phases/indexing.ts b/src/orchestrator/phases/indexing.ts index bd822fc..2368da9 100644 --- a/src/orchestrator/phases/indexing.ts +++ b/src/orchestrator/phases/indexing.ts @@ -1,170 +1,97 @@ -import type { Provider, IndexingProgress } from "../../types/provider" -import type { RunCheckpoint, QuestionCheckpoint } from "../../types/checkpoint" -import { CheckpointManager } from "../checkpoint" +import type { BuildAttemptMetrics, RunCheckpoint } from "../../types/checkpoint" +import type { IndexingProgress, Provider } from "../../types/provider" +import { resolveConcurrency } from "../../types/concurrency" import { logger } from "../../utils/logger" +import { CheckpointManager } from "../checkpoint" import { ConcurrentExecutor } from "../concurrent" -import { resolveConcurrency } from "../../types/concurrency" -function getEpisodeCount(question: QuestionCheckpoint): number { - const ingestResult = question.phases.ingest.ingestResult - if (!ingestResult) return 0 - return (ingestResult.documentIds?.length || 0) + (ingestResult.taskIds?.length || 0) +function totalAttemptDuration(attempts: BuildAttemptMetrics[]): number { + return attempts.reduce((sum, attempt) => sum + (attempt.durationMs ?? 0), 0) } -class IndexingProgressTracker { - private progressByQuestion: Map = - new Map() - private totalEpisodes: number = 0 - private lastDisplayed: string = "" - - constructor(questions: QuestionCheckpoint[]) { - for (const q of questions) { - const count = getEpisodeCount(q) - this.totalEpisodes += count - this.progressByQuestion.set(q.questionId, { completed: 0, failed: 0, total: count }) - } +export function validateCompletedIndexingProgress( + expectedIds: readonly string[], + progress: IndexingProgress +): void { + const expected = new Set(expectedIds) + if (expected.size !== expectedIds.length) { + throw new Error("Indexing input contains duplicate document/task IDs") } - - update(questionId: string, progress: IndexingProgress): void { - const current = this.progressByQuestion.get(questionId) - if (current) { - this.progressByQuestion.set(questionId, { - completed: progress.completedIds.length, - failed: progress.failedIds.length, - total: progress.total, - }) - } - this.display() + if (!Number.isInteger(progress.total) || progress.total !== expected.size) { + throw new Error( + `Indexing provider reported total ${progress.total}; expected ${expected.size} unique IDs` + ) } - - markQuestionDone(questionId: string): void { - const current = this.progressByQuestion.get(questionId) - if (current) { - this.progressByQuestion.set(questionId, { - completed: current.total, - failed: current.failed, - total: current.total, - }) - } + if (new Set(progress.completedIds).size !== progress.completedIds.length) { + throw new Error("Indexing provider reported duplicate completed IDs") } - - getAggregated(): { completed: number; failed: number; total: number } { - let completed = 0 - let failed = 0 - for (const p of this.progressByQuestion.values()) { - completed += p.completed - failed += p.failed - } - return { completed, failed, total: this.totalEpisodes } + if (new Set(progress.failedIds).size !== progress.failedIds.length) { + throw new Error("Indexing provider reported duplicate failed IDs") } - - display(): void { - const agg = this.getAggregated() - const displayStr = `${agg.completed}/${agg.total}` - if (displayStr !== this.lastDisplayed) { - this.lastDisplayed = displayStr - const percent = agg.total > 0 ? Math.round((agg.completed / agg.total) * 100) : 0 - const bar = "█".repeat(Math.floor(percent / 5)) + "░".repeat(20 - Math.floor(percent / 5)) - const failedStr = agg.failed > 0 ? ` (${agg.failed} failed)` : "" - process.stdout.write( - `\r\x1b[36m[${bar}]\x1b[0m ${percent}% Indexing: ${agg.completed}/${agg.total} episodes${failedStr}` - ) - } + const completed = new Set(progress.completedIds) + const failed = new Set(progress.failedIds) + for (const id of [...completed, ...failed]) { + if (!expected.has(id)) throw new Error(`Indexing provider reported unknown ID ${id}`) } - - finish(): void { - const agg = this.getAggregated() - const failedStr = agg.failed > 0 ? ` (${agg.failed} failed)` : "" - process.stdout.write( - `\r\x1b[36m[${"█".repeat(20)}]\x1b[0m 100% Indexing: ${agg.completed}/${agg.total} episodes${failedStr}\n` - ) + for (const id of completed) { + if (failed.has(id)) throw new Error(`Indexing provider reported ${id} as completed and failed`) } - - getTotalEpisodes(): number { - return this.totalEpisodes + if (failed.size > 0) { + throw new Error(`${failed.size} indexing items failed: ${[...failed].join(", ")}`) + } + const missing = [...expected].filter((id) => !completed.has(id)) + if (missing.length > 0) { + throw new Error( + `Indexing provider returned before ${missing.length} IDs completed: ${missing.join(", ")}` + ) } } export async function runIndexingPhase( provider: Provider, checkpoint: RunCheckpoint, - checkpointManager: CheckpointManager, - questionIds?: string[] + checkpointManager: CheckpointManager ): Promise { - const allQuestions = Object.values(checkpoint.questions) - const targetQuestions = questionIds - ? allQuestions.filter((q) => questionIds.includes(q.questionId)) - : allQuestions - - const toIndex = targetQuestions.filter( - (q) => q.phases.ingest.status === "completed" && q.phases.indexing.status !== "completed" + const pendingBuilds = Object.values(checkpoint.builds).filter( + (build) => build.ingest.status === "completed" && build.indexing.status !== "completed" ) - - if (toIndex.length === 0) { - logger.info("No questions pending indexing") + if (pendingBuilds.length === 0) { + logger.info("No builds pending indexing") return } const concurrency = resolveConcurrency("indexing", checkpoint.concurrency, provider.concurrency) - - const questionsByContainer = new Map() - for (const question of toIndex) { - const questions = questionsByContainer.get(question.containerTag) || [] - questions.push(question) - questionsByContainer.set(question.containerTag, questions) - } - - const indexingGroups = Array.from(questionsByContainer.entries()).map( - ([containerTag, questions]) => ({ - containerTag, - questions, - representative: questions[0], - }) - ) - - const tracker = new IndexingProgressTracker(indexingGroups.map((group) => group.representative)) - const totalEpisodes = tracker.getTotalEpisodes() - logger.info( - `Awaiting indexing for ${toIndex.length} questions across ${indexingGroups.length} containers, ${totalEpisodes} episodes (concurrency: ${concurrency})...` + `Awaiting indexing for ${pendingBuilds.length} builds (concurrency: ${concurrency})...` ) - tracker.display() - await ConcurrentExecutor.execute( - indexingGroups, + pendingBuilds, concurrency, checkpoint.runId, "indexing", - async ({ item: group }) => { - const question = group.representative - const groupQuestionIds = group.questions.map((q) => q.questionId) - const ingestResult = question.phases.ingest.ingestResult - const episodeCount = getEpisodeCount(question) - - if (!ingestResult || episodeCount === 0) { - for (const questionId of groupQuestionIds) { - checkpointManager.updatePhase(checkpoint, questionId, "indexing", { - status: "completed", - completedIds: [], - failedIds: [], - completedAt: new Date().toISOString(), - durationMs: 0, - }) - } - tracker.markQuestionDone(question.questionId) - return { questionId: question.questionId, durationMs: 0 } + async ({ item: build, index, total }) => { + const ingestResult = { + documentIds: [...build.ingest.documentIds], + ...(build.ingest.taskIds.length > 0 ? { taskIds: [...build.ingest.taskIds] } : {}), } - - const startTime = Date.now() - for (const questionId of groupQuestionIds) { - checkpointManager.updatePhase(checkpoint, questionId, "indexing", { - status: "in_progress", - completedIds: [], - failedIds: [], - startedAt: new Date().toISOString(), - }) + const expectedIds = [...ingestResult.documentIds, ...(ingestResult.taskIds ?? [])] + const episodeCount = expectedIds.length + const startedAt = new Date().toISOString() + const startedMs = Date.now() + const attempt: BuildAttemptMetrics = { + phase: "indexing", + attempt: build.indexing.attempts.length + 1, + startedAt, + status: "in_progress", + costUsd: null, } + checkpointManager.updateBuild(checkpoint, build.buildId, (current) => { + current.indexing.status = "in_progress" + current.indexing.startedAt ??= startedAt + current.indexing.error = undefined + current.indexing.attempts.push(attempt) + }) try { let lastProgress: IndexingProgress = { @@ -172,48 +99,52 @@ export async function runIndexingPhase( failedIds: [], total: episodeCount, } - - await provider.awaitIndexing(ingestResult, group.containerTag, (progress) => { + await provider.awaitIndexing(ingestResult, build.containerTag, (progress) => { lastProgress = progress - tracker.update(question.questionId, progress) - - for (const questionId of groupQuestionIds) { - checkpointManager.updatePhase(checkpoint, questionId, "indexing", { - status: "in_progress", - completedIds: progress.completedIds, - failedIds: progress.failedIds, - }) - } - }) - - const durationMs = Date.now() - startTime - for (const questionId of groupQuestionIds) { - checkpointManager.updatePhase(checkpoint, questionId, "indexing", { - status: "completed", - completedIds: lastProgress.completedIds, - failedIds: lastProgress.failedIds, - completedAt: new Date().toISOString(), - durationMs, + checkpointManager.updateBuild(checkpoint, build.buildId, (current) => { + current.indexing.completedIds = [...progress.completedIds] + current.indexing.failedIds = [...progress.failedIds] }) - } + }) - return { questionId: question.questionId, durationMs } - } catch (e) { - const error = e instanceof Error ? e.message : String(e) - for (const questionId of groupQuestionIds) { - checkpointManager.updatePhase(checkpoint, questionId, "indexing", { + validateCompletedIndexingProgress(expectedIds, lastProgress) + + const completedAt = new Date().toISOString() + const durationMs = Date.now() - startedMs + checkpointManager.updateBuild(checkpoint, build.buildId, (current) => { + const currentAttempt = current.indexing.attempts.at(-1)! + Object.assign(currentAttempt, { status: "completed", completedAt, durationMs }) + current.indexing.status = "completed" + current.indexing.completedAt = completedAt + current.indexing.durationMs = totalAttemptDuration(current.indexing.attempts) + current.indexing.completedIds = [...lastProgress.completedIds] + current.indexing.failedIds = [] + current.indexing.error = undefined + }) + logger.progress(index + 1, total, `Indexed ${build.ingestionGroupId} (${durationMs}ms)`) + return { buildId: build.buildId, durationMs } + } catch (error) { + const message = error instanceof Error ? error.message : String(error) + const completedAt = new Date().toISOString() + const durationMs = Date.now() - startedMs + checkpointManager.updateBuild(checkpoint, build.buildId, (current) => { + const currentAttempt = current.indexing.attempts.at(-1)! + Object.assign(currentAttempt, { status: "failed", - error, + completedAt, + durationMs, + error: message, }) - } - logger.error(`\nFailed to index ${group.containerTag}: ${error}`) + current.indexing.status = "failed" + current.indexing.error = message + current.indexing.durationMs = totalAttemptDuration(current.indexing.attempts) + }) throw new Error( - `Indexing failed at ${group.containerTag}: ${error}. Fix the issue and resume with the same run ID.` + `Indexing failed at ${build.containerTag}: ${message}. Fix the issue and resume with the same run ID.` ) } } ) - tracker.finish() logger.success("Indexing phase complete") } diff --git a/src/orchestrator/phases/ingest.ts b/src/orchestrator/phases/ingest.ts index 2f93f70..11d3a4f 100644 --- a/src/orchestrator/phases/ingest.ts +++ b/src/orchestrator/phases/ingest.ts @@ -1,158 +1,595 @@ -import type { Provider, IngestResult } from "../../types/provider" -import type { Benchmark } from "../../types/benchmark" -import type { RunCheckpoint } from "../../types/checkpoint" -import { CheckpointManager } from "../checkpoint" +import type { Provider } from "../../types/provider" +import type { BuildAttemptMetrics, RunCheckpoint } from "../../types/checkpoint" import { logger } from "../../utils/logger" -import { ConcurrentExecutor } from "../concurrent" import { resolveConcurrency } from "../../types/concurrency" +import { CheckpointManager } from "../checkpoint" +import { assertCompletedSessionsAreOrderedPrefix, type ValidatedBuildPlan } from "../builds" +import { ConcurrentExecutor } from "../concurrent" +import { validateCompletedIndexingProgress } from "./indexing" +import type { CanonicalIngestionDocument } from "../../types/unified" +import type { IndexingProgress, IngestResult } from "../../types/provider" const RATE_LIMIT_MS = 1000 +export const DEFAULT_INGEST_READINESS_TIMEOUT_MS = 5 * 60 * 1000 + +function totalAttemptDuration(attempts: BuildAttemptMetrics[]): number { + return attempts.reduce((sum, attempt) => sum + (attempt.durationMs ?? 0), 0) +} + +function attributeBatchResult( + documents: CanonicalIngestionDocument[], + result: IngestResult +): Array<{ customId: string; documentIds: string[]; taskIds: string[]; error?: string }> { + if (result.items) { + if (result.items.length !== documents.length) { + throw new Error( + `Provider returned ${result.items.length} item outcomes for ${documents.length} documents` + ) + } + const byCustomId = new Map(result.items.map((item) => [item.customId, item])) + if (byCustomId.size !== result.items.length) { + throw new Error("Provider returned duplicate custom IDs in batch outcomes") + } + const attributed = documents.map((document) => { + const item = byCustomId.get(document.customId) + if (!item) throw new Error(`Provider omitted batch outcome for ${document.customId}`) + return { + customId: item.customId, + documentIds: [...item.documentIds], + taskIds: [...(item.taskIds ?? [])], + ...(item.error ? { error: item.error } : {}), + } + }) + const attributedDocumentIds = attributed.flatMap((item) => item.documentIds) + const attributedTaskIds = attributed.flatMap((item) => item.taskIds) + if ( + new Set(attributedDocumentIds).size !== attributedDocumentIds.length || + new Set(attributedTaskIds).size !== attributedTaskIds.length || + attributedDocumentIds.length !== result.documentIds.length || + attributedTaskIds.length !== (result.taskIds ?? []).length || + attributedDocumentIds.some((id) => !result.documentIds.includes(id)) || + attributedTaskIds.some((id) => !(result.taskIds ?? []).includes(id)) + ) { + throw new Error("Provider aggregate ingest IDs do not match its per-item outcomes") + } + return attributed + } + + if (documents.length === 1) { + return [ + { + customId: documents[0]!.customId, + documentIds: [...result.documentIds], + taskIds: [...(result.taskIds ?? [])], + }, + ] + } + + const documentIds = result.documentIds ?? [] + const taskIds = result.taskIds ?? [] + const count = documents.length + const documentIdsAreAttributable = documentIds.length === 0 || documentIds.length === count + const taskIdsAreAttributable = taskIds.length === 0 || taskIds.length === count + if ( + !documentIdsAreAttributable || + !taskIdsAreAttributable || + (documentIds.length === 0 && taskIds.length === 0) + ) { + throw new Error( + `Provider returned ${documentIds.length} document IDs and ${taskIds.length} task IDs for an ordered batch of ${count} sessions` + ) + } + + return documents.map((_, index) => ({ + customId: documents[index]!.customId, + documentIds: documentIds.length === count ? [documentIds[index]!] : [], + taskIds: taskIds.length === count ? [taskIds[index]!] : [], + })) +} + +function errorMessage(error: unknown): string { + return error instanceof Error ? error.message : String(error) +} + +function sameIds(left: readonly string[], right: readonly string[]): boolean { + return ( + left.length === right.length && + left.every((id) => right.includes(id)) && + right.every((id) => left.includes(id)) + ) +} + +function unresolvedPhysicalIds(build: RunCheckpoint["builds"][string]): string[] { + return [ + ...new Set( + (build.ingest.deferredSessions ?? []).flatMap((deferred) => [ + ...deferred.documentIds, + ...deferred.taskIds, + ]) + ), + ].filter((id) => !build.indexing.completedIds.includes(id)) +} export async function runIngestPhase( provider: Provider, - benchmark: Benchmark, checkpoint: RunCheckpoint, checkpointManager: CheckpointManager, - questionIds?: string[] + plans: ValidatedBuildPlan[] ): Promise { - const questions = benchmark.getQuestions() - const targetQuestions = questionIds - ? questions.filter((q) => questionIds.includes(q.questionId)) - : questions - - const pendingQuestions = targetQuestions.filter((q) => { - const status = checkpointManager.getPhaseStatus(checkpoint, q.questionId, "ingest") - return status !== "completed" - }) + const planByBuildId = new Map(plans.map((plan) => [plan.buildId, plan])) + const pendingBuilds = Object.values(checkpoint.builds).filter( + (build) => build.ingest.status !== "completed" + ) - if (pendingQuestions.length === 0) { - logger.info("No questions pending ingestion") + if (pendingBuilds.length === 0) { + logger.info("No builds pending ingestion") return } - const concurrency = resolveConcurrency("ingest", checkpoint.concurrency, provider.concurrency) - - const questionsByGroup = new Map() - for (const question of pendingQuestions) { - const groupId = benchmark.getIngestionGroupId?.(question.questionId) || question.questionId - const questions = questionsByGroup.get(groupId) || [] - questions.push(question) - questionsByGroup.set(groupId, questions) - } - - const ingestionGroups = Array.from(questionsByGroup.entries()).map(([groupId, questions]) => ({ - groupId, - questions, - representative: questions[0], - })) - + const missingDates = pendingBuilds.reduce((sum, build) => sum + build.missingDocumentDateCount, 0) + const ingestBatchSizes = [...new Set(pendingBuilds.map((build) => build.ingestBatchSize ?? 1))] logger.info( - `Ingesting ${pendingQuestions.length} questions across ${ingestionGroups.length} containers (concurrency: ${concurrency})...` + `Ingesting ${pendingBuilds.length} builds in ordered session batches of ${ingestBatchSizes.join(",")}; ${missingDates} sessions have no valid source date` ) + const ingestConcurrency = resolveConcurrency( + "ingest", + checkpoint.concurrency, + provider.concurrency + ) + const hasPerDocumentBarrier = pendingBuilds.some( + (build) => build.ingestionExecutionPolicy.readinessBarrier === "after-each-document" + ) + const concurrency = hasPerDocumentBarrier + ? Math.min( + ingestConcurrency, + resolveConcurrency("indexing", checkpoint.concurrency, provider.concurrency) + ) + : ingestConcurrency await ConcurrentExecutor.executeBatched({ - items: ingestionGroups, + items: pendingBuilds, concurrency, rateLimitMs: RATE_LIMIT_MS, runId: checkpoint.runId, phaseName: "ingest", - executeTask: async ({ item: group, index, total }) => { - const question = group.representative - const groupQuestionIds = group.questions.map((q) => q.questionId) - const containerTag = checkpoint.questions[question.questionId].containerTag - const sessions = benchmark.getHaystackSessions(question.questionId) - - const sessionsMetadata = sessions.map((s) => ({ - sessionId: s.sessionId, - date: s.metadata?.date as string | undefined, - messageCount: s.messages.length, - })) - for (const questionId of groupQuestionIds) { - checkpointManager.updateSessions(checkpoint, questionId, sessionsMetadata) + executeTask: async ({ item: build, index, total }) => { + const plan = planByBuildId.get(build.buildId) + if (!plan) throw new Error(`Missing validated ingestion plan for ${build.buildId}`) + assertCompletedSessionsAreOrderedPrefix(build) + + const startedAt = new Date().toISOString() + const requiresSessionBarrier = + build.ingestionExecutionPolicy.readinessBarrier === "after-each-document" + const attempt: BuildAttemptMetrics = { + phase: "ingest", + attempt: build.ingest.attempts.length + 1, + startedAt, + status: "in_progress", + costUsd: null, } + const indexingAttempt: BuildAttemptMetrics | undefined = requiresSessionBarrier + ? { + phase: "indexing", + attempt: build.indexing.attempts.length + 1, + startedAt, + status: "in_progress", + costUsd: null, + } + : undefined + let ingestDurationMs = 0 + let indexingDurationMs = 0 + const readinessTimeoutMs = + checkpoint.ingestReadinessTimeoutMs ?? DEFAULT_INGEST_READINESS_TIMEOUT_MS - const startTime = Date.now() - for (const questionId of groupQuestionIds) { - checkpointManager.updatePhase(checkpoint, questionId, "ingest", { - status: "in_progress", - startedAt: new Date().toISOString(), - }) + const awaitReadiness = async ( + result: IngestResult + ): Promise<{ progress: IndexingProgress; error?: string }> => { + const expectedIds = [...result.documentIds, ...(result.taskIds ?? [])] + let progress: IndexingProgress = { + completedIds: [], + failedIds: [], + total: expectedIds.length, + } + const indexingStartedMs = Date.now() + let failure: string | undefined + try { + await provider.awaitIndexing( + result, + build.containerTag, + (current) => { + progress = { + completedIds: [...current.completedIds], + failedIds: [...current.failedIds], + total: current.total, + } + }, + { timeoutMs: readinessTimeoutMs } + ) + try { + validateCompletedIndexingProgress(expectedIds, progress) + } catch (error) { + failure = errorMessage(error) + } + } catch (error) { + failure = errorMessage(error) + } finally { + indexingDurationMs += Date.now() - indexingStartedMs + } + return { progress, ...(failure ? { error: failure } : {}) } } - try { - const completedSessions = - checkpoint.questions[question.questionId].phases.ingest.completedSessions - const combinedResult: IngestResult = { documentIds: [], taskIds: [] } + checkpointManager.updateBuild(checkpoint, build.buildId, (current) => { + current.ingest.status = "in_progress" + current.ingest.startedAt ??= startedAt + current.ingest.error = undefined + current.ingest.attempts.push(attempt) + if (indexingAttempt) { + current.indexing.status = "in_progress" + current.indexing.startedAt ??= startedAt + current.indexing.error = undefined + current.indexing.failedIds = [] + current.indexing.attempts.push(indexingAttempt) + } + }) - for (const session of sessions) { - if (completedSessions.includes(session.sessionId)) { - continue + try { + const firstIncompleteIndex = build.ingest.completedSessionIds.length + const ingestBatchSize = build.ingestBatchSize ?? 1 + for ( + let documentIndex = firstIncompleteIndex; + documentIndex < plan.documents.length; + documentIndex += ingestBatchSize + ) { + const documents = plan.documents.slice(documentIndex, documentIndex + ingestBatchSize) + const ingestStartedMs = Date.now() + let attributedResults: ReturnType + try { + const result = await provider.ingest(documents, { + containerTag: build.containerTag, + ...(build.ingestionExecutionPolicy.processingMode === "instant" + ? { processingMode: "instant" as const } + : {}), + }) + attributedResults = attributeBatchResult(documents, result) + } catch (error) { + const message = errorMessage(error) + attributedResults = documents.map((document) => ({ + customId: document.customId, + documentIds: [], + taskIds: [], + error: message, + })) + } finally { + ingestDurationMs += Date.now() - ingestStartedMs } - const result = await provider.ingest([session], { containerTag }) - - combinedResult.documentIds.push(...result.documentIds) - if (result.taskIds) { - combinedResult.taskIds!.push(...result.taskIds) + let readiness: { progress: IndexingProgress; error?: string } | undefined + if (requiresSessionBarrier) { + const submitted = attributedResults.filter((item) => !item.error) + const submittedResult: IngestResult = { + documentIds: submitted.flatMap((item) => item.documentIds), + taskIds: submitted.flatMap((item) => item.taskIds), + } + if (submittedResult.documentIds.length + (submittedResult.taskIds?.length ?? 0) > 0) { + readiness = await awaitReadiness(submittedResult) + } } - completedSessions.push(session.sessionId) - for (const questionId of groupQuestionIds) { - checkpointManager.updatePhase(checkpoint, questionId, "ingest", { - completedSessions: [...completedSessions], + for (const [offset, document] of documents.entries()) { + const attributed = attributedResults[offset]! + const physicalIds = [...attributed.documentIds, ...attributed.taskIds] + let deferredFailure: + | { customId: string; stage: "submission" | "readiness"; error: string } + | undefined + if (attributed.error) { + deferredFailure = { + customId: document.customId, + stage: "submission", + error: attributed.error, + } + } else if (requiresSessionBarrier) { + const completedIds = new Set(readiness?.progress.completedIds ?? []) + const failedIds = physicalIds.filter((id) => + (readiness?.progress.failedIds ?? []).includes(id) + ) + const pendingIds = physicalIds.filter( + (id) => !completedIds.has(id) && !failedIds.includes(id) + ) + if (physicalIds.length === 0 || failedIds.length > 0 || pendingIds.length > 0) { + deferredFailure = { + customId: document.customId, + stage: "readiness", + error: + failedIds.length > 0 + ? `Provider reported failed IDs: ${failedIds.join(", ")}` + : (readiness?.error ?? + (physicalIds.length === 0 + ? "Provider returned no physical document/task ID" + : `Readiness is still pending for: ${pendingIds.join(", ")}`)), + } + } + } + checkpointManager.recordIngestProgress(checkpoint, build.buildId, { + sequence: documentIndex + offset, + sessionId: document.metadata.sessionId, + documentIds: attributed.documentIds, + taskIds: attributed.taskIds, + readyForNextSession: requiresSessionBarrier && !deferredFailure, + ...(deferredFailure ? { deferredFailure } : {}), }) + if (deferredFailure) { + logger.warn( + `Deferred ${document.customId} in ${build.containerTag} (${deferredFailure.stage}): ${deferredFailure.error}` + ) + } } } - if (combinedResult.taskIds && combinedResult.taskIds.length === 0) { - delete combinedResult.taskIds + const deferredFirstPass = [...(build.ingest.deferredSessions ?? [])].sort( + (left, right) => left.sequence - right.sequence + ) + if (deferredFirstPass.length > 0) { + logger.warn( + `Retrying ${deferredFirstPass.length} deferred sessions in ${build.containerTag}` + ) } + for (const deferredSnapshot of deferredFirstPass) { + const document = plan.documents[deferredSnapshot.sequence] + if ( + !document || + document.metadata.sessionId !== deferredSnapshot.sessionId || + document.customId !== deferredSnapshot.customId + ) { + throw new Error( + `Deferred session ${deferredSnapshot.sessionId} no longer matches its validated plan` + ) + } - const existingResult = checkpoint.questions[question.questionId].phases.ingest.ingestResult - if (existingResult) { - combinedResult.documentIds = [ - ...existingResult.documentIds, - ...combinedResult.documentIds, - ] - if (existingResult.taskIds || combinedResult.taskIds) { - combinedResult.taskIds = [ - ...(existingResult.taskIds || []), - ...(combinedResult.taskIds || []), - ] + let attributed: ReturnType[number] | undefined + let retryStage: "submission" | "readiness" = "submission" + let retryError: string | undefined + let changedPhysicalIds = false + const retryIngestStartedMs = Date.now() + try { + const result = await provider.ingest([document], { + containerTag: build.containerTag, + ...(build.ingestionExecutionPolicy.processingMode === "instant" + ? { processingMode: "instant" as const } + : {}), + }) + attributed = attributeBatchResult([document], result)[0]! + retryError = attributed.error + } catch (error) { + retryError = errorMessage(error) + } finally { + ingestDurationMs += Date.now() - retryIngestStartedMs } - } - const durationMs = Date.now() - startTime - for (const questionId of groupQuestionIds) { - checkpointManager.updatePhase(checkpoint, questionId, "ingest", { - status: "completed", - ingestResult: { - documentIds: [...combinedResult.documentIds], - ...(combinedResult.taskIds ? { taskIds: [...combinedResult.taskIds] } : {}), - }, - completedAt: new Date().toISOString(), - durationMs, + if (attributed && !retryError) { + if ( + deferredSnapshot.documentIds.length + deferredSnapshot.taskIds.length > 0 && + (!sameIds(deferredSnapshot.documentIds, attributed.documentIds) || + !sameIds(deferredSnapshot.taskIds, attributed.taskIds)) + ) { + changedPhysicalIds = true + retryError = `Retry changed physical IDs for ${deferredSnapshot.customId}` + } else if (attributed.documentIds.length + attributed.taskIds.length === 0) { + retryError = "Provider returned no physical document/task ID" + } + } + + if (attributed && !retryError && requiresSessionBarrier) { + retryStage = "readiness" + const retryResult: IngestResult = { + documentIds: [...attributed.documentIds], + taskIds: [...attributed.taskIds], + } + const retryReadiness = await awaitReadiness(retryResult) + const expectedIds = [...attributed.documentIds, ...attributed.taskIds] + const completedIds = new Set(retryReadiness.progress.completedIds) + const failedIds = expectedIds.filter((id) => + retryReadiness.progress.failedIds.includes(id) + ) + const pendingIds = expectedIds.filter( + (id) => !completedIds.has(id) && !failedIds.includes(id) + ) + if (failedIds.length > 0 || pendingIds.length > 0) { + retryError = + failedIds.length > 0 + ? `Provider reported failed IDs: ${failedIds.join(", ")}` + : (retryReadiness.error ?? + `Readiness is still pending for: ${pendingIds.join(", ")}`) + } + } + + if (retryError || !attributed) { + const failedAt = new Date().toISOString() + checkpointManager.updateBuild(checkpoint, build.buildId, (current) => { + const deferred = (current.ingest.deferredSessions ?? []).find( + (candidate) => candidate.sequence === deferredSnapshot.sequence + ) + if (!deferred) { + throw new Error(`Missing deferred session ${deferredSnapshot.sessionId}`) + } + if (attributed && !changedPhysicalIds) { + deferred.documentIds = [...attributed.documentIds] + deferred.taskIds = [...attributed.taskIds] + current.ingest.documentIds = [ + ...new Set([...current.ingest.documentIds, ...attributed.documentIds]), + ] + current.ingest.taskIds = [ + ...new Set([...current.ingest.taskIds, ...attributed.taskIds]), + ] + } + deferred.stage = retryStage + deferred.attempts += 1 + deferred.lastFailedAt = failedAt + deferred.lastError = retryError ?? "Unknown retry failure" + }) + await checkpointManager.flush(checkpoint.runId) + logger.warn( + `Retry failed for ${deferredSnapshot.customId} in ${build.containerTag}: ${retryError ?? "unknown error"}` + ) + continue + } + + checkpointManager.updateBuild(checkpoint, build.buildId, (current) => { + current.ingest.documentIds = [ + ...new Set([...current.ingest.documentIds, ...attributed!.documentIds]), + ] + current.ingest.taskIds = [ + ...new Set([...current.ingest.taskIds, ...attributed!.taskIds]), + ] + current.ingest.deferredSessions = (current.ingest.deferredSessions ?? []).filter( + (candidate) => candidate.sequence !== deferredSnapshot.sequence + ) + if (requiresSessionBarrier) { + current.indexing.completedIds = [ + ...new Set([ + ...current.indexing.completedIds, + ...attributed!.documentIds, + ...attributed!.taskIds, + ]), + ] + current.indexing.failedIds = current.indexing.failedIds.filter( + (id) => !attributed!.documentIds.includes(id) && !attributed!.taskIds.includes(id) + ) + } }) + await checkpointManager.flush(checkpoint.runId) + logger.info(`Recovered ${deferredSnapshot.customId} in ${build.containerTag}`) } - logger.progress(index + 1, total, `Ingested ${group.groupId} (${durationMs}ms)`) + const remainingDeferred = build.ingest.deferredSessions ?? [] + if (remainingDeferred.length > 0) { + const completedAt = new Date().toISOString() + const message = `${remainingDeferred.length} sessions remain deferred after end-of-build retry: ${remainingDeferred + .slice(0, 5) + .map((deferred) => deferred.customId) + .join(", ")}` + checkpointManager.updateBuild(checkpoint, build.buildId, (current) => { + const currentAttempt = current.ingest.attempts.at(-1)! + Object.assign(currentAttempt, { + status: "failed", + completedAt, + durationMs: ingestDurationMs, + error: message, + }) + current.ingest.status = "failed" + current.ingest.error = message + current.ingest.durationMs = totalAttemptDuration(current.ingest.attempts) + if (requiresSessionBarrier) { + const currentIndexingAttempt = current.indexing.attempts.at(-1)! + Object.assign(currentIndexingAttempt, { + status: "failed", + completedAt, + durationMs: indexingDurationMs, + error: message, + }) + current.indexing.status = "failed" + current.indexing.error = message + current.indexing.durationMs = totalAttemptDuration(current.indexing.attempts) + current.indexing.failedIds = unresolvedPhysicalIds(current) + } + }) + await checkpointManager.flush(checkpoint.runId) + logger.warn(`Build ${build.ingestionGroupId} finished its first pass with ${message}`) + return { + buildId: build.buildId, + durationMs: ingestDurationMs + indexingDurationMs, + unresolved: remainingDeferred.length, + } + } - return { questionId: question.questionId, durationMs } - } catch (e) { - const error = e instanceof Error ? e.message : String(e) - for (const questionId of groupQuestionIds) { - checkpointManager.updatePhase(checkpoint, questionId, "ingest", { + const completedAt = new Date().toISOString() + checkpointManager.updateBuild(checkpoint, build.buildId, (current) => { + const currentAttempt = current.ingest.attempts.at(-1)! + Object.assign(currentAttempt, { + status: "completed", + completedAt, + durationMs: ingestDurationMs, + }) + current.ingest.status = "completed" + current.ingest.completedAt = completedAt + current.ingest.durationMs = totalAttemptDuration(current.ingest.attempts) + current.ingest.error = undefined + if (requiresSessionBarrier) { + const currentIndexingAttempt = current.indexing.attempts.at(-1)! + Object.assign(currentIndexingAttempt, { + status: "completed", + completedAt, + durationMs: indexingDurationMs, + }) + current.indexing.status = "completed" + current.indexing.completedAt = completedAt + current.indexing.durationMs = totalAttemptDuration(current.indexing.attempts) + current.indexing.completedIds = [ + ...new Set([...current.ingest.documentIds, ...current.ingest.taskIds]), + ] + current.indexing.failedIds = [] + current.indexing.error = undefined + } + }) + // Compact the fsynced per-session journal only after the full completed + // build checkpoint has itself become durable. + await checkpointManager.flush(checkpoint.runId) + checkpointManager.clearIngestProgressJournal(checkpoint.runId, build.buildId) + const totalDurationMs = ingestDurationMs + indexingDurationMs + logger.progress( + index + 1, + total, + requiresSessionBarrier + ? `Ingested and indexed ${build.ingestionGroupId} (${totalDurationMs}ms)` + : `Ingested ${build.ingestionGroupId} (${ingestDurationMs}ms)` + ) + return { buildId: build.buildId, durationMs: totalDurationMs } + } catch (error) { + const message = errorMessage(error) + const completedAt = new Date().toISOString() + checkpointManager.updateBuild(checkpoint, build.buildId, (current) => { + const currentAttempt = current.ingest.attempts.at(-1)! + Object.assign(currentAttempt, { status: "failed", - error, + completedAt, + durationMs: ingestDurationMs, + error: message, }) - } - logger.error(`Failed to ingest ${group.groupId}: ${error}`) + current.ingest.status = "failed" + current.ingest.error = message + current.ingest.durationMs = totalAttemptDuration(current.ingest.attempts) + if (requiresSessionBarrier) { + const currentIndexingAttempt = current.indexing.attempts.at(-1)! + Object.assign(currentIndexingAttempt, { + status: "failed", + completedAt, + durationMs: indexingDurationMs, + error: message, + }) + current.indexing.status = "failed" + current.indexing.error = message + current.indexing.durationMs = totalAttemptDuration(current.indexing.attempts) + current.indexing.failedIds = unresolvedPhysicalIds(current) + } + }) throw new Error( - `Ingest failed at ${group.groupId}: ${error}. Fix the issue and resume with the same run ID.` + `Ingest failed at ${build.ingestionGroupId}: ${message}. Fix the issue and resume with the same run ID.` ) } }, }) + const unresolvedBuilds = Object.values(checkpoint.builds).filter( + (build) => (build.ingest.deferredSessions ?? []).length > 0 + ) + if (unresolvedBuilds.length > 0) { + const unresolvedSessions = unresolvedBuilds.reduce( + (sum, build) => sum + (build.ingest.deferredSessions?.length ?? 0), + 0 + ) + throw new Error( + `Ingest first pass completed, but ${unresolvedSessions} sessions across ${unresolvedBuilds.length} builds still need retry. Resume with the same run ID.` + ) + } + logger.success("Ingest phase complete") } diff --git a/src/orchestrator/phases/report.ts b/src/orchestrator/phases/report.ts index 43b2102..7812fc2 100644 --- a/src/orchestrator/phases/report.ts +++ b/src/orchestrator/phases/report.ts @@ -1,370 +1,513 @@ -import { writeFileSync, mkdirSync, existsSync } from "fs" -import { join } from "path" +import { existsSync, mkdirSync, writeFileSync } from "node:fs" +import { join } from "node:path" import type { Benchmark } from "../../types/benchmark" import type { RunCheckpoint } from "../../types/checkpoint" +import type { QuestionEvaluation } from "../../types/protocol" import type { BenchmarkResult, + BuildMetrics, + CostCoverageMetrics, EvaluationResult, LatencyStats, + QuestionMetrics, QuestionTypeStats, - RetrievalMetrics, RetrievalAggregates, + RetrievalMetrics, TokenMetrics, + UsageMetrics, } from "../../types/unified" import { logger } from "../../utils/logger" +import { stableSha256 } from "../../utils/stable" +import { canonicalizeSelectedQuestionIds } from "../input-identity" const REPORTS_DIR = "./data/runs" function aggregateRetrievalMetrics(metrics: RetrievalMetrics[]): RetrievalAggregates | undefined { if (metrics.length === 0) return undefined - const sum = metrics.reduce( - (acc, m) => ({ - hitAtK: acc.hitAtK + m.hitAtK, - precisionAtK: acc.precisionAtK + m.precisionAtK, - recallAtK: acc.recallAtK + m.recallAtK, - f1AtK: acc.f1AtK + m.f1AtK, - mrr: acc.mrr + m.mrr, - ndcg: acc.ndcg + m.ndcg, - k: m.k, + (accumulator, metric) => ({ + hitAtK: accumulator.hitAtK + metric.hitAtK, + precisionAtK: accumulator.precisionAtK + metric.precisionAtK, + recallAtK: accumulator.recallAtK + metric.recallAtK, + f1AtK: accumulator.f1AtK + metric.f1AtK, + mrr: accumulator.mrr + metric.mrr, + ndcg: accumulator.ndcg + metric.ndcg, + k: accumulator.k + metric.k, }), - { hitAtK: 0, precisionAtK: 0, recallAtK: 0, f1AtK: 0, mrr: 0, ndcg: 0, k: 10 } + { hitAtK: 0, precisionAtK: 0, recallAtK: 0, f1AtK: 0, mrr: 0, ndcg: 0, k: 0 } ) - - const n = metrics.length + const count = metrics.length return { - hitAtK: sum.hitAtK / n, - precisionAtK: sum.precisionAtK / n, - recallAtK: sum.recallAtK / n, - f1AtK: sum.f1AtK / n, - mrr: sum.mrr / n, - ndcg: sum.ndcg / n, - k: sum.k, + hitAtK: sum.hitAtK / count, + precisionAtK: sum.precisionAtK / count, + recallAtK: sum.recallAtK / count, + f1AtK: sum.f1AtK / count, + mrr: sum.mrr / count, + ndcg: sum.ndcg / count, + k: sum.k / count, } } -function calculateLatencyStats(durations: number[]): LatencyStats { +export function calculateLatencyStats(durations: number[]): LatencyStats { if (durations.length === 0) { return { min: 0, max: 0, mean: 0, median: 0, p95: 0, p99: 0, stdDev: 0, count: 0 } } + const sorted = [...durations].sort((left, right) => left - right) + const count = sorted.length + const mean = sorted.reduce((sum, value) => sum + value, 0) / count + const variance = sorted.reduce((sum, value) => sum + (value - mean) ** 2, 0) / count + const percentile = (fraction: number) => + sorted[Math.min(count - 1, Math.max(0, Math.ceil(count * fraction) - 1))] + return { + min: sorted[0], + max: sorted[count - 1], + mean, + median: percentile(0.5), + p95: percentile(0.95), + p99: percentile(0.99), + stdDev: Math.sqrt(variance), + count, + } +} - const sorted = [...durations].sort((a, b) => a - b) - const n = sorted.length - const sum = sorted.reduce((a, b) => a + b, 0) - const mean = sum / n +function nullableSum(values: Array): number | null { + if (values.length === 0) return null + if (values.some((value) => value == null)) return null + return values.reduce((sum, value) => sum + (value as number), 0) +} - const variance = sorted.reduce((acc, val) => acc + Math.pow(val - mean, 2), 0) / n - const stdDev = Math.sqrt(variance) +function aggregateCostCoverage(values: Array): CostCoverageMetrics { + const knownCosts = values.filter((value): value is number => value != null) + return { + totalCostUsd: + values.length > 0 && knownCosts.length === values.length + ? knownCosts.reduce((sum, value) => sum + value, 0) + : null, + knownCostCount: knownCosts.length, + totalCostCount: values.length, + } +} +function aggregateUsage(values: Array): UsageMetrics | undefined { + const present = values.filter((value): value is UsageMetrics => value !== undefined) + if (present.length === 0) return undefined + const sumField = (field: keyof UsageMetrics): number | undefined => { + const numbers = present + .map((value) => value[field]) + .filter((value): value is number => value !== undefined) + return numbers.length > 0 ? numbers.reduce((sum, value) => sum + value, 0) : undefined + } return { - min: sorted[0], - max: sorted[n - 1], - mean: Math.round(mean), - median: sorted[Math.floor(n / 2)], - p95: sorted[Math.floor(n * 0.95)] || sorted[n - 1], - p99: sorted[Math.floor(n * 0.99)] || sorted[n - 1], - stdDev: Math.round(stdDev), - count: n, + ...(sumField("requestCount") !== undefined ? { requestCount: sumField("requestCount") } : {}), + ...(sumField("tokenUsageCompleteRequestCount") !== undefined + ? { tokenUsageCompleteRequestCount: sumField("tokenUsageCompleteRequestCount") } + : {}), + ...(sumField("tokenUsagePartialRequestCount") !== undefined + ? { tokenUsagePartialRequestCount: sumField("tokenUsagePartialRequestCount") } + : {}), + ...(sumField("tokenUsageUnknownRequestCount") !== undefined + ? { tokenUsageUnknownRequestCount: sumField("tokenUsageUnknownRequestCount") } + : {}), + ...(sumField("inputTokens") !== undefined ? { inputTokens: sumField("inputTokens") } : {}), + ...(sumField("outputTokens") !== undefined ? { outputTokens: sumField("outputTokens") } : {}), + ...(sumField("reasoningTokens") !== undefined + ? { reasoningTokens: sumField("reasoningTokens") } + : {}), + ...(sumField("totalTokens") !== undefined ? { totalTokens: sumField("totalTokens") } : {}), } } +function elapsedMs(startedAt?: string, completedAt?: string): number { + if (!startedAt || !completedAt) return 0 + return Math.max(0, Date.parse(completedAt) - Date.parse(startedAt)) +} + export function generateReport(benchmark: Benchmark, checkpoint: RunCheckpoint): BenchmarkResult { - const questions = benchmark.getQuestions() + if (!checkpoint.benchmarkInputFingerprint?.trim()) { + throw new Error("Cannot generate report without a benchmark input fingerprint") + } + if (!checkpoint.answeringRuntimeIdentity) { + throw new Error("Cannot generate report without an answering runtime identity") + } + if (!Number.isInteger(checkpoint.retrievalTopK) || checkpoint.retrievalTopK <= 0) { + throw new Error("Cannot generate report without an effective retrieval Top-K") + } + const allQuestions = benchmark.getQuestions() + const questionById = new Map(allQuestions.map((question) => [question.questionId, question])) + const requestedQuestionIds = + checkpoint.targetQuestionIds && checkpoint.targetQuestionIds.length > 0 + ? [...checkpoint.targetQuestionIds] + : Object.keys(checkpoint.questions) + const selectedQuestionIds = canonicalizeSelectedQuestionIds(allQuestions, requestedQuestionIds) + if (stableSha256(selectedQuestionIds) !== checkpoint.selectedQuestionIdsDigest) { + throw new Error("Cannot report a checkpoint with non-canonical selected-question identity") + } + const checkpointOnlyIds = Object.keys(checkpoint.questions).filter( + (questionId) => !selectedQuestionIds.includes(questionId) + ) + if (checkpointOnlyIds.length > 0) { + throw new Error( + `Cannot report unselected checkpoint questions: ${checkpointOnlyIds.join(", ")}` + ) + } + const questions = selectedQuestionIds.map((questionId) => questionById.get(questionId)!) + const incompleteQuestionIds = questions + .filter( + (question) => + checkpoint.questions[question.questionId]?.phases.evaluate.status !== "completed" || + !checkpoint.questions[question.questionId]?.phases.evaluate.evaluation + ) + .map((question) => question.questionId) + if (incompleteQuestionIds.length > 0) { + throw new Error( + `Cannot generate a scored report with incomplete evaluations (${incompleteQuestionIds.length}), starting with ${incompleteQuestionIds[0]}` + ) + } const evaluations: EvaluationResult[] = [] - - const ingestDurations: number[] = [] - const indexingDurations: number[] = [] - const searchDurations: number[] = [] - const answerDurations: number[] = [] - const evaluateDurations: number[] = [] - const totalDurations: number[] = [] - + const protocolEvaluations: QuestionEvaluation[] = [] + const evaluatedQuestions = [] as typeof questions + const questionMetrics: QuestionMetrics[] = [] const allRetrievalMetrics: RetrievalMetrics[] = [] - - const byType: Record< + const byType = new Map< string, { total: number - correct: number - searchDurations: number[] - answerDurations: number[] - totalDurations: number[] - retrievalMetrics: RetrievalMetrics[] + passed: number + search: number[] + answer: number[] + online: number[] + retrieval: RetrievalMetrics[] } - > = {} + >() + const completedQuestionCountByBuild = new Map() for (const question of questions) { - const qCheckpoint = checkpoint.questions[question.questionId] - if (!qCheckpoint) continue - - const evalPhase = qCheckpoint.phases.evaluate - if (evalPhase.status !== "completed") continue - - const ingestPhase = qCheckpoint.phases.ingest - const indexingPhase = qCheckpoint.phases.indexing - const searchPhase = qCheckpoint.phases.search - const answerPhase = qCheckpoint.phases.answer - - const ingestDurationMs = ingestPhase.durationMs || 0 - const indexingDurationMs = indexingPhase.durationMs || 0 - const searchDurationMs = searchPhase.durationMs || 0 - const answerDurationMs = answerPhase.durationMs || 0 - const evaluateDurationMs = evalPhase.durationMs || 0 - const totalDurationMs = - ingestDurationMs + - indexingDurationMs + - searchDurationMs + - answerDurationMs + - evaluateDurationMs + const questionCheckpoint = checkpoint.questions[question.questionId] + if (questionCheckpoint?.phases.evaluate.status === "completed") { + completedQuestionCountByBuild.set( + questionCheckpoint.buildId, + (completedQuestionCountByBuild.get(questionCheckpoint.buildId) || 0) + 1 + ) + } + } - const retrievalMetrics = evalPhase.retrievalMetrics + const buildMetrics: BuildMetrics[] = Object.values(checkpoint.builds).map((build) => { + const allAttempts = [...build.ingest.attempts, ...build.indexing.attempts] + const reusedPhases = build.reusedPhases ?? { + ingest: build.reused, + indexing: build.reused, + } + const currentAttempts = allAttempts.filter((attempt) => !reusedPhases[attempt.phase]) + const currentWork = currentAttempts.reduce((sum, attempt) => sum + (attempt.durationMs ?? 0), 0) + const costUsd = nullableSum(currentAttempts.map((attempt) => attempt.costUsd)) + const currentStart = currentAttempts.map((attempt) => attempt.startedAt).sort()[0] + const currentEnd = currentAttempts + .flatMap((attempt) => (attempt.completedAt ? [attempt.completedAt] : [])) + .sort() + .at(-1) + return { + buildId: build.buildId, + containerTag: build.containerTag, + providerIngestionConfigFingerprint: build.providerIngestionConfigFingerprint, + sourceRunId: build.sourceRunId, + reused: build.reused, + reusedPhases, + ingestLatencyMs: reusedPhases.ingest ? 0 : (build.ingest.durationMs ?? 0), + indexingLatencyMs: reusedPhases.indexing ? 0 : (build.indexing.durationMs ?? 0), + buildWallClockMs: elapsedMs(currentStart, currentEnd), + buildWorkMs: currentWork, + attemptCount: currentAttempts.length, + attempts: currentAttempts.map((attempt) => ({ ...attempt })), + ...(aggregateUsage(currentAttempts.map((attempt) => attempt.usage)) + ? { usage: aggregateUsage(currentAttempts.map((attempt) => attempt.usage)) } + : {}), + costUsd, + sessionCount: new Set(build.haystack.orderedSessionIds).size, + documentCount: new Set(build.ingest.documentIds).size, + taskCount: new Set(build.ingest.taskIds).size, + completedIndexingCount: new Set(build.indexing.completedIds).size, + failedIndexingCount: new Set(build.indexing.failedIds).size, + } + }) + const buildMetricsById = new Map(buildMetrics.map((metrics) => [metrics.buildId, metrics])) + for (const question of questions) { + const questionCheckpoint = checkpoint.questions[question.questionId] + if (!questionCheckpoint || questionCheckpoint.phases.evaluate.status !== "completed") continue + const evaluation = questionCheckpoint.phases.evaluate.evaluation + if (!evaluation) { + throw new Error(`Completed question ${question.questionId} is missing protocol evaluation`) + } + if (typeof evaluation.passed !== "boolean") { + throw new Error(`Completed question ${question.questionId} is missing protocol pass state`) + } + const search = questionCheckpoint.phases.search + const answer = questionCheckpoint.phases.answer + const evaluate = questionCheckpoint.phases.evaluate + const searchLatencyMs = search.durationMs ?? 0 + const answerLatencyMs = answer.durationMs ?? 0 + const evaluationLatencyMs = evaluate.durationMs ?? 0 + const onlineQueryLatencyMs = searchLatencyMs + answerLatencyMs + const buildMetric = buildMetricsById.get(questionCheckpoint.buildId) + const denominator = completedQuestionCountByBuild.get(questionCheckpoint.buildId) || 0 + const allocatedBuildWorkMs = + buildMetric && denominator > 0 ? buildMetric.buildWorkMs / denominator : undefined + const passed = evaluation.passed + + protocolEvaluations.push(evaluation) + evaluatedQuestions.push(question) + if (evaluate.retrievalMetrics) allRetrievalMetrics.push(evaluate.retrievalMetrics) evaluations.push({ questionId: question.questionId, questionType: question.questionType, question: question.question, - score: evalPhase.score || 0, - label: evalPhase.label || "incorrect", - explanation: evalPhase.explanation || "", - hypothesis: answerPhase.hypothesis || "", + score: evaluation.primaryScore, + primaryScore: evaluation.primaryScore, + passed, + label: passed ? "correct" : "incorrect", + explanation: evaluation.explanation, + metrics: evaluation.metrics, + hypothesis: answer.hypothesis || "", groundTruth: question.groundTruth, - searchResults: searchPhase.results || [], - searchDurationMs, - answerDurationMs, - totalDurationMs, - retrievalMetrics, - details: evalPhase.details, + searchResults: search.results || [], + searchDurationMs: searchLatencyMs, + answerDurationMs: answerLatencyMs, + totalDurationMs: onlineQueryLatencyMs, + retrievalMetrics: evaluate.retrievalMetrics, + details: evaluation.details, + }) + questionMetrics.push({ + questionId: question.questionId, + buildId: questionCheckpoint.buildId, + searchLatencyMs, + answerLatencyMs, + onlineQueryLatencyMs, + evaluationLatencyMs, + ...(aggregateUsage([search.usage, answer.usage]) + ? { queryUsage: aggregateUsage([search.usage, answer.usage]) } + : {}), + ...(evaluate.usage ? { evaluationUsage: { ...evaluate.usage } } : {}), + queryCostUsd: nullableSum([search.costUsd, answer.costUsd]), + evaluationCostUsd: evaluate.costUsd ?? null, + configuredTopK: search.retrievalPlan?.requestedTopK ?? search.requestedCount ?? 0, + providerRequestLimit: + search.providerRequests?.reduce((sum, request) => sum + request.limit, 0) ?? + search.requestedCount ?? + 0, + rawReturnedCount: + search.rawReturnedCount ?? search.returnedCount ?? search.results?.length ?? 0, + returnedCount: search.normalizedCount ?? search.returnedCount ?? search.results?.length ?? 0, + normalizedCount: search.normalizedCount ?? search.results?.length ?? 0, + droppedCount: search.droppedCount ?? 0, + droppedResults: search.droppedResults ?? [], + answerCutoff: search.retrievalPlan?.answerCutoff ?? search.answerCutoff ?? 0, + answerEvidenceCount: search.answerEvidenceCount ?? answer.evidenceCount ?? 0, + contextTokens: answer.contextTokens ?? 0, + ...(search.retrievalPlan?.searchMode ? { searchMode: search.retrievalPlan.searchMode } : {}), + ...(search.retrievalPlan?.threshold !== undefined + ? { threshold: search.retrievalPlan.threshold } + : {}), + providerRequests: search.providerRequests ?? [], + buildAllocationQuestionCount: denominator, + ...(allocatedBuildWorkMs != null + ? { + allocatedBuildWorkMs, + amortizedOnlinePlusBuildWorkMs: onlineQueryLatencyMs + allocatedBuildWorkMs, + } + : {}), }) - if (retrievalMetrics) { - allRetrievalMetrics.push(retrievalMetrics) - } - - if (ingestPhase.durationMs) ingestDurations.push(ingestPhase.durationMs) - if (indexingPhase.durationMs) indexingDurations.push(indexingPhase.durationMs) - if (searchPhase.durationMs) searchDurations.push(searchPhase.durationMs) - if (answerPhase.durationMs) answerDurations.push(answerPhase.durationMs) - if (evalPhase.durationMs) evaluateDurations.push(evalPhase.durationMs) - if (totalDurationMs > 0) totalDurations.push(totalDurationMs) - - const qType = question.questionType - if (!byType[qType]) { - byType[qType] = { - total: 0, - correct: 0, - searchDurations: [], - answerDurations: [], - totalDurations: [], - retrievalMetrics: [], - } - } - const typeStats = byType[qType]! - typeStats.total++ - if (evalPhase.score === 1) { - typeStats.correct++ + const type = byType.get(question.questionType) || { + total: 0, + passed: 0, + search: [], + answer: [], + online: [], + retrieval: [], } - if (searchDurationMs) typeStats.searchDurations.push(searchDurationMs) - if (answerDurationMs) typeStats.answerDurations.push(answerDurationMs) - if (totalDurationMs > 0) typeStats.totalDurations.push(totalDurationMs) - if (retrievalMetrics) typeStats.retrievalMetrics.push(retrievalMetrics) + type.total++ + if (passed) type.passed++ + type.search.push(searchLatencyMs) + type.answer.push(answerLatencyMs) + type.online.push(onlineQueryLatencyMs) + if (evaluate.retrievalMetrics) type.retrieval.push(evaluate.retrievalMetrics) + byType.set(question.questionType, type) } + const quality = benchmark.protocol.aggregateQuality({ + questions: evaluatedQuestions, + evaluations: protocolEvaluations, + }) const byQuestionType: Record = {} - for (const type of Object.keys(byType)) { - const raw = byType[type]! - byQuestionType[type] = { - total: raw.total, - correct: raw.correct, - accuracy: raw.total > 0 ? raw.correct / raw.total : 0, + for (const [questionType, values] of byType) { + byQuestionType[questionType] = { + total: values.total, + correct: values.passed, + accuracy: values.total > 0 ? values.passed / values.total : 0, latency: { - search: calculateLatencyStats(raw.searchDurations), - answer: calculateLatencyStats(raw.answerDurations), - total: calculateLatencyStats(raw.totalDurations), + search: calculateLatencyStats(values.search), + answer: calculateLatencyStats(values.answer), + total: calculateLatencyStats(values.online), }, - retrieval: aggregateRetrievalMetrics(raw.retrievalMetrics), + retrieval: aggregateRetrievalMetrics(values.retrieval), } } - const overallRetrieval = aggregateRetrievalMetrics(allRetrievalMetrics) - - // Aggregate token metrics — only from evaluated questions (same population as quality/latency) + const promptTokens = evaluations.map( + (evaluation) => checkpoint.questions[evaluation.questionId].phases.answer.promptTokens + ) let tokenMetrics: TokenMetrics | undefined - const allPromptTokens: number[] = [] - const allBasePromptTokens: number[] = [] - const allContextTokens: number[] = [] - - for (const question of questions) { - const qCheckpoint = checkpoint.questions[question.questionId] - if (!qCheckpoint) continue - // Only consider questions that were evaluated (same filter as the quality/latency loop above) - if (qCheckpoint.phases.evaluate.status !== "completed") continue - const answerPhase = qCheckpoint.phases.answer - if (answerPhase.promptTokens != null) allPromptTokens.push(answerPhase.promptTokens) - if (answerPhase.basePromptTokens != null) - allBasePromptTokens.push(answerPhase.basePromptTokens) - if (answerPhase.contextTokens != null) allContextTokens.push(answerPhase.contextTokens) - } - - if (allPromptTokens.length > 0) { - const totalTokens = allPromptTokens.reduce((a, b) => a + b, 0) - const totalBasePromptTokens = allBasePromptTokens.reduce((a, b) => a + b, 0) - const totalContextTokens = allContextTokens.reduce((a, b) => a + b, 0) - - // Use the number of questions with token data as the denominator for averages. - // This is accurate because we already filtered to evaluated questions above. + if (promptTokens.length > 0 && promptTokens.every((value) => value != null)) { + const base = evaluations.map( + (evaluation) => + checkpoint.questions[evaluation.questionId].phases.answer.basePromptTokens ?? 0 + ) + const context = evaluations.map( + (evaluation) => checkpoint.questions[evaluation.questionId].phases.answer.contextTokens ?? 0 + ) + const totalTokens = (promptTokens as number[]).reduce((sum, value) => sum + value, 0) + const basePromptTokens = base.reduce((sum, value) => sum + value, 0) + const contextTokens = context.reduce((sum, value) => sum + value, 0) tokenMetrics = { totalTokens, - basePromptTokens: totalBasePromptTokens, - contextTokens: totalContextTokens, - avgTokensPerQuestion: Math.round(totalTokens / allPromptTokens.length), - avgBasePromptTokens: - allBasePromptTokens.length > 0 - ? Math.round(totalBasePromptTokens / allBasePromptTokens.length) - : 0, - avgContextTokens: - allContextTokens.length > 0 ? Math.round(totalContextTokens / allContextTokens.length) : 0, + basePromptTokens, + contextTokens, + avgTokensPerQuestion: totalTokens / promptTokens.length, + avgBasePromptTokens: basePromptTokens / promptTokens.length, + avgContextTokens: contextTokens / promptTokens.length, } } const totalQuestions = evaluations.length - const correctCount = evaluations.filter((e) => e.score === 1).length + const correctCount = protocolEvaluations.filter((evaluation) => evaluation.passed).length const accuracy = totalQuestions > 0 ? correctCount / totalQuestions : 0 + const averageScore = + totalQuestions > 0 + ? protocolEvaluations.reduce((sum, evaluation) => sum + evaluation.primaryScore, 0) / + totalQuestions + : 0 + const currentBuilds = buildMetrics.filter((build) => !build.reused) + const knownBuildCosts = currentBuilds.filter((build) => build.costUsd != null) + const totalBuildCostUsd = + currentBuilds.length > 0 && knownBuildCosts.length === currentBuilds.length + ? knownBuildCosts.reduce((sum, build) => sum + (build.costUsd as number), 0) + : null + const buildPhaseWallClockMs = checkpoint.buildPhaseAttempts.reduce( + (sum, attempt) => sum + (attempt.durationMs ?? 0), + 0 + ) + const searchDurations = questionMetrics.map((metrics) => metrics.searchLatencyMs) + const answerDurations = questionMetrics.map((metrics) => metrics.answerLatencyMs) + const onlineDurations = questionMetrics.map((metrics) => metrics.onlineQueryLatencyMs) + const evaluateDurations = questionMetrics.map((metrics) => metrics.evaluationLatencyMs) + const queryCosts = aggregateCostCoverage(questionMetrics.map((metrics) => metrics.queryCostUsd)) + const evaluationCosts = aggregateCostCoverage( + questionMetrics.map((metrics) => metrics.evaluationCostUsd) + ) + const ingestDurations = currentBuilds.flatMap((build) => + build.reusedPhases?.ingest ? [] : [build.ingestLatencyMs] + ) + const indexingDurations = currentBuilds.flatMap((build) => + build.reusedPhases?.indexing ? [] : [build.indexingLatencyMs] + ) + const qualityPct = quality.primaryMetric + ? Math.round(quality.primaryMetric.value * 100) + : undefined + const searchLatency = calculateLatencyStats(searchDurations) + const memscore = + tokenMetrics && qualityPct != null + ? `${qualityPct}% / ${Math.round(searchLatency.mean)}ms / ${Math.round(tokenMetrics.avgContextTokens)}tok` + : undefined - const searchLatencyStats = calculateLatencyStats(searchDurations) - const qualityPct = Math.round(accuracy * 100) - const avgLatency = searchLatencyStats.mean - - let memscore: string | undefined - let memscoreComponents: { quality: number; latencyMs: number; contextTokens: number } | undefined - // Only emit MemScore when token data covers all evaluated questions, - // so quality, latency, and tokens are derived from the same population. - if (tokenMetrics && allPromptTokens.length === totalQuestions) { - memscoreComponents = { - quality: qualityPct, - latencyMs: avgLatency, - contextTokens: tokenMetrics.avgContextTokens, - } - memscore = `${qualityPct}% / ${avgLatency}ms / ${tokenMetrics.avgContextTokens}tok` - } - - const result: BenchmarkResult = { + return { provider: checkpoint.provider, + providerPromptFingerprint: checkpoint.providerPromptFingerprint, benchmark: checkpoint.benchmark, runId: checkpoint.runId, dataSourceRunId: checkpoint.dataSourceRunId, judge: checkpoint.judge, answeringModel: checkpoint.answeringModel, + answeringRuntimeIdentity: checkpoint.answeringRuntimeIdentity, timestamp: new Date().toISOString(), - summary: { - totalQuestions, - correctCount, - accuracy, + selectedQuestionIdsDigest: checkpoint.selectedQuestionIdsDigest, + retrievalTopK: checkpoint.retrievalTopK, + benchmarkScope: checkpoint.benchmarkScope, + datasetIdentity: checkpoint.datasetIdentity as unknown as Record | undefined, + benchmarkInputFingerprint: checkpoint.benchmarkInputFingerprint, + protocolIdentity: checkpoint.protocolIdentity as unknown as Record, + quality, + summary: { totalQuestions, correctCount, accuracy, averageScore }, + builds: { + uniqueBuildCount: buildMetrics.length, + sumContainerBuildWorkMs: currentBuilds.reduce((sum, build) => sum + build.buildWorkMs, 0), + buildPhaseWallClockMs, + totalBuildCostUsd, + knownCostBuildCount: knownBuildCosts.length, + totalCostBuildCount: currentBuilds.length, + items: buildMetrics, }, + costs: { + query: queryCosts, + evaluation: evaluationCosts, + }, + questionMetrics, latency: { ingest: calculateLatencyStats(ingestDurations), indexing: calculateLatencyStats(indexingDurations), - search: searchLatencyStats, + search: searchLatency, answer: calculateLatencyStats(answerDurations), evaluate: calculateLatencyStats(evaluateDurations), - total: calculateLatencyStats(totalDurations), + total: calculateLatencyStats(onlineDurations), }, tokens: tokenMetrics, memscore, - memscoreComponents, - retrieval: overallRetrieval, + memscoreComponents: + tokenMetrics && qualityPct != null + ? { + quality: qualityPct, + latencyMs: searchLatency.mean, + contextTokens: tokenMetrics.avgContextTokens, + } + : undefined, + retrieval: aggregateRetrievalMetrics(allRetrievalMetrics), byQuestionType, questionTypeRegistry: benchmark.getQuestionTypes(), evaluations, } - - return result } export function saveReport(result: BenchmarkResult): string { const reportsDir = join(REPORTS_DIR, result.runId) - if (!existsSync(reportsDir)) { - mkdirSync(reportsDir, { recursive: true }) - } - + if (!existsSync(reportsDir)) mkdirSync(reportsDir, { recursive: true }) const reportPath = join(reportsDir, "report.json") writeFileSync(reportPath, JSON.stringify(result, null, 2)) - logger.success(`Report saved to ${reportPath}`) return reportPath } -function formatLatencyRow(stats: LatencyStats): string { - const pad = (n: number) => n.toString().padStart(7) - return `${pad(stats.min)} ${pad(stats.max)} ${pad(stats.mean)} ${pad(stats.median)} ${pad(stats.p95)} ${pad(stats.p99)}` -} - export function printReport(result: BenchmarkResult): void { console.log("\n" + "=".repeat(60)) console.log("MEMORYBENCH RESULTS") console.log("=".repeat(60)) console.log(`Provider: ${result.provider}`) - console.log(`Benchmark: ${result.benchmark}`) + console.log(`Benchmark: ${result.benchmarkScope.displayName} (${result.benchmark})`) console.log(`Run ID: ${result.runId}`) - console.log(`Data Source: ${result.dataSourceRunId}`) - console.log(`Judge: ${result.judge}`) - console.log(`Answering Model: ${result.answeringModel}`) + console.log(`Protocol: ${result.protocolIdentity.id}@${result.protocolIdentity.version}`) console.log("-".repeat(60)) - console.log("\nSUMMARY:") - console.log(` Total Questions: ${result.summary.totalQuestions}`) - console.log(` Correct: ${result.summary.correctCount}`) - console.log(` Accuracy: ${(result.summary.accuracy * 100).toFixed(2)}%`) - - if (result.memscore && result.tokens) { - const qualityPct = Math.round(result.summary.accuracy * 100) - const avgLatency = result.latency.search.mean - console.log("") - console.log(` Quality: ${qualityPct}%`) - console.log(` Latency: ${avgLatency}ms (avg)`) + if (result.quality.primaryMetric) { console.log( - ` Tokens: ${result.tokens.avgContextTokens.toLocaleString()} (avg context sent to answering model)` + `Primary ${result.quality.primaryMetric.key}: ${result.quality.primaryMetric.value.toFixed(4)}` ) - console.log("") - console.log(` MemScore: ${result.memscore}`) - } - - console.log("-".repeat(60)) - console.log("\nLATENCY (ms):") - console.log(" min max mean median p95 p99") - console.log(` Ingest: ${formatLatencyRow(result.latency.ingest)}`) - console.log(` Indexing: ${formatLatencyRow(result.latency.indexing)}`) - console.log(` Search: ${formatLatencyRow(result.latency.search)}`) - console.log(` Answer: ${formatLatencyRow(result.latency.answer)}`) - console.log(` Evaluate: ${formatLatencyRow(result.latency.evaluate)}`) - console.log(` Total: ${formatLatencyRow(result.latency.total)}`) - - if (result.retrieval) { - console.log("-".repeat(60)) - console.log("\nRETRIEVAL QUALITY (K=" + result.retrieval.k + "):") - console.log(` Hit@K: ${(result.retrieval.hitAtK * 100).toFixed(1)}%`) - console.log(` Precision: ${(result.retrieval.precisionAtK * 100).toFixed(1)}%`) - console.log(` Recall: ${(result.retrieval.recallAtK * 100).toFixed(1)}%`) - console.log(` F1: ${(result.retrieval.f1AtK * 100).toFixed(1)}%`) - console.log(` MRR: ${result.retrieval.mrr.toFixed(3)}`) - console.log(` NDCG: ${result.retrieval.ndcg.toFixed(3)}`) - } - - console.log("-".repeat(60)) - console.log("\nBY QUESTION TYPE:") - for (const [type, stats] of Object.entries(result.byQuestionType)) { - const typeInfo = result.questionTypeRegistry?.[type] - const description = typeInfo?.description ? ` (${typeInfo.description})` : "" - console.log(` ${type}${description}:`) - console.log( - ` Total: ${stats.total}, Correct: ${stats.correct}, Accuracy: ${(stats.accuracy * 100).toFixed(2)}%` - ) - console.log( - ` Latency: search=${stats.latency.search.median}ms, answer=${stats.latency.answer.median}ms, total=${stats.latency.total.median}ms (median)` - ) - if (stats.retrieval) { - console.log( - ` Retrieval: Hit@${stats.retrieval.k}=${(stats.retrieval.hitAtK * 100).toFixed(0)}%, P=${(stats.retrieval.precisionAtK * 100).toFixed(0)}%, R=${(stats.retrieval.recallAtK * 100).toFixed(0)}%, MRR=${stats.retrieval.mrr.toFixed(2)}` - ) - } + } else { + console.log("Primary metric: none (official tier scores are reported separately)") } + console.log(`Average question score: ${result.summary.averageScore.toFixed(4)}`) + console.log(`Pass accuracy: ${(result.summary.accuracy * 100).toFixed(2)}%`) + console.log(`Unique builds: ${result.builds.uniqueBuildCount}`) + console.log(`Build work: ${result.builds.sumContainerBuildWorkMs}ms`) + console.log(`Build phase wall-clock: ${result.builds.buildPhaseWallClockMs}ms`) + console.log(`Online query latency (mean): ${result.latency.total.mean.toFixed(1)}ms`) + console.log(`Offline evaluation latency (mean): ${result.latency.evaluate.mean.toFixed(1)}ms`) + if (result.memscore) console.log(`MemScore: ${result.memscore}`) console.log("=".repeat(60) + "\n") } diff --git a/src/orchestrator/phases/retrieval-eval.ts b/src/orchestrator/phases/retrieval-eval.ts index 344197c..ca5c9e7 100644 --- a/src/orchestrator/phases/retrieval-eval.ts +++ b/src/orchestrator/phases/retrieval-eval.ts @@ -1,29 +1,65 @@ -import type { RetrievalMetrics } from "../../types/unified" -import type { LanguageModel } from "ai" -import { generateText } from "ai" +import { z } from "zod" +import type { AuxiliaryRetrievalEvaluationPolicy, EvaluationRuntime } from "../../types/protocol" +import type { RetrievalMetrics, UnifiedSearchResult } from "../../types/unified" interface RelevanceResult { id: string relevant: 0 | 1 } +function formatResultForRelevance(result: UnifiedSearchResult, id: string): string { + return [ + `=== ${id} ===`, + `PROVIDER: ${result.provider}`, + `RESULT_TYPE: ${result.resultType}`, + `RANK: ${result.rank}`, + ...(result.score === undefined ? [] : [`SCORE: ${result.score}`]), + ...(result.sessionId ? [`SESSION_ID: ${result.sessionId}`] : []), + ...(result.documentDate ? [`DOCUMENT_DATE: ${result.documentDate}`] : []), + "TEXT:", + result.text, + ].join("\n") +} + async function evaluateAllChunks( - model: LanguageModel, + runtime: EvaluationRuntime, question: string, groundTruth: string, - searchResults: unknown[] + searchResults: UnifiedSearchResult[] ): Promise { if (searchResults.length === 0) return [] - const formattedResults = searchResults - .map((result, index) => { - const id = `result_${index + 1}` - const content = JSON.stringify(result, null, 2) - return `=== ${id} ===\n${content}` + const expectedIds = searchResults.map((_, index) => `result_${index + 1}`) + const responseSchema = z + .object({ + results: z + .array( + z + .object({ + id: z.string().min(1), + relevant: z.union([z.literal(0), z.literal(1)]), + }) + .strict() + ) + .length(searchResults.length), + }) + .strict() + .superRefine((output, context) => { + for (let index = 0; index < expectedIds.length; index++) { + if (output.results[index]?.id !== expectedIds[index]) { + context.addIssue({ + code: "custom", + path: ["results", index, "id"], + message: `Retrieval relevance output ID mismatch: expected ${expectedIds[index]}, got ${output.results[index]?.id ?? ""}`, + }) + } + } }) + const formattedResults = searchResults + .map((result, index) => formatResultForRelevance(result, expectedIds[index]!)) .join("\n\n") - const prompt = `You are evaluating search results for relevance to a question. + const prompt = `Evaluate each normalized search result for relevance to the question. QUESTION: ${question} @@ -31,42 +67,31 @@ ${question} EXPECTED ANSWER: ${groundTruth} -SEARCH RESULTS: +NORMALIZED SEARCH RESULTS: ${formattedResults} -TASK: -For each search result, determine if it contains information relevant to answering the question. -A result is relevant if it contains content that helps answer the question or supports the expected answer. - -Return a JSON array with your evaluation for each result: -[ - {"id": "result_1", "relevant": 1}, - {"id": "result_2", "relevant": 0}, - ... -] - -Where: -- "id" is the result identifier (result_1, result_2, etc.) -- "relevant" is 1 if relevant, 0 if not relevant - -Return ONLY the JSON array, no other text.` - - try { - const response = await generateText({ - model, - messages: [{ role: "user", content: prompt }], - }) +A result is relevant when its TEXT contains information that helps answer the question or supports the expected answer. Return one result for every supplied ID, in the same order.` + + const output = await runtime.generateStructured({ + system: + "You are a retrieval relevance evaluator. Judge only the normalized result text supplied by the harness and follow the response schema exactly.", + prompt, + schema: responseSchema, + schemaName: "legacy_retrieval_relevance", + temperature: 0, + maxOutputTokens: Math.max(256, searchResults.length * 32), + maxAttempts: 3, + timeoutMs: 120_000, + }) - const jsonMatch = response.text.match(/\[[\s\S]*\]/) - if (!jsonMatch) { - return searchResults.map((_, i) => ({ id: `result_${i + 1}`, relevant: 0 as const })) + for (let index = 0; index < expectedIds.length; index++) { + if (output.results[index]?.id !== expectedIds[index]) { + throw new Error( + `Retrieval relevance output ID mismatch at index ${index}: expected ${expectedIds[index]}, got ${output.results[index]?.id ?? ""}` + ) } - - const parsed = JSON.parse(jsonMatch[0]) as RelevanceResult[] - return parsed - } catch { - return searchResults.map((_, i) => ({ id: `result_${i + 1}`, relevant: 0 as const })) } + return output.results } function calculateNDCG(relevanceScores: number[], idealRelevant: number): number { @@ -86,10 +111,10 @@ function calculateNDCG(relevanceScores: number[], idealRelevant: number): number } export async function calculateRetrievalMetrics( - model: LanguageModel, + runtime: EvaluationRuntime, question: string, groundTruth: string, - searchResults: unknown[], + searchResults: UnifiedSearchResult[], k: number = 10 ): Promise { const resultsToEval = searchResults.slice(0, k) @@ -108,29 +133,18 @@ export async function calculateRetrievalMetrics( } } - const relevanceResults = await evaluateAllChunks(model, question, groundTruth, resultsToEval) - - const relevanceScores = resultsToEval.map((_, i) => { - const id = `result_${i + 1}` - const result = relevanceResults.find((r) => r.id === id) - return result?.relevant === 1 ? 1 : 0 - }) + const relevanceResults = await evaluateAllChunks(runtime, question, groundTruth, resultsToEval) + const relevanceScores = relevanceResults.map((result) => result.relevant) - const relevantRetrieved = relevanceScores.filter((r) => r === 1).length + const relevantRetrieved = relevanceScores.filter((relevance) => relevance === 1).length const totalRelevant = Math.max(1, relevantRetrieved) - const hitAtK = relevantRetrieved > 0 ? 1 : 0 - - const precisionAtK = resultsToEval.length > 0 ? relevantRetrieved / resultsToEval.length : 0 - + const precisionAtK = relevantRetrieved / resultsToEval.length const recallAtK = relevantRetrieved > 0 ? 1 : 0 - const f1AtK = - precisionAtK + recallAtK > 0 ? (2 * (precisionAtK * recallAtK)) / (precisionAtK + recallAtK) : 0 - - const firstRelevantIndex = relevanceScores.findIndex((r) => r === 1) + precisionAtK + recallAtK > 0 ? (2 * precisionAtK * recallAtK) / (precisionAtK + recallAtK) : 0 + const firstRelevantIndex = relevanceScores.findIndex((relevance) => relevance === 1) const mrr = firstRelevantIndex >= 0 ? 1 / (firstRelevantIndex + 1) : 0 - const ndcg = calculateNDCG(relevanceScores, totalRelevant) return { @@ -145,3 +159,23 @@ export async function calculateRetrievalMetrics( totalRelevant, } } + +export async function calculateProtocolRetrievalMetrics( + policy: AuxiliaryRetrievalEvaluationPolicy, + runtime: EvaluationRuntime, + question: string, + groundTruth: string, + searchResults: UnifiedSearchResult[], + k: number +): Promise { + switch (policy) { + case "disabled": + return undefined + case "legacy-llm-relevance-v1": + return calculateRetrievalMetrics(runtime, question, groundTruth, searchResults, k) + default: { + const unsupported: never = policy + throw new Error(`Unsupported auxiliary retrieval evaluation policy: ${unsupported}`) + } + } +} diff --git a/src/orchestrator/phases/search.ts b/src/orchestrator/phases/search.ts index c774ab8..3fcaef8 100644 --- a/src/orchestrator/phases/search.ts +++ b/src/orchestrator/phases/search.ts @@ -1,12 +1,134 @@ -import { writeFileSync, mkdirSync, existsSync } from "fs" -import { join } from "path" -import type { Provider } from "../../types/provider" +import { existsSync, mkdirSync, writeFileSync } from "node:fs" import type { Benchmark } from "../../types/benchmark" import type { RunCheckpoint } from "../../types/checkpoint" -import { CheckpointManager } from "../checkpoint" +import type { Provider, ProviderSearchResponse } from "../../types/provider" +import { UNIFIED_SEARCH_RESULT_TYPES, type UnifiedSearchResult } from "../../types/unified" +import { resolveConcurrency } from "../../types/concurrency" import { logger } from "../../utils/logger" +import { CheckpointManager } from "../checkpoint" import { ConcurrentExecutor } from "../concurrent" -import { resolveConcurrency } from "../../types/concurrency" + +function validateProviderResults( + results: UnifiedSearchResult[], + provider: Pick, + requestedTopK: number +): void { + if (results.length > requestedTopK) { + throw new Error( + `${provider.name} returned ${results.length} results for requested Top-K ${requestedTopK}` + ) + } + const ranks = new Set() + for (const [index, result] of results.entries()) { + if (!result.text.trim()) + throw new Error(`${provider.name} returned empty text at index ${index}`) + if (!(UNIFIED_SEARCH_RESULT_TYPES as readonly string[]).includes(result.resultType)) { + throw new Error( + `${provider.name} returned unsupported result type ${JSON.stringify(result.resultType)} at index ${index}` + ) + } + if (result.provider !== provider.name) { + throw new Error( + `${provider.name} returned a result attributed to ${result.provider} at index ${index}` + ) + } + if (!Number.isInteger(result.rank) || result.rank < 1 || ranks.has(result.rank)) { + throw new Error(`${provider.name} returned an invalid or duplicate rank at index ${index}`) + } + ranks.add(result.rank) + } +} + +export function validateProviderSearchResponse( + response: ProviderSearchResponse, + provider: Pick, + requestedTopK: number +): void { + const { results, diagnostics } = response + const inconsistent = (reason: string): never => { + throw new Error(`${provider.name} returned inconsistent retrieval diagnostics: ${reason}`) + } + + if (diagnostics.requestedLimit !== requestedTopK) { + inconsistent( + `adapter requestedLimit ${diagnostics.requestedLimit} does not equal benchmark Top-K ${requestedTopK}` + ) + } + + for (const [index, request] of diagnostics.providerRequests.entries()) { + if (!request.operation.trim()) inconsistent(`provider request ${index} has no operation`) + if (!Number.isInteger(request.limit) || request.limit < 1) { + inconsistent(`provider request ${index} has invalid limit ${request.limit}`) + } + } + + const providerRequestLimit = diagnostics.providerRequests.reduce( + (sum, request) => sum + request.limit, + 0 + ) + if (provider.searchRequestStructure.kind === "single") { + if (diagnostics.providerRequests.length !== 1) { + inconsistent( + `single-request adapter reported ${diagnostics.providerRequests.length} provider requests` + ) + } + if (diagnostics.providerRequests[0].limit !== requestedTopK) { + inconsistent( + `single provider request limit ${diagnostics.providerRequests[0].limit} does not equal benchmark Top-K ${requestedTopK}` + ) + } + } else if (providerRequestLimit !== requestedTopK) { + inconsistent( + `split provider request limits total ${providerRequestLimit}, expected benchmark Top-K ${requestedTopK}` + ) + } + + const counts = [ + ["rawReturnedCount", diagnostics.rawReturnedCount], + ["normalizedCount", diagnostics.normalizedCount], + ["droppedCount", diagnostics.droppedCount], + ] as const + for (const [name, value] of counts) { + if (!Number.isInteger(value) || value < 0) inconsistent(`${name} is invalid: ${value}`) + } + if (diagnostics.rawReturnedCount > requestedTopK) { + inconsistent( + `rawReturnedCount ${diagnostics.rawReturnedCount} exceeds benchmark Top-K ${requestedTopK}` + ) + } + if (diagnostics.normalizedCount !== results.length) { + inconsistent( + `normalizedCount ${diagnostics.normalizedCount} does not equal evidence count ${results.length}` + ) + } + if (diagnostics.droppedCount !== diagnostics.rawReturnedCount - diagnostics.normalizedCount) { + inconsistent( + `droppedCount ${diagnostics.droppedCount} does not equal raw minus normalized count` + ) + } + if (!Array.isArray(diagnostics.droppedResults)) { + inconsistent("droppedResults is missing") + } + if (diagnostics.droppedResults.length !== diagnostics.droppedCount) { + inconsistent( + `recorded ${diagnostics.droppedResults.length} drop reasons for droppedCount ${diagnostics.droppedCount}` + ) + } + const droppedIndices = new Set() + for (const dropped of diagnostics.droppedResults) { + if ( + !Number.isInteger(dropped.index) || + dropped.index < 0 || + dropped.index >= diagnostics.rawReturnedCount || + droppedIndices.has(dropped.index) + ) { + inconsistent(`invalid or duplicate dropped result index ${dropped.index}`) + } + droppedIndices.add(dropped.index) + } + + validateProviderResults(results, provider, requestedTopK) +} export async function runSearchPhase( provider: Provider, @@ -17,13 +139,33 @@ export async function runSearchPhase( ): Promise { const questions = benchmark.getQuestions() const targetQuestions = questionIds - ? questions.filter((q) => questionIds.includes(q.questionId)) + ? questions.filter((question) => questionIds.includes(question.questionId)) : questions - - const pendingQuestions = targetQuestions.filter((q) => { - const status = checkpointManager.getPhaseStatus(checkpoint, q.questionId, "search") - const indexingStatus = checkpointManager.getPhaseStatus(checkpoint, q.questionId, "indexing") - return status !== "completed" && indexingStatus === "completed" + for (const question of targetQuestions) { + const questionCheckpoint = checkpoint.questions[question.questionId] + if (!questionCheckpoint) { + throw new Error(`Question ${question.questionId} has no checkpoint record`) + } + const build = checkpoint.builds[questionCheckpoint.buildId] + if (!build) { + throw new Error( + `Question ${question.questionId} references missing build ${questionCheckpoint.buildId}` + ) + } + if (build.ingest.status !== "completed") { + throw new Error( + `Question ${question.questionId} cannot search because build ${build.buildId} ingestion is ${build.ingest.status}` + ) + } + if (build.indexing.status !== "completed" || build.indexing.failedIds.length > 0) { + throw new Error( + `Question ${question.questionId} cannot search because build ${build.buildId} indexing is ${build.indexing.status}${build.indexing.failedIds.length > 0 ? ` with ${build.indexing.failedIds.length} failed IDs` : ""}` + ) + } + } + const pendingQuestions = targetQuestions.filter((question) => { + const questionCheckpoint = checkpoint.questions[question.questionId] + return questionCheckpoint.phases.search.status !== "completed" }) if (pendingQuestions.length === 0) { @@ -32,12 +174,8 @@ export async function runSearchPhase( } const resultsDir = checkpointManager.getResultsDir(checkpoint.runId) - if (!existsSync(resultsDir)) { - mkdirSync(resultsDir, { recursive: true }) - } - + if (!existsSync(resultsDir)) mkdirSync(resultsDir, { recursive: true }) const concurrency = resolveConcurrency("search", checkpoint.concurrency, provider.concurrency) - logger.info(`Searching ${pendingQuestions.length} questions (concurrency: ${concurrency})...`) await ConcurrentExecutor.execute( @@ -46,55 +184,108 @@ export async function runSearchPhase( checkpoint.runId, "search", async ({ item: question, index, total }) => { - const containerTag = checkpoint.questions[question.questionId].containerTag + const questionCheckpoint = checkpoint.questions[question.questionId] + const build = checkpoint.builds[questionCheckpoint.buildId] + if (!build) throw new Error(`Question ${question.questionId} references missing build`) + const retrievalPlan = benchmark.protocol.createRetrievalPlan({ question }) + if ( + !Number.isInteger(retrievalPlan.requestedTopK) || + retrievalPlan.requestedTopK < 1 || + retrievalPlan.requestedTopK > 100 + ) { + throw new Error(`Invalid requestedTopK ${retrievalPlan.requestedTopK}`) + } + if ( + !Number.isInteger(retrievalPlan.answerCutoff) || + retrievalPlan.answerCutoff < 0 || + retrievalPlan.answerCutoff > retrievalPlan.requestedTopK + ) { + throw new Error(`Invalid answerCutoff ${retrievalPlan.answerCutoff}`) + } - const startTime = Date.now() + const startedAt = new Date().toISOString() + const startedMs = Date.now() checkpointManager.updatePhase(checkpoint, question.questionId, "search", { status: "in_progress", - startedAt: new Date().toISOString(), + retrievalPlan, + requestedCount: retrievalPlan.requestedTopK, + answerCutoff: retrievalPlan.answerCutoff, + startedAt, + costUsd: null, + error: undefined, }) try { - const results = await provider.search(question.question, { - containerTag, - limit: 10, - threshold: 0.3, + const response = await provider.search(retrievalPlan.query, { + containerTag: build.containerTag, + limit: retrievalPlan.requestedTopK, + threshold: retrievalPlan.threshold, + searchMode: retrievalPlan.searchMode, + filters: retrievalPlan.filters, }) + const { results, diagnostics } = response + validateProviderSearchResponse(response, provider, retrievalPlan.requestedTopK) - const durationMs = Date.now() - startTime - const resultFile = join(resultsDir, `${question.questionId}.json`) + const completedAt = new Date().toISOString() + const durationMs = Date.now() - startedMs + const resultFile = checkpointManager.getQuestionResultsPath( + checkpoint.runId, + question.questionId + ) const resultData = { + benchmark: checkpoint.benchmark, + benchmarkScope: checkpoint.benchmarkScope, + datasetIdentity: checkpoint.datasetIdentity, + benchmarkInputFingerprint: checkpoint.benchmarkInputFingerprint, + selectedQuestionIdsDigest: checkpoint.selectedQuestionIdsDigest, questionId: question.questionId, question: question.question, questionType: question.questionType, groundTruth: question.groundTruth, - containerTag, - timestamp: new Date().toISOString(), + buildId: build.buildId, + containerTag: build.containerTag, + protocolIdentity: checkpoint.protocolIdentity, + retrievalPlan, + requestedCount: retrievalPlan.requestedTopK, + rawReturnedCount: diagnostics.rawReturnedCount, + returnedCount: diagnostics.normalizedCount, + normalizedCount: diagnostics.normalizedCount, + droppedCount: diagnostics.droppedCount, + droppedResults: diagnostics.droppedResults, + providerRequests: diagnostics.providerRequests, + timestamp: completedAt, durationMs, results, } - writeFileSync(resultFile, JSON.stringify(resultData, null, 2)) - checkpointManager.updatePhase(checkpoint, question.questionId, "search", { status: "completed", resultFile, results, - completedAt: new Date().toISOString(), + rawReturnedCount: diagnostics.rawReturnedCount, + returnedCount: diagnostics.normalizedCount, + normalizedCount: diagnostics.normalizedCount, + droppedCount: diagnostics.droppedCount, + droppedResults: diagnostics.droppedResults, + providerRequests: diagnostics.providerRequests, + completedAt, durationMs, + error: undefined, }) - - logger.progress(index + 1, total, `Searched ${question.questionId} (${durationMs}ms)`) + logger.progress( + index + 1, + total, + `Searched ${question.questionId}: ${results.length}/${retrievalPlan.requestedTopK} (${durationMs}ms)` + ) return { questionId: question.questionId, durationMs } - } catch (e) { - const error = e instanceof Error ? e.message : String(e) + } catch (error) { + const message = error instanceof Error ? error.message : String(error) checkpointManager.updatePhase(checkpoint, question.questionId, "search", { status: "failed", - error, + error: message, }) - logger.error(`Failed to search ${question.questionId}: ${error}`) throw new Error( - `Search failed at ${question.questionId}: ${error}. Fix the issue and resume with the same run ID.` + `Search failed at ${question.questionId}: ${message}. Fix the issue and resume with the same run ID.` ) } } diff --git a/src/prompts/beam.ts b/src/prompts/beam.ts index b91f113..b749c4e 100644 --- a/src/prompts/beam.ts +++ b/src/prompts/beam.ts @@ -1,21 +1,4 @@ -export interface BeamRubricJudgeResult { - score: number - reason: string -} - -/** - * Number of memories to retrieve per BEAM search. Matches mem0's evaluation - * cutoff exactly (their --top-k-cutoffs default) so the answering model sees - * the same context budget. mem0 retrieves 200 and trims to 100; we retrieve - * 100 directly since we evaluate at a single cutoff. - */ -export const BEAM_SEARCH_TOP_K = 100 - -/** - * How many of the retrieved memories to expose to the answering model. - * Mirrors BEAM_SEARCH_TOP_K — every retrieved memory is shown. - */ -export const BEAM_ANSWER_TOP_K = 100 +import { sha256Text } from "../utils/stable" const MONTH_NAMES = [ "January", @@ -70,20 +53,31 @@ interface BeamMemoryLike { metadata?: { sessionId?: string } } +export type BeamAnswerFormat = "default" | "event-ordering-lines" + +/** + * The pinned BEAM scorer treats every non-empty answer line as one predicted + * event. This format-only rule makes that scorer input explicit without adding + * a second semantic extraction step. + */ +export const BEAM_EVENT_ORDERING_ANSWER_FORMAT_VERSION = "authors-newline-scorer-compatible-v1" + /** - * Ported verbatim from mem0's get_beam_answer_generation_prompt - * (mem0ai/memory-benchmarks/benchmarks/beam/prompts.py). Memories are sliced - * to BEAM_ANSWER_TOP_K, sorted chronologically (oldest first), numbered, and - * prefixed with their session date when available — exactly the format mem0's - * answering LLM sees on their published BEAM 1M / 10M numbers. + * BEAM answer prompt. The protocol applies its configured answer cutoff before + * calling this formatter; this function only orders and renders normalized + * evidence and never sees provider-specific raw JSON. */ export function buildBeamAnswerPrompt( question: string, memories: unknown[], - sessionDateMap: Map + sessionDateMap: Map, + answerFormat: BeamAnswerFormat = "default" ): string { - const sliced = memories.slice(0, BEAM_ANSWER_TOP_K) as BeamMemoryLike[] - const memoriesText = formatBeamMemories(sliced, sessionDateMap) + const memoriesText = formatBeamMemories(memories as BeamMemoryLike[], sessionDateMap) + const eventOrderingRule = + answerFormat === "event-ordering-lines" + ? "\n10. For this event-ordering question, output exactly one event per line in chronological order. Do not use bullets, numbering, headings, or explanations." + : "" return `You are an AI assistant with access to stored memories from prior conversations with a user. Use these memories to answer the following question as accurately and completely as possible. @@ -96,7 +90,7 @@ IMPORTANT RULES: 6. For ordering questions: present events in chronological order. 7. For preference questions: use the most recently stated preference. 8. Be specific and direct — include exact names, dates, numbers, and details from the memories. -9. Do NOT invent or assume information that isn't in the memories. +9. Do NOT invent or assume information that isn't in the memories.${eventOrderingRule} QUESTION: ${question} @@ -114,14 +108,11 @@ function formatBeamMemories( // Resolve text + date per memory. const items = memories.map((m) => { - const text = - typeof m?.memory === "string" - ? m.memory - : typeof m?.content === "string" - ? m.content - : JSON.stringify(m) - const sessionId = - typeof m?.metadata?.sessionId === "string" ? m.metadata.sessionId : "" + const text = typeof m?.memory === "string" ? m.memory : m.content + if (typeof text !== "string" || !text.trim()) { + throw new Error("BEAM prompt received evidence without normalized text") + } + const sessionId = typeof m?.metadata?.sessionId === "string" ? m.metadata.sessionId : "" const date = sessionId ? sessionDateMap.get(sessionId) : undefined return { text, sessionId, date } }) @@ -144,124 +135,7 @@ function formatBeamMemories( .join("\n") } -/** - * Ported verbatim from mem0's BEAM benchmark setup - * (mem0ai/memory-benchmarks/benchmarks/beam/prompts.py) so our judging is - * apples-to-apples with their published BEAM 1M / 10M numbers. - */ -export const BEAM_JUDGE_SYSTEM_PROMPT = - "You are an expert evaluator assessing whether an AI assistant's response satisfies " + - "specific rubric criteria. You must be objective, fair, and consistent. " + - "Return ONLY valid JSON with the exact format requested." - -function parseJsonResponse(response: string): Record { - const trimmed = response.trim() - - if (trimmed.startsWith("```")) { - const codeFenceMatch = trimmed.match(/```(?:json)?\s*([\s\S]*?)\s*```/) - if (codeFenceMatch?.[1]) { - return JSON.parse(codeFenceMatch[1]) - } - } - - const jsonMatch = trimmed.match(/\{[\s\S]*\}/) - if (jsonMatch) { - return JSON.parse(jsonMatch[0]) - } - - return JSON.parse(trimmed) -} - -/** - * Ported from mem0's get_beam_nugget_judge_prompt. Each rubric "nugget" is - * scored independently on a 3-point scale by the judge LLM. - */ -export function buildBeamRubricJudgePrompt( - question: string, - nugget: string, - llmResponse: string -): string { - return `Evaluate whether the following LLM response demonstrates compliance with the specified RUBRIC CRITERION. - -QUESTION: -${question} - -LLM RESPONSE: -${llmResponse} - -RUBRIC CRITERION: -${nugget} - -SCORING GUIDELINES: - -First, determine whether the rubric criterion is a POSITIVE requirement (the response SHOULD include something) or a NEGATIVE constraint (the response SHOULD NOT include something). - -**For POSITIVE requirements** (response should contain, mention, or demonstrate something): -- **1.0 (Complete Compliance)**: The required element is present, accurate, and complete. The response fully and clearly satisfies the rubric criterion. -- **0.5 (Partial Compliance)**: The required element is partially present, has minor inaccuracies, or is incomplete. The core intent is present but not fully realized. -- **0.0 (No Compliance)**: The required element is missing, incorrect, or the response is entirely off-topic / non-responsive. - -**For NEGATIVE constraints** (response should NOT contain or should avoid something): -- **1.0 (Complete Compliance)**: The response is responsive to the question AND the prohibited element is absent. -- **0.5 (Partial Compliance)**: The response is responsive but contains a borderline or ambiguous reference to the prohibited element. -- **0.0 (No Compliance)**: The prohibited element is present in the response, OR the response is non-responsive (off-topic, refusal, empty). - -**Compound statement handling**: If the rubric criterion contains "and" or commas connecting multiple required elements: -- All elements present and correct = 1.0 -- Some (but not all) elements present and correct = 0.5 -- No elements present or correct = 0.0 - -EVALUATION RULES: -1. **Semantic tolerance**: Paraphrases and synonyms are acceptable. The response does not need to use the exact same words as the rubric. -2. **Numeric and date equivalence**: Treat equivalent representations as identical. "$68,000" = "68k" = "sixty-eight thousand dollars". "2 years" = "24 months". Prefer normalized comparison for numbers, currencies, dates, and durations. -3. **Case / punctuation / whitespace tolerance**: Differences in capitalization, punctuation, and whitespace must be ignored when comparing content. -4. **Hedging tolerance**: Do not penalize hedging language ("I think", "probably", "it seems"), passive voice, or verbosity if the substantive content satisfies the rubric criterion. -5. **Style neutrality**: Do not penalize for tone, formatting, or length unless the rubric criterion specifically requires a particular format. -6. **Responsiveness**: If the LLM response is completely off-topic or refuses to answer, score 0.0 for all criteria. -7. **Independence**: Evaluate this criterion in isolation — do not consider other rubric items. -8. **Specificity matters**: Vague or generic answers that could apply to any question score lower than specific, detailed answers. - -STEP-BY-STEP EVALUATION: -Follow these steps in order: -1. **Understand the Requirement**: Read the rubric criterion and classify it as a positive requirement or a negative constraint. -2. **Parse Compound Statements**: If the criterion contains multiple sub-requirements joined by "and" or commas, identify each element separately. -3. **Check Compliance**: Compare the LLM response against each element, applying the tolerance rules above (semantic, numeric, case, hedging). -4. **Assign Score**: Use the appropriate scoring table (positive or negative) and compound-statement rule to determine the score. -5. **Provide Reasoning**: Write a concise explanation referencing which elements were or were not satisfied. - -Return your evaluation as a JSON object with exactly two fields: -{"score": <0.0 or 0.5 or 1.0>, "reason": ""}` -} - -/** - * Ported from mem0's _clamp_nugget_score. Snaps any numeric score the judge - * returns to the nearest of {0.0, 0.5, 1.0}, instead of penalizing on exact - * match failure. - */ -export function clampNuggetScore(raw: number): 0 | 0.5 | 1 { - if (!Number.isFinite(raw)) return 0 - if (raw >= 0.75) return 1 - if (raw >= 0.25) return 0.5 - return 0 -} - -/** - * Ported from mem0's judge_single_nugget response handling. If JSON parsing - * works, clamp the score; otherwise fall back to scanning the raw text for - * "1.0" / "0.5" markers. - */ -export function parseBeamRubricJudgeResponse(response: string): BeamRubricJudgeResult { - try { - const parsed = parseJsonResponse(response) - const raw = typeof parsed.score === "number" ? parsed.score : Number(parsed.score) - return { - score: clampNuggetScore(raw), - reason: typeof parsed.reason === "string" ? parsed.reason : "", - } - } catch { - const snippet = response.slice(0, 200) - if (snippet.includes("1.0")) return { score: 1, reason: snippet } - if (snippet.includes("0.5")) return { score: 0.5, reason: snippet } - return { score: 0, reason: `Parse error: ${snippet}` } - } -} +/** Changes whenever either the public prompt builder or its evidence/date renderer changes. */ +export const BEAM_ANSWER_FORMATTER_IMPLEMENTATION_HASH = sha256Text( + [buildBeamAnswerPrompt.toString(), formatBeamMemories.toString()].join("\n\n") +) diff --git a/src/prompts/extraction.ts b/src/prompts/extraction.ts index daca2b9..5a933ae 100644 --- a/src/prompts/extraction.ts +++ b/src/prompts/extraction.ts @@ -1,9 +1,12 @@ import { createOpenAI } from "@ai-sdk/openai" import { generateText } from "ai" import type { UnifiedSession } from "../types/unified" +import { sha256Text, stableSha256 } from "../utils/stable" /** Model used for memory extraction (fast, cheap, sufficient for extraction) */ const EXTRACTION_MODEL = "gpt-4o-mini" +const EXTRACTION_MAX_OUTPUT_TOKENS = 2000 +const EXTRACTION_TEMPERATURE = 0 /** * Build an extraction prompt that instructs the LLM to extract structured @@ -64,6 +67,17 @@ Rules: - Resolve relative date references ("yesterday", "last week") to absolute dates using the conversation date when possible` } +/** Provider build identity for the shared filesystem/RAG extraction stage. */ +export function getMemoryExtractionConfigFingerprint(): string { + return stableSha256({ + schemaVersion: 1, + model: EXTRACTION_MODEL, + maxOutputTokens: EXTRACTION_MAX_OUTPUT_TOKENS, + temperature: EXTRACTION_TEMPERATURE, + promptBuilderSourceSha256: sha256Text(Function.prototype.toString.call(buildExtractionPrompt)), + }) +} + /** * Call LLM to extract structured memories from a conversation session. * Returns MEMORY.md-style markdown with categorized facts, events, preferences. @@ -77,8 +91,8 @@ export async function extractMemories( const params: Record = { model: openai(EXTRACTION_MODEL), prompt, - maxTokens: 2000, - temperature: 0, + maxTokens: EXTRACTION_MAX_OUTPUT_TOKENS, + temperature: EXTRACTION_TEMPERATURE, } const { text } = await generateText(params as Parameters[0]) diff --git a/src/protocols/beam-mem0/index.ts b/src/protocols/beam-mem0/index.ts new file mode 100644 index 0000000..5ff16d1 --- /dev/null +++ b/src/protocols/beam-mem0/index.ts @@ -0,0 +1,514 @@ +import { z } from "zod" +import { + ANSWER_RUNTIME_EXECUTION_VERSION, + generateAnswerWithRetries, + getLanguageModel, + runAnswerPhase, + shouldRunAnswerPhase, +} from "../../orchestrator/phases/answer" +import { + STRUCTURED_RUNTIME_EXECUTION_VERSION, + executeStructuredWithRetries, + generateStructuredObject, +} from "../../orchestrator/evaluation-runtime" +import { hasEvaluableAnswer, runEvaluatePhase } from "../../orchestrator/phases/evaluate" +import { buildBeamAnswerPrompt } from "../../prompts/beam" +import type { BenchmarkProtocol, ProtocolIdentity, QuestionEvaluation } from "../../types/protocol" +import type { UnifiedQuestion, UnifiedSearchResult, UnifiedSession } from "../../types/unified" +import { sha256Text, stableSha256 } from "../../utils/stable" +import { + BEAM_ABILITY_IDS, + BEAM_PASS_THRESHOLD, + BeamPaperProtocol, + type BeamAbilityId, +} from "../beam-paper" +import { + BEAM_MEM0_JUDGE_SYSTEM_PROMPT, + BEAM_MEM0_NUGGET_PROMPT_VERSION, + buildBeamMem0NuggetPrompt, +} from "./prompts" + +export * from "./prompts" + +export const BEAM_MEM0_NUGGET_PROTOCOL_ID = "beam-mem0-nugget" +export const BEAM_MEM0_NUGGET_PROTOCOL_VERSION = "1.2.0" +export const BEAM_MEM0_NUGGET_PROFILE = "mem0-nugget" + +const MAX_DIRECT_RETRIEVAL_TOP_K = 100 +const MEM0_GPT5_MAX_OUTPUT_TOKENS = 4096 +const MEM0_GPT5_MAX_ATTEMPTS = 5 +const MEM0_GPT5_INNER_MAX_RETRIES = 2 +const MEM0_GPT5_TIMEOUT_MS = 120_000 +const MEM0_GPT5_RETRY_BACKOFF_MS = 2_000 +const MEM0_GPT5_TRANSPORT = "openai-chat-completions" as const +const MEM0_TERMINAL_EMPTY_OUTPUT_POLICY = "accept-and-evaluate" as const + +const NUGGET_JUDGMENT_SCHEMA = z + .object({ + score: z.number(), + reason: z.string().trim().min(1), + }) + .strict() + +const NUGGET_PROGRESS_SCHEMA = z + .object({ + kind: z.literal("beam-mem0-nuggets-v1"), + questionId: z.string().min(1), + rubricHash: z.string().regex(/^[a-f0-9]{64}$/), + judgments: z + .array( + z + .object({ + nugget: z.string().min(1), + score: z.union([z.literal(0), z.literal(0.5), z.literal(1)]), + reason: z.string().min(1), + }) + .strict() + ) + .default([]), + }) + .strict() + +interface NuggetJudgment { + nugget: string + score: 0 | 0.5 | 1 + reason: string +} + +/** Match mem0's `_clamp_nugget_score` thresholds. */ +export function clampMem0NuggetScore(score: number): 0 | 0.5 | 1 { + if (!Number.isFinite(score)) return 0 + if (score >= 0.75) return 1 + if (score >= 0.25) return 0.5 + return 0 +} + +export interface BeamMem0NuggetProtocolConfig { + retrievalTopK?: number + answerCutoff?: number +} + +function mean(values: readonly number[]): number { + if (values.length === 0) throw new Error("Cannot average an empty list") + return values.reduce((sum, value) => sum + value, 0) / values.length +} + +function getRubric(question: UnifiedQuestion): string[] { + const rubric = question.metadata?.rubric + if ( + !Array.isArray(rubric) || + rubric.length === 0 || + rubric.some((item) => typeof item !== "string" || !item.trim()) + ) { + throw new Error( + `BEAM question ${question.questionId || ""} must have a non-empty string rubric` + ) + } + return rubric as string[] +} + +function getDocumentDate(session: UnifiedSession): string | undefined { + const documentDate = session.metadata?.documentDate + const legacyDate = session.metadata?.date + if (documentDate !== undefined && legacyDate !== undefined && documentDate !== legacyDate) { + throw new Error(`BEAM session ${session.sessionId} has conflicting document dates`) + } + const value = documentDate ?? legacyDate + return typeof value === "string" ? value : undefined +} + +function createSessionDateMap( + sessions: readonly UnifiedSession[], + results: readonly UnifiedSearchResult[] +): Map { + const dates = new Map() + for (const session of sessions) { + const date = getDocumentDate(session) + if (date) dates.set(session.sessionId, date) + } + for (const result of results) { + if (result.sessionId && result.documentDate && !dates.has(result.sessionId)) { + dates.set(result.sessionId, result.documentDate) + } + } + return dates +} + +function toPromptEvidence(results: readonly UnifiedSearchResult[], dates: Map) { + return results.map((result) => ({ + content: result.text, + metadata: (() => { + const promptSessionId = + result.sessionId ?? (result.documentDate ? `normalized-result:${result.id}` : undefined) + if (!promptSessionId) return undefined + if (result.documentDate) dates.set(promptSessionId, result.documentDate) + return { sessionId: promptSessionId } + })(), + })) +} + +const MEM0_EVALUATOR_IDENTITY = { + profile: BEAM_MEM0_NUGGET_PROFILE, + sourceRepository: "mem0ai/memory-benchmarks", + sourceCommit: "4b61c5d31b9c668a12b4f5e78064248a02c82d2b", + judgeProvider: "openai", + judgeModel: "gpt-5", + nuggetPromptVersion: BEAM_MEM0_NUGGET_PROMPT_VERSION, + nuggetPromptSha256: sha256Text( + buildBeamMem0NuggetPrompt({ + question: "", + nugget: "", + answer: "", + }) + ), + systemPromptSha256: sha256Text(BEAM_MEM0_JUDGE_SYSTEM_PROMPT), + structuredOutputSchemaSha256: stableSha256({ + score: "number-clamped-at-0.25-and-0.75", + reason: "non-empty-string", + additionalProperties: false, + }), + structuredOutputMode: "ai-sdk-generate-object-json-schema-v1", + parseFallback: "none-fail-closed-deviation-from-mem0-raw-text-marker-fallback", + eventOrderingPolicy: "ordinary-nugget-average-primary", + temperature: null, + maxOutputTokens: MEM0_GPT5_MAX_OUTPUT_TOKENS, + maxAttempts: MEM0_GPT5_MAX_ATTEMPTS, + innerMaxRetries: MEM0_GPT5_INNER_MAX_RETRIES, + timeoutMs: MEM0_GPT5_TIMEOUT_MS, + retryBackoffMs: MEM0_GPT5_RETRY_BACKOFF_MS, + transport: MEM0_GPT5_TRANSPORT, + runtimeExecutionVersion: STRUCTURED_RUNTIME_EXECUTION_VERSION, + runtimeExecutionSha256: stableSha256({ + executeStructuredWithRetries: executeStructuredWithRetries.toString(), + generateStructuredObject: generateStructuredObject.toString(), + hasEvaluableAnswer: hasEvaluableAnswer.toString(), + runEvaluatePhase: runEvaluatePhase.toString(), + }), +} as const + +function createIdentity( + protocol: BeamMem0NuggetProtocol, + ingestionDelegate: BeamPaperProtocol +): ProtocolIdentity { + const retrieval = { + policy: "direct-top-k", + requestedTopK: protocol.retrievalTopK, + answerCutoff: protocol.answerCutoff, + threshold: 0, + maximumDirectTopK: MAX_DIRECT_RETRIEVAL_TOP_K, + } + const answerProbe = buildBeamAnswerPrompt("", [], new Map()) + const answer = { + formatter: "mem0-public-beam-answer-prompt", + formatterVersion: "normalized-evidence-v1", + transport: MEM0_GPT5_TRANSPORT, + maxOutputTokens: MEM0_GPT5_MAX_OUTPUT_TOKENS, + maxAttempts: MEM0_GPT5_MAX_ATTEMPTS, + innerMaxRetries: MEM0_GPT5_INNER_MAX_RETRIES, + timeoutMs: MEM0_GPT5_TIMEOUT_MS, + retryBackoffMs: MEM0_GPT5_RETRY_BACKOFF_MS, + emptyOutputPolicy: "retry", + terminalEmptyOutputPolicy: MEM0_TERMINAL_EMPTY_OUTPUT_POLICY, + runtimeExecutionVersion: ANSWER_RUNTIME_EXECUTION_VERSION, + runtimeExecutionSha256: stableSha256({ + getLanguageModel: getLanguageModel.toString(), + generateAnswerWithRetries: generateAnswerWithRetries.toString(), + shouldRunAnswerPhase: shouldRunAnswerPhase.toString(), + runAnswerPhase: runAnswerPhase.toString(), + }), + promptSha256: sha256Text(answerProbe), + implementationSha256: stableSha256({ + createAnswerPlan: protocol.createAnswerPlan.toString(), + promptBuilder: buildBeamAnswerPrompt.toString(), + documentDate: getDocumentDate.toString(), + sessionDateMap: createSessionDateMap.toString(), + promptEvidence: toPromptEvidence.toString(), + }), + } + const evaluator = { + ...MEM0_EVALUATOR_IDENTITY, + implementationSha256: stableSha256({ + evaluateQuestion: protocol.evaluateQuestion.toString(), + promptBuilder: buildBeamMem0NuggetPrompt.toString(), + rubricValidation: getRubric.toString(), + scoreClamp: clampMem0NuggetScore.toString(), + mean: mean.toString(), + }), + } + const aggregation = { + primaryMetric: "mem0NuggetAverage", + questionWeighting: "equal-micro", + eventOrderingPolicy: "ordinary-nugget-average-primary", + implementationSha256: stableSha256({ + aggregateQuality: protocol.aggregateQuality.toString(), + mean: mean.toString(), + }), + } + const implementation = { + protocol: BEAM_MEM0_NUGGET_PROTOCOL_ID, + version: BEAM_MEM0_NUGGET_PROTOCOL_VERSION, + ingestionPolicyHash: ingestionDelegate.identity.ingestionPolicyHash, + retrieval, + answer, + evaluator, + aggregation, + } + return { + id: BEAM_MEM0_NUGGET_PROTOCOL_ID, + version: BEAM_MEM0_NUGGET_PROTOCOL_VERSION, + configFingerprint: stableSha256({ retrieval, answer, evaluator, aggregation }), + implementationFingerprint: stableSha256(implementation), + ingestionPolicyHash: ingestionDelegate.identity.ingestionPolicyHash, + retrievalPolicyHash: stableSha256(retrieval), + answerPromptHash: stableSha256(answer), + evaluatorHash: stableSha256(evaluator), + aggregationHash: stableSha256(aggregation), + details: { + comparisonProfile: BEAM_MEM0_NUGGET_PROFILE, + ingestionPolicy: ingestionDelegate.identity.details?.ingestionPolicy, + retrievalPolicy: retrieval, + answerPrompt: answer, + evaluatorIdentity: evaluator, + aggregation, + comparabilityNotice: + "Experimental direct-retrieval mem0-style profile; unlike mem0's runner, schema-invalid judge output fails closed without raw-text marker fallback. This is not the BEAM paper protocol.", + }, + } +} + +export class BeamMem0NuggetProtocol implements BenchmarkProtocol { + readonly auxiliaryRetrievalEvaluation = "disabled" as const + readonly ingestionExecutionPolicy = { + readinessBarrier: "after-each-document", + processingMode: "instant", + } as const + readonly requiredJudge = { + provider: "openai", + modelId: "gpt-5", + modelAlias: "gpt-5", + } + readonly retrievalTopK: number + readonly answerCutoff: number + readonly identity: ProtocolIdentity + private readonly ingestionDelegate = new BeamPaperProtocol() + + constructor(config: BeamMem0NuggetProtocolConfig = {}) { + const retrievalTopK = config.retrievalTopK ?? 50 + const answerCutoff = config.answerCutoff ?? retrievalTopK + if ( + !Number.isInteger(retrievalTopK) || + retrievalTopK < 1 || + retrievalTopK > MAX_DIRECT_RETRIEVAL_TOP_K + ) { + throw new Error( + `BEAM mem0 comparison retrieval Top-K must be an integer from 1 to ${MAX_DIRECT_RETRIEVAL_TOP_K}; got ${retrievalTopK}` + ) + } + if (!Number.isInteger(answerCutoff) || answerCutoff < 1 || answerCutoff > retrievalTopK) { + throw new Error( + `BEAM mem0 comparison answer cutoff must be an integer from 1 to Top-K ${retrievalTopK}; got ${answerCutoff}` + ) + } + this.retrievalTopK = retrievalTopK + this.answerCutoff = answerCutoff + this.identity = createIdentity(this, this.ingestionDelegate) + } + + validateQuestion(question: UnifiedQuestion): void { + this.ingestionDelegate.validateQuestion(question) + } + + createIngestionPlan(input: Parameters[0]) { + return this.ingestionDelegate.createIngestionPlan(input) + } + + createRetrievalPlan({ question }: Parameters[0]) { + this.validateQuestion(question) + return { + query: question.question, + requestedTopK: this.retrievalTopK, + answerCutoff: this.answerCutoff, + threshold: 0, + } + } + + createAnswerPlan({ + question, + sessions, + results, + retrieval, + }: Parameters[0]) { + this.validateQuestion(question) + if ( + retrieval.requestedTopK !== this.retrievalTopK || + retrieval.answerCutoff !== this.answerCutoff + ) { + throw new Error("BEAM mem0 comparison retrieval plan drifted from its configured budget") + } + const evidence = results.slice(0, retrieval.answerCutoff) + const dates = createSessionDateMap(sessions, evidence) + const promptEvidence = toPromptEvidence(evidence, dates) + return { + request: { + prompt: buildBeamAnswerPrompt(question.question, promptEvidence, dates), + maxOutputTokens: MEM0_GPT5_MAX_OUTPUT_TOKENS, + transport: MEM0_GPT5_TRANSPORT, + maxAttempts: MEM0_GPT5_MAX_ATTEMPTS, + innerMaxRetries: MEM0_GPT5_INNER_MAX_RETRIES, + timeoutMs: MEM0_GPT5_TIMEOUT_MS, + retryBackoffMs: MEM0_GPT5_RETRY_BACKOFF_MS, + terminalEmptyOutputPolicy: MEM0_TERMINAL_EMPTY_OUTPUT_POLICY, + }, + baseRequest: { + prompt: buildBeamAnswerPrompt(question.question, [], dates), + maxOutputTokens: MEM0_GPT5_MAX_OUTPUT_TOKENS, + transport: MEM0_GPT5_TRANSPORT, + maxAttempts: MEM0_GPT5_MAX_ATTEMPTS, + innerMaxRetries: MEM0_GPT5_INNER_MAX_RETRIES, + timeoutMs: MEM0_GPT5_TIMEOUT_MS, + retryBackoffMs: MEM0_GPT5_RETRY_BACKOFF_MS, + terminalEmptyOutputPolicy: MEM0_TERMINAL_EMPTY_OUTPUT_POLICY, + }, + answerEvidenceCount: evidence.length, + } + } + + async evaluateQuestion( + { + question, + hypothesis, + protocolProgress, + onProtocolProgress, + }: Parameters[0], + runtime: Parameters[1] + ): Promise { + this.validateQuestion(question) + const rubric = getRubric(question) + const rubricHash = stableSha256(rubric) + const progress = protocolProgress + ? NUGGET_PROGRESS_SCHEMA.parse(protocolProgress) + : { + kind: "beam-mem0-nuggets-v1" as const, + questionId: question.questionId, + rubricHash, + judgments: [] as NuggetJudgment[], + } + if ( + progress.questionId !== question.questionId || + progress.rubricHash !== rubricHash || + progress.judgments.length > rubric.length || + progress.judgments.some((judgment, index) => judgment.nugget !== rubric[index]) + ) { + throw new Error(`BEAM mem0 nugget progress identity mismatch for ${question.questionId}`) + } + + const judgments: NuggetJudgment[] = [...progress.judgments] + for (let index = judgments.length; index < rubric.length; index++) { + const nugget = rubric[index]! + const result = await runtime.generateStructured({ + system: BEAM_MEM0_JUDGE_SYSTEM_PROMPT, + prompt: buildBeamMem0NuggetPrompt({ + question: question.question, + nugget, + answer: hypothesis, + }), + schema: NUGGET_JUDGMENT_SCHEMA, + schemaName: "beam_mem0_nugget_judgment", + // GPT-5 does not accept a temperature override. Leaving it absent is + // part of this profile's recorded evaluator identity. + maxOutputTokens: MEM0_EVALUATOR_IDENTITY.maxOutputTokens, + maxAttempts: MEM0_EVALUATOR_IDENTITY.maxAttempts, + innerMaxRetries: MEM0_EVALUATOR_IDENTITY.innerMaxRetries, + timeoutMs: MEM0_EVALUATOR_IDENTITY.timeoutMs, + retryBackoffMs: MEM0_EVALUATOR_IDENTITY.retryBackoffMs, + transport: MEM0_EVALUATOR_IDENTITY.transport, + }) + judgments.push({ + nugget, + score: clampMem0NuggetScore(result.score), + reason: result.reason, + }) + await onProtocolProgress?.({ ...progress, judgments: [...judgments] }) + } + + const primaryScore = mean(judgments.map((judgment) => judgment.score)) + const passed = primaryScore >= BEAM_PASS_THRESHOLD + return { + questionId: question.questionId, + questionType: question.questionType, + primaryScore, + passed, + label: passed ? "pass" : "fail", + explanation: `Mem0-style BEAM nugget average: ${primaryScore.toFixed(4)}`, + metrics: { + nuggetAverage: primaryScore, + nuggetCount: judgments.length, + }, + details: { + evaluatorIdentity: MEM0_EVALUATOR_IDENTITY, + nuggetJudgments: judgments, + eventOrderingScoreUsed: 0, + }, + } + } + + aggregateQuality({ + questions, + evaluations, + }: Parameters[0]) { + if (evaluations.length === 0) { + throw new Error("Cannot aggregate an empty BEAM mem0 comparison run") + } + const questionById = new Map(questions.map((question) => [question.questionId, question])) + const seen = new Set() + const scoresByAbility = new Map() + for (const evaluation of evaluations) { + const question = questionById.get(evaluation.questionId) + if (!question || seen.has(evaluation.questionId)) { + throw new Error(`Invalid or duplicate evaluation ${evaluation.questionId}`) + } + if (evaluation.questionType !== question.questionType) { + throw new Error(`Evaluation type mismatch for ${evaluation.questionId}`) + } + if (!BEAM_ABILITY_IDS.includes(evaluation.questionType as BeamAbilityId)) { + throw new Error(`Unsupported BEAM ability ${evaluation.questionType}`) + } + if ( + !Number.isFinite(evaluation.primaryScore) || + evaluation.primaryScore < 0 || + evaluation.primaryScore > 1 + ) { + throw new Error(`Invalid nugget score for ${evaluation.questionId}`) + } + seen.add(evaluation.questionId) + const ability = evaluation.questionType as BeamAbilityId + const scores = scoresByAbility.get(ability) ?? [] + scores.push(evaluation.primaryScore) + scoresByAbility.set(ability, scores) + } + if (seen.size !== questions.length) { + throw new Error(`Missing BEAM mem0 evaluations: received ${seen.size}/${questions.length}`) + } + + const score = mean(evaluations.map((evaluation) => evaluation.primaryScore)) + const bySlice: Record> = {} + for (const ability of BEAM_ABILITY_IDS) { + const abilityScores = scoresByAbility.get(ability) + if (!abilityScores?.length) continue + bySlice[ability] = { + averageScore: mean(abilityScores), + questionCount: abilityScores.length, + } + } + return { + primaryMetric: { key: "mem0NuggetAverage", value: score, higherIsBetter: true }, + metrics: { + mem0NuggetAverage: score, + totalQuestions: evaluations.length, + }, + bySlice, + } + } +} + +export const beamMem0NuggetProtocol = new BeamMem0NuggetProtocol() diff --git a/src/protocols/beam-mem0/prompts.ts b/src/protocols/beam-mem0/prompts.ts new file mode 100644 index 0000000..a230e19 --- /dev/null +++ b/src/protocols/beam-mem0/prompts.ts @@ -0,0 +1,68 @@ +/** + * Prompts used by mem0's public BEAM comparison runner. This is deliberately + * separate from the paper profile so an experimental comparison cannot change + * the benchmark-author evaluator. + */ +export const BEAM_MEM0_JUDGE_SYSTEM_PROMPT = + "You are an expert evaluator assessing whether an AI assistant's response satisfies " + + "specific rubric criteria. You must be objective, fair, and consistent. " + + "Return ONLY valid JSON with the exact format requested." + +export const BEAM_MEM0_NUGGET_PROMPT_VERSION = "mem0-public-beam-nugget-v1" + +export function buildBeamMem0NuggetPrompt(input: { + question: string + nugget: string + answer: string +}): string { + return `Evaluate whether the following LLM response demonstrates compliance with the specified RUBRIC CRITERION. + +QUESTION: +${input.question} + +LLM RESPONSE: +${input.answer} + +RUBRIC CRITERION: +${input.nugget} + +SCORING GUIDELINES: + +First, determine whether the rubric criterion is a POSITIVE requirement (the response SHOULD include something) or a NEGATIVE constraint (the response SHOULD NOT include something). + +**For POSITIVE requirements** (response should contain, mention, or demonstrate something): +- **1.0 (Complete Compliance)**: The required element is present, accurate, and complete. The response fully and clearly satisfies the rubric criterion. +- **0.5 (Partial Compliance)**: The required element is partially present, has minor inaccuracies, or is incomplete. The core intent is present but not fully realized. +- **0.0 (No Compliance)**: The required element is missing, incorrect, or the response is entirely off-topic / non-responsive. + +**For NEGATIVE constraints** (response should NOT contain or should avoid something): +- **1.0 (Complete Compliance)**: The response is responsive to the question AND the prohibited element is absent. +- **0.5 (Partial Compliance)**: The response is responsive but contains a borderline or ambiguous reference to the prohibited element. +- **0.0 (No Compliance)**: The prohibited element is present in the response, OR the response is non-responsive (off-topic, refusal, empty). + +**Compound statement handling**: If the rubric criterion contains "and" or commas connecting multiple required elements: +- All elements present and correct = 1.0 +- Some (but not all) elements present and correct = 0.5 +- No elements present or correct = 0.0 + +EVALUATION RULES: +1. **Semantic tolerance**: Paraphrases and synonyms are acceptable. The response does not need to use the exact same words as the rubric. +2. **Numeric and date equivalence**: Treat equivalent representations as identical. "$68,000" = "68k" = "sixty-eight thousand dollars". "2 years" = "24 months". Prefer normalized comparison for numbers, currencies, dates, and durations. +3. **Case / punctuation / whitespace tolerance**: Differences in capitalization, punctuation, and whitespace must be ignored when comparing content. +4. **Hedging tolerance**: Do not penalize hedging language ("I think", "probably", "it seems"), passive voice, or verbosity if the substantive content satisfies the rubric criterion. +5. **Style neutrality**: Do not penalize for tone, formatting, or length unless the rubric criterion specifically requires a particular format. +6. **Responsiveness**: If the LLM response is completely off-topic or refuses to answer, score 0.0 for all criteria. +7. **Independence**: Evaluate this criterion in isolation — do not consider other rubric items. +8. **Specificity matters**: Vague or generic answers that could apply to any question score lower than specific, detailed answers. + +STEP-BY-STEP EVALUATION: +Follow these steps in order: +1. **Understand the Requirement**: Read the rubric criterion and classify it as a positive requirement or a negative constraint. +2. **Parse Compound Statements**: If the criterion contains multiple sub-requirements joined by "and" or commas, identify each element separately. +3. **Check Compliance**: Compare the LLM response against each element, applying the tolerance rules above (semantic, numeric, case, hedging). +4. **Assign Score**: Use the appropriate scoring table (positive or negative) and compound-statement rule to determine the score. +5. **Provide Reasoning**: Write a concise explanation referencing which elements were or were not satisfied. + +Return your evaluation as a JSON object with exactly two fields: +{"score": <0.0 or 0.5 or 1.0>, "reason": ""}` +} diff --git a/src/protocols/beam-paper/event-ordering.ts b/src/protocols/beam-paper/event-ordering.ts new file mode 100644 index 0000000..06f1f7b --- /dev/null +++ b/src/protocols/beam-paper/event-ordering.ts @@ -0,0 +1,378 @@ +export const BEAM_KENDALL_TAU_IMPLEMENTATION = "scipy-compatible-tau-b-fail-closed-v1" +export const BEAM_EVENT_EXTRACTION_VERSION = "authors-literal-newline-split-v2" +export const BEAM_EVENT_ORDERING_SCORING_VERSION = + "paper-table-tau-norm-primary-authors-f1-product-diagnostic-v1" + +export interface BeamEventAlignmentAttempt { + predictedIndex: number + referenceIndex: number + predictedEvent: string + referenceEvent: string + equivalent: boolean +} + +export interface BeamEventAlignment { + predictedIndex: number + predictedEvent: string + referenceIndex?: number + referenceEvent?: string +} + +export interface BeamEventRankItem { + id: string + kind: "reference" | "unmatched-prediction" + referenceIndex?: number + predictedIndex?: number + event: string +} + +export interface BeamEventRankVectors { + union: BeamEventRankItem[] + referenceRanks: number[] + predictedRanks: number[] + bottomTieRank: number +} + +export interface KendallTauBResult { + concordantPairs: number + discordantPairs: number + tiesOnlyInReference: number + tiesOnlyInPrediction: number + tiesInBoth: number + denominator: number + tauB: number + degenerate: boolean +} + +export interface BeamEventOrderingScore { + referenceEvents: string[] + predictedEvents: string[] + canonicalPredictedEvents: string[] + alignmentAttempts: BeamEventAlignmentAttempt[] + alignments: BeamEventAlignment[] + missingReferenceEvents: Array<{ referenceIndex: number; event: string }> + unmatchedPredictedEvents: Array<{ predictedIndex: number; event: string }> + rankVectors: BeamEventRankVectors + kendall: KendallTauBResult + normalizedKendallTauB: number + matchedCount: number + precision: number + recall: number + f1: number + finalScore: number +} + +export type BeamEventEquivalence = (input: { + predictedIndex: number + referenceIndex: number + predictedEvent: string + referenceEvent: string +}) => Promise + +export function extractBeamPredictedEvents(answer: string): string[] { + // Exact authors' scorer semantics: llm_response.split("\n"). Blank lines, + // surrounding whitespace, duplicates, and a trailing empty line all remain + // score-bearing events. + return answer.split("\n") +} + +export async function alignBeamEvents( + referenceEvents: readonly string[], + predictedEvents: readonly string[], + equivalent: BeamEventEquivalence +): Promise<{ + attempts: BeamEventAlignmentAttempt[] + alignments: BeamEventAlignment[] +}> { + const usedReferenceIndices = new Set() + const attempts: BeamEventAlignmentAttempt[] = [] + const alignments: BeamEventAlignment[] = [] + + for (let predictedIndex = 0; predictedIndex < predictedEvents.length; predictedIndex++) { + const predictedEvent = predictedEvents[predictedIndex]! + let matchedReferenceIndex: number | undefined + + for (let referenceIndex = 0; referenceIndex < referenceEvents.length; referenceIndex++) { + if (usedReferenceIndices.has(referenceIndex)) continue + + const referenceEvent = referenceEvents[referenceIndex]! + const isEquivalent = await equivalent({ + predictedIndex, + referenceIndex, + predictedEvent, + referenceEvent, + }) + + attempts.push({ + predictedIndex, + referenceIndex, + predictedEvent, + referenceEvent, + equivalent: isEquivalent, + }) + + if (isEquivalent) { + matchedReferenceIndex = referenceIndex + usedReferenceIndices.add(referenceIndex) + break + } + } + + alignments.push({ + predictedIndex, + predictedEvent, + ...(matchedReferenceIndex === undefined + ? {} + : { + referenceIndex: matchedReferenceIndex, + referenceEvent: referenceEvents[matchedReferenceIndex]!, + }), + }) + } + + return { attempts, alignments } +} + +export function buildBeamEventRankVectors( + referenceEvents: readonly string[], + predictedEvents: readonly string[], + alignments: readonly BeamEventAlignment[] +): BeamEventRankVectors { + if (alignments.length !== predictedEvents.length) { + throw new Error( + `Expected one event alignment per prediction (${predictedEvents.length}), got ${alignments.length}` + ) + } + + const matchedPredictionByReference = new Map() + const alignmentByPrediction = new Map() + + for (const alignment of alignments) { + if (alignment.predictedIndex < 0 || alignment.predictedIndex >= predictedEvents.length) { + throw new Error(`Invalid predicted event index: ${alignment.predictedIndex}`) + } + if (alignmentByPrediction.has(alignment.predictedIndex)) { + throw new Error(`Predicted event ${alignment.predictedIndex} was aligned more than once`) + } + if (predictedEvents[alignment.predictedIndex] !== alignment.predictedEvent) { + throw new Error( + `Predicted event ${alignment.predictedIndex} content does not match alignment` + ) + } + alignmentByPrediction.set(alignment.predictedIndex, alignment) + + if (alignment.referenceIndex !== undefined) { + if (alignment.referenceIndex < 0 || alignment.referenceIndex >= referenceEvents.length) { + throw new Error(`Invalid reference event index: ${alignment.referenceIndex}`) + } + if (referenceEvents[alignment.referenceIndex] !== alignment.referenceEvent) { + throw new Error( + `Reference event ${alignment.referenceIndex} content does not match alignment` + ) + } + if (matchedPredictionByReference.has(alignment.referenceIndex)) { + throw new Error(`Reference event ${alignment.referenceIndex} was matched more than once`) + } + matchedPredictionByReference.set(alignment.referenceIndex, alignment.predictedIndex) + } + } + + for (let predictedIndex = 0; predictedIndex < predictedEvents.length; predictedIndex++) { + if (!alignmentByPrediction.has(predictedIndex)) { + throw new Error(`Missing alignment for predicted event ${predictedIndex}`) + } + } + + const canonicalPredictedEvents = predictedEvents.map((predictedEvent, predictedIndex) => { + const referenceIndex = alignmentByPrediction.get(predictedIndex)!.referenceIndex + return referenceIndex === undefined ? predictedEvent : referenceEvents[referenceIndex]! + }) + + // Exact authors' event_ordering_score semantics: + // union = list(dict.fromkeys(reference_canon + system_canon)) + // ranks = {item: i + 1 for i, item in enumerate(sequence)} + // The union keeps the first duplicate, while the rank map keeps the last. + const unionEvents = [...new Set([...referenceEvents, ...canonicalPredictedEvents])] + const union: BeamEventRankItem[] = unionEvents.map((event) => { + const referenceIndex = referenceEvents.indexOf(event) + if (referenceIndex >= 0) { + return { id: `reference:${referenceIndex}`, kind: "reference", referenceIndex, event } + } + const predictedIndex = canonicalPredictedEvents.indexOf(event) + return { + id: `prediction:${predictedIndex}`, + kind: "unmatched-prediction", + predictedIndex, + event, + } + }) + + const bottomTieRank = union.length + 1 + const toRanks = (sequence: readonly string[]) => { + const ranks = new Map() + sequence.forEach((event, index) => ranks.set(event, index + 1)) + return unionEvents.map((event) => ranks.get(event) ?? bottomTieRank) + } + + return { + union, + referenceRanks: toRanks(referenceEvents), + predictedRanks: toRanks(canonicalPredictedEvents), + bottomTieRank, + } +} + +export function computeKendallTauB( + referenceRanks: readonly number[], + predictedRanks: readonly number[] +): KendallTauBResult { + if (referenceRanks.length !== predictedRanks.length) { + throw new Error( + `Kendall tau rank vectors must have equal length (${referenceRanks.length} !== ${predictedRanks.length})` + ) + } + if ( + referenceRanks.some((rank) => !Number.isFinite(rank)) || + predictedRanks.some((rank) => !Number.isFinite(rank)) + ) { + throw new Error("Kendall tau rank vectors must contain only finite numbers") + } + + let concordantPairs = 0 + let discordantPairs = 0 + let tiesOnlyInReference = 0 + let tiesOnlyInPrediction = 0 + let tiesInBoth = 0 + + for (let left = 0; left < referenceRanks.length; left++) { + for (let right = left + 1; right < referenceRanks.length; right++) { + const referenceSign = Math.sign(referenceRanks[left]! - referenceRanks[right]!) + const predictedSign = Math.sign(predictedRanks[left]! - predictedRanks[right]!) + + if (referenceSign === 0 && predictedSign === 0) { + tiesInBoth++ + } else if (referenceSign === 0) { + tiesOnlyInReference++ + } else if (predictedSign === 0) { + tiesOnlyInPrediction++ + } else if (referenceSign === predictedSign) { + concordantPairs++ + } else { + discordantPairs++ + } + } + } + + const denominator = Math.sqrt( + (concordantPairs + discordantPairs + tiesOnlyInReference) * + (concordantPairs + discordantPairs + tiesOnlyInPrediction) + ) + + const degenerate = denominator === 0 + // scipy.stats.kendalltau (used by the pinned authors' evaluator) returns NaN + // when tau-b has no denominator. Preserve that semantic here; the caller + // fails closed instead of inventing a perfect or fully-incorrect score. + const tauB = degenerate ? Number.NaN : (concordantPairs - discordantPairs) / denominator + + return { + concordantPairs, + discordantPairs, + tiesOnlyInReference, + tiesOnlyInPrediction, + tiesInBoth, + denominator, + tauB, + degenerate, + } +} + +export function scoreAlignedBeamEvents(input: { + referenceEvents: readonly string[] + predictedEvents: readonly string[] + attempts?: readonly BeamEventAlignmentAttempt[] + alignments: readonly BeamEventAlignment[] +}): BeamEventOrderingScore { + const rankVectors = buildBeamEventRankVectors( + input.referenceEvents, + input.predictedEvents, + input.alignments + ) + const kendall = computeKendallTauB(rankVectors.referenceRanks, rankVectors.predictedRanks) + if (!Number.isFinite(kendall.tauB)) { + throw new Error( + "BEAM paper Kendall tau-b is undefined for degenerate rank vectors; refusing to invent a score" + ) + } + const normalizedKendallTauB = (kendall.tauB + 1) / 2 + const canonicalPredictedEvents = input.predictedEvents.map((predictedEvent, predictedIndex) => { + const alignment = input.alignments.find( + (candidate) => candidate.predictedIndex === predictedIndex + ) + return alignment?.referenceIndex === undefined + ? predictedEvent + : input.referenceEvents[alignment.referenceIndex]! + }) + const referenceSet = new Set(input.referenceEvents) + const predictedSet = new Set(canonicalPredictedEvents) + const matchedCount = [...referenceSet].filter((event) => predictedSet.has(event)).length + const falsePositiveCount = canonicalPredictedEvents.filter( + (event) => !referenceSet.has(event) + ).length + const falseNegativeCount = input.referenceEvents.filter( + (event) => !predictedSet.has(event) + ).length + const missingReferenceEvents = input.referenceEvents.flatMap((event, referenceIndex) => + predictedSet.has(event) ? [] : [{ referenceIndex, event }] + ) + const unmatchedPredictedEvents = input.alignments.flatMap((alignment) => + alignment.referenceIndex === undefined + ? [{ predictedIndex: alignment.predictedIndex, event: alignment.predictedEvent }] + : [] + ) + const precision = + matchedCount + falsePositiveCount > 0 ? matchedCount / (matchedCount + falsePositiveCount) : 0 + const recall = + matchedCount + falseNegativeCount > 0 ? matchedCount / (matchedCount + falseNegativeCount) : 0 + const f1 = precision + recall > 0 ? (2 * precision * recall) / (precision + recall) : 0 + // The pinned authors helper calculates this product, but report_results.py + // uses tau_norm for the published Table 1 value. Retain the product only as + // an auditable diagnostic; the protocol selects the paper-table score. + const finalScore = normalizedKendallTauB * f1 + + return { + referenceEvents: [...input.referenceEvents], + predictedEvents: [...input.predictedEvents], + canonicalPredictedEvents, + alignmentAttempts: [...(input.attempts ?? [])], + alignments: [...input.alignments], + missingReferenceEvents, + unmatchedPredictedEvents, + rankVectors, + kendall, + normalizedKendallTauB, + matchedCount, + precision, + recall, + f1, + finalScore, + } +} + +export async function evaluateBeamEventOrdering(input: { + referenceEvents: readonly string[] + predictedEvents: readonly string[] + equivalent: BeamEventEquivalence +}): Promise { + const { attempts, alignments } = await alignBeamEvents( + input.referenceEvents, + input.predictedEvents, + input.equivalent + ) + + return scoreAlignedBeamEvents({ + referenceEvents: input.referenceEvents, + predictedEvents: input.predictedEvents, + attempts, + alignments, + }) +} diff --git a/src/protocols/beam-paper/index.ts b/src/protocols/beam-paper/index.ts new file mode 100644 index 0000000..b399cc2 --- /dev/null +++ b/src/protocols/beam-paper/index.ts @@ -0,0 +1,916 @@ +import { z } from "zod" +import { + BEAM_ANSWER_FORMATTER_IMPLEMENTATION_HASH, + BEAM_EVENT_ORDERING_ANSWER_FORMAT_VERSION, + buildBeamAnswerPrompt, +} from "../../prompts/beam" +import type { BenchmarkProtocol, ProtocolIdentity, QuestionEvaluation } from "../../types/protocol" +import type { UnifiedQuestion, UnifiedSearchResult, UnifiedSession } from "../../types/unified" +import { sha256Text, stableSha256 } from "../../utils/stable" +import { + BEAM_EVENT_EXTRACTION_VERSION, + BEAM_EVENT_ORDERING_SCORING_VERSION, + BEAM_KENDALL_TAU_IMPLEMENTATION, + alignBeamEvents, + buildBeamEventRankVectors, + computeKendallTauB, + evaluateBeamEventOrdering, + extractBeamPredictedEvents, + scoreAlignedBeamEvents, +} from "./event-ordering" +import { + BEAM_EVENT_EQUIVALENCE_PROMPT_VERSION, + BEAM_EVENT_EQUIVALENCE_SYSTEM_PROMPT, + BEAM_EVENT_EQUIVALENCE_USER_PROMPT, + BEAM_NUGGET_JUDGE_PROMPT, + BEAM_NUGGET_JUDGE_PROMPT_VERSION, + buildBeamEventEquivalencePrompt, + buildBeamPaperNuggetPrompt, +} from "./prompts" + +export * from "./event-ordering" +export * from "./prompts" + +export const BEAM_PAPER_PROTOCOL_ID = "beam-paper" +export const BEAM_PAPER_PROTOCOL_VERSION = "1.5.0" +export const BEAM_PAPER_ID = "arXiv:2510.27246" +export const BEAM_PAPER_REVISION = "v2" +export const BEAM_PAPER_PDF_SHA256 = + "8ae85b00eb0f93f0717edb082f5471716f6c757670d7157dc5ba94df01fbb303" +export const BEAM_REFERENCE_REPOSITORY = "mohammadtavakoli78/BEAM" +export const BEAM_REFERENCE_COMMIT = "3e12035532eb85768f1a7cd779832b650c4b2ef9" +export const BEAM_PASS_THRESHOLD = 0.5 +export const BEAM_RETRIEVAL_TOP_K_VALUES = [5, 10, 15, 20] as const +export type BeamRetrievalTopK = (typeof BEAM_RETRIEVAL_TOP_K_VALUES)[number] + +export const BEAM_ABILITY_IDS = [ + "abstention", + "contradiction_resolution", + "event_ordering", + "information_extraction", + "instruction_following", + "knowledge_update", + "multi_session_reasoning", + "preference_following", + "summarization", + "temporal_reasoning", +] as const +export type BeamAbilityId = (typeof BEAM_ABILITY_IDS)[number] + +export const BEAM_OFFICIAL_TIER_COUNTS = { + "1M": { questions: 700, questionsPerAbility: 70 }, + "10M": { questions: 200, questionsPerAbility: 20 }, +} as const + +const BEAM_ABILITY_ID_SET = new Set(BEAM_ABILITY_IDS) + +function createBeamNuggetJudgmentSchema() { + return z + .object({ + score: z.union([z.literal(0), z.literal(0.5), z.literal(1)]), + reason: z.string().trim().min(1), + }) + .strict() +} + +function createBeamEventEquivalenceSchema() { + return z + .object({ + answer: z.enum(["YES", "NO"]), + }) + .strict() +} + +function createBeamNuggetProgressSchema() { + return z + .object({ + kind: z.literal("beam-nuggets-v1"), + questionId: z.string().min(1), + rubricHash: z.string().regex(/^[a-f0-9]{64}$/), + judgments: z + .array( + z + .object({ + nugget: z.string().min(1), + score: z.union([z.literal(0), z.literal(0.5), z.literal(1)]), + reason: z.string().min(1), + }) + .strict() + ) + .default([]), + }) + .strict() +} + +function createBeamEventProgressSchema() { + return z + .object({ + kind: z.literal("beam-event-equivalence-v1"), + questionId: z.string().min(1), + rubricHash: z.string().regex(/^[a-f0-9]{64}$/), + predictedEventsHash: z.string().regex(/^[a-f0-9]{64}$/), + judgments: z.array( + z + .object({ + predictedIndex: z.number().int().nonnegative(), + referenceIndex: z.number().int().nonnegative(), + // The pinned authors scorer preserves blank lines from split("\n"). + predictedEvent: z.string(), + referenceEvent: z.string().min(1), + equivalent: z.boolean(), + }) + .strict() + ), + }) + .strict() +} + +export const BEAM_NUGGET_JUDGMENT_SCHEMA = createBeamNuggetJudgmentSchema() +export const BEAM_EVENT_EQUIVALENCE_SCHEMA = createBeamEventEquivalenceSchema() +const BEAM_NUGGET_PROGRESS_SCHEMA = createBeamNuggetProgressSchema() +const BEAM_EVENT_PROGRESS_SCHEMA = createBeamEventProgressSchema() + +export type BeamNuggetJudgmentOutput = z.infer + +export interface BeamNuggetJudgment extends BeamNuggetJudgmentOutput { + nugget: string +} + +export interface BeamEvaluatorIdentity { + protocolId: typeof BEAM_PAPER_PROTOCOL_ID + paperId: typeof BEAM_PAPER_ID + paperRevision: typeof BEAM_PAPER_REVISION + paperPdfSha256: typeof BEAM_PAPER_PDF_SHA256 + referenceRepository: typeof BEAM_REFERENCE_REPOSITORY + referenceCommit: typeof BEAM_REFERENCE_COMMIT + nuggetPromptVersion: string + nuggetPromptSha256: string + eventEquivalencePromptVersion: string + eventEquivalencePromptSha256: string + evaluatorImplementationSha256: string + judgeProvider: "openai" + judgeModel: "gpt-4.1-mini" + structuredOutputSchemaVersion: string + structuredOutputSchemaSha256: string + structuredOutputMode: string + kendallTauImplementation: string + eventExtractionVersion: string + eventOrderingScoringVersion: string + temperature: 0 + maxOutputTokens: number + maxAttempts: number + timeoutMs: number + retryPolicy: string +} + +export interface BeamPaperProtocolConfig { + retrievalTopK?: number +} + +function isBeamRetrievalTopK(value: number): value is BeamRetrievalTopK { + return (BEAM_RETRIEVAL_TOP_K_VALUES as readonly number[]).includes(value) +} + +function getRubric(question: UnifiedQuestion): string[] { + const rubric = question.metadata?.rubric + if ( + !Array.isArray(rubric) || + rubric.length === 0 || + rubric.some((item) => typeof item !== "string" || item.trim().length === 0) + ) { + throw new Error( + `BEAM question ${question.questionId || ""} must have a non-empty string rubric` + ) + } + return rubric as string[] +} + +function getDocumentDate(session: UnifiedSession): string | undefined { + const documentDate = session.metadata?.documentDate + const legacyDate = session.metadata?.date + if (documentDate !== undefined && legacyDate !== undefined && documentDate !== legacyDate) { + throw new Error( + `BEAM session ${session.sessionId} has conflicting document dates ${String(documentDate)} and ${String(legacyDate)}` + ) + } + const value = documentDate ?? legacyDate + if (value === undefined) return undefined + if (typeof value !== "string") { + throw new Error(`BEAM session ${session.sessionId} document date must be a YYYY-MM-DD string`) + } + const match = value.match(/^(\d{4})-(\d{2})-(\d{2})$/) + if (!match) { + throw new Error(`BEAM session ${session.sessionId} has invalid document date ${value}`) + } + const year = Number(match[1]) + const month = Number(match[2]) + const day = Number(match[3]) + const date = new Date(Date.UTC(year, month - 1, day)) + if ( + date.getUTCFullYear() !== year || + date.getUTCMonth() + 1 !== month || + date.getUTCDate() !== day + ) { + throw new Error(`BEAM session ${session.sessionId} has invalid document date ${value}`) + } + return value +} + +function renderBeamTranscript(session: UnifiedSession): string { + return session.messages + .map((message) => `[${message.role.toUpperCase()}]\n${message.content}`) + .join("\n\n") +} + +function createSessionDateMap( + sessions: readonly UnifiedSession[], + results: readonly UnifiedSearchResult[] +): Map { + const dates = new Map() + for (const session of sessions) { + const date = getDocumentDate(session) + if (date) dates.set(session.sessionId, date) + } + for (const result of results) { + if (result.sessionId && result.documentDate && !dates.has(result.sessionId)) { + dates.set(result.sessionId, result.documentDate) + } + } + return dates +} + +function toBeamPromptEvidence(results: readonly UnifiedSearchResult[], dates: Map) { + return results.map((result) => ({ + content: result.text, + metadata: (() => { + const promptSessionId = + result.sessionId ?? (result.documentDate ? `normalized-result:${result.id}` : undefined) + if (!promptSessionId) return undefined + if (result.documentDate) dates.set(promptSessionId, result.documentDate) + return { sessionId: promptSessionId } + })(), + })) +} + +function mean(values: readonly number[]): number { + if (values.length === 0) throw new Error("Cannot average an empty list") + return values.reduce((sum, value) => sum + value, 0) / values.length +} + +function createProtocolIdentity( + topK: BeamRetrievalTopK, + protocol: BeamPaperProtocol +): ProtocolIdentity { + const methodHashes = { + validateQuestion: sha256Text(protocol.validateQuestion.toString()), + createIngestionPlan: sha256Text(protocol.createIngestionPlan.toString()), + createRetrievalPlan: sha256Text(protocol.createRetrievalPlan.toString()), + createAnswerPlan: sha256Text(protocol.createAnswerPlan.toString()), + evaluateQuestion: sha256Text(protocol.evaluateQuestion.toString()), + aggregateQuality: sha256Text(protocol.aggregateQuality.toString()), + } + const helperHashes = { + getRubric: sha256Text(getRubric.toString()), + getDocumentDate: sha256Text(getDocumentDate.toString()), + renderTranscript: sha256Text(renderBeamTranscript.toString()), + createSessionDateMap: sha256Text(createSessionDateMap.toString()), + toPromptEvidence: sha256Text(toBeamPromptEvidence.toString()), + arithmeticMean: sha256Text(mean.toString()), + } + const orchestration = { methodHashes, helperHashes } + const ingestion = { + policy: "one-document-per-session-date-prefixed-transcript-v1", + customId: "sessionId", + metadata: "sessionId-and-optional-documentDate", + messageProjection: "ordered-role-and-content", + missingDatePolicy: "omit-document-date-prefix-and-metadata", + executionPolicy: protocol.ingestionExecutionPolicy, + implementationSha256: stableSha256({ + createIngestionPlan: methodHashes.createIngestionPlan, + getDocumentDate: helperHashes.getDocumentDate, + renderTranscript: helperHashes.renderTranscript, + }), + } + const retrieval = { + policy: "paper-top-k", + requestedTopK: topK, + answerCutoff: topK, + threshold: 0, + allowedValues: BEAM_RETRIEVAL_TOP_K_VALUES, + implementationSha256: methodHashes.createRetrievalPlan, + } + const answerProbe = buildBeamAnswerPrompt("", [], new Map()) + const eventOrderingAnswerProbe = buildBeamAnswerPrompt( + "", + [], + new Map(), + "event-ordering-lines" + ) + const answer = { + formatter: "buildBeamAnswerPrompt", + formatterVersion: "beam-normalized-evidence-v2", + basePromptSha256: sha256Text(answerProbe), + eventOrderingAnswerFormatVersion: BEAM_EVENT_ORDERING_ANSWER_FORMAT_VERSION, + eventOrderingPromptSha256: sha256Text(eventOrderingAnswerProbe), + formatterImplementationSha256: BEAM_ANSWER_FORMATTER_IMPLEMENTATION_HASH, + answerPlanImplementationSha256: stableSha256({ + createAnswerPlan: methodHashes.createAnswerPlan, + createSessionDateMap: helperHashes.createSessionDateMap, + toPromptEvidence: helperHashes.toPromptEvidence, + }), + datedEvidenceProbeSha256: sha256Text( + buildBeamAnswerPrompt( + "", + [{ memory: "", metadata: { sessionId: "" } }], + new Map([["", "2000-01-01"]]) + ) + ), + } + const auxiliaryRetrievalEvaluation = { + policy: protocol.auxiliaryRetrievalEvaluation, + protocolStatus: "disabled-because-not-defined-by-beam-paper", + } + const evaluator = { + paperEvaluator: BEAM_EVALUATOR_IDENTITY, + auxiliaryRetrievalEvaluation, + } + const aggregation = { + primaryMetric: "beamScore", + abilityIds: BEAM_ABILITY_IDS, + abilityWeighting: "equal-macro", + officialTierCounts: BEAM_OFFICIAL_TIER_COUNTS, + passThreshold: BEAM_PASS_THRESHOLD, + partialMetricName: "beamScorePartial", + combinedTierPolicy: "equal-tier-macro-secondary-not-paper-score-v1", + implementationSha256: stableSha256({ + aggregateQuality: protocol.aggregateQuality.toString(), + arithmeticMean: mean.toString(), + }), + } + const implementation = { + protocol: BEAM_PAPER_PROTOCOL_ID, + version: BEAM_PAPER_PROTOCOL_VERSION, + ingestion, + evaluator, + answer, + aggregation, + orchestration, + } + + return { + id: BEAM_PAPER_PROTOCOL_ID, + version: BEAM_PAPER_PROTOCOL_VERSION, + configFingerprint: stableSha256({ ingestion, retrieval, answer, evaluator, aggregation }), + implementationFingerprint: stableSha256(implementation), + ingestionPolicyHash: stableSha256(ingestion), + retrievalPolicyHash: stableSha256(retrieval), + answerPromptHash: stableSha256(answer), + evaluatorHash: stableSha256(evaluator), + aggregationHash: stableSha256(aggregation), + details: { + paperProfile: { + paperId: BEAM_PAPER_ID, + paperRevision: BEAM_PAPER_REVISION, + paperPdfSha256: BEAM_PAPER_PDF_SHA256, + referenceRepository: BEAM_REFERENCE_REPOSITORY, + referenceCommit: BEAM_REFERENCE_COMMIT, + }, + ingestionPolicy: ingestion, + retrievalPolicy: retrieval, + answerPrompt: answer, + evaluatorIdentity: BEAM_EVALUATOR_IDENTITY, + auxiliaryRetrievalEvaluation, + aggregation, + orchestrationImplementation: orchestration, + }, + } +} + +export class BeamPaperProtocol implements BenchmarkProtocol { + readonly auxiliaryRetrievalEvaluation = "disabled" as const + readonly ingestionExecutionPolicy = { + readinessBarrier: "after-each-document", + processingMode: "instant", + } as const + readonly requiredJudge = { + provider: "openai", + modelId: "gpt-4.1-mini", + modelAlias: "gpt-4.1-mini", + } + readonly retrievalTopK: BeamRetrievalTopK + readonly identity: ProtocolIdentity + + constructor(config: BeamPaperProtocolConfig = {}) { + const retrievalTopK = config.retrievalTopK ?? 5 + if (!isBeamRetrievalTopK(retrievalTopK)) { + throw new Error( + `BEAM retrieval Top-K must be one of ${BEAM_RETRIEVAL_TOP_K_VALUES.join(", ")}; got ${retrievalTopK}` + ) + } + this.retrievalTopK = retrievalTopK + this.identity = createProtocolIdentity(retrievalTopK, this) + } + + validateQuestion(question: UnifiedQuestion): void { + if (!question.questionId || !question.questionId.trim()) { + throw new Error("BEAM question must have a non-empty question ID") + } + if (!question.question || !question.question.trim()) { + throw new Error(`BEAM question ${question.questionId} must have non-empty question text`) + } + if (!BEAM_ABILITY_ID_SET.has(question.questionType)) { + throw new Error( + `BEAM question ${question.questionId} has unsupported ability ${JSON.stringify(question.questionType)}` + ) + } + const rubric = getRubric(question) + if (question.questionType === "event_ordering" && rubric.length < 2) { + throw new Error( + `BEAM event-ordering question ${question.questionId} must have at least two reference events for Kendall tau-b` + ) + } + } + + createIngestionPlan({ + question, + sessions, + }: Parameters[0]) { + this.validateQuestion(question) + return sessions.map((session) => { + if ( + typeof session.sessionId !== "string" || + !session.sessionId.trim() || + !Array.isArray(session.messages) || + session.messages.length !== 2 || + session.messages[0]?.role !== "user" || + session.messages[1]?.role !== "assistant" || + session.messages.some( + (message) => typeof message.content !== "string" || !message.content.trim() + ) + ) { + throw new Error( + `BEAM question ${question.questionId} session ${session.sessionId || ""} must contain exactly one non-empty user message followed by one non-empty assistant message` + ) + } + const transcript = renderBeamTranscript(session) + const documentDate = getDocumentDate(session) + return { + customId: session.sessionId, + content: documentDate ? `DOCUMENT_DATE: ${documentDate}\n\n${transcript}` : transcript, + metadata: { + sessionId: session.sessionId, + ...(documentDate ? { documentDate } : {}), + }, + // All provider adapters receive the same BEAM evidence. Per-message + // source anchors/speaker labels remain dataset provenance, but must not + // leak into extraction-based providers when the approved document is + // only the dated role/content transcript. + messages: session.messages.map(({ role, content }) => ({ role, content })), + } + }) + } + + createRetrievalPlan({ question }: Parameters[0]) { + this.validateQuestion(question) + return { + query: question.question, + requestedTopK: this.retrievalTopK, + answerCutoff: this.retrievalTopK, + threshold: 0, + } + } + + createAnswerPlan({ + question, + sessions, + results, + retrieval, + }: Parameters[0]) { + this.validateQuestion(question) + if (retrieval.answerCutoff !== this.retrievalTopK) { + throw new Error( + `BEAM answer cutoff drifted from configured Top-K (${retrieval.answerCutoff} !== ${this.retrievalTopK})` + ) + } + const evidence = results.slice(0, retrieval.answerCutoff) + const dates = createSessionDateMap(sessions, evidence) + const promptEvidence = toBeamPromptEvidence(evidence, dates) + const answerFormat = + question.questionType === "event_ordering" ? "event-ordering-lines" : "default" + + return { + request: { + prompt: buildBeamAnswerPrompt(question.question, promptEvidence, dates, answerFormat), + }, + baseRequest: { + prompt: buildBeamAnswerPrompt(question.question, [], dates, answerFormat), + }, + answerEvidenceCount: evidence.length, + } + } + + async evaluateQuestion( + { + question, + hypothesis, + protocolProgress, + onProtocolProgress, + }: Parameters[0], + runtime: Parameters[1] + ): Promise { + this.validateQuestion(question) + const rubric = getRubric(question) + + if (question.questionType === "event_ordering") { + const predictedEvents = extractBeamPredictedEvents(hypothesis) + const rubricHash = stableSha256(rubric) + const predictedEventsHash = stableSha256(predictedEvents) + const eventProgress = protocolProgress + ? BEAM_EVENT_PROGRESS_SCHEMA.parse(protocolProgress) + : { + kind: "beam-event-equivalence-v1" as const, + questionId: question.questionId, + rubricHash, + predictedEventsHash, + judgments: [], + } + if ( + eventProgress.questionId !== question.questionId || + eventProgress.rubricHash !== rubricHash || + eventProgress.predictedEventsHash !== predictedEventsHash + ) { + throw new Error(`BEAM event progress identity mismatch for ${question.questionId}`) + } + const seenEventPairs = new Set() + for (const judgment of eventProgress.judgments) { + if ( + rubric[judgment.referenceIndex] !== judgment.referenceEvent || + predictedEvents[judgment.predictedIndex] !== judgment.predictedEvent + ) { + throw new Error(`BEAM event progress content mismatch for ${question.questionId}`) + } + const key = `${judgment.predictedIndex}:${judgment.referenceIndex}` + if (seenEventPairs.has(key)) { + throw new Error(`BEAM event progress contains duplicate pair ${key}`) + } + seenEventPairs.add(key) + } + const eventScore = await evaluateBeamEventOrdering({ + referenceEvents: rubric, + predictedEvents, + equivalent: async ({ predictedIndex, referenceIndex, referenceEvent, predictedEvent }) => { + const existing = eventProgress.judgments.find( + (judgment) => + judgment.predictedIndex === predictedIndex && + judgment.referenceIndex === referenceIndex + ) + if (existing) return existing.equivalent + + const result = await runtime.generateStructured({ + system: BEAM_EVENT_EQUIVALENCE_SYSTEM_PROMPT, + prompt: buildBeamEventEquivalencePrompt({ referenceEvent, predictedEvent }), + schema: BEAM_EVENT_EQUIVALENCE_SCHEMA, + schemaName: "beam_event_equivalence", + temperature: BEAM_EVALUATOR_IDENTITY.temperature, + maxOutputTokens: BEAM_EVALUATOR_IDENTITY.maxOutputTokens, + maxAttempts: BEAM_EVALUATOR_IDENTITY.maxAttempts, + timeoutMs: BEAM_EVALUATOR_IDENTITY.timeoutMs, + }) + const equivalent = result.answer === "YES" + eventProgress.judgments.push({ + predictedIndex, + referenceIndex, + predictedEvent, + referenceEvent, + equivalent, + }) + await onProtocolProgress?.({ ...eventProgress, judgments: [...eventProgress.judgments] }) + return equivalent + }, + }) + // Paper Section 2.4 names Kendall tau-b, and the pinned authors' + // report_results.py reads tau_norm for Table 1. Their helper's + // tau_norm * F1 value remains diagnostic only. + const primaryScore = eventScore.normalizedKendallTauB + const passed = primaryScore >= BEAM_PASS_THRESHOLD + + return { + questionId: question.questionId, + questionType: question.questionType, + primaryScore, + passed, + label: passed ? "pass" : "fail", + explanation: `BEAM event-ordering score (normalized Kendall tau-b): ${primaryScore.toFixed(4)}`, + metrics: { + kendallTauB: eventScore.kendall.tauB, + normalizedKendallTauB: eventScore.normalizedKendallTauB, + eventPrecision: eventScore.precision, + eventRecall: eventScore.recall, + eventF1: eventScore.f1, + authorsHelperFinalScore: eventScore.finalScore, + }, + details: { + evaluatorIdentity: BEAM_EVALUATOR_IDENTITY, + eventOrdering: eventScore, + }, + } + } + + const rubricHash = stableSha256(rubric) + const nuggetProgress = protocolProgress + ? BEAM_NUGGET_PROGRESS_SCHEMA.parse(protocolProgress) + : { + kind: "beam-nuggets-v1" as const, + questionId: question.questionId, + rubricHash, + judgments: [] as BeamNuggetJudgment[], + } + if ( + nuggetProgress.questionId !== question.questionId || + nuggetProgress.rubricHash !== rubricHash || + nuggetProgress.judgments.length > rubric.length || + nuggetProgress.judgments.some((judgment, index) => judgment.nugget !== rubric[index]) + ) { + throw new Error(`BEAM nugget progress identity mismatch for ${question.questionId}`) + } + + const nuggetJudgments: BeamNuggetJudgment[] = [...nuggetProgress.judgments] + for (let index = nuggetJudgments.length; index < rubric.length; index++) { + const nugget = rubric[index]! + const result = await runtime.generateStructured({ + prompt: buildBeamPaperNuggetPrompt({ + question: question.question, + nugget, + answer: hypothesis, + }), + schema: BEAM_NUGGET_JUDGMENT_SCHEMA, + schemaName: "beam_nugget_judgment", + temperature: BEAM_EVALUATOR_IDENTITY.temperature, + maxOutputTokens: BEAM_EVALUATOR_IDENTITY.maxOutputTokens, + maxAttempts: BEAM_EVALUATOR_IDENTITY.maxAttempts, + timeoutMs: BEAM_EVALUATOR_IDENTITY.timeoutMs, + }) + nuggetJudgments.push({ nugget, score: result.score, reason: result.reason }) + await onProtocolProgress?.({ + ...nuggetProgress, + judgments: [...nuggetJudgments], + }) + } + + const primaryScore = mean(nuggetJudgments.map((judgment) => judgment.score)) + const passed = primaryScore >= BEAM_PASS_THRESHOLD + return { + questionId: question.questionId, + questionType: question.questionType, + primaryScore, + passed, + label: passed ? "pass" : "fail", + explanation: `BEAM nugget average: ${primaryScore.toFixed(4)}`, + metrics: { + nuggetAverage: primaryScore, + nuggetCount: nuggetJudgments.length, + }, + details: { + evaluatorIdentity: BEAM_EVALUATOR_IDENTITY, + nuggetJudgments, + }, + } + } + + aggregateQuality({ + questions, + evaluations, + }: Parameters[0]) { + if (evaluations.length === 0) { + throw new Error("Cannot aggregate a BEAM run with no completed evaluations") + } + + const questionById = new Map(questions.map((question) => [question.questionId, question])) + const evaluationIds = new Set() + const entries: Array<{ + ability: BeamAbilityId + score: number + scale?: "1M" | "10M" + }> = [] + + for (const evaluation of evaluations) { + if (evaluationIds.has(evaluation.questionId)) { + throw new Error(`Duplicate BEAM evaluation for question ${evaluation.questionId}`) + } + evaluationIds.add(evaluation.questionId) + const question = questionById.get(evaluation.questionId) + if (!question) { + throw new Error(`BEAM evaluation references unknown question ${evaluation.questionId}`) + } + if (evaluation.questionType !== question.questionType) { + throw new Error(`BEAM evaluation type mismatch for question ${evaluation.questionId}`) + } + if (!BEAM_ABILITY_ID_SET.has(evaluation.questionType)) { + throw new Error(`Unsupported BEAM ability in evaluation: ${evaluation.questionType}`) + } + if ( + !Number.isFinite(evaluation.primaryScore) || + evaluation.primaryScore < 0 || + evaluation.primaryScore > 1 + ) { + throw new Error(`Invalid BEAM score for question ${evaluation.questionId}`) + } + + const scale = question.metadata?.scale + if (scale !== undefined && scale !== "1M" && scale !== "10M") { + throw new Error( + `BEAM question ${question.questionId} has unsupported tier ${JSON.stringify(scale)}` + ) + } + entries.push({ + ability: evaluation.questionType as BeamAbilityId, + score: evaluation.primaryScore, + ...(scale ? { scale } : {}), + }) + } + + const missingQuestionEvaluations = questions.filter( + (question) => !evaluationIds.has(question.questionId) + ) + if (missingQuestionEvaluations.length > 0) { + throw new Error( + `Missing BEAM evaluations for ${missingQuestionEvaluations.length} question(s), starting with ${missingQuestionEvaluations[0]!.questionId}` + ) + } + + const aggregateEntries = (subset: typeof entries) => { + const scoresByAbility = new Map() + for (const entry of subset) { + const scores = scoresByAbility.get(entry.ability) ?? [] + scores.push(entry.score) + scoresByAbility.set(entry.ability, scores) + } + const abilitySlices: Record> = {} + const abilityScores: number[] = [] + let passedQuestions = 0 + for (const ability of BEAM_ABILITY_IDS) { + const scores = scoresByAbility.get(ability) + if (!scores?.length) continue + const averageScore = mean(scores) + const passed = scores.filter((score) => score >= BEAM_PASS_THRESHOLD).length + abilityScores.push(averageScore) + passedQuestions += passed + abilitySlices[ability] = { + averageScore, + passAccuracy: passed / scores.length, + questionCount: scores.length, + passedQuestions: passed, + } + } + return { + score: mean(abilityScores), + passAccuracy: passedQuestions / subset.length, + passedQuestions, + questionCount: subset.length, + coveredAbilities: abilityScores.length, + allAbilitiesCovered: abilityScores.length === BEAM_ABILITY_IDS.length, + abilitySlices, + } + } + + const pooled = aggregateEntries(entries) + const bySlice: Record> = { ...pooled.abilitySlices } + const metrics: Record = { + passAccuracy: pooled.passAccuracy, + passedQuestions: pooled.passedQuestions, + totalQuestions: pooled.questionCount, + coveredAbilities: pooled.coveredAbilities, + } + + const scales = (["1M", "10M"] as const).filter((scale) => + entries.some((entry) => entry.scale === scale) + ) + const tierScores: number[] = [] + const tierPassAccuracies: number[] = [] + let allTiersOfficial = true + let singleTierOfficial = false + for (const scale of scales) { + const tierEntries = entries.filter((entry) => entry.scale === scale) + const tier = aggregateEntries(tierEntries) + const expected = BEAM_OFFICIAL_TIER_COUNTS[scale] + const officialQuestionSet = + tierEntries.length === expected.questions && + BEAM_ABILITY_IDS.every( + (ability) => + tierEntries.filter((entry) => entry.ability === ability).length === + expected.questionsPerAbility + ) + const metricSuffix = officialQuestionSet ? scale : `${scale}Partial` + metrics[`beamScore${metricSuffix}`] = tier.score + metrics[`passAccuracy${metricSuffix}`] = tier.passAccuracy + metrics[`questionCount${scale}`] = tier.questionCount + metrics[`officialQuestionSet${scale}`] = officialQuestionSet ? 1 : 0 + tierScores.push(tier.score) + tierPassAccuracies.push(tier.passAccuracy) + allTiersOfficial &&= officialQuestionSet + singleTierOfficial = officialQuestionSet + bySlice[`tier:${scale}`] = { + averageScore: tier.score, + passAccuracy: tier.passAccuracy, + questionCount: tier.questionCount, + passedQuestions: tier.passedQuestions, + coveredAbilities: tier.coveredAbilities, + officialQuestionSet: officialQuestionSet ? 1 : 0, + } + for (const [ability, values] of Object.entries(tier.abilitySlices)) { + bySlice[`tier:${scale}/ability:${ability}`] = values + } + } + + if (scales.length > 1) { + const score = mean(tierScores) + const secondaryKey = + allTiersOfficial && entries.every((entry) => entry.scale !== undefined) + ? "beamTierMacroAverageSecondary" + : "beamTierMacroAverageSecondaryPartial" + metrics[secondaryKey] = score + metrics.passAccuracyTierMacro = mean(tierPassAccuracies) + metrics.beamAbilityPooledAverage = pooled.score + metrics.beamQuestionMicroAverage = mean(entries.map((entry) => entry.score)) + return { + metrics, + bySlice, + } + } + + const primaryKey = + scales.length === 1 && + singleTierOfficial && + entries.every((entry) => entry.scale === scales[0]) + ? "beamScore" + : "beamScorePartial" + metrics[primaryKey] = pooled.score + + return { + primaryMetric: { key: primaryKey, value: pooled.score, higherIsBetter: true }, + metrics, + bySlice, + } + } +} + +const BEAM_EVALUATOR_IMPLEMENTATION_SHA256 = stableSha256({ + protocolEvaluateQuestion: BeamPaperProtocol.prototype.evaluateQuestion.toString(), + rubricValidation: getRubric.toString(), + arithmeticMean: mean.toString(), + nuggetPromptBuilder: buildBeamPaperNuggetPrompt.toString(), + eventPromptBuilder: buildBeamEventEquivalencePrompt.toString(), + nuggetJudgmentSchema: createBeamNuggetJudgmentSchema.toString(), + eventEquivalenceSchema: createBeamEventEquivalenceSchema.toString(), + nuggetProgressSchema: createBeamNuggetProgressSchema.toString(), + eventProgressSchema: createBeamEventProgressSchema.toString(), + eventExtraction: extractBeamPredictedEvents.toString(), + eventAlignment: alignBeamEvents.toString(), + eventRankVectors: buildBeamEventRankVectors.toString(), + kendallTauB: computeKendallTauB.toString(), + eventScore: scoreAlignedBeamEvents.toString(), + eventEvaluation: evaluateBeamEventOrdering.toString(), +}) + +export const BEAM_STRUCTURED_OUTPUT_SCHEMA_SHA256 = stableSha256({ + nuggetJudgmentSchema: createBeamNuggetJudgmentSchema.toString(), + eventEquivalenceSchema: createBeamEventEquivalenceSchema.toString(), + nuggetProgressSchema: createBeamNuggetProgressSchema.toString(), + eventProgressSchema: createBeamEventProgressSchema.toString(), +}) + +export const BEAM_AGGREGATION_IMPLEMENTATION_SHA256 = stableSha256({ + aggregateQuality: BeamPaperProtocol.prototype.aggregateQuality.toString(), + arithmeticMean: mean.toString(), +}) + +export const BEAM_EVALUATOR_IDENTITY: BeamEvaluatorIdentity = { + protocolId: BEAM_PAPER_PROTOCOL_ID, + paperId: BEAM_PAPER_ID, + paperRevision: BEAM_PAPER_REVISION, + paperPdfSha256: BEAM_PAPER_PDF_SHA256, + referenceRepository: BEAM_REFERENCE_REPOSITORY, + referenceCommit: BEAM_REFERENCE_COMMIT, + nuggetPromptVersion: BEAM_NUGGET_JUDGE_PROMPT_VERSION, + nuggetPromptSha256: sha256Text(BEAM_NUGGET_JUDGE_PROMPT), + eventEquivalencePromptVersion: BEAM_EVENT_EQUIVALENCE_PROMPT_VERSION, + eventEquivalencePromptSha256: sha256Text( + `${BEAM_EVENT_EQUIVALENCE_SYSTEM_PROMPT}\n\n${BEAM_EVENT_EQUIVALENCE_USER_PROMPT}` + ), + evaluatorImplementationSha256: BEAM_EVALUATOR_IMPLEMENTATION_SHA256, + judgeProvider: "openai", + judgeModel: "gpt-4.1-mini", + structuredOutputSchemaVersion: "beam-paper-structured-output-v1", + structuredOutputSchemaSha256: BEAM_STRUCTURED_OUTPUT_SCHEMA_SHA256, + structuredOutputMode: "ai-sdk-generate-object-json-schema-v1", + kendallTauImplementation: BEAM_KENDALL_TAU_IMPLEMENTATION, + eventExtractionVersion: BEAM_EVENT_EXTRACTION_VERSION, + eventOrderingScoringVersion: BEAM_EVENT_ORDERING_SCORING_VERSION, + temperature: 0, + maxOutputTokens: 512, + maxAttempts: 3, + timeoutMs: 120_000, + retryPolicy: "immediate-transport-or-schema-retry-v1", +} + +export const beamPaperProtocol = new BeamPaperProtocol() diff --git a/src/protocols/beam-paper/prompts.ts b/src/protocols/beam-paper/prompts.ts new file mode 100644 index 0000000..90261f2 --- /dev/null +++ b/src/protocols/beam-paper/prompts.ts @@ -0,0 +1,112 @@ +/** + * BEAM paper evaluator prompts. + * + * Source: Appendix H, Listings 20 and 21, as implemented at + * mohammadtavakoli78/BEAM@3e12035532eb85768f1a7cd779832b650c4b2ef9. + * Keep these strings stable: their hashes are part of the protocol identity. + */ + +export const BEAM_NUGGET_JUDGE_PROMPT_VERSION = "beam-paper-listing-20-v1" +export const BEAM_EVENT_EQUIVALENCE_PROMPT_VERSION = "beam-paper-listing-21-v1" + +export const BEAM_NUGGET_JUDGE_PROMPT = `You are an expert evaluator tasked with judging whether the LLM's response demonstrates compliance with the specified RUBRIC CRITERION. + +## EVALUATION INPUTS +- QUESTION (what the user asked): +- RUBRIC CRITERION (what to check): +- RESPONSE TO EVALUATE: +## EVALUATION RUBRIC: +The rubric defines a specific requirement, constraint, or expected behavior that the LLM response should demonstrate. + +**IMPORTANT**: Pay careful attention to whether the rubric specifies: +- **Positive requirements** (things the response SHOULD include/do) +- **Negative constraints** (things the response SHOULD NOT include/do, often indicated by "no", "not", "avoid", "absent") +## RESPONSIVENESS REQUIREMENT (anchored to the QUESTION) +A compliant response must be **on-topic with respect to the QUESTION** and attempt to answer it. +- If the response does not address the QUESTION, score **0.0** and stop. +- For negative constraints, both must hold: (a) the response is responsive to the QUESTION, and (b) the prohibited element is absent. +## SEMANTIC TOLERANCE RULES: +Judge by meaning, not exact wording. +- Accept **paraphrases** and **synonyms** that preserve intent. +- **Case/punctuation/whitespace** differences must be ignored. +- **Numbers/currencies/dates** may appear in equivalent forms (e.g., "$68,000", "68k", "68,000 USD", or "sixty-eight thousand dollars"). Treat them as equal when numerically equivalent. +- If the rubric expects a number or duration, prefer **normalized comparison** (extract and compare values) over string matching. +## STYLE NEUTRALITY (prevents style contamination): +Ignore tone, politeness, length, and flourish unless the rubric explicitly requires a format/structure (e.g., "itemized list", "no citations", "one sentence"). +- Do **not** penalize hedging, voice, or verbosity if content satisfies the rubric. +- Only evaluate format when the rubric **explicitly** mandates it. +## SCORING SCALE: +- **1.0 (Complete Compliance)**: Fully complies with the rubric criterion. + - Positive: required element present, accurate, properly executed (allowing semantic equivalents). + - Negative: prohibited element **absent** AND response is **responsive**. + +- **0.5 (Partial Compliance)**: Partially complies. + - Positive: element present but minor inaccuracies/incomplete execution. + - Negative: generally responsive and mostly avoids the prohibited element but with minor/edge violations. +- **0.0 (No Compliance)**: Fails to comply. + - Positive: required element missing or incorrect. + - Negative: prohibited element present **or** response is non-responsive/evasive even if the element is absent. +## EVALUATION INSTRUCTIONS: +1. **Understand the Requirement**: Determine if the rubric is asking for something to be present (positive) or absent (negative/constraint). +2. **Parse Compound Statements**: If the rubric contains multiple elements connected by "and" or commas, evaluate whether: + - **All elements** must be present for full compliance (1.0) + - **Some elements** present indicates partial compliance (0.5) + - **No elements** present indicates no compliance (0.0) + +3. **Check Compliance**: + - For positive requirements: Look for the presence and quality of the required element + - For negative constraints: Look for the absence of the prohibited element +4. **Assign Score**: Based on compliance with the specific rubric criterion according to the scoring scale above. + +5. **Provide Reasoning**: Explain whether the rubric criterion was satisfied and justify the score. +## OUTPUT FORMAT: +Return your evaluation in JSON format with two fields: + +{ + "score": [your score: 1.0, 0.5, or 0.0], + "reason": "[detailed explanation of whether the rubric criterion was satisfied and why this justified the assigned score]" +} + +NOTE: ONLY output the json object, without any explanation before or after that` + +export const BEAM_EVENT_EQUIVALENCE_SYSTEM_PROMPT = `You are a binary classifier. +If the TWO snippets describe the SAME event/fact, reply **YES** +Otherwise reply **NO**. No extra words. DO NOT provide any exaplanation.` + +export const BEAM_EVENT_EQUIVALENCE_USER_PROMPT = + "First snippet: \n Second snippet: " + +function replaceAllLiteral(value: string, token: string, replacement: string): string { + return value.split(token).join(replacement) +} + +export function buildBeamPaperNuggetPrompt(input: { + question: string + nugget: string + answer: string +}): string { + return replaceAllLiteral( + replaceAllLiteral( + replaceAllLiteral(BEAM_NUGGET_JUDGE_PROMPT, "", input.question), + "", + input.nugget + ), + "", + input.answer + ) +} + +export function buildBeamEventEquivalencePrompt(input: { + referenceEvent: string + predictedEvent: string +}): string { + return replaceAllLiteral( + replaceAllLiteral( + BEAM_EVENT_EQUIVALENCE_USER_PROMPT, + "", + input.referenceEvent + ), + "", + input.predictedEvent + ) +} diff --git a/src/protocols/index.ts b/src/protocols/index.ts new file mode 100644 index 0000000..e570606 --- /dev/null +++ b/src/protocols/index.ts @@ -0,0 +1,2 @@ +export { LegacyBenchmarkProtocol, legacyBenchmarkProtocol } from "./legacy" +export { BeamPaperProtocol } from "./beam-paper" diff --git a/src/protocols/legacy.ts b/src/protocols/legacy.ts new file mode 100644 index 0000000..4367411 --- /dev/null +++ b/src/protocols/legacy.ts @@ -0,0 +1,179 @@ +import type { BenchmarkProtocol, ProtocolIdentity } from "../types/protocol" +import type { ProviderPrompts } from "../types/prompts" +import { buildContextString } from "../types/prompts" +import { buildDefaultAnswerPrompt } from "../prompts/defaults" +import { sha256Text, stableSha256 } from "../utils/stable" + +const LEGACY_VERSION = "1.1.0" + +function renderLegacyAnswerPrompt( + question: string, + context: unknown[], + questionDate: string | undefined, + prompts?: ProviderPrompts +): string { + if (prompts?.answerPrompt) { + if (typeof prompts.answerPrompt === "function") { + return prompts.answerPrompt(question, context, questionDate) + } + + return prompts.answerPrompt + .replace("{{question}}", question) + .replace("{{questionDate}}", questionDate || "Not specified") + .replace("{{context}}", buildContextString(context)) + } + + return buildDefaultAnswerPrompt(question, context, questionDate) +} + +function identity(protocol: LegacyBenchmarkProtocol): ProtocolIdentity { + const ingestion = { + policy: "legacy-stringified-session-json-v1", + customId: "sessionId", + metadata: "sessionId-and-optional-documentDate-from-metadata-date", + formattedDatePolicy: "optional-human-readable-content-prefix", + messageProjection: "stringified-session-messages", + executionPolicy: protocol.ingestionExecutionPolicy, + implementationSha256: sha256Text(protocol.createIngestionPlan.toString()), + } + const retrieval = { requestedTopK: 10, answerCutoff: 10, threshold: 0.3 } + const answer = { formatter: "legacy-provider-prompt-or-default", version: 1 } + const evaluator = { + evaluator: "configured-generic-judge", + version: 1, + auxiliaryRetrievalEvaluation: protocol.auxiliaryRetrievalEvaluation, + } + const aggregation = { primary: "accuracy", passThreshold: 1 } + + return { + id: "memorybench.legacy", + version: LEGACY_VERSION, + configFingerprint: stableSha256({ ingestion, retrieval, answer, evaluator, aggregation }), + implementationFingerprint: stableSha256({ + protocol: "legacy", + version: LEGACY_VERSION, + ingestion, + }), + ingestionPolicyHash: stableSha256(ingestion), + retrievalPolicyHash: stableSha256(retrieval), + answerPromptHash: stableSha256(answer), + evaluatorHash: stableSha256(evaluator), + aggregationHash: stableSha256(aggregation), + details: { ingestionPolicy: ingestion }, + } +} + +export class LegacyBenchmarkProtocol implements BenchmarkProtocol { + readonly auxiliaryRetrievalEvaluation = "legacy-llm-relevance-v1" as const + readonly ingestionExecutionPolicy = { + readinessBarrier: "after-build", + processingMode: "provider-default", + } as const + readonly identity: ProtocolIdentity + + constructor() { + this.identity = identity(this) + } + + validateQuestion(question: Parameters[0]): void { + if (!question.questionId || !question.question.trim()) { + throw new Error("Legacy benchmark question must have a non-empty ID and question") + } + } + + createIngestionPlan({ sessions }: Parameters[0]) { + return sessions.map((session) => { + const formattedDate = session.metadata?.formattedDate + const sourceDate = session.metadata?.date + const documentDate = typeof sourceDate === "string" ? sourceDate : undefined + const serialized = JSON.stringify(session.messages) + .replace(//g, ">") + const content = + typeof formattedDate === "string" && formattedDate + ? `Here is the date the following session took place: ${formattedDate}\n\nHere is the session as a stringified JSON:\n${serialized}` + : `Here is the session as a stringified JSON:\n${serialized}` + + return { + customId: session.sessionId, + content, + metadata: { + sessionId: session.sessionId, + ...(documentDate ? { documentDate } : {}), + }, + messages: session.messages, + } + }) + } + + createRetrievalPlan({ question }: Parameters[0]) { + return { + query: question.question, + requestedTopK: 10, + answerCutoff: 10, + threshold: 0.3, + } + } + + createAnswerPlan({ + question, + results, + questionDate, + providerPrompts, + }: Parameters[0]) { + const prompt = renderLegacyAnswerPrompt( + question.question, + results, + questionDate, + providerPrompts + ) + const basePrompt = renderLegacyAnswerPrompt( + question.question, + [], + questionDate, + providerPrompts + ) + + return { + request: { prompt }, + baseRequest: { prompt: basePrompt }, + answerEvidenceCount: results.length, + } + } + + async evaluateQuestion( + { question, hypothesis, providerPrompts }: Parameters[0], + runtime: Parameters[1] + ) { + const result = await runtime.evaluateLegacy({ + question: question.question, + questionType: question.questionType, + groundTruth: question.groundTruth, + hypothesis, + providerPrompts, + }) + + return { + questionId: question.questionId, + questionType: question.questionType, + primaryScore: result.score, + passed: result.label === "correct", + label: result.label, + explanation: result.explanation, + details: result.details, + } + } + + aggregateQuality({ evaluations }: Parameters[0]) { + const total = evaluations.length + const passed = evaluations.filter((evaluation) => evaluation.passed).length + const accuracy = total > 0 ? passed / total : 0 + + return { + primaryMetric: { key: "accuracy", value: accuracy, higherIsBetter: true }, + metrics: { accuracy, passed, total }, + } + } +} + +export const legacyBenchmarkProtocol = new LegacyBenchmarkProtocol() diff --git a/src/providers/filesystem/index.ts b/src/providers/filesystem/index.ts index 0c51f7a..cecea85 100644 --- a/src/providers/filesystem/index.ts +++ b/src/providers/filesystem/index.ts @@ -8,14 +8,74 @@ import type { IngestResult, SearchOptions, IndexingProgressCallback, + ProviderSearchResponse, } from "../../types/provider" -import type { UnifiedSession } from "../../types/unified" +import type { + CanonicalIngestionDocument, + ProviderResultDropDiagnostic, + UnifiedSearchResult, +} from "../../types/unified" import { logger } from "../../utils/logger" -import { extractMemories } from "../../prompts/extraction" +import { extractMemories, getMemoryExtractionConfigFingerprint } from "../../prompts/extraction" +import { stableSha256 } from "../../utils/stable" import { FILESYSTEM_PROMPTS } from "./prompts" +import { + assertResultBudget, + canonicalDocumentToSession, + rankResults, + recordResultDrop, + requireSearchLimit, + resolveDocumentDate, + createProviderSearchResponse, +} from "../normalization" const BASE_DIR = join(process.cwd(), "data", "providers", "filesystem") +interface FilesystemSearchCandidate { + id: string + sessionId: string + content: string + score: number + documentDate?: string +} + +export function renderFilesystemMemoryFile( + sessionId: string, + documentDate: string | undefined, + extractedMemories: string +): string { + const header = documentDate + ? `# Memory: ${sessionId}\n**Date:** ${documentDate}\n\n` + : `# Memory: ${sessionId}\n\n` + return header + extractedMemories +} + +export function normalizeFilesystemSearchResults( + rawResults: FilesystemSearchCandidate[], + limit: number, + droppedResults: ProviderResultDropDiagnostic[] = [] +): UnifiedSearchResult[] { + requireSearchLimit(limit, "filesystem") + assertResultBudget(rawResults.length, limit, "filesystem") + const normalized: Omit[] = [] + for (const [index, result] of rawResults.entries()) { + if (!result.content.trim()) { + recordResultDrop(droppedResults, index, "empty-text") + continue + } + normalized.push({ + id: result.id, + text: result.content, + score: result.score, + sessionId: result.sessionId, + ...(result.documentDate ? { documentDate: result.documentDate } : {}), + provider: "filesystem", + resultType: "document", + }) + } + return rankResults(normalized) +} + /** * Simple tokenizer: lowercase, split on non-alphanumeric, filter short tokens. * Deliberately kept simple to represent the filesystem-based approach. @@ -33,7 +93,10 @@ function tokenize(text: string): string[] { * Returns a score between 0 and 1 representing the fraction of query terms found, * with a small frequency bonus for repeated matches. */ -function scoreDocument(queryTerms: string[], docText: string): { score: number; matchCount: number } { +function scoreDocument( + queryTerms: string[], + docText: string +): { score: number; matchCount: number } { if (queryTerms.length === 0) return { score: 0, matchCount: 0 } const docLower = docText.toLowerCase() @@ -77,6 +140,8 @@ function scoreDocument(queryTerms: string[], docText: string): { score: number; */ export class FilesystemProvider implements Provider { name = "filesystem" + adapterVersion = "2.0.0" + searchRequestStructure = { kind: "single" } as const prompts = FILESYSTEM_PROMPTS concurrency = { default: 50, @@ -85,6 +150,17 @@ export class FilesystemProvider implements Provider { private openai: ReturnType | null = null + getIngestionConfigFingerprint(_config: ProviderConfig): string { + return stableSha256({ + schemaVersion: 1, + provider: this.name, + adapterVersion: this.adapterVersion, + extractionConfigFingerprint: getMemoryExtractionConfigFingerprint(), + storage: "memory-markdown-plus-json-sidecar-v1", + documentId: "sanitized-customId", + }) + } + async initialize(config: ProviderConfig): Promise { if (!config.apiKey || config.apiKey === "none") { throw new Error("Filesystem provider requires OPENAI_API_KEY for memory extraction") @@ -94,7 +170,10 @@ export class FilesystemProvider implements Provider { logger.info("Initialized Filesystem memory provider (MEMORY.md-style with LLM extraction)") } - async ingest(sessions: UnifiedSession[], options: IngestOptions): Promise { + async ingest( + documents: CanonicalIngestionDocument[], + options: IngestOptions + ): Promise { if (!this.openai) throw new Error("Provider not initialized") const containerDir = join(BASE_DIR, sanitizePath(options.containerTag)) @@ -103,22 +182,31 @@ export class FilesystemProvider implements Provider { const documentIds: string[] = [] - for (const session of sessions) { + for (const document of documents) { + const session = canonicalDocumentToSession(document) const extractedMemories = await extractMemories(this.openai, session) // Build a memory file with date header + extracted content - const date = - (session.metadata?.formattedDate as string) || - (session.metadata?.date as string) || - "Unknown date" - const header = `# Memory: ${session.sessionId}\n**Date:** ${date}\n\n` - const content = header + extractedMemories - - const safeId = sanitizePath(session.sessionId) + const date = document.metadata.documentDate + const content = renderFilesystemMemoryFile( + document.metadata.sessionId, + date, + extractedMemories + ) + + const safeId = sanitizePath(document.customId) const filePath = join(memoriesDir, `${safeId}.md`) await writeFile(filePath, content, "utf-8") + await writeFile( + join(memoriesDir, `${safeId}.json`), + JSON.stringify({ + sessionId: document.metadata.sessionId, + ...(date ? { documentDate: date } : {}), + }), + "utf-8" + ) documentIds.push(safeId) - logger.debug(`Extracted and stored memories for session ${session.sessionId}`) + logger.debug(`Extracted and stored memories for session ${document.metadata.sessionId}`) } return { documentIds } @@ -137,7 +225,24 @@ export class FilesystemProvider implements Provider { }) } - async search(query: string, options: SearchOptions): Promise { + async search(query: string, options: SearchOptions): Promise { + const limit = requireSearchLimit(options.limit, this.name) + const respond = (raw: FilesystemSearchCandidate[]) => { + const droppedResults: ProviderResultDropDiagnostic[] = [] + return createProviderSearchResponse({ + results: normalizeFilesystemSearchResults(raw, limit, droppedResults), + requestedLimit: limit, + rawReturnedCount: raw.length, + droppedResults, + providerRequests: [ + { + operation: "filesystem.scan", + limit, + parameters: { scoring: "term-coverage-frequency-v1" }, + }, + ], + }) + } const containerDir = join(BASE_DIR, sanitizePath(options.containerTag)) const memoriesDir = join(containerDir, "memories") @@ -146,46 +251,71 @@ export class FilesystemProvider implements Provider { files = await readdir(memoriesDir) } catch { logger.warn(`No memories directory found for ${options.containerTag}`) - return [] + return respond([]) } const mdFiles = files.filter((f) => f.endsWith(".md")) - if (mdFiles.length === 0) return [] + if (mdFiles.length === 0) return respond([]) const queryTerms = tokenize(query) const scored: Array<{ + id: string sessionId: string content: string score: number matchCount: number + documentDate?: string }> = [] for (const file of mdFiles) { const content = await readFile(join(memoriesDir, file), "utf-8") const { score, matchCount } = scoreDocument(queryTerms, content) + const safeId = file.replace(".md", "") + let sessionId = safeId + let documentDate: string | undefined + try { + const metadata = JSON.parse( + await readFile(join(memoriesDir, `${safeId}.json`), "utf-8") + ) as { sessionId?: unknown; documentDate?: unknown } + if (typeof metadata.sessionId === "string" && metadata.sessionId) { + sessionId = metadata.sessionId + } + documentDate = resolveDocumentDate(metadata) + } catch { + const sessionMatch = content.match(/^# Memory: (.+)$/m) + const dateMatch = content.match(/^\*\*Date:\*\* (.+)$/m) + if (sessionMatch?.[1]) sessionId = sessionMatch[1] + const legacyDate = dateMatch?.[1]?.trim() + if ( + legacyDate && + !["unknown", "unknown date", "not specified"].includes(legacyDate.toLowerCase()) + ) { + documentDate = legacyDate + } + } scored.push({ - sessionId: file.replace(".md", ""), + id: safeId, + sessionId, content, score, matchCount, + ...(documentDate ? { documentDate } : {}), }) } // Sort by score (desc), then by matchCount (desc) as tiebreaker scored.sort((a, b) => b.score - a.score || b.matchCount - a.matchCount) - const limit = options.limit || 10 - // Return top results; include score=0 results only if we have fewer than limit scored results const scoredResults = scored.filter((r) => r.score > 0) if (scoredResults.length >= limit) { - return scoredResults.slice(0, limit) + return respond(scoredResults.slice(0, limit)) } // Fill remaining slots with unscored results (chronological order fallback) const unscoredResults = scored.filter((r) => r.score === 0) - return [...scoredResults, ...unscoredResults].slice(0, limit) + return respond([...scoredResults, ...unscoredResults].slice(0, limit)) } async clear(containerTag: string): Promise { diff --git a/src/providers/filesystem/prompts.ts b/src/providers/filesystem/prompts.ts index b90791b..86a968a 100644 --- a/src/providers/filesystem/prompts.ts +++ b/src/providers/filesystem/prompts.ts @@ -1,14 +1,8 @@ import type { ProviderPrompts } from "../../types/prompts" - -interface FilesystemResult { - sessionId: string - content: string - score: number - matchCount: number -} +import type { UnifiedSearchResult } from "../../types/unified" function buildFilesystemContext(context: unknown[]): string { - const results = context as FilesystemResult[] + const results = context as UnifiedSearchResult[] if (results.length === 0) { return "No relevant memory files were found." @@ -16,8 +10,8 @@ function buildFilesystemContext(context: unknown[]): string { return results .map((result, i) => { - const header = `=== Memory File ${i + 1}: ${result.sessionId} (relevance: ${(result.score * 100).toFixed(0)}%) ===` - return `${header}\n${result.content}` + const date = result.documentDate ? ` [${result.documentDate}]` : "" + return `=== Memory File ${i + 1}${date} ===\n${result.text}` }) .join("\n\n---\n\n") } diff --git a/src/providers/mem0/index.ts b/src/providers/mem0/index.ts index e01d343..dfc8097 100644 --- a/src/providers/mem0/index.ts +++ b/src/providers/mem0/index.ts @@ -16,10 +16,75 @@ import type { IngestResult, SearchOptions, IndexingProgressCallback, + ProviderSearchResponse, } from "../../types/provider" -import type { UnifiedSession } from "../../types/unified" +import type { + CanonicalIngestionDocument, + ProviderResultDropDiagnostic, + UnifiedSearchResult, +} from "../../types/unified" import { logger } from "../../utils/logger" +import { stableSha256 } from "../../utils/stable" import { MEM0_PROMPTS } from "./prompts" +import { + asFiniteNumber, + asNonEmptyString, + asRecord, + assertResultBudget, + canonicalDocumentToSession, + rankResults, + recordResultDrop, + requireSearchLimit, + resolveDocumentDate, + resolveSessionId, + createProviderSearchResponse, +} from "../normalization" + +export function normalizeMem0SearchResults( + rawResults: unknown[], + limit: number, + droppedResults: ProviderResultDropDiagnostic[] = [] +): UnifiedSearchResult[] { + requireSearchLimit(limit, "mem0") + assertResultBudget(rawResults.length, limit, "mem0") + + const normalized: Omit[] = [] + for (const [index, rawResult] of rawResults.entries()) { + const result = asRecord(rawResult) + if (!result) { + recordResultDrop(droppedResults, index, "malformed-result") + continue + } + + const data = asRecord(result.data) + const id = asNonEmptyString(result.id) ?? asNonEmptyString(data?.id) + const text = asNonEmptyString(result.memory) ?? asNonEmptyString(data?.memory) + if (!id) { + recordResultDrop(droppedResults, index, "missing-id") + continue + } + if (!text) { + recordResultDrop(droppedResults, index, "empty-text") + continue + } + + const metadata = asRecord(result.metadata) ?? asRecord(data?.metadata) + const score = asFiniteNumber(result.score) ?? asFiniteNumber(data?.score) + const sessionId = resolveSessionId(metadata) + const documentDate = resolveDocumentDate(metadata) + normalized.push({ + id, + text, + ...(score !== undefined ? { score } : {}), + ...(sessionId ? { sessionId } : {}), + ...(documentDate ? { documentDate } : {}), + provider: "mem0", + resultType: "memory", + }) + } + + return rankResults(normalized) +} /** * Custom instructions from Mem0's official evaluation. @@ -51,36 +116,77 @@ const CUSTOM_INSTRUCTIONS = `Generate personal memories that follow these guidel 5. Format each memory as a paragraph with a clear narrative structure that captures the person's experience, challenges, and aspirations` +export async function configureMem0Project(client: { + updateProject(input: { custom_instructions: string }): Promise +}): Promise { + await client.updateProject({ custom_instructions: CUSTOM_INSTRUCTIONS }) +} + export class Mem0Provider implements Provider { name = "mem0" + adapterVersion = "2.1.0" + searchRequestStructure = { kind: "single" } as const prompts = MEM0_PROMPTS concurrency = { default: 50, } private client: MemoryClient | null = null - private apiKey: string = "" + + getIngestionConfigFingerprint(_config: ProviderConfig): string { + return stableSha256({ + schemaVersion: 1, + provider: this.name, + adapterVersion: this.adapterVersion, + customInstructions: CUSTOM_INSTRUCTIONS, + version: "v2", + enableGraph: false, + asyncMode: false, + reconciliation: "user_id-run_id-v1", + }) + } async initialize(config: ProviderConfig): Promise { - this.apiKey = config.apiKey this.client = new MemoryClient({ apiKey: config.apiKey }) - try { - await this.client.updateProject({ - custom_instructions: CUSTOM_INSTRUCTIONS, - }) - } catch (e) { - logger.warn(`Could not set custom instructions: ${e}`) - } + await configureMem0Project(this.client) logger.info(`Initialized Mem0 provider`) } - async ingest(sessions: UnifiedSession[], options: IngestOptions): Promise { + async ingest( + documents: CanonicalIngestionDocument[], + options: IngestOptions + ): Promise { if (!this.client) throw new Error("Provider not initialized") const eventIds: string[] = [] - for (const session of sessions) { + for (const document of documents) { + const existingResponse = await this.client.getAll({ + user_id: options.containerTag, + run_id: document.customId, + page_size: 100, + output_format: "v1.1", + }) + const existingRecord = asRecord(existingResponse) + const existing = Array.isArray(existingResponse) + ? existingResponse + : Array.isArray(existingRecord?.results) + ? existingRecord.results + : [] + const existingIds = existing.flatMap((item) => { + const record = asRecord(item) + const data = asRecord(record?.data) + const id = asNonEmptyString(record?.id) ?? asNonEmptyString(data?.id) + return id ? [id] : [] + }) + if (existingIds.length > 0) { + eventIds.push(...existingIds) + logger.debug(`Reconciled previously ingested session ${document.customId}`) + continue + } + + const session = canonicalDocumentToSession(document) const messages = session.messages.map((m) => ({ role: m.role, content: m.content, @@ -88,36 +194,32 @@ export class Mem0Provider implements Provider { const addOptions: MemoryOptions = { user_id: options.containerTag, + run_id: document.customId, version: "v2", enable_graph: false, - async_mode: true, + // Synchronous completion makes the run_id marker queryable before this call returns. + // Mem0's default infer-mode conflict resolution supplies duplicate protection if the + // process dies after remote success but before the local checkpoint becomes durable. + async_mode: false, metadata: { - sessionId: session.sessionId, - timestamp: session.metadata?.date, - ...session.metadata, + ...document.metadata, + ...(document.metadata.documentDate ? { timestamp: document.metadata.documentDate } : {}), ...options.metadata, }, } const result = (await this.client.add(messages, addOptions)) as Array<{ + id?: string event_id?: string }> for (const event of result) { - if (event.event_id) eventIds.push(event.event_id) + const id = event.id ?? event.event_id + if (id) eventIds.push(id) } } return { documentIds: eventIds } } - private async getEventStatus(eventId: string): Promise { - const response = await fetch(`https://api.mem0.ai/v1/event/${eventId}/`, { - headers: { Authorization: `Token ${this.apiKey}` }, - }) - if (!response.ok) return "UNKNOWN" - const data = await response.json() - return data.status || "UNKNOWN" - } - async awaitIndexing( result: IngestResult, _containerTag: string, @@ -129,63 +231,50 @@ export class Mem0Provider implements Provider { return } - const total = eventIds.length - const pending = new Set(eventIds) - const completedIds: string[] = [] - const failedIds: string[] = [] - let backoffMs = 500 - - onProgress?.({ completedIds: [], failedIds: [], total }) - - while (pending.size > 0) { - const pendingArray = Array.from(pending) - const results = await Promise.allSettled( - pendingArray.map(async (eventId) => { - const status = await this.getEventStatus(eventId) - return { eventId, status } - }) - ) - - for (const res of results) { - if (res.status === "fulfilled") { - const { eventId, status } = res.value - if (status === "SUCCEEDED") { - pending.delete(eventId) - completedIds.push(eventId) - } else if (status === "FAILED") { - pending.delete(eventId) - failedIds.push(eventId) - } - } - } - - onProgress?.({ completedIds: [...completedIds], failedIds: [...failedIds], total }) - - if (pending.size > 0) { - await new Promise((r) => setTimeout(r, backoffMs)) - backoffMs = Math.min(backoffMs * 1.5, 5000) - } - } - - if (failedIds.length > 0) { - logger.warn(`${failedIds.length} events failed indexing`) - } + // async_mode=false returns fully processed memory IDs. + onProgress?.({ completedIds: [...eventIds], failedIds: [], total: eventIds.length }) } - async search(query: string, options: SearchOptions): Promise { + async search(query: string, options: SearchOptions): Promise { if (!this.client) throw new Error("Provider not initialized") + const limit = requireSearchLimit(options.limit, this.name) const searchOptions: Mem0SearchOptions = { user_id: options.containerTag, - top_k: options.limit || 30, + top_k: limit, enable_graph: false, output_format: "v1.1", + ...(options.threshold !== undefined ? { threshold: options.threshold } : {}), + ...(options.filters ? { filters: options.filters } : {}), } const response = await this.client.search(query, searchOptions) - - const res = response as { results?: unknown[] } - return res.results ?? [] + const responseRecord = asRecord(response) + const rawResults = Array.isArray(response) + ? response + : Array.isArray(responseRecord?.results) + ? responseRecord.results + : [] + const droppedResults: ProviderResultDropDiagnostic[] = [] + return createProviderSearchResponse({ + results: normalizeMem0SearchResults(rawResults, limit, droppedResults), + requestedLimit: limit, + rawReturnedCount: rawResults.length, + droppedResults, + providerRequests: [ + { + operation: "search.memories", + limit, + parameters: { + topK: limit, + enableGraph: false, + outputFormat: "v1.1", + thresholdProvided: options.threshold !== undefined, + ...(options.threshold !== undefined ? { threshold: options.threshold } : {}), + }, + }, + ], + }) } async clear(containerTag: string): Promise { diff --git a/src/providers/mem0/prompts.ts b/src/providers/mem0/prompts.ts index 526cfd4..f87a0e4 100644 --- a/src/providers/mem0/prompts.ts +++ b/src/providers/mem0/prompts.ts @@ -1,20 +1,12 @@ import type { ProviderPrompts } from "../../types/prompts" - -interface Mem0Memory { - memory?: string - metadata?: Record -} +import type { UnifiedSearchResult } from "../../types/unified" export function buildMem0AnswerPrompt(question: string, context: unknown[]): string { const memoriesStr = context .map((r, i) => { - const mem = r as Mem0Memory - const metadata = mem.metadata - const timestampInfo = - metadata?.date || metadata?.timestamp - ? ` [Timestamp: ${metadata.date || metadata.timestamp}]` - : "" - return `[${i + 1}]${timestampInfo} ${mem.memory || JSON.stringify(r)}` + const memory = r as UnifiedSearchResult + const timestampInfo = memory.documentDate ? ` [Document date: ${memory.documentDate}]` : "" + return `[${i + 1}]${timestampInfo} ${memory.text}` }) .join("\n\n") diff --git a/src/providers/normalization.ts b/src/providers/normalization.ts new file mode 100644 index 0000000..05d070a --- /dev/null +++ b/src/providers/normalization.ts @@ -0,0 +1,142 @@ +import type { + CanonicalIngestionDocument, + UnifiedSearchResult, + UnifiedSession, +} from "../types/unified" +import type { ProviderSearchResponse } from "../types/provider" +import type { ProviderRequestDiagnostic } from "../types/unified" +import type { ProviderResultDropDiagnostic } from "../types/unified" + +export type UnknownRecord = Record + +export function asRecord(value: unknown): UnknownRecord | undefined { + return typeof value === "object" && value !== null && !Array.isArray(value) + ? (value as UnknownRecord) + : undefined +} + +export function asNonEmptyString(value: unknown): string | undefined { + if (typeof value !== "string") return undefined + const trimmed = value.trim() + return trimmed.length > 0 ? trimmed : undefined +} + +export function asFiniteNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) ? value : undefined +} + +export function requireSearchLimit(limit: number, provider: string): number { + if (!Number.isInteger(limit) || limit <= 0) { + throw new Error(`${provider} search limit must be a positive integer, received ${limit}`) + } + return limit +} + +export function assertResultBudget(rawCount: number, limit: number, provider: string): void { + if (rawCount > limit) { + throw new Error( + `${provider} returned ${rawCount} results for requested Top-K ${limit}; refusing to silently truncate evidence` + ) + } +} + +export function rankResults(results: Omit[]): UnifiedSearchResult[] { + return results.map((result, index) => ({ ...result, rank: index + 1 })) +} + +export function recordResultDrop( + drops: ProviderResultDropDiagnostic[], + index: number, + reason: ProviderResultDropDiagnostic["reason"] +): void { + drops.push({ index, reason }) +} + +export function resolveSessionId(...values: unknown[]): string | undefined { + for (const value of values) { + const record = asRecord(value) + const sessionId = asNonEmptyString(record?.sessionId) + if (sessionId) return sessionId + } + return undefined +} + +export function resolveDocumentDate(...values: unknown[]): string | undefined { + const sourceDate = (value: unknown): string | undefined => { + const date = asNonEmptyString(value) + if (!date) return undefined + const normalized = date.toLowerCase() + return normalized === "unknown" || + normalized === "unknown date" || + normalized === "not specified" + ? undefined + : date + } + for (const value of values) { + const record = asRecord(value) + if (!record) continue + + const direct = sourceDate(record.documentDate) + if (direct) return direct + + const temporal = asRecord(record.temporalContext) + const temporalDate = sourceDate(temporal?.documentDate) + if (temporalDate) return temporalDate + + const legacy = sourceDate(record.date) + if (legacy) return legacy + } + return undefined +} + +export function canonicalDocumentToSession(document: CanonicalIngestionDocument): UnifiedSession { + return { + sessionId: document.metadata.sessionId, + messages: + document.messages && document.messages.length > 0 + ? document.messages + : [{ role: "user", content: document.content }], + metadata: { + ...document.metadata, + ...(document.metadata.documentDate ? { date: document.metadata.documentDate } : {}), + }, + } +} + +export function assertContainerTag(containerTag: string): void { + if (containerTag.length === 0 || containerTag.length > 100) { + throw new Error("containerTag must be between 1 and 100 characters") + } + if (!/^[a-zA-Z0-9_:-]+$/.test(containerTag)) { + throw new Error( + "containerTag may only contain alphanumeric characters, hyphens, underscores, and colons" + ) + } +} + +export function createProviderSearchResponse(input: { + results: UnifiedSearchResult[] + requestedLimit: number + rawReturnedCount: number + providerRequests: ProviderRequestDiagnostic[] + droppedResults?: ProviderResultDropDiagnostic[] +}): ProviderSearchResponse { + const droppedCount = input.rawReturnedCount - input.results.length + const droppedResults = input.droppedResults ?? [] + if (droppedResults.length !== droppedCount) { + throw new Error( + `Provider normalized ${input.results.length}/${input.rawReturnedCount} results but recorded ${droppedResults.length}/${droppedCount} drop reasons` + ) + } + return { + results: input.results, + diagnostics: { + requestedLimit: input.requestedLimit, + providerRequests: input.providerRequests, + rawReturnedCount: input.rawReturnedCount, + normalizedCount: input.results.length, + droppedCount, + droppedResults, + }, + } +} diff --git a/src/providers/prompt-identity.ts b/src/providers/prompt-identity.ts new file mode 100644 index 0000000..4a9281c --- /dev/null +++ b/src/providers/prompt-identity.ts @@ -0,0 +1,25 @@ +import type { ProviderPrompts } from "../types/prompts" +import { stableSha256 } from "../utils/stable" + +function promptValueIdentity( + value: ProviderPrompts["answerPrompt"] | ProviderPrompts["judgePrompt"] +) { + if (value === undefined) return null + if (typeof value === "string") return { kind: "string", value } + return { + kind: "function", + source: Function.prototype.toString.call(value), + } +} + +/** + * Identifies the provider-selected legacy answer/judge prompt implementation. + * BEAM never consumes these prompts, but legacy runs must not resume across drift. + */ +export function fingerprintProviderPrompts(prompts?: ProviderPrompts): string { + return stableSha256({ + schemaVersion: 1, + answerPrompt: promptValueIdentity(prompts?.answerPrompt), + judgePrompt: promptValueIdentity(prompts?.judgePrompt), + }) +} diff --git a/src/providers/rag/index.ts b/src/providers/rag/index.ts index c90b723..b084e01 100644 --- a/src/providers/rag/index.ts +++ b/src/providers/rag/index.ts @@ -1,5 +1,7 @@ import { embedMany, embed } from "ai" import { createOpenAI } from "@ai-sdk/openai" +import { mkdir, readFile, rename, rm, writeFile } from "node:fs/promises" +import { join } from "node:path" import type { Provider, ProviderConfig, @@ -7,13 +9,28 @@ import type { IngestResult, SearchOptions, IndexingProgressCallback, + ProviderSearchResponse, } from "../../types/provider" -import type { UnifiedSession } from "../../types/unified" +import type { + CanonicalIngestionDocument, + ProviderResultDropDiagnostic, + UnifiedSearchResult, +} from "../../types/unified" import { logger } from "../../utils/logger" import { HybridSearchEngine } from "./search" -import type { Chunk } from "./search" +import type { Chunk, SearchResult as RagSearchResult } from "./search" import { RAG_PROMPTS } from "./prompts" -import { extractMemories } from "../../prompts/extraction" +import { extractMemories, getMemoryExtractionConfigFingerprint } from "../../prompts/extraction" +import { stableSha256 } from "../../utils/stable" +import { + assertResultBudget, + canonicalDocumentToSession, + rankResults, + recordResultDrop, + requireSearchLimit, + resolveDocumentDate, + createProviderSearchResponse, +} from "../normalization" /** Target chunk size in characters (~400 tokens) */ const CHUNK_SIZE = 1600 @@ -23,6 +40,144 @@ const CHUNK_OVERLAP = 320 const EMBEDDING_BATCH_SIZE = 100 /** Embedding model to use */ const EMBEDDING_MODEL = "text-embedding-3-small" +const RAG_INDEX_SCHEMA_VERSION = 1 +const DEFAULT_INDEX_ROOT = join(process.cwd(), "data", "providers", "rag") + +interface PersistedRagIndex { + schemaVersion: typeof RAG_INDEX_SCHEMA_VERSION + embeddingModel: typeof EMBEDDING_MODEL + chunks: Chunk[] +} + +function sanitizePath(input: string): string { + return input.replace(/[^a-zA-Z0-9_.-]/g, "_") +} + +function validatePersistedChunk(value: unknown, index: number): Chunk { + if (!value || typeof value !== "object" || Array.isArray(value)) { + throw new Error(`RAG index chunk ${index} is not an object`) + } + const chunk = value as Partial + if ( + typeof chunk.id !== "string" || + chunk.id.length === 0 || + typeof chunk.content !== "string" || + chunk.content.length === 0 || + typeof chunk.sessionId !== "string" || + chunk.sessionId.length === 0 || + !Number.isInteger(chunk.chunkIndex) || + (chunk.chunkIndex as number) < 0 || + !Array.isArray(chunk.embedding) || + chunk.embedding.length === 0 || + chunk.embedding.some((number) => typeof number !== "number" || !Number.isFinite(number)) + ) { + throw new Error(`RAG index chunk ${index} is malformed`) + } + if (chunk.date !== undefined && typeof chunk.date !== "string") { + throw new Error(`RAG index chunk ${index} has an invalid date`) + } + if ( + chunk.metadata !== undefined && + (!chunk.metadata || typeof chunk.metadata !== "object" || Array.isArray(chunk.metadata)) + ) { + throw new Error(`RAG index chunk ${index} has invalid metadata`) + } + return chunk as Chunk +} + +function indexPath(root: string, containerTag: string): string { + return join(root, sanitizePath(containerTag), "index.json") +} + +export async function loadPersistedRagChunks( + root: string, + containerTag: string +): Promise { + const path = indexPath(root, containerTag) + let content: string + try { + content = await readFile(path, "utf8") + } catch (error) { + if ((error as NodeJS.ErrnoException).code === "ENOENT") return null + throw error + } + + let parsed: unknown + try { + parsed = JSON.parse(content) + } catch (error) { + throw new Error(`RAG index ${path} is unreadable: ${String(error)}`) + } + if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) { + throw new Error(`RAG index ${path} has an invalid root object`) + } + const index = parsed as Partial + if ( + index.schemaVersion !== RAG_INDEX_SCHEMA_VERSION || + index.embeddingModel !== EMBEDDING_MODEL + ) { + throw new Error(`RAG index ${path} uses an incompatible schema or embedding model`) + } + if (!Array.isArray(index.chunks)) throw new Error(`RAG index ${path} has no chunk array`) + const chunks = index.chunks.map(validatePersistedChunk) + if (new Set(chunks.map((chunk) => chunk.id)).size !== chunks.length) { + throw new Error(`RAG index ${path} contains duplicate chunk IDs`) + } + return chunks +} + +export async function persistRagChunks( + root: string, + containerTag: string, + chunks: Chunk[] +): Promise { + const path = indexPath(root, containerTag) + const directory = join(root, sanitizePath(containerTag)) + const temporaryPath = `${path}.${process.pid}.tmp` + await mkdir(directory, { recursive: true }) + const persisted: PersistedRagIndex = { + schemaVersion: RAG_INDEX_SCHEMA_VERSION, + embeddingModel: EMBEDDING_MODEL, + chunks, + } + await writeFile(temporaryPath, JSON.stringify(persisted), "utf8") + await rename(temporaryPath, path) +} + +export function normalizeRagSearchResults( + rawResults: RagSearchResult[], + limit: number, + threshold?: number, + droppedResults: ProviderResultDropDiagnostic[] = [] +): UnifiedSearchResult[] { + requireSearchLimit(limit, "rag") + assertResultBudget(rawResults.length, limit, "rag") + const normalized: Omit[] = [] + for (const [index, result] of rawResults.entries()) { + if (!result.content.trim()) { + recordResultDrop(droppedResults, index, "empty-text") + continue + } + if (threshold !== undefined && result.score < threshold) { + recordResultDrop(droppedResults, index, "below-threshold") + continue + } + const documentDate = + result.date && result.date !== "unknown" + ? resolveDocumentDate({ documentDate: result.date }) + : resolveDocumentDate(result.metadata) + normalized.push({ + id: result.id, + text: result.content, + score: result.score, + sessionId: result.sessionId, + ...(documentDate ? { documentDate } : {}), + provider: "rag", + resultType: "chunk", + }) + } + return rankResults(normalized) +} // ─── Chunking ──────────────────────────────────────────────────────────────── @@ -30,7 +185,11 @@ const EMBEDDING_MODEL = "text-embedding-3-small" * Split text into overlapping chunks, attempting to break on sentence boundaries. * Follows the chunking approach from OpenClaw/QMD: ~400 tokens with overlap. */ -function chunkText(text: string, chunkSize: number = CHUNK_SIZE, overlap: number = CHUNK_OVERLAP): string[] { +function chunkText( + text: string, + chunkSize: number = CHUNK_SIZE, + overlap: number = CHUNK_OVERLAP +): string[] { if (text.length <= chunkSize) { return [text.trim()] } @@ -86,6 +245,8 @@ function chunkText(text: string, chunkSize: number = CHUNK_SIZE, overlap: number */ export class RAGProvider implements Provider { name = "rag" + adapterVersion = "2.2.0" + searchRequestStructure = { kind: "single" } as const prompts = RAG_PROMPTS concurrency = { default: 20, @@ -96,6 +257,40 @@ export class RAGProvider implements Provider { private searchEngine = new HybridSearchEngine() private openai: ReturnType | null = null private apiKey: string = "" + private loadedContainers = new Set() + + constructor(private readonly indexRoot: string = DEFAULT_INDEX_ROOT) {} + + getIngestionConfigFingerprint(_config: ProviderConfig): string { + return stableSha256({ + schemaVersion: 1, + provider: this.name, + adapterVersion: this.adapterVersion, + extractionConfigFingerprint: getMemoryExtractionConfigFingerprint(), + chunkSize: CHUNK_SIZE, + chunkOverlap: CHUNK_OVERLAP, + embeddingBatchSize: EMBEDDING_BATCH_SIZE, + embeddingModel: EMBEDDING_MODEL, + indexSchemaVersion: RAG_INDEX_SCHEMA_VERSION, + sessionReplacement: "replace-complete-session-v1", + }) + } + + private async ensureContainerLoaded(containerTag: string, required: boolean): Promise { + if (this.loadedContainers.has(containerTag)) return + const chunks = await loadPersistedRagChunks(this.indexRoot, containerTag) + if (!chunks) { + if (required) { + throw new Error( + `Durable RAG index is missing for ${containerTag}; rebuild this run from ingestion instead of reusing an empty in-memory build` + ) + } + this.loadedContainers.add(containerTag) + return + } + this.searchEngine.replaceChunks(containerTag, chunks) + this.loadedContainers.add(containerTag) + } async initialize(config: ProviderConfig): Promise { this.apiKey = config.apiKey @@ -103,30 +298,38 @@ export class RAGProvider implements Provider { throw new Error("RAG provider requires OPENAI_API_KEY for memory extraction and embeddings") } this.openai = createOpenAI({ apiKey: this.apiKey }) - logger.info("Initialized RAG memory provider (OpenClaw/QMD-style with LLM extraction + hybrid search)") + await mkdir(this.indexRoot, { recursive: true }) + logger.info( + "Initialized RAG memory provider (OpenClaw/QMD-style with LLM extraction + hybrid search)" + ) } - async ingest(sessions: UnifiedSession[], options: IngestOptions): Promise { + async ingest( + documents: CanonicalIngestionDocument[], + options: IngestOptions + ): Promise { if (!this.openai) throw new Error("Provider not initialized") + await this.ensureContainerLoaded(options.containerTag, false) const allChunks: Array<{ text: string sessionId: string chunkIndex: number - date: string + date?: string metadata?: Record }> = [] // Step 1: Extract memories from each session via LLM, then chunk - for (const session of sessions) { + for (const document of documents) { + const session = canonicalDocumentToSession(document) const extracted = await extractMemories(this.openai, session) // Extract ISO date for OpenClaw-style date organization - const isoDate = (session.metadata?.date as string) || "unknown" - const dateStr = isoDate !== "unknown" ? isoDate.split("T")[0] : "unknown" + const documentDate = document.metadata.documentDate + const dateStr = documentDate?.split("T")[0] // Prepend date context (like OpenClaw's memory/YYYY-MM-DD.md) - const dateHeader = `# Memories from ${dateStr}\n\n` + const dateHeader = dateStr ? `# Memories from ${dateStr}\n\n` : "" const content = dateHeader + extracted const textChunks = chunkText(content) @@ -134,18 +337,25 @@ export class RAGProvider implements Provider { for (let i = 0; i < textChunks.length; i++) { allChunks.push({ text: textChunks[i], - sessionId: session.sessionId, + sessionId: document.metadata.sessionId, chunkIndex: i, - date: dateStr, + date: documentDate, metadata: { - ...session.metadata, - memoryDate: dateStr, + ...document.metadata, + ...(dateStr ? { memoryDate: dateStr } : {}), }, }) } } + const replacedSessionIds = documents.map((document) => document.metadata.sessionId) if (allChunks.length === 0) { + this.searchEngine.replaceSessionChunks(options.containerTag, replacedSessionIds, []) + await persistRagChunks( + this.indexRoot, + options.containerTag, + this.searchEngine.getChunks(options.containerTag) + ) return { documentIds: [] } } @@ -182,11 +392,16 @@ export class RAGProvider implements Provider { } // Step 3: Add to search engine - this.searchEngine.addChunks(options.containerTag, embeddedChunks) + this.searchEngine.replaceSessionChunks(options.containerTag, replacedSessionIds, embeddedChunks) + await persistRagChunks( + this.indexRoot, + options.containerTag, + this.searchEngine.getChunks(options.containerTag) + ) const documentIds = embeddedChunks.map((c) => c.id) logger.debug( - `Ingested ${sessions.length} session(s) as ${embeddedChunks.length} extracted memory chunks for ${options.containerTag}` + `Ingested ${documents.length} session(s) as ${embeddedChunks.length} extracted memory chunks for ${options.containerTag}` ) return { documentIds } @@ -205,8 +420,10 @@ export class RAGProvider implements Provider { }) } - async search(query: string, options: SearchOptions): Promise { + async search(query: string, options: SearchOptions): Promise { if (!this.openai) throw new Error("Provider not initialized") + const limit = requireSearchLimit(options.limit, this.name) + await this.ensureContainerLoaded(options.containerTag, true) // Generate query embedding const embeddingModel = this.openai.embedding(EMBEDDING_MODEL) @@ -215,21 +432,40 @@ export class RAGProvider implements Provider { value: query, }) - const limit = options.limit || 10 - // Hybrid search - const results = this.searchEngine.search(options.containerTag, queryEmbedding, query, limit) + const rawResults = this.searchEngine.search(options.containerTag, queryEmbedding, query, limit) + const droppedResults: ProviderResultDropDiagnostic[] = [] + const results = normalizeRagSearchResults(rawResults, limit, options.threshold, droppedResults) logger.debug( `Search returned ${results.length} results for "${query.substring(0, 50)}..." ` + `(${this.searchEngine.getChunkCount(options.containerTag)} total chunks)` ) - return results + return createProviderSearchResponse({ + results, + requestedLimit: limit, + rawReturnedCount: rawResults.length, + droppedResults, + providerRequests: [ + { + operation: "rag.hybrid", + limit, + parameters: { + vectorWeight: 0.7, + bm25Weight: 0.3, + embeddingModel: EMBEDDING_MODEL, + ...(options.threshold !== undefined ? { threshold: options.threshold } : {}), + }, + }, + ], + }) } async clear(containerTag: string): Promise { this.searchEngine.clear(containerTag) + this.loadedContainers.delete(containerTag) + await rm(join(this.indexRoot, sanitizePath(containerTag)), { recursive: true, force: true }) logger.info(`Cleared RAG data for: ${containerTag}`) } } diff --git a/src/providers/rag/prompts.ts b/src/providers/rag/prompts.ts index e775cfe..b892a52 100644 --- a/src/providers/rag/prompts.ts +++ b/src/providers/rag/prompts.ts @@ -1,18 +1,8 @@ import type { ProviderPrompts } from "../../types/prompts" - -interface RAGSearchResult { - content: string - score: number - vectorScore: number - bm25Score: number - sessionId: string - chunkIndex: number - date?: string - metadata?: Record -} +import type { UnifiedSearchResult } from "../../types/unified" function buildRAGContext(context: unknown[]): string { - const results = context as RAGSearchResult[] + const results = context as UnifiedSearchResult[] if (results.length === 0) { return "No relevant memory chunks were retrieved." @@ -20,17 +10,8 @@ function buildRAGContext(context: unknown[]): string { return results .map((result, i) => { - const scoreParts = [ - `hybrid: ${result.score.toFixed(3)}`, - `semantic: ${result.vectorScore.toFixed(3)}`, - `keyword: ${result.bm25Score.toFixed(3)}`, - ].join(", ") - - const date = result.date || (result.metadata?.date as string) || undefined - const dateStr = date ? ` | Date: ${date}` : "" - - return `[Chunk ${i + 1}] (session: ${result.sessionId}, scores: ${scoreParts}${dateStr}) -${result.content}` + const date = result.documentDate ? ` [${result.documentDate}]` : "" + return `[Chunk ${i + 1}]${date}\n${result.text}` }) .join("\n\n---\n\n") } diff --git a/src/providers/rag/search.ts b/src/providers/rag/search.ts index df2a365..66c3a3c 100644 --- a/src/providers/rag/search.ts +++ b/src/providers/rag/search.ts @@ -22,6 +22,7 @@ export interface Chunk { } export interface SearchResult { + id: string content: string score: number vectorScore: number @@ -278,16 +279,58 @@ export class HybridSearchEngine { return this.containers.get(containerTag)! } - addChunks(containerTag: string, chunks: Chunk[]): void { - const container = this.getContainer(containerTag) + private replaceContainer(containerTag: string, chunks: Iterable): void { + const container = { + chunks: new Map(), + bm25Index: createBM25Index(), + } for (const chunk of chunks) { container.chunks.set(chunk.id, chunk) + } + for (const chunk of container.chunks.values()) { addToBM25Index(container.bm25Index, chunk.id, chunk.content) } + + this.containers.set(containerTag, container) + } + + addChunks(containerTag: string, chunks: Chunk[]): void { + const container = this.getContainer(containerTag) + const merged = new Map(container.chunks) + for (const chunk of chunks) { + merged.set(chunk.id, chunk) + } + // Rebuild the lexical index so retrying a deterministic chunk ID is an + // idempotent replacement rather than a second BM25 document. + this.replaceContainer(containerTag, merged.values()) + } + + /** + * Atomically replace all chunks owned by the supplied sessions. A retry may + * produce fewer chunks than the first attempt, so merging IDs alone would + * leave stale evidence from the abandoned attempt. + */ + replaceSessionChunks(containerTag: string, sessionIds: string[], chunks: Chunk[]): void { + const replaced = new Set(sessionIds) + const retained = this.getChunks(containerTag).filter((chunk) => !replaced.has(chunk.sessionId)) + this.replaceContainer(containerTag, [...retained, ...chunks]) + } + + replaceChunks(containerTag: string, chunks: Chunk[]): void { + this.replaceContainer(containerTag, chunks) + } + + getChunks(containerTag: string): Chunk[] { + return [...(this.containers.get(containerTag)?.chunks.values() ?? [])] } - search(containerTag: string, queryEmbedding: number[], query: string, limit: number): SearchResult[] { + search( + containerTag: string, + queryEmbedding: number[], + query: string, + limit: number + ): SearchResult[] { const container = this.containers.get(containerTag) if (!container || container.chunks.size === 0) return [] @@ -340,6 +383,7 @@ export class HybridSearchEngine { return hybridScores.slice(0, limit).map((result) => { const chunk = container.chunks.get(result.chunkId)! return { + id: chunk.id, content: chunk.content, score: result.score, vectorScore: result.vectorScore, diff --git a/src/providers/supermemory/index.ts b/src/providers/supermemory/index.ts index 027bc32..cf5a93d 100644 --- a/src/providers/supermemory/index.ts +++ b/src/providers/supermemory/index.ts @@ -6,13 +6,237 @@ import type { IngestResult, SearchOptions, IndexingProgressCallback, + ProviderSearchResponse, + AwaitIndexingOptions, } from "../../types/provider" -import type { UnifiedSession } from "../../types/unified" +import type { + CanonicalIngestionDocument, + ProviderResultDropDiagnostic, + UnifiedSearchResult, +} from "../../types/unified" import { logger } from "../../utils/logger" +import { stableSha256 } from "../../utils/stable" import { SUPERMEMORY_PROMPTS } from "./prompts" +import { + asFiniteNumber, + asNonEmptyString, + asRecord, + assertContainerTag, + assertResultBudget, + rankResults, + recordResultDrop, + requireSearchLimit, + resolveDocumentDate, + resolveSessionId, + createProviderSearchResponse, +} from "../normalization" + +type SupermemoryAddBody = Parameters[0] & { + dreaming?: "instant" +} + +type SupermemoryBatchAddBody = { + documents: Array<{ + content: string + customId: string + metadata: Record + }> + containerTag: string + dreaming?: "instant" +} + +type SupermemoryDocumentStatus = { + status?: string + dreamingStatus?: string +} + +export type SupermemoryReadiness = "pending" | "completed" | "failed" + +export function resolveSupermemoryBatchIngestResult( + response: unknown, + documents: CanonicalIngestionDocument[] +): IngestResult { + const responseRecord = asRecord(response) + const rawItems = Array.isArray(response) + ? response + : responseRecord && Array.isArray(responseRecord.results) + ? responseRecord.results + : null + if (!rawItems) throw new Error("Supermemory batch ingest returned an invalid response") + if (rawItems.length !== documents.length) { + throw new Error( + `Supermemory batch ingest returned ${rawItems.length} results for ${documents.length} documents` + ) + } + + const failedByCustomId = new Map() + const successfulDocumentIds: string[] = [] + for (const [index, rawItem] of rawItems.entries()) { + const item = asRecord(rawItem) + const id = asNonEmptyString(item?.id) + const status = asNonEmptyString(item?.status) + if (!item || !id) { + throw new Error( + `Supermemory batch ingest returned an invalid item at index ${index}: missing document ID` + ) + } + if (status === "error") { + if (!documents.some((document) => document.customId === id)) { + throw new Error(`Supermemory batch ingest returned an error for unknown custom ID ${id}`) + } + if (failedByCustomId.has(id)) { + throw new Error(`Supermemory batch ingest returned duplicate failure for ${id}`) + } + failedByCustomId.set( + id, + asNonEmptyString(item.error) ?? asNonEmptyString(item.details) ?? "batch validation failed" + ) + continue + } + successfulDocumentIds.push(id) + } + + const unclaimedFailures = new Set(failedByCustomId.keys()) + let successIndex = 0 + const items = documents.map((document) => { + const error = failedByCustomId.get(document.customId) + if (error) { + unclaimedFailures.delete(document.customId) + return { customId: document.customId, documentIds: [], error } + } + const documentId = successfulDocumentIds[successIndex++] + if (!documentId) { + throw new Error(`Supermemory batch ingest did not return a result for ${document.customId}`) + } + return { customId: document.customId, documentIds: [documentId] } + }) + if (successIndex !== successfulDocumentIds.length || unclaimedFailures.size > 0) { + throw new Error("Supermemory batch ingest response could not be attributed to every input") + } + + return { documentIds: successfulDocumentIds, items } +} + +/** Mono exposes memory inference separately from the normal document lifecycle. */ +export function classifySupermemoryReadiness( + document: SupermemoryDocumentStatus +): SupermemoryReadiness { + if (document.status === "failed") return "failed" + if (document.status === "done" && document.dreamingStatus === "done") return "completed" + return "pending" +} + +function resolveSupermemoryMetadataField( + resultId: string, + fieldName: "sessionId" | "documentDate", + resultMetadata: unknown, + documentMetadata: unknown[], + resolve: (...values: unknown[]) => string | undefined +): string | undefined { + const resultValue = resolve(resultMetadata) + if (resultValue) return resultValue + + const documentValues = documentMetadata + .map((metadata) => resolve(metadata)) + .filter((value): value is string => value !== undefined) + const distinctValues = [...new Set(documentValues)] + if (distinctValues.length > 1) { + throw new Error( + `Supermemory result ${resultId} has conflicting document ${fieldName} values: ${distinctValues.join(", ")}` + ) + } + return documentValues[0] +} + +async function withDeadline( + operation: Promise, + deadlineMs: number, + timeoutMessage: string +): Promise { + const remainingMs = deadlineMs - Date.now() + if (remainingMs <= 0) throw new Error(timeoutMessage) + + let timeout: ReturnType | undefined + try { + return await Promise.race([ + operation, + new Promise((_, reject) => { + timeout = setTimeout(() => reject(new Error(timeoutMessage)), remainingMs) + }), + ]) + } finally { + if (timeout) clearTimeout(timeout) + } +} + +export function normalizeSupermemorySearchResults( + rawResults: unknown[], + limit: number, + droppedResults: ProviderResultDropDiagnostic[] = [] +): UnifiedSearchResult[] { + requireSearchLimit(limit, "supermemory") + assertResultBudget(rawResults.length, limit, "supermemory") + + const normalized: Omit[] = [] + for (const [index, rawResult] of rawResults.entries()) { + const result = asRecord(rawResult) + if (!result) { + recordResultDrop(droppedResults, index, "malformed-result") + continue + } + + const id = asNonEmptyString(result.id) + const memory = asNonEmptyString(result.memory) + const chunk = asNonEmptyString(result.chunk) + if (!id) { + recordResultDrop(droppedResults, index, "missing-id") + continue + } + if (!memory && !chunk) { + recordResultDrop(droppedResults, index, "empty-text") + continue + } + + const metadata = asRecord(result.metadata) + const documents = Array.isArray(result.documents) ? result.documents : [] + const documentMetadata = documents + .map((document) => asRecord(document)) + .map((document) => asRecord(document?.metadata)) + .filter((value): value is NonNullable => value !== undefined) + const score = asFiniteNumber(result.similarity) + const sessionId = resolveSupermemoryMetadataField( + id, + "sessionId", + metadata, + documentMetadata, + resolveSessionId + ) + const documentDate = resolveSupermemoryMetadataField( + id, + "documentDate", + metadata, + documentMetadata, + resolveDocumentDate + ) + + normalized.push({ + id, + text: memory ?? chunk!, + ...(score !== undefined ? { score } : {}), + ...(sessionId ? { sessionId } : {}), + ...(documentDate ? { documentDate } : {}), + provider: "supermemory", + resultType: memory ? "memory" : "chunk", + }) + } + + return rankResults(normalized) +} export class SupermemoryProvider implements Provider { name = "supermemory" + adapterVersion = "2.4.0" + searchRequestStructure = { kind: "single" } as const prompts = SUPERMEMORY_PROMPTS concurrency = { default: 50, @@ -21,48 +245,106 @@ export class SupermemoryProvider implements Provider { } private client: Supermemory | null = null + constructor( + private readonly indexingTimeoutMs = 30 * 60 * 1000, + private readonly initialPollingIntervalMs = 1000 + ) {} + + getIngestionConfigFingerprint(config: ProviderConfig): string { + return stableSha256({ + schemaVersion: 1, + provider: this.name, + adapterVersion: this.adapterVersion, + baseUrl: config.baseUrl ?? "https://api.supermemory.ai", + addContract: "single-or-batch-content-containerTag-customId-metadata-dreaming-idempotency-v3", + readinessContract: "document-status-done-and-dreaming-status-done-v1", + indexingTimeoutMs: this.indexingTimeoutMs, + initialPollingIntervalMs: this.initialPollingIntervalMs, + }) + } + async initialize(config: ProviderConfig): Promise { this.client = new Supermemory({ apiKey: config.apiKey, + ...(config.baseUrl ? { baseURL: config.baseUrl } : {}), }) logger.info(`Initialized Supermemory provider`) } - async ingest(sessions: UnifiedSession[], options: IngestOptions): Promise { + async ingest( + documents: CanonicalIngestionDocument[], + options: IngestOptions + ): Promise { if (!this.client) throw new Error("Provider not initialized") + assertContainerTag(options.containerTag) - const documentIds: string[] = [] - - for (const session of sessions) { - const sessionStr = JSON.stringify(session.messages) - .replace(//g, ">") + if (documents.length > 1) { + const request: SupermemoryBatchAddBody = { + documents: documents.map((document) => ({ + content: document.content, + customId: document.customId, + metadata: document.metadata as Record, + })), + containerTag: options.containerTag, + ...(options.processingMode === "instant" ? { dreaming: "instant" as const } : {}), + } + const response = await this.client.documents.batchAdd( + request as Parameters[0], + { + idempotencyKey: stableSha256({ + provider: this.name, + operation: "documents.batchAdd", + containerTag: options.containerTag, + customIds: documents.map((document) => document.customId), + }), + } + ) + const result = resolveSupermemoryBatchIngestResult(response, documents) + for (const item of result.items ?? []) { + if (item.error) { + logger.warn(`Deferred session ${item.customId}: ${item.error}`) + } else { + logger.debug(`Ingested session ${item.customId}`) + } + } + return result + } - const formattedDate = session.metadata?.formattedDate as string - const isoDate = session.metadata?.date as string - const content = formattedDate - ? `Here is the date the following session took place: ${formattedDate}\n\nHere is the session as a stringified JSON:\n${sessionStr}` - : `Here is the session as a stringified JSON:\n${sessionStr}` + const documentIds: string[] = [] - const response = await this.client.add({ - content, + for (const document of documents) { + const request: SupermemoryAddBody = { + content: document.content, containerTag: options.containerTag, - metadata: { - sessionId: session.sessionId, - ...(isoDate ? { date: isoDate } : {}), - }, + customId: document.customId, + metadata: document.metadata as Record, + ...(options.processingMode === "instant" ? { dreaming: "instant" as const } : {}), + } + const response = await this.client.add(request, { + idempotencyKey: stableSha256({ + provider: this.name, + containerTag: options.containerTag, + customId: document.customId, + }), }) documentIds.push(response.id) - logger.debug(`Ingested session ${session.sessionId}`) + logger.debug(`Ingested session ${document.metadata.sessionId}`) } - return { documentIds } + return { + documentIds, + items: documents.map((document, index) => ({ + customId: document.customId, + documentIds: documentIds[index] ? [documentIds[index]!] : [], + })), + } } async awaitIndexing( result: IngestResult, _containerTag: string, - onProgress?: IndexingProgressCallback + onProgress?: IndexingProgressCallback, + options?: AwaitIndexingOptions ): Promise { if (!this.client) throw new Error("Provider not initialized") if (result.documentIds.length === 0) { @@ -70,34 +352,43 @@ export class SupermemoryProvider implements Provider { return } + const timeoutMs = options?.timeoutMs ?? this.indexingTimeoutMs + if (!Number.isFinite(timeoutMs) || timeoutMs < 0) { + throw new Error(`Supermemory indexing timeout cannot be negative; received ${timeoutMs}`) + } const total = result.documentIds.length const pending = new Set(result.documentIds) const completedIds: string[] = [] const failedIds: string[] = [] - let backoffMs = 1000 + let backoffMs = this.initialPollingIntervalMs + const indexingStartedMs = Date.now() + const indexingDeadlineMs = indexingStartedMs + timeoutMs onProgress?.({ completedIds: [], failedIds: [], total }) while (pending.size > 0) { + const timeoutMessage = `Supermemory indexing timed out after ${timeoutMs}ms with ${pending.size} documents pending` + if (Date.now() >= indexingDeadlineMs) throw new Error(timeoutMessage) + const pendingArray = Array.from(pending) - const results = await Promise.allSettled( - pendingArray.map(async (docId) => { - const doc = await this.client!.documents.get(docId) - if (doc.status === "done" || doc.status === "failed") { - const memory = await this.client!.memories.get(docId) - return { docId, docStatus: doc.status, memStatus: memory.status } - } - return { docId, docStatus: doc.status, memStatus: "pending" } - }) + const results = await withDeadline( + Promise.allSettled( + pendingArray.map(async (docId) => { + const document = (await this.client!.documents.get(docId)) as SupermemoryDocumentStatus + return { docId, readiness: classifySupermemoryReadiness(document) } + }) + ), + indexingDeadlineMs, + timeoutMessage ) for (const res of results) { if (res.status === "fulfilled") { - const { docId, docStatus, memStatus } = res.value - if (docStatus === "failed" || memStatus === "failed") { + const { docId, readiness } = res.value + if (readiness === "failed") { pending.delete(docId) failedIds.push(docId) - } else if (docStatus === "done" && memStatus === "done") { + } else if (readiness === "completed") { pending.delete(docId) completedIds.push(docId) } @@ -107,8 +398,14 @@ export class SupermemoryProvider implements Provider { onProgress?.({ completedIds: [...completedIds], failedIds: [...failedIds], total }) if (pending.size > 0) { - await new Promise((r) => setTimeout(r, backoffMs)) - backoffMs = Math.min(backoffMs * 1.2, 5000) + const remainingMs = indexingDeadlineMs - Date.now() + if (remainingMs <= 0) { + throw new Error( + `Supermemory indexing timed out after ${timeoutMs}ms with ${pending.size} documents pending` + ) + } + await new Promise((r) => setTimeout(r, Math.min(backoffMs, remainingMs))) + backoffMs = Math.min(Math.max(backoffMs * 1.2, this.initialPollingIntervalMs), 5000) } } @@ -117,22 +414,50 @@ export class SupermemoryProvider implements Provider { } } - async search(query: string, options: SearchOptions): Promise { + async search(query: string, options: SearchOptions): Promise { if (!this.client) throw new Error("Provider not initialized") + const limit = requireSearchLimit(options.limit, this.name) + assertContainerTag(options.containerTag) + const searchMode = options.searchMode ?? "hybrid" + if (searchMode !== "memories" && searchMode !== "hybrid") { + throw new Error(`Supermemory adapter does not support search mode: ${searchMode}`) + } + + const threshold = options.threshold ?? 0.6 const response = await this.client.search.memories({ q: query, containerTag: options.containerTag, - limit: 30, - threshold: options.threshold || 0.3, - searchMode: "hybrid", - include: { - summaries: true, - chunks: true - } + limit, + searchMode, + include: { documents: true }, + threshold, + rerank: false, + rewriteQuery: false, }) - return response.results || [] + const rawResults = response.results ?? [] + const droppedResults: ProviderResultDropDiagnostic[] = [] + return createProviderSearchResponse({ + results: normalizeSupermemorySearchResults(rawResults, limit, droppedResults), + requestedLimit: limit, + rawReturnedCount: rawResults.length, + droppedResults, + providerRequests: [ + { + operation: `search.${searchMode}`, + limit, + parameters: { + searchMode, + threshold, + includeDocuments: true, + includeChunks: false, + rerank: false, + rewriteQuery: false, + }, + }, + ], + }) } async clear(containerTag: string): Promise { diff --git a/src/providers/supermemory/prompts.ts b/src/providers/supermemory/prompts.ts index 5113b89..45aebf4 100644 --- a/src/providers/supermemory/prompts.ts +++ b/src/providers/supermemory/prompts.ts @@ -1,83 +1,16 @@ import type { ProviderPrompts } from "../../types/prompts" - -interface SupermemoryChunk { - content: string - position: number -} - -interface SupermemoryResult { - memory?: string - chunk?: string - chunks?: SupermemoryChunk[] - metadata?: { - temporalContext?: { - documentDate?: string - eventDate?: string | string[] - } - } -} - -function deduplicateAndSortChunks(chunks: SupermemoryChunk[]): SupermemoryChunk[] { - const uniqueChunks = chunks.filter( - (chunk, index, self) => index === self.findIndex((c) => c.content === chunk.content) - ) - return uniqueChunks.sort((a, b) => a.position - b.position) -} +import type { UnifiedSearchResult } from "../../types/unified" function buildSupermemoryContext(context: unknown[]): string { - const results = context as SupermemoryResult[] - const allChunks: SupermemoryChunk[] = [] - - for (let i = 0; i < results.length; i++) { - const result = results[i] - - const chunks = result.chunks || [] - for (const chunk of chunks) { - allChunks.push({ - content: chunk.content, - position: chunk.position ?? 0, - }) - } - - if (result.chunk && typeof result.chunk === "string" && result.chunk.trim()) { - allChunks.push({ - content: result.chunk, - position: i, - }) - } - } - - const deduplicatedChunks = deduplicateAndSortChunks(allChunks) - - const memoriesSection = results - .map((result, i) => { - const memory = result.memory || "" - const temporalContext = result.metadata?.temporalContext - const documentDate = temporalContext?.documentDate - const eventDate = temporalContext?.eventDate - - const memoryParts = [`Result ${i + 1}:`, memory] - - if (documentDate || eventDate) { - const temporalInfo: string[] = [] - if (documentDate) temporalInfo.push(`documentDate: ${documentDate}`) - if (eventDate) { - const eventDates = Array.isArray(eventDate) ? eventDate : [eventDate] - temporalInfo.push(`eventDate: ${eventDates.join(", ")}`) - } - memoryParts.push(`Temporal Context: ${temporalInfo.join(" | ")}`) - } + const results = context as UnifiedSearchResult[] + if (results.length === 0) return "No relevant evidence was retrieved." - return memoryParts.join("\n") + return results + .map((result, index) => { + const date = result.documentDate ? `[${result.documentDate}] ` : "" + return `${index + 1}. ${date}${result.text}` }) - .join("\n\n---\n\n") - - const chunksSection = - deduplicatedChunks.length > 0 - ? `\n\n=== DEDUPLICATED CHUNKS ===\n${deduplicatedChunks.map((chunk) => chunk.content).join("\n\n---\n\n")}` - : "" - - return memoriesSection + chunksSection + .join("\n") } export function buildSupermemoryAnswerPrompt( @@ -85,74 +18,21 @@ export function buildSupermemoryAnswerPrompt( context: unknown[], questionDate?: string ): string { - const results = context as SupermemoryResult[] - const retrievedContext = buildSupermemoryContext(context) - - // console.log(`\n=== DEBUG: Processing ${results.length} search results ===`) - // for (let i = 0; i < Math.min(results.length, 3); i++) { - // const r = results[i] - // console.log(`Result ${i + 1}:`) - // console.log(` - memory: ${r.memory?.substring(0, 80)}...`) - // console.log(` - chunk (singular): ${r.chunk ? r.chunk.substring(0, 80) + "..." : "EMPTY"}`) - // console.log(` - chunks (array): ${r.chunks?.length || 0} items`) - // if (r.chunks && r.chunks.length > 0) { - // console.log(` First chunk: ${r.chunks[0].content?.substring(0, 80)}...`) - // } - // } - // console.log(`\n=== Total chunks extracted: ${retrievedContext.includes("DEDUPLICATED CHUNKS") ? "YES" : "NO CHUNKS"} ===`) - // console.log("Retrieved context preview:", retrievedContext) - - return `You are a question-answering system. Based on the retrieved context below, answer the question. + return `You are a question-answering system. Based only on the retrieved evidence below, answer the question. Question: ${question} Question Date: ${questionDate || "Not specified"} -Retrieved Context: -${retrievedContext} - -**Understanding the Context:** -The context contains search results from a memory system. Each result has multiple components you can use: - -1. **Memory**: A high-level summary/atomic fact (e.g., "Alex loves hiking in mountains", "John reports to Maria") - - This is the searchable title/summary of what was stored - -2. **Chunks**: The actual detailed raw content where the memory was extracted from - - Contains conversations, documents, messages, or text excerpts - - **This is your primary source for detailed information and facts** - - Look here for specifics, context, quotes, and evidence - -3. **Temporal Context** (if present): - - **Question Date**: The date when the question was asked (provided above). Use this to understand the temporal perspective of the question. - - **documentDate**: ISO date string for when the content was originally authored/written/said by the user (NOT the system createdAt timestamp). This is the reference point for calculating relative dates. Extract from document metadata, timestamps, or context. - - **eventDate**: Array of ISO date strings for when the event/fact being referenced actually occurred or will occur. Always provided as an array, even for single dates. For past events use past dates, for future events use future dates. Calculate relative dates (today, yesterday, last week) based on documentDate, NOT the current date. - - Useful for time-based questions (what happened when, recent vs old info) - - **Important**: When you see relative terms like "today", "yesterday", calculate them relative to the documentDate, NOT the current date. The question date helps you understand the temporal context of what the user is asking about. - -4. **Version**: Shows if a memory has been updated/extended over time - -**How to Answer:** -1. Start by scanning memory titles to find relevant results -2. **Read the chunks carefully** - they contain the actual details you need -3. Use temporal context to understand when things happened -4. Synthesize information from multiple results if needed +Retrieved Evidence: +${buildSupermemoryContext(context)} Instructions: -- First, think through the problem step by step. Show your reasoning process. -- Identify which parts of the context are relevant to answering the question -- Consider temporal relationships, sequences of events, and any updates to information over time -- If the context contains enough information to answer the question, provide a clear, concise answer -- If the context does not contain enough information, respond with "I don't know" or explain what information is missing -- Base your answer ONLY on the provided context -- **Prioritize information from chunks** - they're the raw source material - -**Response Format:** -Think step by step, then provide your answer. - -Reasoning: -[Your step-by-step reasoning process here] +- Read all evidence before answering. +- Use document dates to resolve temporal relationships and prefer newer evidence when facts conflict. +- If the evidence is insufficient, respond with "I don't know". +- Give a clear, concise answer without exposing internal result metadata. -Answer: -[Your final answer here]` +Answer:` } export const SUPERMEMORY_PROMPTS: ProviderPrompts = { diff --git a/src/providers/zep/index.ts b/src/providers/zep/index.ts index 083b3db..97e9255 100644 --- a/src/providers/zep/index.ts +++ b/src/providers/zep/index.ts @@ -1,4 +1,4 @@ -import { ZepClient, Zep } from "@getzep/zep-cloud" +import { ZepClient } from "@getzep/zep-cloud" import type { Provider, ProviderConfig, @@ -6,13 +6,130 @@ import type { IngestResult, SearchOptions, IndexingProgressCallback, + ProviderSearchResponse, } from "../../types/provider" -import type { UnifiedSession } from "../../types/unified" +import type { + CanonicalIngestionDocument, + ProviderResultDropDiagnostic, + UnifiedSearchResult, +} from "../../types/unified" import { logger } from "../../utils/logger" +import { stableSha256 } from "../../utils/stable" import { ZEP_PROMPTS } from "./prompts" +import { + asFiniteNumber, + asNonEmptyString, + asRecord, + assertResultBudget, + rankResults, + recordResultDrop, + requireSearchLimit, + createProviderSearchResponse, +} from "../normalization" const MAX_DATA_SIZE = 9500 +async function withDeadline( + operation: Promise, + deadlineMs: number, + timeoutMessage: string +): Promise { + const remainingMs = deadlineMs - Date.now() + if (remainingMs <= 0) throw new Error(timeoutMessage) + + let timeout: ReturnType | undefined + try { + return await Promise.race([ + operation, + new Promise((_, reject) => { + timeout = setTimeout(() => reject(new Error(timeoutMessage)), remainingMs) + }), + ]) + } finally { + if (timeout) clearTimeout(timeout) + } +} + +export function allocateZepSearchBudget(limit: number): { edgeLimit: number; nodeLimit: number } { + requireSearchLimit(limit, "zep") + return { + edgeLimit: Math.ceil(limit / 2), + nodeLimit: Math.floor(limit / 2), + } +} + +export function normalizeZepSearchResults( + rawResults: unknown[], + limit: number, + threshold?: number, + droppedResults: ProviderResultDropDiagnostic[] = [] +): UnifiedSearchResult[] { + requireSearchLimit(limit, "zep") + assertResultBudget(rawResults.length, limit, "zep") + + const normalized: Array & { score?: number }> = [] + for (const [index, rawResult] of rawResults.entries()) { + const result = asRecord(rawResult) + if (!result) { + recordResultDrop(droppedResults, index, "malformed-result") + continue + } + + const resultType = result._type + const id = asNonEmptyString(result.uuid) + const score = asFiniteNumber(result.relevance) ?? asFiniteNumber(result.score) + if (!id) { + recordResultDrop(droppedResults, index, "missing-id") + continue + } + if (threshold !== undefined && score !== undefined && score < threshold) { + recordResultDrop(droppedResults, index, "below-threshold") + continue + } + + if (resultType === "edge") { + const text = asNonEmptyString(result.fact) + if (!text) { + recordResultDrop(droppedResults, index, "empty-text") + continue + } + normalized.push({ + id, + text, + ...(score !== undefined ? { score } : {}), + provider: "zep", + resultType: "graph-edge", + }) + continue + } + + if (resultType === "node") { + const name = asNonEmptyString(result.name) + const summary = asNonEmptyString(result.summary) + const text = name && summary ? `${name}: ${summary}` : (summary ?? name) + if (!text) { + recordResultDrop(droppedResults, index, "empty-text") + continue + } + normalized.push({ + id, + text, + ...(score !== undefined ? { score } : {}), + provider: "zep", + resultType: "graph-node", + }) + continue + } + + recordResultDrop(droppedResults, index, "unsupported-result-type") + } + + normalized.sort( + (a, b) => (b.score ?? Number.NEGATIVE_INFINITY) - (a.score ?? Number.NEGATIVE_INFINITY) + ) + return rankResults(normalized) +} + function splitIntoChunks(text: string, maxSize: number): string[] { if (text.length <= maxSize) return [text] @@ -81,6 +198,8 @@ const ZEP_ENTITY_TYPES = { export class ZepProvider implements Provider { name = "zep" + adapterVersion = "2.1.0" + searchRequestStructure = { kind: "split", budget: "shared-total" } as const prompts = ZEP_PROMPTS concurrency = { default: 10, @@ -90,12 +209,31 @@ export class ZepProvider implements Provider { private graphIds: Map = new Map() private ontologySet: Set = new Set() + constructor(private readonly indexingTimeoutMs = 30 * 60 * 1000) {} + + getIngestionConfigFingerprint(_config: ProviderConfig): string { + return stableSha256({ + schemaVersion: 1, + provider: this.name, + adapterVersion: this.adapterVersion, + maxDataSize: MAX_DATA_SIZE, + episodeType: "message", + sourceDescription: "memorybench::/", + addOrder: "sequential", + entityTypes: ZEP_ENTITY_TYPES, + indexingTimeoutMs: this.indexingTimeoutMs, + }) + } + async initialize(config: ProviderConfig): Promise { this.client = new ZepClient({ apiKey: config.apiKey }) logger.info(`Initialized Zep provider`) } - async ingest(sessions: UnifiedSession[], options: IngestOptions): Promise { + async ingest( + documents: CanonicalIngestionDocument[], + options: IngestOptions + ): Promise { if (!this.client) throw new Error("Provider not initialized") const graphId = `memorybench_${options.containerTag.replace(/[^a-zA-Z0-9_-]/g, "_")}` @@ -113,64 +251,53 @@ export class ZepProvider implements Provider { } if (!this.ontologySet.has(graphId)) { - try { - await this.client.graph.setOntology(ZEP_ENTITY_TYPES, {}, { graphIds: [graphId] }) - this.ontologySet.add(graphId) - logger.debug(`Set ontology for graph: ${graphId}`) - } catch (e) { - logger.debug(`Ontology may already be set: ${e}`) - } + await this.client.graph.setOntology(ZEP_ENTITY_TYPES, {}, { graphIds: [graphId] }) + this.ontologySet.add(graphId) + logger.debug(`Set ontology for graph: ${graphId}`) } - const episodes: Zep.EpisodeData[] = [] - - for (const session of sessions) { - const rawDate = session.metadata?.date as string | undefined + const documentIds: string[] = [] + for (const document of documents) { + const rawDate = document.metadata.documentDate const isoDate = rawDate && /^\d{4}-\d{2}-\d{2}$/.test(rawDate) ? `${rawDate}T00:00:00Z` : rawDate - - for (const message of session.messages) { - const speaker = message.speaker || message.role - const messageData = `${speaker}: ${message.content}` - - if (messageData.length > MAX_DATA_SIZE) { - const chunks = splitIntoChunks(messageData, MAX_DATA_SIZE) - for (const chunk of chunks) { - episodes.push({ - type: "message", - data: chunk, - createdAt: isoDate, - }) - } - } else { - episodes.push({ - type: "message", - data: messageData, - createdAt: isoDate, - }) - } - } - logger.debug(`Ingested session ${session.sessionId}`) - } - - const taskIds: string[] = [] - - const BATCH_SIZE = 20 - for (let i = 0; i < episodes.length; i += BATCH_SIZE) { - const batch = episodes.slice(i, i + BATCH_SIZE) - const result = await this.client.graph.addBatch({ - graphId, - episodes: batch, + const chunks = splitIntoChunks(document.content, MAX_DATA_SIZE) + const sourceDescriptions = chunks.map( + (_chunk, index) => `memorybench:${document.customId}:${index + 1}/${chunks.length}` + ) + const recent = await this.client.graph.episode.getByGraphId(graphId, { + lastn: Math.min(1000, Math.max(50, chunks.length * 2)), }) + const existingBySource = new Map( + (recent.episodes ?? []).flatMap((episode) => + episode.sourceDescription && sourceDescriptions.includes(episode.sourceDescription) + ? [[episode.sourceDescription, episode] as const] + : [] + ) + ) - for (const episode of result) { - if (episode.taskId) { - taskIds.push(episode.taskId) + for (let index = 0; index < chunks.length; index++) { + const sourceDescription = sourceDescriptions[index]! + const existing = existingBySource.get(sourceDescription) + if (existing) { + documentIds.push(existing.uuid) + continue } + + // Sequential adds preserve the exact chunk order; addBatch explicitly does not. + const episode = await this.client.graph.add({ + graphId, + type: "message", + data: chunks[index]!, + createdAt: isoDate, + sourceDescription, + }) + documentIds.push(episode.uuid) } + logger.debug(`Ingested or reconciled session ${document.metadata.sessionId}`) } - return { documentIds: [], taskIds: [...new Set(taskIds)] } + return { documentIds: [...new Set(documentIds)] } } async awaitIndexing( @@ -181,45 +308,77 @@ export class ZepProvider implements Provider { if (!this.client) throw new Error("Provider not initialized") const taskIds = result.taskIds || [] - if (taskIds.length === 0) { + const episodeIds = result.documentIds || [] + if (taskIds.length === 0 && episodeIds.length === 0) { onProgress?.({ completedIds: [], failedIds: [], total: 0 }) return } - const total = taskIds.length - const pending = new Set(taskIds) + const total = taskIds.length + episodeIds.length + const pendingTasks = new Set(taskIds) + const pendingEpisodes = new Set(episodeIds) const completedIds: string[] = [] const failedIds: string[] = [] let backoffMs = 500 + const indexingStartedMs = Date.now() + const indexingDeadlineMs = indexingStartedMs + this.indexingTimeoutMs onProgress?.({ completedIds: [], failedIds: [], total }) - while (pending.size > 0) { - const pendingArray = Array.from(pending) - const results = await Promise.allSettled( - pendingArray.map((taskId) => this.client!.task.get(taskId)) + while (pendingTasks.size > 0 || pendingEpisodes.size > 0) { + const timeoutMessage = `Zep indexing timed out after ${this.indexingTimeoutMs}ms with ${pendingTasks.size + pendingEpisodes.size} items pending` + if (Date.now() >= indexingDeadlineMs) throw new Error(timeoutMessage) + + const pendingTaskArray = Array.from(pendingTasks) + const taskResults = await withDeadline( + Promise.allSettled(pendingTaskArray.map((taskId) => this.client!.task.get(taskId))), + indexingDeadlineMs, + timeoutMessage ) - for (let i = 0; i < results.length; i++) { - const taskId = pendingArray[i] - const res = results[i] + for (let i = 0; i < taskResults.length; i++) { + const taskId = pendingTaskArray[i] + const res = taskResults[i] if (res.status === "fulfilled") { const task = res.value if (task.status === "succeeded" || task.status === "completed") { - pending.delete(taskId) + pendingTasks.delete(taskId) completedIds.push(taskId) } else if (task.status === "failed") { - pending.delete(taskId) + pendingTasks.delete(taskId) failedIds.push(taskId) } } } + const pendingEpisodeArray = Array.from(pendingEpisodes) + const episodeResults = await withDeadline( + Promise.allSettled( + pendingEpisodeArray.map((episodeId) => this.client!.graph.episode.get(episodeId)) + ), + indexingDeadlineMs, + timeoutMessage + ) + for (let index = 0; index < episodeResults.length; index++) { + const episodeId = pendingEpisodeArray[index]! + const response = episodeResults[index]! + if (response.status === "fulfilled" && response.value.processed) { + pendingEpisodes.delete(episodeId) + completedIds.push(episodeId) + } + } + onProgress?.({ completedIds: [...completedIds], failedIds: [...failedIds], total }) - if (pending.size > 0) { - await new Promise((r) => setTimeout(r, backoffMs)) + if (pendingTasks.size > 0 || pendingEpisodes.size > 0) { + const remainingMs = indexingDeadlineMs - Date.now() + if (remainingMs <= 0) { + throw new Error( + `Zep indexing timed out after ${this.indexingTimeoutMs}ms with ${pendingTasks.size + pendingEpisodes.size} items pending` + ) + } + await new Promise((r) => setTimeout(r, Math.min(backoffMs, remainingMs))) backoffMs = Math.min(backoffMs * 1.5, 5000) } } @@ -229,8 +388,9 @@ export class ZepProvider implements Provider { } } - async search(query: string, options: SearchOptions): Promise { + async search(query: string, options: SearchOptions): Promise { if (!this.client) throw new Error("Provider not initialized") + const limit = requireSearchLimit(options.limit, this.name) const graphId = this.graphIds.get(options.containerTag) if (!graphId) { @@ -240,8 +400,7 @@ export class ZepProvider implements Provider { } const finalGraphId = this.graphIds.get(options.containerTag)! - const edgeLimit = options.limit || 20 - const nodeLimit = Math.min(edgeLimit, 10) + const { edgeLimit, nodeLimit } = allocateZepSearchBudget(limit) const [edgesResponse, nodesResponse] = await Promise.all([ this.client.graph.search({ @@ -251,13 +410,15 @@ export class ZepProvider implements Provider { scope: "edges", reranker: "cross_encoder", }), - this.client.graph.search({ - graphId: finalGraphId, - query, - limit: nodeLimit, - scope: "nodes", - reranker: "cross_encoder", - }), + nodeLimit > 0 + ? this.client.graph.search({ + graphId: finalGraphId, + query, + limit: nodeLimit, + scope: "nodes", + reranker: "cross_encoder", + }) + : Promise.resolve({ nodes: [] }), ]) const results: unknown[] = [] @@ -274,7 +435,37 @@ export class ZepProvider implements Provider { } } - return results + const droppedResults: ProviderResultDropDiagnostic[] = [] + return createProviderSearchResponse({ + results: normalizeZepSearchResults(results, limit, options.threshold, droppedResults), + requestedLimit: limit, + rawReturnedCount: results.length, + droppedResults, + providerRequests: [ + { + operation: "graph.edges", + limit: edgeLimit, + parameters: { + scope: "edges", + reranker: "cross_encoder", + ...(options.threshold !== undefined ? { threshold: options.threshold } : {}), + }, + }, + ...(nodeLimit > 0 + ? [ + { + operation: "graph.nodes", + limit: nodeLimit, + parameters: { + scope: "nodes", + reranker: "cross_encoder", + ...(options.threshold !== undefined ? { threshold: options.threshold } : {}), + }, + }, + ] + : []), + ], + }) } async clear(containerTag: string): Promise { diff --git a/src/providers/zep/prompts.ts b/src/providers/zep/prompts.ts index 5b9f581..a11a8ea 100644 --- a/src/providers/zep/prompts.ts +++ b/src/providers/zep/prompts.ts @@ -1,30 +1,17 @@ import type { ProviderPrompts } from "../../types/prompts" - -interface ZepResult { - _type?: string - fact?: string - name?: string - summary?: string - valid_at?: string - invalid_at?: string -} +import type { UnifiedSearchResult } from "../../types/unified" function buildZepContext(context: unknown[]): string { const facts: string[] = [] const entities: string[] = [] for (const r of context) { - const result = r as ZepResult - const type = result._type + const result = r as UnifiedSearchResult - if (type === "node") { - const name = result.name || "Unknown" - const summary = result.summary || "" - entities.push(` - ${name}: ${summary}`) + if (result.resultType === "graph-node") { + entities.push(` - ${result.text}`) } else { - const content = result.fact || JSON.stringify(r) - const validAt = result.valid_at - facts.push(` - ${content} (event_time: ${validAt || "unknown"})`) + facts.push(` - ${result.text}`) } } diff --git a/src/server/db/index.ts b/src/server/db/index.ts index b05b540..6e128d6 100644 --- a/src/server/db/index.ts +++ b/src/server/db/index.ts @@ -30,6 +30,17 @@ export function initDatabase() { provider TEXT NOT NULL, benchmark TEXT NOT NULL, version TEXT NOT NULL DEFAULT 'baseline', + benchmark_scope TEXT, + dataset_identity TEXT, + dataset_fingerprint TEXT, + question_set_fingerprint TEXT, + protocol_identity TEXT, + protocol_fingerprint TEXT, + retrieval_top_k INTEGER, + primary_metric_key TEXT, + primary_metric_value REAL, + primary_metric_higher_is_better INTEGER, + comparison_cohort_key TEXT, accuracy REAL NOT NULL, total_questions INTEGER NOT NULL, correct_count INTEGER NOT NULL, @@ -45,9 +56,37 @@ export function initDatabase() { ) `) + const columns = new Set( + ( + sqlite.query("PRAGMA table_info(leaderboard_entries)").all() as Array<{ + name: string + }> + ).map((column) => column.name) + ) + const identityColumns: Array<[string, string]> = [ + ["benchmark_scope", "TEXT"], + ["dataset_identity", "TEXT"], + ["dataset_fingerprint", "TEXT"], + ["question_set_fingerprint", "TEXT"], + ["protocol_identity", "TEXT"], + ["protocol_fingerprint", "TEXT"], + ["retrieval_top_k", "INTEGER"], + ["primary_metric_key", "TEXT"], + ["primary_metric_value", "REAL"], + ["primary_metric_higher_is_better", "INTEGER"], + ["comparison_cohort_key", "TEXT"], + ] + for (const [name, type] of identityColumns) { + if (!columns.has(name)) { + sqlite.exec(`ALTER TABLE leaderboard_entries ADD COLUMN ${name} ${type}`) + } + } + + // The old index collapsed different datasets/protocols/retrieval policies. + sqlite.exec("DROP INDEX IF EXISTS provider_benchmark_version_idx") sqlite.exec(` - CREATE UNIQUE INDEX IF NOT EXISTS provider_benchmark_version_idx - ON leaderboard_entries (provider, benchmark, version) + CREATE UNIQUE INDEX IF NOT EXISTS provider_benchmark_version_cohort_idx + ON leaderboard_entries (provider, benchmark, version, comparison_cohort_key) `) } diff --git a/src/server/db/schema.ts b/src/server/db/schema.ts index 6a19a6b..3e2f245 100644 --- a/src/server/db/schema.ts +++ b/src/server/db/schema.ts @@ -11,6 +11,22 @@ export const leaderboardEntries = sqliteTable( benchmark: text("benchmark").notNull(), version: text("version").notNull().default("baseline"), + // Immutable like-for-like comparison identity. JSON columns preserve the + // complete source identities; scalar fingerprints make cohorts auditable. + benchmarkScope: text("benchmark_scope"), + datasetIdentity: text("dataset_identity"), + datasetFingerprint: text("dataset_fingerprint"), + questionSetFingerprint: text("question_set_fingerprint"), + protocolIdentity: text("protocol_identity"), + protocolFingerprint: text("protocol_fingerprint"), + retrievalTopK: integer("retrieval_top_k"), + primaryMetricKey: text("primary_metric_key"), + primaryMetricValue: real("primary_metric_value"), + primaryMetricHigherIsBetter: integer("primary_metric_higher_is_better", { + mode: "boolean", + }), + comparisonCohortKey: text("comparison_cohort_key"), + // Results snapshot accuracy: real("accuracy").notNull(), totalQuestions: integer("total_questions").notNull(), @@ -36,11 +52,12 @@ export const leaderboardEntries = sqliteTable( notes: text("notes"), }, (table) => ({ - // Unique constraint: same provider+benchmark+version replaces existing entry - providerBenchmarkVersion: uniqueIndex("provider_benchmark_version_idx").on( + // A display version only replaces a result from the exact same cohort. + providerBenchmarkVersionCohort: uniqueIndex("provider_benchmark_version_cohort_idx").on( table.provider, table.benchmark, - table.version + table.version, + table.comparisonCohortKey ), }) ) diff --git a/src/server/leaderboard-identity.ts b/src/server/leaderboard-identity.ts new file mode 100644 index 0000000..f777c92 --- /dev/null +++ b/src/server/leaderboard-identity.ts @@ -0,0 +1,500 @@ +import { sha256Text, stableSha256 } from "../utils/stable" +import type { RunCheckpoint } from "../types/checkpoint" +import type { BenchmarkResult } from "../types/unified" +import type { AnsweringRuntimeIdentity } from "../types/model" +import type { BenchmarkProtocol, QuestionEvaluation } from "../types/protocol" +import type { UnifiedQuestion } from "../types/unified" +import { resolveAnsweringRuntimeIdentity } from "../utils/models" + +export interface LeaderboardPrimaryMetric { + key: string + value: number + higherIsBetter: boolean +} + +export interface LeaderboardComparisonIdentity { + schemaVersion: 3 + benchmark: string + benchmarkScope: Record + datasetIdentity: Record + datasetFingerprint: string + questionSetFingerprint: string + benchmarkInputFingerprint: string + protocolIdentity: Record + protocolFingerprint: string + retrievalTopK: number | null + judgeModel: string + answeringModel: string + answeringRuntimeFingerprint: string + providerPromptFingerprint: string | null + primaryMetric: LeaderboardPrimaryMetric + cohortKey: string + legacy: boolean +} + +export interface LeaderboardIdentitySource { + benchmark: string + benchmarkScope?: Record + datasetIdentity?: Record + selectedQuestionIdsDigest?: string + benchmarkInputFingerprint?: string + protocolIdentity?: Record + retrievalTopK?: number + questionMetrics?: Array<{ configuredTopK?: number }> + judgeModel?: string + answeringModel?: string + answeringRuntimeIdentity?: AnsweringRuntimeIdentity + providerPromptFingerprint?: string | null + primaryMetric?: Partial + accuracy: number +} + +export interface RankableLeaderboardEntry { + id: number + benchmark: string + accuracy: number + judgeModel?: string + answeringModel?: string + providerPromptFingerprint?: string | null + addedAt?: string + comparisonIdentity?: LeaderboardComparisonIdentity + quality?: { + primaryMetric?: Partial + } +} + +export type RankedLeaderboardEntry = T & { + comparisonIdentity: LeaderboardComparisonIdentity + cohortKey: string + cohortRank: number + cohortSize: number +} + +export interface LeaderboardAggregationContext { + protocol: Pick + questions: UnifiedQuestion[] +} + +function finiteNumber(value: unknown, fallback: number): number { + return typeof value === "number" && Number.isFinite(value) ? value : fallback +} + +function nonEmptyString(value: unknown, fallback: string): string { + return typeof value === "string" && value.trim() ? value : fallback +} + +function resolveRetrievalTopK(source: LeaderboardIdentitySource): number | null { + const reported = [ + ...new Set( + (source.questionMetrics ?? []) + .map((metric) => metric.configuredTopK) + .filter( + (value): value is number => + typeof value === "number" && Number.isInteger(value) && value > 0 + ) + ), + ] + + if (reported.length > 1) { + throw new Error( + `Cannot publish a run with mixed retrieval Top-K values: ${reported.join(", ")}` + ) + } + + if ( + reported.length === 1 && + Number.isInteger(source.retrievalTopK) && + source.retrievalTopK !== reported[0] + ) { + throw new Error( + `Cannot publish a run whose configured retrieval Top-K (${source.retrievalTopK}) differs from its recorded question Top-K (${reported[0]})` + ) + } + + if (reported.length === 1) return reported[0] + return Number.isInteger(source.retrievalTopK) && (source.retrievalTopK ?? 0) > 0 + ? source.retrievalTopK! + : null +} + +/** + * Build the immutable comparison identity stored with a leaderboard snapshot. + * The primary metric value is deliberately excluded from cohortKey: scores vary + * inside a cohort, while every input and metric semantic must remain identical. + */ +export function createLeaderboardComparisonIdentity( + source: LeaderboardIdentitySource +): LeaderboardComparisonIdentity { + const benchmarkScope = source.benchmarkScope ?? { + displayName: source.benchmark, + includedTiers: [], + coverage: "subset", + } + const questionSetFingerprint = nonEmptyString( + source.selectedQuestionIdsDigest, + stableSha256({ benchmark: source.benchmark, questionSet: "unknown" }) + ) + const benchmarkInputFingerprint = nonEmptyString( + source.benchmarkInputFingerprint, + `derived:${stableSha256({ benchmark: source.benchmark, questionSetFingerprint })}` + ) + const datasetIdentity = source.datasetIdentity ?? { + identityKind: "derived-from-question-set", + benchmark: source.benchmark, + questionSetFingerprint, + } + const datasetFingerprint = nonEmptyString( + datasetIdentity.datasetFingerprint, + stableSha256(datasetIdentity) + ) + const protocolIdentity = source.protocolIdentity ?? { + id: "memorybench.legacy", + version: "unknown", + } + const protocolFingerprint = stableSha256(protocolIdentity) + const retrievalTopK = resolveRetrievalTopK(source) + const judgeModel = nonEmptyString(source.judgeModel, "unknown") + const answeringModel = nonEmptyString(source.answeringModel, "unknown") + const answeringRuntimeIdentity = + source.answeringRuntimeIdentity ?? resolveAnsweringRuntimeIdentity(answeringModel) + const answeringRuntimeFingerprint = stableSha256(answeringRuntimeIdentity) + const providerPromptFingerprint = + protocolIdentity.id === "memorybench.legacy" + ? nonEmptyString(source.providerPromptFingerprint, "unknown") + : null + const primaryMetric: LeaderboardPrimaryMetric = { + key: nonEmptyString(source.primaryMetric?.key, "accuracy"), + value: finiteNumber(source.primaryMetric?.value, source.accuracy), + higherIsBetter: source.primaryMetric?.higherIsBetter !== false, + } + + const cohortKey = stableSha256({ + schemaVersion: 3, + benchmark: source.benchmark, + benchmarkScope, + datasetFingerprint, + questionSetFingerprint, + benchmarkInputFingerprint, + protocolFingerprint, + retrievalTopK, + judgeModel, + answeringModel, + answeringRuntimeFingerprint, + providerPromptFingerprint, + primaryMetric: { + key: primaryMetric.key, + higherIsBetter: primaryMetric.higherIsBetter, + }, + }) + + return { + schemaVersion: 3, + benchmark: source.benchmark, + benchmarkScope, + datasetIdentity, + datasetFingerprint, + questionSetFingerprint, + benchmarkInputFingerprint, + protocolIdentity, + protocolFingerprint, + retrievalTopK, + judgeModel, + answeringModel, + answeringRuntimeFingerprint, + providerPromptFingerprint, + primaryMetric, + cohortKey, + legacy: + !source.datasetIdentity || + !source.protocolIdentity || + !source.benchmarkInputFingerprint || + !source.judgeModel || + !source.answeringModel || + !source.answeringRuntimeIdentity || + (protocolIdentity.id === "memorybench.legacy" && !source.providerPromptFingerprint), + } +} + +function approximatelyEqual(left: number, right: number): boolean { + return Math.abs(left - right) <= 1e-12 +} + +function checkpointDatasetIdentity(checkpoint: RunCheckpoint): Record { + if (checkpoint.datasetIdentity) { + return checkpoint.datasetIdentity as unknown as Record + } + return { + identityKind: "selected-provider-visible-haystacks-v1", + benchmark: checkpoint.benchmark, + selectedQuestionIdsDigest: checkpoint.selectedQuestionIdsDigest, + builds: Object.values(checkpoint.builds) + .map((build) => ({ + ingestionGroupId: build.ingestionGroupId, + memberQuestionIds: build.memberQuestionIds, + haystackFingerprint: build.haystack.fingerprint, + })) + .sort((left, right) => left.ingestionGroupId.localeCompare(right.ingestionGroupId)), + } +} + +/** Fail closed before a report snapshot is admitted to a ranked leaderboard. */ +export function validateLeaderboardReportForPublication( + checkpoint: RunCheckpoint, + report: BenchmarkResult, + evaluatedQuestionCount: number, + aggregation: LeaderboardAggregationContext +): LeaderboardComparisonIdentity { + const errors: string[] = [] + if (report.runId !== checkpoint.runId) errors.push("run ID") + if (report.provider !== checkpoint.provider) errors.push("provider") + if (report.providerPromptFingerprint !== checkpoint.providerPromptFingerprint) { + errors.push("provider-prompt fingerprint") + } + if (report.benchmark !== checkpoint.benchmark) errors.push("benchmark") + if (report.judge !== checkpoint.judge) errors.push("judge model") + if (report.answeringModel !== checkpoint.answeringModel) errors.push("answering model") + if ( + !report.answeringRuntimeIdentity || + !checkpoint.answeringRuntimeIdentity || + stableSha256(report.answeringRuntimeIdentity ?? null) !== + stableSha256(checkpoint.answeringRuntimeIdentity ?? null) + ) { + errors.push("answering runtime") + } + if (stableSha256(report.benchmarkScope) !== stableSha256(checkpoint.benchmarkScope)) { + errors.push("benchmark scope") + } + const effectiveDatasetIdentity = checkpointDatasetIdentity(checkpoint) + if (checkpoint.datasetIdentity && !report.datasetIdentity) { + errors.push("missing dataset identity") + } else if ( + report.datasetIdentity && + stableSha256(report.datasetIdentity) !== stableSha256(effectiveDatasetIdentity) + ) { + errors.push("dataset identity") + } + if (stableSha256(report.protocolIdentity) !== stableSha256(checkpoint.protocolIdentity)) { + errors.push("protocol identity") + } + if (stableSha256(aggregation.protocol.identity) !== stableSha256(checkpoint.protocolIdentity)) { + errors.push("loaded aggregation protocol identity") + } + if (report.selectedQuestionIdsDigest !== checkpoint.selectedQuestionIdsDigest) { + errors.push("selected-question fingerprint") + } + if (report.benchmarkInputFingerprint !== checkpoint.benchmarkInputFingerprint) { + errors.push("benchmark-input fingerprint") + } + + const selectedQuestionIds = checkpoint.targetQuestionIds?.length + ? checkpoint.targetQuestionIds + : Object.keys(checkpoint.questions) + const evaluationQuestionIds = report.evaluations.map((evaluation) => evaluation.questionId) + const metricQuestionIds = report.questionMetrics.map((metric) => metric.questionId) + if ( + selectedQuestionIds.length !== evaluatedQuestionCount || + report.summary.totalQuestions !== evaluatedQuestionCount || + evaluationQuestionIds.length !== evaluatedQuestionCount || + metricQuestionIds.length !== evaluatedQuestionCount + ) { + errors.push("complete question count") + } + if (stableSha256(evaluationQuestionIds) !== checkpoint.selectedQuestionIdsDigest) { + errors.push("evaluation question set/order") + } + if (stableSha256(metricQuestionIds) !== checkpoint.selectedQuestionIdsDigest) { + errors.push("question-metric set/order") + } + + const aggregationQuestionIds = aggregation.questions.map((question) => question.questionId) + if (stableSha256(aggregationQuestionIds) !== checkpoint.selectedQuestionIdsDigest) { + errors.push("aggregation question set/order") + } + + const checkpointEvaluations: QuestionEvaluation[] = [] + for (const questionId of selectedQuestionIds) { + const evaluation = checkpoint.questions[questionId]?.phases?.evaluate?.evaluation + if (!evaluation) { + errors.push(`checkpoint evaluation ${questionId}`) + continue + } + checkpointEvaluations.push(evaluation) + } + + const reportAggregationEvaluations = report.evaluations.map((evaluation) => ({ + questionId: evaluation.questionId, + questionType: evaluation.questionType, + primaryScore: evaluation.primaryScore ?? evaluation.score, + passed: evaluation.passed, + explanation: evaluation.explanation, + metrics: evaluation.metrics, + details: evaluation.details, + })) + const checkpointAggregationEvaluations = checkpointEvaluations.map((evaluation) => ({ + questionId: evaluation.questionId, + questionType: evaluation.questionType, + primaryScore: evaluation.primaryScore, + passed: evaluation.passed, + explanation: evaluation.explanation, + metrics: evaluation.metrics, + details: evaluation.details, + })) + if ( + stableSha256(reportAggregationEvaluations) !== stableSha256(checkpointAggregationEvaluations) + ) { + errors.push("report/checkpoint evaluations") + } + + let expectedQuality: BenchmarkResult["quality"] | undefined + try { + expectedQuality = aggregation.protocol.aggregateQuality({ + questions: aggregation.questions, + evaluations: checkpointEvaluations, + }) + if (stableSha256(report.quality) !== stableSha256(expectedQuality)) { + errors.push("protocol quality aggregation") + } + } catch (error) { + errors.push( + `protocol quality aggregation failed: ${error instanceof Error ? error.message : String(error)}` + ) + } + + const primaryMetric = expectedQuality?.primaryMetric + if (!primaryMetric || !primaryMetric.key.trim() || !Number.isFinite(primaryMetric.value)) { + errors.push("finite scalar primary metric") + } + if (checkpoint.protocolIdentity.id === "beam-paper") { + if (primaryMetric?.key !== "beamScore") { + errors.push("official complete-tier BEAM primary metric") + } + const includedTiers = checkpoint.benchmarkScope.includedTiers + const orderedQuestionIdsDigest = effectiveDatasetIdentity.orderedQuestionIdsDigest + const tier = includedTiers.length === 1 ? includedTiers[0] : undefined + const expectedQuestionCount = tier === "1M" ? 700 : tier === "10M" ? 200 : undefined + const officialTierDigest = + tier && orderedQuestionIdsDigest && typeof orderedQuestionIdsDigest === "object" + ? (orderedQuestionIdsDigest as Record)[tier] + : undefined + if ( + (tier !== "1M" && tier !== "10M") || + selectedQuestionIds.length !== expectedQuestionCount || + typeof officialTierDigest !== "string" || + sha256Text(selectedQuestionIds.join("\n")) !== officialTierDigest + ) { + errors.push("official complete-tier BEAM question set") + } + } + const invalidEvaluation = report.evaluations.some( + (evaluation) => + typeof evaluation.passed !== "boolean" || !Number.isFinite(evaluation.primaryScore) + ) + if (invalidEvaluation) errors.push("complete protocol evaluations") + + const correctCount = checkpointEvaluations.filter((evaluation) => evaluation.passed).length + const averageScore = + checkpointEvaluations.length > 0 + ? checkpointEvaluations.reduce((sum, evaluation) => sum + evaluation.primaryScore, 0) / + checkpointEvaluations.length + : 0 + const accuracy = + checkpointEvaluations.length > 0 ? correctCount / checkpointEvaluations.length : 0 + if (report.summary.correctCount !== correctCount) errors.push("correct-count aggregation") + if (!approximatelyEqual(report.summary.accuracy, accuracy)) errors.push("accuracy aggregation") + if (!approximatelyEqual(report.summary.averageScore, averageScore)) { + errors.push("average-score aggregation") + } + + let identity: LeaderboardComparisonIdentity | undefined + try { + identity = createLeaderboardComparisonIdentity({ + benchmark: report.benchmark, + benchmarkScope: report.benchmarkScope, + datasetIdentity: effectiveDatasetIdentity, + selectedQuestionIdsDigest: report.selectedQuestionIdsDigest, + benchmarkInputFingerprint: report.benchmarkInputFingerprint, + protocolIdentity: report.protocolIdentity, + retrievalTopK: report.retrievalTopK, + questionMetrics: report.questionMetrics, + judgeModel: report.judge, + answeringModel: report.answeringModel, + answeringRuntimeIdentity: report.answeringRuntimeIdentity, + providerPromptFingerprint: report.providerPromptFingerprint, + primaryMetric: primaryMetric!, + accuracy: report.summary.accuracy, + }) + if (checkpoint.retrievalTopK != null && identity.retrievalTopK !== checkpoint.retrievalTopK) { + errors.push("retrieval Top-K") + } + if (identity.retrievalTopK == null) errors.push("recorded retrieval Top-K") + } catch (error) { + errors.push(error instanceof Error ? error.message : "comparison identity") + } + + if (errors.length > 0 || !identity) { + throw new Error(`Report is not publishable: ${[...new Set(errors)].join(", ")}`) + } + return identity +} + +function fallbackIdentity(entry: RankableLeaderboardEntry): LeaderboardComparisonIdentity { + const primary = entry.quality?.primaryMetric + return createLeaderboardComparisonIdentity({ + benchmark: entry.benchmark, + accuracy: entry.accuracy, + judgeModel: entry.judgeModel, + answeringModel: entry.answeringModel, + providerPromptFingerprint: entry.providerPromptFingerprint, + primaryMetric: { + key: primary?.key, + value: primary?.value, + higherIsBetter: primary?.higherIsBetter, + }, + }) +} + +/** Sort and rank only inside exact comparison cohorts. Cross-cohort order is lexical. */ +export function rankLeaderboardEntries( + entries: T[] +): Array> { + const cohorts = new Map>() + + for (const entry of entries) { + const identity = entry.comparisonIdentity ?? fallbackIdentity(entry) + const values = cohorts.get(identity.cohortKey) ?? [] + values.push({ entry, identity }) + cohorts.set(identity.cohortKey, values) + } + + const result: Array> = [] + for (const cohortKey of [...cohorts.keys()].sort()) { + const cohort = cohorts.get(cohortKey)! + const higherIsBetter = cohort[0].identity.primaryMetric.higherIsBetter + cohort.sort((left, right) => { + const leftValue = left.identity.primaryMetric.value + const rightValue = right.identity.primaryMetric.value + const scoreOrder = higherIsBetter ? rightValue - leftValue : leftValue - rightValue + if (scoreOrder !== 0) return scoreOrder + const dateOrder = (left.entry.addedAt ?? "").localeCompare(right.entry.addedAt ?? "") + return dateOrder || left.entry.id - right.entry.id + }) + + let priorValue: number | undefined + let priorRank = 0 + cohort.forEach(({ entry, identity }, index) => { + const value = identity.primaryMetric.value + const rank = priorValue !== undefined && value === priorValue ? priorRank : index + 1 + result.push({ + ...entry, + comparisonIdentity: identity, + cohortKey, + cohortRank: rank, + cohortSize: cohort.length, + }) + priorValue = value + priorRank = rank + }) + } + + return result +} diff --git a/src/server/routes/benchmarks.ts b/src/server/routes/benchmarks.ts index af0b377..618de14 100644 --- a/src/server/routes/benchmarks.ts +++ b/src/server/routes/benchmarks.ts @@ -28,11 +28,16 @@ export async function handleBenchmarksRoutes(req: Request, url: URL): Promise ({ - name, - displayName: getBenchmarkDisplayName(name), - description: getBenchmarkDescription(name), - })), + benchmarks: benchmarks.map((name) => { + const benchmark = createBenchmark(name) + return { + name, + displayName: getBenchmarkDisplayName(name), + description: getBenchmarkDescription(name), + scope: benchmark.scope, + requiredJudge: benchmark.protocol.requiredJudge, + } + }), }) } @@ -85,9 +90,22 @@ export async function handleBenchmarksRoutes(req: Request, url: URL): Promise ({ questionId: q.questionId, question: q.question, @@ -123,8 +144,11 @@ export async function handleBenchmarksRoutes(req: Request, url: URL): Promise() +function getEvaluationPassState(value: any): boolean | undefined { + const protocolEvaluation = value?.evaluation + if (typeof protocolEvaluation?.passed === "boolean") return protocolEvaluation.passed + if (typeof value?.passed === "boolean") return value.passed + + for (const label of [protocolEvaluation?.label, value?.label]) { + if (typeof label !== "string") continue + const normalized = label.toLowerCase() + if (normalized === "pass" || normalized === "correct") return true + if (normalized === "fail" || normalized === "incorrect" || normalized === "wrong") { + return false + } + } + + const legacyScore = value?.score ?? protocolEvaluation?.primaryScore + return typeof legacyScore === "number" ? legacyScore === 1 : undefined +} + function json(data: unknown, status = 200): Response { return new Response(JSON.stringify(data), { status, @@ -126,6 +144,10 @@ export async function handleCompareRoutes(req: Request, url: URL): Promise q.phases?.evaluate?.status === "completed" ) const correctCount = evaluatedQuestions.filter( - (q: any) => q.phases?.evaluate?.score === 1 + (q: any) => getEvaluationPassState(q.phases?.evaluate) === true ).length const accuracy = evaluatedQuestions.length > 0 ? correctCount / evaluatedQuestions.length : null @@ -254,10 +289,21 @@ export async function handleCompareRoutes(req: Request, url: URL): Promise ({ provider: r.provider, report: r.report, @@ -345,20 +399,31 @@ function getRunStatus(checkpoint: any, summary: any): string { return "failed" } - // Check if any question has a failed phase const questions = Object.values(checkpoint.questions || {}) as any[] - const hasFailed = questions.some((q: any) => { + const builds = Object.values(checkpoint.builds || {}) as any[] + const usesSharedBuilds = builds.length > 0 || questions.some((question) => question.buildId) + + // Build phases are owned once per shared haystack, not by each question. + const hasFailedBuild = builds.some( + (build) => build.ingest?.status === "failed" || build.indexing?.status === "failed" + ) + + // Search, answer, and evaluation remain question-owned. + const hasMissingReferencedBuild = questions.some( + (question) => question.buildId && !checkpoint.builds?.[question.buildId] + ) + const hasFailedQuestion = questions.some((q: any) => { const phases = q.phases || {} return ( - phases.ingest?.status === "failed" || - phases.indexing?.status === "failed" || + (!usesSharedBuilds && + (phases.ingest?.status === "failed" || phases.indexing?.status === "failed")) || phases.search?.status === "failed" || phases.answer?.status === "failed" || phases.evaluate?.status === "failed" ) }) - if (hasFailed) { + if (hasFailedBuild || hasMissingReferencedBuild || hasFailedQuestion) { return "failed" } @@ -388,11 +453,22 @@ async function initializeComparison(options: { answeringModel: string sampling?: SamplingConfig force?: boolean + dataPath?: string + datasetRevision?: string + retrievalTopK?: number }): Promise<{ compareId: string }> { - // Only await manifest creation - this is fast const manifest = await batchManager.createManifest(options) const compareId = manifest.compareId + try { + // Every provider must pass the same durable dataset/protocol/build/resume + // preflight before this endpoint can return a successful start response. + await batchManager.preflightRuns(manifest) + } catch (error) { + batchManager.delete(compareId) + throw error + } + startCompare( compareId, options.benchmark, diff --git a/src/server/routes/leaderboard.ts b/src/server/routes/leaderboard.ts index 7a7a507..f825aee 100644 --- a/src/server/routes/leaderboard.ts +++ b/src/server/routes/leaderboard.ts @@ -1,21 +1,176 @@ -import { eq, desc, and } from "drizzle-orm" +import { eq, and } from "drizzle-orm" import { existsSync, readFileSync } from "fs" import { join } from "path" import { db, schema } from "../db" import { CheckpointManager } from "../../orchestrator/checkpoint" import { createBenchmark } from "../../benchmarks" import type { BenchmarkName } from "../../types/benchmark" +import type { BenchmarkResult } from "../../types/unified" +import { + canonicalizeSelectedQuestionIds, + fingerprintSelectedBenchmarkInput, +} from "../../orchestrator/input-identity" +import { stableSha256 } from "../../utils/stable" +import { + createLeaderboardComparisonIdentity, + rankLeaderboardEntries, + validateLeaderboardReportForPublication, + type LeaderboardComparisonIdentity, +} from "../leaderboard-identity" const checkpointManager = new CheckpointManager() const benchmarkRegistryCache: Record = {} +const REPORT_METADATA_KEY = "__memorybenchReport" + +function extractReportMetadata(evaluations: unknown[]): Record | undefined { + const first = evaluations[0] + if (!first || typeof first !== "object" || Array.isArray(first)) return undefined + const metadata = (first as Record)[REPORT_METADATA_KEY] + return metadata && typeof metadata === "object" && !Array.isArray(metadata) + ? (metadata as Record) + : undefined +} + +function attachReportMetadata(evaluations: any[], report: any): any[] { + if (evaluations.length === 0 || !report) return evaluations + const metadata = { + quality: report.quality, + summary: report.summary, + builds: report.builds, + costs: report.costs, + questionMetrics: report.questionMetrics, + protocolIdentity: report.protocolIdentity, + benchmarkScope: report.benchmarkScope, + datasetIdentity: report.datasetIdentity, + selectedQuestionIdsDigest: report.selectedQuestionIdsDigest, + benchmarkInputFingerprint: report.benchmarkInputFingerprint, + retrievalTopK: report.retrievalTopK, + retrieval: report.retrieval, + judge: report.judge, + answeringModel: report.answeringModel, + answeringRuntimeIdentity: report.answeringRuntimeIdentity, + providerPromptFingerprint: report.providerPromptFingerprint, + } + return evaluations.map((evaluation, index) => + index === 0 ? { ...evaluation, [REPORT_METADATA_KEY]: metadata } : evaluation + ) +} + +function getEvaluationPassState(value: any): boolean | undefined { + const protocolEvaluation = value?.evaluation + if (typeof protocolEvaluation?.passed === "boolean") return protocolEvaluation.passed + if (typeof value?.passed === "boolean") return value.passed + + for (const label of [protocolEvaluation?.label, value?.label]) { + if (typeof label !== "string") continue + const normalized = label.toLowerCase() + if (normalized === "pass" || normalized === "correct") return true + if (normalized === "fail" || normalized === "incorrect" || normalized === "wrong") { + return false + } + } + + const legacyScore = value?.score ?? protocolEvaluation?.primaryScore + return typeof legacyScore === "number" ? legacyScore === 1 : undefined +} function getQuestionTypeRegistry(benchmarkName: string) { - if (!benchmarkRegistryCache[benchmarkName]) { - const benchmark = createBenchmark(benchmarkName as BenchmarkName) - benchmarkRegistryCache[benchmarkName] = benchmark.getQuestionTypes() + try { + if (!benchmarkRegistryCache[benchmarkName]) { + const benchmark = createBenchmark(benchmarkName as BenchmarkName) + benchmarkRegistryCache[benchmarkName] = benchmark.getQuestionTypes() + } + return benchmarkRegistryCache[benchmarkName] + } catch { + // Historical rows can reference benchmark ids removed from the registry. + return null + } +} + +function parseJson(value: string | null | undefined, fallback: T): T { + if (!value) return fallback + try { + return JSON.parse(value) as T + } catch { + return fallback + } +} + +function parseLeaderboardEntry(entry: typeof schema.leaderboardEntries.$inferSelect) { + const byQuestionType = parseJson>(entry.byQuestionType, {}) + const latencyStats = parseJson(entry.latencyStats, null) + const evaluations = parseJson(entry.evaluations, []) + const promptsUsed = parseJson | null>(entry.promptsUsed, null) + const reportMetadata = extractReportMetadata(evaluations) + const benchmarkScope = + parseJson | null>(entry.benchmarkScope, null) ?? + reportMetadata?.benchmarkScope + const datasetIdentity = + parseJson | null>(entry.datasetIdentity, null) ?? + reportMetadata?.datasetIdentity + const protocolIdentity = + parseJson | null>(entry.protocolIdentity, null) ?? + reportMetadata?.protocolIdentity + const primaryMetric = { + key: entry.primaryMetricKey ?? reportMetadata?.quality?.primaryMetric?.key ?? "accuracy", + value: + entry.primaryMetricValue ?? reportMetadata?.quality?.primaryMetric?.value ?? entry.accuracy, + higherIsBetter: + entry.primaryMetricHigherIsBetter ?? + reportMetadata?.quality?.primaryMetric?.higherIsBetter ?? + true, + } + const comparisonIdentity = createLeaderboardComparisonIdentity({ + benchmark: entry.benchmark, + benchmarkScope, + datasetIdentity, + selectedQuestionIdsDigest: entry.questionSetFingerprint ?? undefined, + benchmarkInputFingerprint: reportMetadata?.benchmarkInputFingerprint, + protocolIdentity, + retrievalTopK: entry.retrievalTopK ?? undefined, + questionMetrics: reportMetadata?.questionMetrics, + judgeModel: entry.judgeModel, + answeringModel: entry.answeringModel, + answeringRuntimeIdentity: reportMetadata?.answeringRuntimeIdentity, + providerPromptFingerprint: reportMetadata?.providerPromptFingerprint, + primaryMetric, + accuracy: entry.accuracy, + }) + const persistedIdentity: LeaderboardComparisonIdentity = { + ...comparisonIdentity, + datasetFingerprint: entry.datasetFingerprint ?? comparisonIdentity.datasetFingerprint, + protocolFingerprint: entry.protocolFingerprint ?? comparisonIdentity.protocolFingerprint, + // Recompute the current cohort so historical keys that omitted parts of + // the effective benchmark/model identity cannot merge unlike runs. + cohortKey: comparisonIdentity.cohortKey, + } + + return { + ...entry, + byQuestionType, + questionTypeRegistry: getQuestionTypeRegistry(entry.benchmark), + latencyStats, + evaluations, + promptsUsed, + benchmarkScope: persistedIdentity.benchmarkScope, + datasetIdentity: persistedIdentity.datasetIdentity, + datasetFingerprint: persistedIdentity.datasetFingerprint, + questionSetFingerprint: persistedIdentity.questionSetFingerprint, + benchmarkInputFingerprint: persistedIdentity.benchmarkInputFingerprint, + protocolIdentity: persistedIdentity.protocolIdentity, + protocolFingerprint: persistedIdentity.protocolFingerprint, + retrievalTopK: persistedIdentity.retrievalTopK, + providerPromptFingerprint: persistedIdentity.providerPromptFingerprint, + primaryMetric: persistedIdentity.primaryMetric, + comparisonIdentity: persistedIdentity, + quality: reportMetadata?.quality ?? { primaryMetric, metrics: {} }, + averageScore: reportMetadata?.summary?.averageScore, + builds: reportMetadata?.builds, + costs: reportMetadata?.costs, + questionMetrics: reportMetadata?.questionMetrics, + retrieval: reportMetadata?.retrieval, } - return benchmarkRegistryCache[benchmarkName] } function json(data: unknown, status = 200): Response { @@ -32,49 +187,9 @@ export async function handleLeaderboardRoutes(req: Request, url: URL): Promise { - let byQuestionType = {} - let latencyStats = null - let evaluations: unknown[] = [] - let promptsUsed = null - - try { - byQuestionType = JSON.parse(entry.byQuestionType) - } catch { - /* ignore */ - } - try { - latencyStats = entry.latencyStats ? JSON.parse(entry.latencyStats) : null - } catch { - /* ignore */ - } - try { - evaluations = entry.evaluations ? JSON.parse(entry.evaluations) : [] - } catch { - /* ignore */ - } - try { - promptsUsed = entry.promptsUsed ? JSON.parse(entry.promptsUsed) : null - } catch { - /* ignore */ - } + const entries = db.select().from(schema.leaderboardEntries).all() - return { - ...entry, - byQuestionType, - questionTypeRegistry: getQuestionTypeRegistry(entry.benchmark), - latencyStats, - evaluations, - promptsUsed, - } - }) + const parsed = rankLeaderboardEntries(entries.map(parseLeaderboardEntry)) return json({ entries: parsed }) } catch (e) { @@ -107,7 +222,76 @@ export async function handleLeaderboardRoutes(req: Request, url: URL): Promise [question.questionId, question]) + ) + const selectedQuestions = selectedQuestionIds.map((questionId) => { + const question = questionById.get(questionId) + if (!question) throw new Error(`Leaderboard question is missing: ${questionId}`) + return question + }) + const benchmarkInputFingerprint = fingerprintSelectedBenchmarkInput( + benchmark, + selectedQuestions + ) + if (benchmarkInputFingerprint !== checkpoint.benchmarkInputFingerprint) { + throw new Error("Loaded benchmark input does not match the run checkpoint") + } + const loadedDatasetIdentity = benchmark.getDatasetIdentity?.() + if ( + checkpoint.datasetIdentity && + stableSha256(loadedDatasetIdentity ?? null) !== stableSha256(checkpoint.datasetIdentity) + ) { + throw new Error("Loaded dataset identity does not match the run checkpoint") + } + comparisonIdentity = validateLeaderboardReportForPublication( + checkpoint, + report, + summary.total, + { protocol: benchmark.protocol, questions: selectedQuestions } + ) + } catch (error) { + return json( + { + error: + error instanceof Error + ? error.message + : "Leaderboard report is malformed or incomplete", + }, + 400 + ) + } + + const questions = Object.values(checkpoint.questions) + const correctCount = report.summary.correctCount + const accuracy = report.summary.accuracy + + // Upsert only within the exact dataset/protocol/retrieval/metric cohort. const existing = db .select() .from(schema.leaderboardEntries) @@ -115,23 +299,12 @@ export async function handleLeaderboardRoutes(req: Request, url: URL): Promise q.phases?.evaluate?.score === 1).length - const accuracy = report?.summary?.accuracy ?? correctCount / summary.total - // Get provider code const providerCode = getProviderCode(checkpoint.provider) @@ -139,8 +312,17 @@ export async function handleLeaderboardRoutes(req: Request, url: URL): Promise = - {} + const byQuestionType: Record< + string, + { + total: number + correct: number + accuracy: number + averageScore?: number + passAccuracy?: number + retrieval?: unknown + } + > = {} for (const q of questions) { const qData = q as any const type = qData.questionType || "unknown" @@ -148,40 +330,44 @@ export async function handleLeaderboardRoutes(req: Request, url: URL): Promise ({ - questionId: q.questionId, - questionType: q.questionType, - question: q.question, - groundTruth: q.groundTruth, - hypothesis: q.phases?.answer?.hypothesis || "", - score: q.phases?.evaluate?.score || 0, - label: q.phases?.evaluate?.label || "incorrect", - explanation: q.phases?.evaluate?.explanation || "", - searchResults: q.phases?.search?.results || [], - })) - } + let evaluations = report.evaluations.map((evaluation) => ({ ...evaluation })) + evaluations = attachReportMetadata(evaluations, report) const entryData = { runId, provider: checkpoint.provider, benchmark: checkpoint.benchmark, version: entryVersion, + benchmarkScope: JSON.stringify(comparisonIdentity.benchmarkScope), + datasetIdentity: JSON.stringify(comparisonIdentity.datasetIdentity), + datasetFingerprint: comparisonIdentity.datasetFingerprint, + questionSetFingerprint: comparisonIdentity.questionSetFingerprint, + protocolIdentity: JSON.stringify(comparisonIdentity.protocolIdentity), + protocolFingerprint: comparisonIdentity.protocolFingerprint, + retrievalTopK: comparisonIdentity.retrievalTopK, + primaryMetricKey: comparisonIdentity.primaryMetric.key, + primaryMetricValue: comparisonIdentity.primaryMetric.value, + primaryMetricHigherIsBetter: comparisonIdentity.primaryMetric.higherIsBetter, + comparisonCohortKey: comparisonIdentity.cohortKey, accuracy, totalQuestions: summary.total, correctCount, byQuestionType: JSON.stringify(byQuestionType), - latencyStats: report?.latency ? JSON.stringify(report.latency) : null, + latencyStats: report.latency ? JSON.stringify(report.latency) : null, evaluations: JSON.stringify(evaluations), providerCode, promptsUsed: promptsUsed ? JSON.stringify(promptsUsed) : null, @@ -213,7 +399,26 @@ export async function handleLeaderboardRoutes(req: Request, url: URL): Promise candidate.id === entry.id)!) } catch (e) { return json({ error: e instanceof Error ? e.message : "Failed to get entry" }, 500) } diff --git a/src/server/routes/runs.ts b/src/server/routes/runs.ts index 1aaab7b..d0f0504 100644 --- a/src/server/routes/runs.ts +++ b/src/server/routes/runs.ts @@ -1,20 +1,53 @@ import { existsSync, readFileSync, readdirSync } from "fs" import { join } from "path" -import { CheckpointManager } from "../../orchestrator/checkpoint" -import { orchestrator } from "../../orchestrator" +import { assertCopyPhaseOverrides, CheckpointManager } from "../../orchestrator/checkpoint" +import { + assertResumeIdentity, + orchestrator, + resolveEffectiveRetrievalTopK, +} from "../../orchestrator" +import { assertResumeBuilds, prepareValidatedBuildPlans } from "../../orchestrator/builds" import { wsManager } from "../index" import { activeRuns, startRun, endRun, requestStop, isRunActive, getRunState } from "../runState" import { createBenchmark } from "../../benchmarks" +import { createProvider } from "../../providers" +import { fingerprintProviderPrompts } from "../../providers/prompt-identity" import type { ProviderName } from "../../types/provider" import type { BenchmarkName } from "../../types/benchmark" -import type { PhaseId, SamplingConfig } from "../../types/checkpoint" +import type { PhaseId, RunCheckpoint, SamplingConfig } from "../../types/checkpoint" import type { ConcurrencyConfig } from "../../types/concurrency" import { getPhasesFromPhase, PHASE_ORDER } from "../../types/checkpoint" +import { resolveAnsweringRuntimeIdentity, resolveModel } from "../../utils/models" +import { stableSha256 } from "../../utils/stable" +import { getRunListIdentity } from "../run-identity" +import { + canonicalizeSelectedQuestionIds, + fingerprintSelectedBenchmarkInput, +} from "../../orchestrator/input-identity" +import { getProviderConfig } from "../../utils/config" const checkpointManager = new CheckpointManager() const benchmarkRegistryCache: Record = {} +function getEvaluationPassState(value: any): boolean | undefined { + const protocolEvaluation = value?.evaluation + if (typeof protocolEvaluation?.passed === "boolean") return protocolEvaluation.passed + if (typeof value?.passed === "boolean") return value.passed + + for (const label of [protocolEvaluation?.label, value?.label]) { + if (typeof label !== "string") continue + const normalized = label.toLowerCase() + if (normalized === "pass" || normalized === "correct") return true + if (normalized === "fail" || normalized === "incorrect" || normalized === "wrong") { + return false + } + } + + const legacyScore = value?.score ?? protocolEvaluation?.primaryScore + return typeof legacyScore === "number" ? legacyScore === 1 : undefined +} + function getQuestionTypeRegistry(benchmarkName: string) { if (!benchmarkRegistryCache[benchmarkName]) { const benchmark = createBenchmark(benchmarkName as BenchmarkName) @@ -30,6 +63,108 @@ function json(data: unknown, status = 200): Response { }) } +async function validateCheckpointCopy(input: { + source: RunCheckpoint + provider: ProviderName + benchmark: BenchmarkName + judgeModel: string + answeringModel: string + dataPath?: string + datasetRevision?: string + retrievalTopK?: number + fromPhase: PhaseId +}): Promise { + const dataPath = input.dataPath ?? input.source.dataPath + const datasetRevision = input.datasetRevision ?? input.source.datasetRevision + const configuredRetrievalTopK = input.retrievalTopK ?? input.source.retrievalTopK + const benchmark = createBenchmark(input.benchmark) + + await benchmark.load({ + dataPath, + datasetRevision, + retrievalTopK: configuredRetrievalTopK, + }) + const resolvedJudge = resolveModel(input.judgeModel) + const requiredJudge = benchmark.protocol.requiredJudge + if ( + requiredJudge && + (resolvedJudge.provider !== requiredJudge.provider || + resolvedJudge.id !== requiredJudge.modelId) + ) { + throw new Error( + `Protocol ${benchmark.protocol.identity.id} requires judge ${requiredJudge.provider}/${requiredJudge.modelId}; received ${resolvedJudge.provider}/${resolvedJudge.id}` + ) + } + + const allQuestions = benchmark.getQuestions() + for (const question of allQuestions) benchmark.protocol.validateQuestion(question) + if (!input.source.targetQuestionIds?.length) { + throw new Error( + `Source run ${input.source.runId} does not record its selected question IDs; start a new run` + ) + } + const targetQuestionIds = canonicalizeSelectedQuestionIds( + allQuestions, + input.source.targetQuestionIds + ) + const questionById = new Map(allQuestions.map((question) => [question.questionId, question])) + + const selectedQuestions = targetQuestionIds.map((questionId) => questionById.get(questionId)!) + const retrievalTopK = resolveEffectiveRetrievalTopK( + benchmark.protocol, + selectedQuestions, + configuredRetrievalTopK + ) + const benchmarkInputFingerprint = fingerprintSelectedBenchmarkInput(benchmark, selectedQuestions) + const provider = createProvider(input.provider) + const providerConfig = getProviderConfig(input.provider) + const providerPromptFingerprint = fingerprintProviderPrompts(provider.prompts) + const providerIngestionConfigFingerprint = provider.getIngestionConfigFingerprint(providerConfig) + assertCopyPhaseOverrides(input.source, input.fromPhase, { + judge: input.judgeModel, + answeringModel: input.answeringModel, + }) + const buildPlans = prepareValidatedBuildPlans({ + benchmark, + questions: selectedQuestions, + provider: provider.name, + providerAdapterVersion: provider.adapterVersion, + providerPromptFingerprint, + providerIngestionConfigFingerprint, + dataSourceRunId: input.source.dataSourceRunId || input.source.runId, + }) + + const copyIdentity = structuredClone(input.source) + copyIdentity.judge = input.judgeModel + copyIdentity.answeringModel = input.answeringModel + const rerunsAnswer = PHASE_ORDER.indexOf(input.fromPhase) <= PHASE_ORDER.indexOf("answer") + const answeringRuntimeIdentity = rerunsAnswer + ? resolveAnsweringRuntimeIdentity(input.answeringModel) + : input.source.answeringRuntimeIdentity + if (!answeringRuntimeIdentity) { + throw new Error( + `Cannot copy ${input.source.runId} from ${input.fromPhase}; source answer runtime identity is missing` + ) + } + copyIdentity.answeringRuntimeIdentity = answeringRuntimeIdentity + assertResumeIdentity(copyIdentity, { + provider: provider.name, + providerAdapterVersion: provider.adapterVersion, + providerPromptFingerprint, + benchmark: benchmark.name, + benchmarkScope: benchmark.scope, + datasetIdentity: benchmark.getDatasetIdentity?.(), + benchmarkInputFingerprint, + selectedQuestionIdsDigest: stableSha256(targetQuestionIds), + protocolIdentity: benchmark.protocol.identity, + retrievalTopK, + judge: input.judgeModel, + answeringModel: input.answeringModel, + answeringRuntimeIdentity, + }) + assertResumeBuilds(input.source, buildPlans) +} + export async function handleRunsRoutes(req: Request, url: URL): Promise { const method = req.method const pathname = url.pathname @@ -49,7 +184,7 @@ export async function handleRunsRoutes(req: Request, url: URL): Promise q.phases?.evaluate?.status === "completed" ) const correctCount = evaluatedQuestions.filter( - (q: any) => q.phases?.evaluate?.score === 1 + (q: any) => getEvaluationPassState(q.phases?.evaluate) === true ).length const accuracy = evaluatedQuestions.length > 0 ? correctCount / evaluatedQuestions.length : null @@ -58,6 +193,7 @@ export async function handleRunsRoutes(req: Request, url: URL): Promise void + let rejectPreflight!: (error: Error) => void + const preflight = new Promise((resolve, reject) => { + resolvePreflight = resolve + rejectPreflight = reject + }) + + void runBenchmark({ provider: provider as ProviderName, benchmark: benchmark as BenchmarkName, runId, @@ -282,10 +454,24 @@ export async function handleRunsRoutes(req: Request, url: URL): Promise { + preflightResolved = true + resolvePreflight() + }, + onFailure: (error) => { + if (!preflightResolved) rejectPreflight(error) + }, }).finally(() => { endRun(runId) }) + // Do not return a successful start response until all fail-closed dataset, + // protocol, build-identity, and resume checks have passed and a checkpoint exists. + await preflight + return json({ message: "Run started", runId }) } catch (e) { return json({ error: e instanceof Error ? e.message : "Invalid request body" }, 400) @@ -330,20 +516,31 @@ function getRunStatus(checkpoint: any, summary: any): string { return "failed" } - // Check if any question has a failed phase const questions = Object.values(checkpoint.questions || {}) as any[] - const hasFailed = questions.some((q: any) => { + const builds = Object.values(checkpoint.builds || {}) as any[] + const usesSharedBuilds = builds.length > 0 || questions.some((question) => question.buildId) + + // Build phases are owned once per shared haystack, not by each question. + const hasFailedBuild = builds.some( + (build) => build.ingest?.status === "failed" || build.indexing?.status === "failed" + ) + + // Search, answer, and evaluation remain question-owned. + const hasMissingReferencedBuild = questions.some( + (question) => question.buildId && !checkpoint.builds?.[question.buildId] + ) + const hasFailedQuestion = questions.some((q: any) => { const phases = q.phases || {} return ( - phases.ingest?.status === "failed" || - phases.indexing?.status === "failed" || + (!usesSharedBuilds && + (phases.ingest?.status === "failed" || phases.indexing?.status === "failed")) || phases.search?.status === "failed" || phases.answer?.status === "failed" || phases.evaluate?.status === "failed" ) }) - if (hasFailed) { + if (hasFailedBuild || hasMissingReferencedBuild || hasFailedQuestion) { return "failed" } @@ -377,6 +574,11 @@ async function runBenchmark(options: { concurrency?: ConcurrencyConfig force?: boolean fromPhase?: PhaseId + dataPath?: string + datasetRevision?: string + retrievalTopK?: number + onPreflightComplete?: () => void + onFailure?: (error: Error) => void }) { try { wsManager.broadcast({ @@ -399,6 +601,10 @@ async function runBenchmark(options: { concurrency: options.concurrency, force: options.force, phases, + dataPath: options.dataPath, + datasetRevision: options.datasetRevision, + retrievalTopK: options.retrievalTopK, + onPreflightComplete: options.onPreflightComplete, }) wsManager.broadcast({ @@ -406,13 +612,22 @@ async function runBenchmark(options: { runId: options.runId, }) } catch (error) { - const message = error instanceof Error ? error.message : "Unknown error" + let message = error instanceof Error ? error.message : "Unknown error" + const normalizedError = error instanceof Error ? error : new Error(message) + options.onFailure?.(normalizedError) const wasStoppedByUser = message.includes("stopped by user") // Update checkpoint status to persist the failure/stopped state - const checkpoint = checkpointManager.load(options.runId) - if (checkpoint) { - checkpointManager.updateStatus(checkpoint, "failed") + try { + const checkpoint = checkpointManager.load(options.runId) + if (checkpoint) { + checkpointManager.updateStatus(checkpoint, "failed") + await checkpointManager.flush(options.runId) + } + } catch (persistenceError) { + const persistenceMessage = + persistenceError instanceof Error ? persistenceError.message : String(persistenceError) + message = `${message}; failed to persist terminal run state: ${persistenceMessage}` } wsManager.broadcast({ diff --git a/src/server/run-identity.ts b/src/server/run-identity.ts new file mode 100644 index 0000000..4b94fbb --- /dev/null +++ b/src/server/run-identity.ts @@ -0,0 +1,13 @@ +import type { RunCheckpoint } from "../types/checkpoint" + +/** Identity fields every run-list row must expose for like-for-like interpretation. */ +export function getRunListIdentity(checkpoint: RunCheckpoint) { + return { + benchmarkScope: checkpoint.benchmarkScope, + datasetIdentity: checkpoint.datasetIdentity, + benchmarkInputFingerprint: checkpoint.benchmarkInputFingerprint, + selectedQuestionIdsDigest: checkpoint.selectedQuestionIdsDigest, + protocolIdentity: checkpoint.protocolIdentity, + providerPromptFingerprint: checkpoint.providerPromptFingerprint, + } +} diff --git a/src/types/benchmark.ts b/src/types/benchmark.ts index 7066f7b..86a6743 100644 --- a/src/types/benchmark.ts +++ b/src/types/benchmark.ts @@ -1,7 +1,42 @@ import type { UnifiedQuestion, UnifiedSession, QuestionTypeRegistry } from "./unified" +import type { BenchmarkProtocol } from "./protocol" export interface BenchmarkConfig { dataPath?: string + datasetRevision?: string + retrievalTopK?: number + answerCutoff?: number + evaluationProfile?: string +} + +export interface BenchmarkScope { + displayName: string + includedTiers: string[] + coverage: "full" | "subset" +} + +export interface DatasetIdentity { + /** Identity of the exact tier(s) consumed by this benchmark instance. */ + datasetFingerprint: string + manifestHash: string + /** Enclosing prepared snapshot provenance; may contain additional unused tiers. */ + snapshotFingerprint?: string + snapshotManifestHash?: string + manifestSchemaVersion: number + canonicalSchemaVersion: number + converterVersion: string + converterImplementationHash: string + includedTiers: string[] + counts: Record + orderedQuestionIdsDigest: Record + sourceFiles: Array<{ path: string; byteSize: number; sha256: string }> + canonicalFiles: Array<{ path: string; byteSize: number; sha256: string }> + sources: Array<{ + repository: string + split: string + revision: string + sourceIdentity: "reviewed-published" | "injected-test-fixture" + }> } export interface QuestionFilter { @@ -13,12 +48,21 @@ export interface QuestionFilter { export interface Benchmark { name: string + scope: BenchmarkScope + protocol: BenchmarkProtocol load(config?: BenchmarkConfig): Promise getQuestions(filter?: QuestionFilter): UnifiedQuestion[] getHaystackSessions(questionId: string): UnifiedSession[] getGroundTruth(questionId: string): string getQuestionTypes(): QuestionTypeRegistry getIngestionGroupId?(questionId: string): string + getDatasetIdentity?(): DatasetIdentity | undefined } -export type BenchmarkName = "locomo" | "longmemeval" | "convomem" | "beam-1m" | "beam-10m" | "beam" +export type BenchmarkName = + | "locomo" + | "longmemeval" + | "convomem" + | "beam-1m" + | "beam-10m" + | "beam-1m-10m" diff --git a/src/types/checkpoint.ts b/src/types/checkpoint.ts index b3ff572..5d9d7c9 100644 --- a/src/types/checkpoint.ts +++ b/src/types/checkpoint.ts @@ -1,6 +1,20 @@ -import type { SearchResult, RetrievalMetrics } from "./unified" -import type { IngestResult } from "./provider" import type { ConcurrencyConfig } from "./concurrency" +import type { BenchmarkScope, DatasetIdentity } from "./benchmark" +import type { + IngestionExecutionPolicy, + ProtocolIdentity, + QuestionEvaluation, + RetrievalPlan, +} from "./protocol" +import type { AnsweringRuntimeIdentity } from "./model" +import type { + ProviderRequestDiagnostic, + ProviderResultDropDiagnostic, + RetrievalMetrics, + UnifiedSearchResult, +} from "./unified" + +export const CHECKPOINT_SCHEMA_VERSION = 4 export type PhaseStatus = "pending" | "in_progress" | "completed" | "failed" @@ -21,51 +35,159 @@ export function getPhasesFromPhase(fromPhase: PhaseId): PhaseId[] { return PHASE_ORDER.slice(startIndex) } -export interface IngestPhaseCheckpoint { - status: PhaseStatus - completedSessions: string[] - ingestResult?: IngestResult - startedAt?: string +export interface ProviderUsage { + requestCount?: number + tokenUsageCompleteRequestCount?: number + tokenUsagePartialRequestCount?: number + tokenUsageUnknownRequestCount?: number + inputTokens?: number + outputTokens?: number + reasoningTokens?: number + totalTokens?: number +} + +export interface BuildAttemptMetrics { + phase: "ingest" | "indexing" + attempt: number + startedAt: string completedAt?: string durationMs?: number + status: "in_progress" | "completed" | "failed" + usage?: ProviderUsage + costUsd: number | null error?: string } -export interface IndexingPhaseCheckpoint { +export interface HaystackIdentity { + schemaVersion: 2 + algorithm: "sha256" + fingerprint: string + orderedSessionIds: string[] + sessionFingerprints: string[] +} + +export interface SessionMetadata { + sessionId: string + documentDate?: string + messageCount: number +} + +export interface DeferredIngestSession { + sequence: number + sessionId: string + customId: string + documentIds: string[] + taskIds: string[] + stage: "submission" | "readiness" + attempts: number + firstFailedAt: string + lastFailedAt: string + lastError: string +} + +export interface BuildCheckpoint { + buildId: string + ingestionGroupId: string + memberQuestionIds: string[] + containerTag: string + haystack: HaystackIdentity + buildFingerprint: string + providerIngestionConfigFingerprint: string + ingestionExecutionPolicy: IngestionExecutionPolicy + /** Number of ordered sessions submitted in one provider request before a readiness barrier. */ + ingestBatchSize?: number + sessions: SessionMetadata[] + missingDocumentDateCount: number + sourceRunId?: string + reused: boolean + reusedPhases?: { + ingest: boolean + indexing: boolean + } + ingest: { + status: PhaseStatus + completedSessionIds: string[] + documentIds: string[] + taskIds: string[] + /** Sessions deferred during the first pass and retried in order at build end. */ + deferredSessions?: DeferredIngestSession[] + startedAt?: string + completedAt?: string + durationMs?: number + attempts: BuildAttemptMetrics[] + error?: string + } + indexing: { + status: PhaseStatus + completedIds: string[] + failedIds: string[] + startedAt?: string + completedAt?: string + durationMs?: number + attempts: BuildAttemptMetrics[] + error?: string + } +} + +export interface SearchPhaseCheckpoint { status: PhaseStatus - completedIds?: string[] - failedIds?: string[] + retrievalPlan?: RetrievalPlan + resultFile?: string + results?: UnifiedSearchResult[] + requestedCount?: number + rawReturnedCount?: number + returnedCount?: number + normalizedCount?: number + droppedCount?: number + droppedResults?: ProviderResultDropDiagnostic[] + providerRequests?: ProviderRequestDiagnostic[] + answerCutoff?: number + answerEvidenceCount?: number startedAt?: string completedAt?: string durationMs?: number + usage?: ProviderUsage + costUsd?: number | null error?: string } -export interface SearchPhaseCheckpoint { - status: PhaseStatus - resultFile?: string - results?: SearchResult[] - resultCount?: number - startedAt?: string +export interface AnswerAttemptMetrics { + attempt: number + startedAt: string completedAt?: string durationMs?: number + status: "in_progress" | "completed" | "failed" + finishReason?: string + reasoningTokens?: number + usage?: ProviderUsage error?: string } export interface AnswerPhaseCheckpoint { status: PhaseStatus hypothesis?: string + /** True only when a benchmark explicitly accepts terminal all-empty model output. */ + terminalEmptyAccepted?: boolean promptTokens?: number basePromptTokens?: number contextTokens?: number + evidenceCount?: number startedAt?: string completedAt?: string durationMs?: number + /** Durable outer attempts, including empty-output and transport retries. */ + attempts?: AnswerAttemptMetrics[] + usage?: ProviderUsage + costUsd?: number | null error?: string } export interface EvaluatePhaseCheckpoint { status: PhaseStatus + /** Benchmark-owned durable state for multi-call evaluators. */ + protocolProgress?: Record + evaluation?: QuestionEvaluation + /** Compatibility mirrors for existing UI/readers. */ label?: "correct" | "incorrect" score?: number explanation?: string @@ -74,26 +196,19 @@ export interface EvaluatePhaseCheckpoint { startedAt?: string completedAt?: string durationMs?: number + usage?: ProviderUsage + costUsd?: number | null error?: string } -export interface SessionMetadata { - sessionId: string - date?: string - messageCount: number -} - export interface QuestionCheckpoint { questionId: string - containerTag: string + buildId: string question: string groundTruth: string questionType: string questionDate?: string - sessions?: SessionMetadata[] phases: { - ingest: IngestPhaseCheckpoint - indexing: IndexingPhaseCheckpoint search: SearchPhaseCheckpoint answer: AnswerPhaseCheckpoint evaluate: EvaluatePhaseCheckpoint @@ -112,19 +227,48 @@ export interface SamplingConfig { limit?: number } +export interface BuildPhaseAttempt { + startedAt: string + completedAt?: string + durationMs?: number + status: "in_progress" | "completed" | "failed" +} + export interface RunCheckpoint { + schemaVersion: typeof CHECKPOINT_SCHEMA_VERSION runId: string dataSourceRunId: string status: RunStatus provider: string + providerAdapterVersion: string + providerPromptFingerprint: string benchmark: string + benchmarkScope: BenchmarkScope + datasetIdentity?: DatasetIdentity + benchmarkInputFingerprint: string + selectedQuestionIdsDigest: string + protocolIdentity: ProtocolIdentity judge: string answeringModel: string + answeringRuntimeIdentity: AnsweringRuntimeIdentity createdAt: string updatedAt: string + dataPath?: string + datasetRevision?: string + retrievalTopK: number + /** Explicit non-default benchmark evaluation profile, when selected. */ + evaluationProfile?: string + /** Maximum retrieved results exposed to the answering model. */ + answerCutoff?: number limit?: number sampling?: SamplingConfig targetQuestionIds?: string[] concurrency?: ConcurrencyConfig + /** Defaults to 1 for checkpoints created before ordered batch ingestion. */ + ingestBatchSize?: number + /** Operational per-readiness-call deadline; defaults to five minutes. */ + ingestReadinessTimeoutMs?: number + buildPhaseAttempts: BuildPhaseAttempt[] + builds: Record questions: Record } diff --git a/src/types/index.ts b/src/types/index.ts index 27aafe8..3c21b3f 100644 --- a/src/types/index.ts +++ b/src/types/index.ts @@ -2,4 +2,6 @@ export * from "./unified" export * from "./provider" export * from "./benchmark" export * from "./judge" +export * from "./model" export * from "./checkpoint" +export * from "./protocol" diff --git a/src/types/judge.ts b/src/types/judge.ts index aef64aa..1df7b65 100644 --- a/src/types/judge.ts +++ b/src/types/judge.ts @@ -1,4 +1,5 @@ import type { ProviderPrompts } from "./prompts" +import type { ModelTransport } from "./protocol" export interface JudgeConfig { apiKey: string @@ -28,7 +29,7 @@ export interface Judge { initialize(config: JudgeConfig): Promise evaluate(input: JudgeInput): Promise getPromptForQuestionType(questionType: string, providerPrompts?: ProviderPrompts): string - getModel(): import("ai").LanguageModel + getModel(transport?: ModelTransport): import("ai").LanguageModel } export type JudgeName = "openai" | "anthropic" | "google" diff --git a/src/types/model.ts b/src/types/model.ts new file mode 100644 index 0000000..e64d32c --- /dev/null +++ b/src/types/model.ts @@ -0,0 +1,10 @@ +export interface AnsweringRuntimeIdentity { + schemaVersion: 1 + transport: "ai-sdk-generate-text-v1" + modelAlias: string + provider: "openai" | "anthropic" | "google" + modelId: string + supportsTemperature: boolean + effectiveDefaultTemperature: number | null + effectiveDefaultMaxOutputTokens: number +} diff --git a/src/types/protocol.ts b/src/types/protocol.ts new file mode 100644 index 0000000..29fc236 --- /dev/null +++ b/src/types/protocol.ts @@ -0,0 +1,175 @@ +import type { z } from "zod" +import type { JudgeInput, JudgeResult } from "./judge" +import type { ProviderPrompts } from "./prompts" +import type { + CanonicalIngestionDocument, + UnifiedQuestion, + UnifiedSearchResult, + UnifiedSession, +} from "./unified" + +export interface ProtocolIdentity { + id: string + version: string + configFingerprint: string + implementationFingerprint: string + /** + * Fingerprint of only the protocol policy and implementation that produce + * canonical ingestion documents. Build reuse must depend on this value, not + * retrieval, answer, evaluator, or aggregation identity. + */ + ingestionPolicyHash: string + retrievalPolicyHash: string + answerPromptHash: string + evaluatorHash: string + aggregationHash: string + /** Auditable pinned profile data whose hashes are recorded above. */ + details?: Record +} + +export interface RetrievalPlan { + query: string + requestedTopK: number + answerCutoff: number + threshold?: number + searchMode?: string + filters?: Record +} + +export type ModelTransport = "provider-default" | "openai-chat-completions" +export type TerminalEmptyOutputPolicy = "fail" | "accept-and-evaluate" + +export interface ModelRequest { + system?: string + prompt: string + maxOutputTokens?: number + temperature?: number + /** Explicit transport when a reference runner does not use the provider default. */ + transport?: ModelTransport + /** Protocol-owned outer attempts, including retries of empty model output. */ + maxAttempts?: number + /** SDK-level retries inside each outer attempt. */ + innerMaxRetries?: number + /** Deadline for each outer attempt. */ + timeoutMs?: number + /** Linear backoff base: attempt N waits N * retryBackoffMs before N+1. */ + retryBackoffMs?: number + /** Benchmark-owned behavior after outer attempts exhaust without non-empty text. */ + terminalEmptyOutputPolicy?: TerminalEmptyOutputPolicy +} + +export interface AnswerPlan { + request: ModelRequest + baseRequest: ModelRequest + answerEvidenceCount: number +} + +export interface QuestionEvaluation { + questionId: string + questionType: string + primaryScore: number + passed: boolean + label?: string + explanation: string + metrics?: Record + details?: Record +} + +export interface BenchmarkQualityReport { + /** Absent when a combined scope has multiple official metrics and no official scalar score. */ + primaryMetric?: { + key: string + value: number + higherIsBetter: boolean + } + metrics: Record + bySlice?: Record> +} + +export interface StructuredModelRequest extends ModelRequest { + schema: z.ZodType + schemaName: string + maxAttempts?: number + timeoutMs?: number +} + +export interface EvaluationRuntime { + evaluateLegacy(input: JudgeInput): Promise + generateStructured(request: StructuredModelRequest): Promise + getUsage?(): ModelUsage | undefined +} + +export interface ModelUsage { + /** Number of model requests attempted, including failed paid attempts. */ + requestCount?: number + tokenUsageCompleteRequestCount?: number + tokenUsagePartialRequestCount?: number + tokenUsageUnknownRequestCount?: number + inputTokens?: number + outputTokens?: number + reasoningTokens?: number + totalTokens?: number +} + +/** + * Optional LLM relevance diagnostics are not part of every benchmark's + * protocol. The benchmark must explicitly own whether they run. + */ +export type AuxiliaryRetrievalEvaluationPolicy = "disabled" | "legacy-llm-relevance-v1" + +/** + * Benchmark-owned build semantics. A causal benchmark can require each + * document to be fully ready before the next ordered document is submitted; + * independent builds may still execute concurrently. + */ +export interface IngestionExecutionPolicy { + readinessBarrier: "after-build" | "after-each-document" + processingMode: "provider-default" | "instant" +} + +export interface BenchmarkProtocol { + identity: ProtocolIdentity + auxiliaryRetrievalEvaluation: AuxiliaryRetrievalEvaluationPolicy + ingestionExecutionPolicy: IngestionExecutionPolicy + requiredJudge?: { + provider: string + modelId: string + modelAlias: string + } + + validateQuestion(question: UnifiedQuestion): void + + createIngestionPlan(input: { + question: UnifiedQuestion + sessions: UnifiedSession[] + }): CanonicalIngestionDocument[] + + createRetrievalPlan(input: { question: UnifiedQuestion }): RetrievalPlan + + createAnswerPlan(input: { + question: UnifiedQuestion + sessions: UnifiedSession[] + results: UnifiedSearchResult[] + retrieval: RetrievalPlan + questionDate?: string + providerPrompts?: ProviderPrompts + }): AnswerPlan + + evaluateQuestion( + input: { + question: UnifiedQuestion + hypothesis: string + results: UnifiedSearchResult[] + retrieval: RetrievalPlan + providerPrompts?: ProviderPrompts + protocolProgress?: Record + onProtocolProgress?: (progress: Record) => Promise + }, + runtime: EvaluationRuntime + ): Promise + + aggregateQuality(input: { + questions: UnifiedQuestion[] + evaluations: QuestionEvaluation[] + }): BenchmarkQualityReport +} diff --git a/src/types/provider.ts b/src/types/provider.ts index cdc0228..5d0c4bf 100644 --- a/src/types/provider.ts +++ b/src/types/provider.ts @@ -1,4 +1,9 @@ -import type { UnifiedSession } from "./unified" +import type { + CanonicalIngestionDocument, + ProviderRequestDiagnostic, + ProviderResultDropDiagnostic, + UnifiedSearchResult, +} from "./unified" import type { ProviderPrompts } from "./prompts" import type { ConcurrencyConfig } from "./concurrency" @@ -11,17 +16,30 @@ export interface ProviderConfig { export interface IngestOptions { containerTag: string metadata?: Record + /** Request immediate provider processing when the benchmark requires a causal barrier. */ + processingMode?: "instant" } export interface SearchOptions { containerTag: string - limit?: number + limit: number threshold?: number + searchMode?: string + filters?: Record } export interface IngestResult { documentIds: string[] taskIds?: string[] + /** Per-input outcomes for providers that can partially accept a batch. */ + items?: IngestItemResult[] +} + +export interface IngestItemResult { + customId: string + documentIds: string[] + taskIds?: string[] + error?: string } export interface IndexingProgress { @@ -30,20 +48,53 @@ export interface IndexingProgress { total: number } +export interface ProviderSearchDiagnostics { + requestedLimit: number + providerRequests: ProviderRequestDiagnostic[] + rawReturnedCount: number + normalizedCount: number + droppedCount: number + droppedResults: ProviderResultDropDiagnostic[] +} + +export interface ProviderSearchResponse { + results: UnifiedSearchResult[] + diagnostics: ProviderSearchDiagnostics +} + +export type ProviderSearchRequestStructure = + | { kind: "single" } + | { kind: "split"; budget: "shared-total" } + export type IndexingProgressCallback = (progress: IndexingProgress) => void +export interface AwaitIndexingOptions { + /** Operational polling deadline. This does not change the built container identity. */ + timeoutMs?: number +} + export interface Provider { name: string + adapterVersion: string + searchRequestStructure: ProviderSearchRequestStructure prompts?: ProviderPrompts concurrency?: ConcurrencyConfig + /** + * Fingerprint every non-secret provider setting that can change the built + * memory container. The orchestrator records this in the build identity + * before initializing the provider so resume cannot cross configuration + * drift silently. + */ + getIngestionConfigFingerprint(config: ProviderConfig): string initialize(config: ProviderConfig): Promise - ingest(sessions: UnifiedSession[], options: IngestOptions): Promise + ingest(documents: CanonicalIngestionDocument[], options: IngestOptions): Promise awaitIndexing( result: IngestResult, containerTag: string, - onProgress?: IndexingProgressCallback + onProgress?: IndexingProgressCallback, + options?: AwaitIndexingOptions ): Promise - search(query: string, options: SearchOptions): Promise + search(query: string, options: SearchOptions): Promise clear(containerTag: string): Promise } diff --git a/src/types/unified.ts b/src/types/unified.ts index e5ef939..2627402 100644 --- a/src/types/unified.ts +++ b/src/types/unified.ts @@ -1,3 +1,6 @@ +import type { AnsweringRuntimeIdentity } from "./model" +import type { ProviderName } from "./provider" + export interface QuestionTypeInfo { id: string alias: string @@ -19,6 +22,18 @@ export interface UnifiedSession { metadata?: Record } +export interface CanonicalIngestionDocument { + customId: string + content: string + metadata: { + sessionId: string + documentDate?: string + [key: string]: unknown + } + /** Provider-neutral source messages for adapters that ingest chat messages. */ + messages?: UnifiedMessage[] +} + export interface UnifiedQuestion { questionId: string question: string @@ -28,7 +43,46 @@ export interface UnifiedQuestion { metadata?: Record } -export type SearchResult = unknown +export const UNIFIED_SEARCH_RESULT_TYPES = [ + "memory", + "chunk", + "graph-edge", + "graph-node", + "document", +] as const + +export type UnifiedSearchResultType = (typeof UNIFIED_SEARCH_RESULT_TYPES)[number] + +export interface UnifiedSearchResult { + id: string + rank: number + text: string + score?: number + sessionId?: string + documentDate?: string + provider: ProviderName + resultType: UnifiedSearchResultType + /** Optional pointer to a separately stored raw artifact; never prompt content. */ + rawArtifactRef?: string +} + +export interface ProviderRequestDiagnostic { + operation: string + limit: number + parameters?: Record +} + +export interface ProviderResultDropDiagnostic { + index: number + reason: + | "malformed-result" + | "missing-id" + | "empty-text" + | "unsupported-result-type" + | "below-threshold" +} + +export type SearchResult = UnifiedSearchResult export interface RetrievalMetrics { hitAtK: number @@ -57,8 +111,11 @@ export interface EvaluationResult { questionType: string question: string score: number + primaryScore: number + passed: boolean label: "correct" | "incorrect" explanation: string + metrics?: Record hypothesis: string groundTruth: string searchResults: SearchResult[] @@ -69,6 +126,91 @@ export interface EvaluationResult { details?: Record } +export interface UsageMetrics { + requestCount?: number + tokenUsageCompleteRequestCount?: number + tokenUsagePartialRequestCount?: number + tokenUsageUnknownRequestCount?: number + inputTokens?: number + outputTokens?: number + reasoningTokens?: number + totalTokens?: number +} + +export interface BuildAttemptReport { + phase: "ingest" | "indexing" + attempt: number + startedAt: string + completedAt?: string + durationMs?: number + status: "in_progress" | "completed" | "failed" + usage?: UsageMetrics + costUsd: number | null + error?: string +} + +export interface BuildMetrics { + buildId: string + containerTag: string + providerIngestionConfigFingerprint: string + sourceRunId?: string + reused: boolean + reusedPhases?: { + ingest: boolean + indexing: boolean + } + ingestLatencyMs: number + indexingLatencyMs: number + buildWallClockMs: number + buildWorkMs: number + attemptCount: number + attempts: BuildAttemptReport[] + usage?: UsageMetrics + costUsd: number | null + sessionCount: number + documentCount: number + taskCount: number + completedIndexingCount: number + failedIndexingCount: number +} + +export interface QuestionMetrics { + questionId: string + buildId: string + searchLatencyMs: number + answerLatencyMs: number + onlineQueryLatencyMs: number + evaluationLatencyMs: number + queryUsage?: UsageMetrics + evaluationUsage?: UsageMetrics + queryCostUsd: number | null + evaluationCostUsd: number | null + configuredTopK: number + providerRequestLimit: number + rawReturnedCount: number + returnedCount: number + normalizedCount: number + droppedCount: number + answerCutoff: number + answerEvidenceCount: number + contextTokens: number + searchMode?: string + threshold?: number + providerRequests: ProviderRequestDiagnostic[] + droppedResults: ProviderResultDropDiagnostic[] + /** Completed questions sharing this build-work allocation. */ + buildAllocationQuestionCount: number + allocatedBuildWorkMs?: number + amortizedOnlinePlusBuildWorkMs?: number +} + +export interface CostCoverageMetrics { + /** Null unless every relevant question has a known cost. */ + totalCostUsd: number | null + knownCostCount: number + totalCostCount: number +} + export interface LatencyStats { min: number max: number @@ -103,17 +245,49 @@ export interface TokenMetrics { export interface BenchmarkResult { provider: string + providerPromptFingerprint: string benchmark: string runId: string dataSourceRunId: string judge: string answeringModel: string + answeringRuntimeIdentity: AnsweringRuntimeIdentity timestamp: string + selectedQuestionIdsDigest: string + retrievalTopK: number + benchmarkScope: { + displayName: string + includedTiers: string[] + coverage: "full" | "subset" + } + datasetIdentity?: Record + benchmarkInputFingerprint: string + protocolIdentity: Record + quality: { + primaryMetric?: { key: string; value: number; higherIsBetter: boolean } + metrics: Record + bySlice?: Record> + } summary: { totalQuestions: number correctCount: number accuracy: number + averageScore: number + } + builds: { + uniqueBuildCount: number + sumContainerBuildWorkMs: number + buildPhaseWallClockMs: number + totalBuildCostUsd: number | null + knownCostBuildCount: number + totalCostBuildCount: number + items: BuildMetrics[] + } + costs: { + query: CostCoverageMetrics + evaluation: CostCoverageMetrics } + questionMetrics: QuestionMetrics[] latency: { ingest: LatencyStats indexing: LatencyStats diff --git a/src/utils/models.ts b/src/utils/models.ts index b29ac80..7d9fd6a 100644 --- a/src/utils/models.ts +++ b/src/utils/models.ts @@ -1,3 +1,5 @@ +import type { AnsweringRuntimeIdentity } from "../types/model" + export interface ModelConfig { id: string provider: "openai" | "anthropic" | "google" @@ -322,6 +324,22 @@ export function resolveModel(alias: string): ModelConfig { return getModelConfig(alias) } +/** Immutable effective defaults used by the generic answering runtime. */ +export function resolveAnsweringRuntimeIdentity(alias: string): AnsweringRuntimeIdentity { + const resolvedAlias = alias || DEFAULT_ANSWERING_MODEL + const model = getModelConfig(resolvedAlias) + return { + schemaVersion: 1, + transport: "ai-sdk-generate-text-v1", + modelAlias: resolvedAlias, + provider: model.provider, + modelId: model.id, + supportsTemperature: model.supportsTemperature, + effectiveDefaultTemperature: model.supportsTemperature ? model.defaultTemperature : null, + effectiveDefaultMaxOutputTokens: model.defaultMaxTokens, + } +} + export function getModelId(alias: string): string { return getModelConfig(alias).id } diff --git a/src/utils/stable.ts b/src/utils/stable.ts new file mode 100644 index 0000000..4216459 --- /dev/null +++ b/src/utils/stable.ts @@ -0,0 +1,33 @@ +import { createHash } from "node:crypto" + +function canonicalize(value: unknown): unknown { + if (Array.isArray(value)) { + return value.map((item) => canonicalize(item)) + } + + if (value && typeof value === "object") { + const entries = Object.entries(value as Record) + .filter(([, item]) => item !== undefined) + .sort(([left], [right]) => left.localeCompare(right)) + + return Object.fromEntries(entries.map(([key, item]) => [key, canonicalize(item)])) + } + + if (typeof value === "number" && !Number.isFinite(value)) { + throw new Error("Cannot fingerprint a non-finite number") + } + + return value +} + +export function stableStringify(value: unknown): string { + return JSON.stringify(canonicalize(value)) +} + +export function sha256Text(value: string): string { + return createHash("sha256").update(value, "utf8").digest("hex") +} + +export function stableSha256(value: unknown): string { + return sha256Text(stableStringify(value)) +} diff --git a/test/answer-runtime.test.ts b/test/answer-runtime.test.ts new file mode 100644 index 0000000..e5300bd --- /dev/null +++ b/test/answer-runtime.test.ts @@ -0,0 +1,172 @@ +import { describe, expect, test } from "bun:test" +import { + aggregateAnswerAttemptUsage, + generateAnswerWithRetries, + normalizeAnsweringUsage, + shouldRunAnswerPhase, +} from "../src/orchestrator/phases/answer" +import type { AnswerAttemptMetrics } from "../src/types/checkpoint" +import { hasEvaluableAnswer } from "../src/orchestrator/phases/evaluate" +import { resolveAnsweringRuntimeIdentity } from "../src/utils/models" + +describe("answering runtime identity and usage", () => { + test("resolves the effective model defaults instead of persisting only an alias", () => { + expect(resolveAnsweringRuntimeIdentity("gpt-4.1-mini")).toEqual({ + schemaVersion: 1, + transport: "ai-sdk-generate-text-v1", + modelAlias: "gpt-4.1-mini", + provider: "openai", + modelId: "gpt-4.1-mini", + supportsTemperature: true, + effectiveDefaultTemperature: 0, + effectiveDefaultMaxOutputTokens: 1000, + }) + + expect(resolveAnsweringRuntimeIdentity("gpt-5").effectiveDefaultTemperature).toBeNull() + }) + + test("retains generateText token usage and derives a missing total", () => { + expect( + normalizeAnsweringUsage({ inputTokens: 120, outputTokens: 30, totalTokens: 150 }) + ).toEqual({ requestCount: 1, inputTokens: 120, outputTokens: 30, totalTokens: 150 }) + expect(normalizeAnsweringUsage({ inputTokens: 10, outputTokens: 2 })).toEqual({ + requestCount: 1, + inputTokens: 10, + outputTokens: 2, + totalTokens: 12, + }) + }) + + test("retries empty and failed answers with durable attempt usage", async () => { + const attempts: AnswerAttemptMetrics[] = [] + const delays: number[] = [] + const scripted = [ + { + text: "", + finishReason: "length", + usage: { inputTokens: 10, outputTokens: 2, reasoningTokens: 2, totalTokens: 12 }, + }, + new Error("transport failed"), + { + text: " final answer ", + finishReason: "stop", + usage: { inputTokens: 11, outputTokens: 3, reasoningTokens: 1, totalTokens: 14 }, + }, + ] + + const outcome = await generateAnswerWithRetries({ + maxAttempts: 5, + timeoutMs: 120_000, + retryBackoffMs: 2_000, + execute: async () => { + const next = scripted.shift()! + if (next instanceof Error) throw next + return next + }, + onAttempt: (attempt) => { + const index = attempts.findIndex((candidate) => candidate.attempt === attempt.attempt) + if (index >= 0) attempts[index] = attempt + else attempts.push(attempt) + }, + sleep: async (delayMs) => { + delays.push(delayMs) + }, + }) + + expect(outcome).toEqual({ hypothesis: "final answer", terminalEmptyAccepted: false }) + expect(delays).toEqual([2_000, 4_000]) + expect(attempts.map(({ status, error }) => ({ status, error }))).toEqual([ + { status: "failed", error: "Answering model returned an empty hypothesis" }, + { status: "failed", error: "transport failed" }, + { status: "completed", error: undefined }, + ]) + expect( + attempts.map(({ finishReason, reasoningTokens }) => ({ finishReason, reasoningTokens })) + ).toEqual([ + { finishReason: "length", reasoningTokens: 2 }, + { finishReason: undefined, reasoningTokens: undefined }, + { finishReason: "stop", reasoningTokens: 1 }, + ]) + expect(aggregateAnswerAttemptUsage(attempts)).toEqual({ + requestCount: 3, + tokenUsageCompleteRequestCount: 2, + tokenUsageUnknownRequestCount: 1, + inputTokens: 21, + outputTokens: 5, + reasoningTokens: 3, + totalTokens: 26, + }) + }) + + test("accepts terminal exhausted output only when the protocol explicitly owns that policy", async () => { + const attempts: AnswerAttemptMetrics[] = [] + const outcome = await generateAnswerWithRetries({ + maxAttempts: 5, + terminalEmptyOutputPolicy: "accept-and-evaluate", + execute: async () => ({ text: "", finishReason: "length" }), + onAttempt: (attempt) => { + const index = attempts.findIndex((candidate) => candidate.attempt === attempt.attempt) + if (index >= 0) attempts[index] = attempt + else attempts.push(attempt) + }, + }) + + expect(outcome).toEqual({ hypothesis: "", terminalEmptyAccepted: true }) + expect(attempts).toHaveLength(5) + expect(attempts.every((attempt) => attempt.status === "failed")).toBe(true) + + await expect( + generateAnswerWithRetries({ + maxAttempts: 2, + execute: async () => ({ text: "" }), + onAttempt: () => {}, + }) + ).rejects.toThrow("failed after 2 attempts") + + await expect( + generateAnswerWithRetries({ + maxAttempts: 2, + terminalEmptyOutputPolicy: "accept-and-evaluate", + execute: async (attempt) => { + if (attempt === 1) throw new Error("transport failure") + return { text: "" } + }, + onAttempt: () => {}, + }) + ).resolves.toEqual({ hypothesis: "", terminalEmptyAccepted: true }) + }) + + test("includes only explicitly accepted empty hypotheses in evaluation", () => { + expect( + hasEvaluableAnswer({ + status: "completed", + hypothesis: "", + terminalEmptyAccepted: true, + }) + ).toBe(true) + expect(hasEvaluableAnswer({ status: "completed", hypothesis: "" })).toBe(false) + expect(hasEvaluableAnswer({ status: "completed", hypothesis: "answer" })).toBe(true) + expect(hasEvaluableAnswer(undefined)).toBe(false) + }) + + test("does not rerun a completed accepted-empty answer on resume", () => { + expect( + shouldRunAnswerPhase({ + search: { status: "completed" }, + answer: { + status: "completed", + hypothesis: "", + terminalEmptyAccepted: true, + }, + evaluate: { status: "pending" }, + }) + ).toBe(false) + expect( + shouldRunAnswerPhase({ + search: { status: "completed" }, + answer: { status: "failed" }, + evaluate: { status: "pending" }, + }) + ).toBe(true) + }) +}) diff --git a/test/batch-report.test.ts b/test/batch-report.test.ts new file mode 100644 index 0000000..38b7850 --- /dev/null +++ b/test/batch-report.test.ts @@ -0,0 +1,368 @@ +import { describe, expect, test } from "bun:test" +import { + BatchManager, + comparePrimaryMetrics, + type CompareManifest, +} from "../src/orchestrator/batch" +import type { BenchmarkResult, LatencyStats } from "../src/types/unified" +import { stableSha256 } from "../src/utils/stable" +import { resolveAnsweringRuntimeIdentity } from "../src/utils/models" + +const ZERO_LATENCY: LatencyStats = { + min: 0, + max: 0, + mean: 0, + median: 0, + p95: 0, + p99: 0, + stdDev: 0, + count: 0, +} + +function report(input: { + provider: string + primaryKey?: string + primaryValue?: number + higherIsBetter?: boolean + passAccuracy: number +}): BenchmarkResult { + const totalQuestions = 10 + const questionIds = Array.from({ length: totalQuestions }, (_, index) => `q${index + 1}`) + const primaryValue = input.primaryValue ?? 0 + return { + provider: input.provider, + providerPromptFingerprint: `prompt-${input.provider}`, + benchmark: "beam-1m", + runId: `run-${input.provider}`, + dataSourceRunId: `run-${input.provider}`, + judge: "gpt-4.1-mini", + answeringModel: "gpt-4.1-mini", + answeringRuntimeIdentity: resolveAnsweringRuntimeIdentity("gpt-4.1-mini"), + timestamp: "2026-08-03T00:00:00.000Z", + selectedQuestionIdsDigest: stableSha256(questionIds), + benchmarkInputFingerprint: "benchmark-input-a", + retrievalTopK: 5, + benchmarkScope: { + displayName: "BEAM 1M", + includedTiers: ["1M"], + coverage: "subset", + }, + datasetIdentity: { datasetFingerprint: "dataset-a" }, + protocolIdentity: { id: "beam-paper", version: "1.1.0" }, + quality: { + primaryMetric: + input.primaryValue == null + ? undefined + : { + key: input.primaryKey ?? "beamScore", + value: input.primaryValue, + higherIsBetter: input.higherIsBetter ?? true, + }, + metrics: { passAccuracy: input.passAccuracy }, + }, + summary: { + totalQuestions, + correctCount: Math.round(input.passAccuracy * totalQuestions), + accuracy: input.passAccuracy, + averageScore: primaryValue, + }, + builds: { + uniqueBuildCount: 0, + sumContainerBuildWorkMs: 0, + buildPhaseWallClockMs: 0, + totalBuildCostUsd: null, + knownCostBuildCount: 0, + totalCostBuildCount: 0, + items: [], + }, + questionMetrics: questionIds.map((questionId) => ({ + questionId, + buildId: "build-1", + searchLatencyMs: 0, + answerLatencyMs: 0, + onlineQueryLatencyMs: 0, + evaluationLatencyMs: 0, + queryCostUsd: null, + evaluationCostUsd: null, + configuredTopK: 5, + providerRequestLimit: 5, + rawReturnedCount: 5, + returnedCount: 5, + normalizedCount: 5, + droppedCount: 0, + answerCutoff: 5, + answerEvidenceCount: 5, + contextTokens: 0, + providerRequests: [], + })), + latency: { + ingest: ZERO_LATENCY, + indexing: ZERO_LATENCY, + search: ZERO_LATENCY, + answer: ZERO_LATENCY, + evaluate: ZERO_LATENCY, + total: ZERO_LATENCY, + }, + byQuestionType: {}, + evaluations: [], + } +} + +function manifest(reports: BenchmarkResult[]): CompareManifest { + return { + compareId: "compare-test", + createdAt: "2026-08-03T00:00:00.000Z", + updatedAt: "2026-08-03T00:00:00.000Z", + benchmark: "beam-1m", + judge: "gpt-4.1-mini", + answeringModel: "gpt-4.1-mini", + answeringRuntimeIdentity: resolveAnsweringRuntimeIdentity("gpt-4.1-mini"), + targetQuestionIds: Array.from({ length: 10 }, (_, index) => `q${index + 1}`), + retrievalTopK: 5, + datasetIdentity: { datasetFingerprint: "dataset-a" } as any, + benchmarkScope: reports[0].benchmarkScope, + protocolIdentity: reports[0].protocolIdentity as any, + selectedQuestionIdsDigest: reports[0].selectedQuestionIdsDigest, + benchmarkInputFingerprint: reports[0].benchmarkInputFingerprint, + runs: reports.map((value) => ({ + provider: value.provider, + runId: value.runId, + providerPromptFingerprint: value.providerPromptFingerprint, + })), + } +} + +function captureConsoleLog(run: () => void): string { + const lines: string[] = [] + const original = console.log + console.log = (...values: unknown[]) => lines.push(values.map(String).join(" ")) + try { + run() + } finally { + console.log = original + } + return lines.join("\n") +} + +describe("batch comparison primary metric semantics", () => { + test("ranks by the comparable primary metric while keeping pass accuracy secondary", () => { + const alpha = report({ provider: "alpha", primaryValue: 0.75, passAccuracy: 0.5 }) + const beta = report({ provider: "beta", primaryValue: 0.7, passAccuracy: 0.9 }) + const reports = [ + { provider: beta.provider, report: beta }, + { provider: alpha.provider, report: alpha }, + ] + + const comparison = comparePrimaryMetrics(reports) + expect(comparison.comparable).toBe(true) + expect(comparison.rows.map(({ provider }) => provider)).toEqual(["alpha", "beta"]) + expect(comparison.rows[0].deltaFromBest).toBe(0) + expect(comparison.rows[1].deltaFromBest).toBeCloseTo(-0.05) + expect(comparison.winners).toEqual(["alpha"]) + + const manager = new BatchManager() + manager.getReports = () => reports + const output = captureConsoleLog(() => manager.printComparisonReport(manifest([beta, alpha]))) + expect(output).toContain("QUALITY — beamScore (higher is better)") + expect(output).toContain("Pass accuracy") + expect(output).toContain("-0.0500") + expect(output).toContain( + "WINNER: alpha (beamScore=0.7500; pass accuracy secondary: alpha 50.0%)" + ) + expect(output).not.toContain("WINNER: beta") + }) + + test("honors lower-is-better primary metrics", () => { + const alpha = report({ + provider: "alpha", + primaryKey: "errorRate", + primaryValue: 0.2, + higherIsBetter: false, + passAccuracy: 0.9, + }) + const beta = report({ + provider: "beta", + primaryKey: "errorRate", + primaryValue: 0.1, + higherIsBetter: false, + passAccuracy: 0.5, + }) + + const comparison = comparePrimaryMetrics([ + { provider: alpha.provider, report: alpha }, + { provider: beta.provider, report: beta }, + ]) + expect(comparison.rows.map(({ provider }) => provider)).toEqual(["beta", "alpha"]) + expect(comparison.rows.map(({ deltaFromBest }) => deltaFromBest)).toEqual([0, 0.1]) + expect(comparison.winners).toEqual(["beta"]) + }) + + test("does not rank or declare a winner when primary metric identities differ", () => { + const beam = report({ provider: "beam", primaryValue: 0.75, passAccuracy: 0.6 }) + const legacy = report({ + provider: "legacy", + primaryKey: "accuracy", + primaryValue: 0.9, + passAccuracy: 0.9, + }) + const reports = [ + { provider: beam.provider, report: beam }, + { provider: legacy.provider, report: legacy }, + ] + + const comparison = comparePrimaryMetrics(reports) + expect(comparison.comparable).toBe(false) + expect(comparison.rows.map(({ provider }) => provider)).toEqual(["beam", "legacy"]) + expect(comparison.rows.every(({ deltaFromBest }) => deltaFromBest === undefined)).toBe(true) + expect(comparison.winners).toEqual([]) + + const manager = new BatchManager() + manager.getReports = () => reports + const output = captureConsoleLog(() => manager.printComparisonReport(manifest([beam, legacy]))) + expect(output).toContain("suppressing ranking, deltas, and winner") + expect(output).toContain("NO WINNER: reports are not like-for-like") + expect(output).not.toContain("WINNER: beam (") + expect(output).not.toContain("WINNER: legacy (") + }) + + test("requires matching dataset, question set, protocol, Top-K, judge, and answering model", () => { + const baseline = report({ provider: "alpha", primaryValue: 0.75, passAccuracy: 0.6 }) + const variants: BenchmarkResult[] = [ + { + ...report({ provider: "beta", primaryValue: 0.7, passAccuracy: 0.6 }), + datasetIdentity: { datasetFingerprint: "dataset-b" }, + }, + { + ...report({ provider: "beta", primaryValue: 0.7, passAccuracy: 0.6 }), + selectedQuestionIdsDigest: "different-questions", + }, + { + ...report({ provider: "beta", primaryValue: 0.7, passAccuracy: 0.6 }), + benchmarkInputFingerprint: "different-transformed-input", + }, + { + ...report({ provider: "beta", primaryValue: 0.7, passAccuracy: 0.6 }), + protocolIdentity: { id: "beam-paper", version: "different" }, + }, + { + ...report({ provider: "beta", primaryValue: 0.7, passAccuracy: 0.6 }), + retrievalTopK: 10, + questionMetrics: baseline.questionMetrics.map((metric) => ({ + ...metric, + configuredTopK: 10, + })), + }, + { + ...report({ provider: "beta", primaryValue: 0.7, passAccuracy: 0.6 }), + judge: "different-judge", + }, + { + ...report({ provider: "beta", primaryValue: 0.7, passAccuracy: 0.6 }), + answeringModel: "different-answering-model", + }, + ] + + for (const variant of variants) { + const comparison = comparePrimaryMetrics([ + { provider: baseline.provider, report: baseline }, + { provider: variant.provider, report: variant }, + ]) + expect(comparison.comparable).toBe(false) + expect(comparison.winners).toEqual([]) + } + }) + + test("keeps combined BEAM non-comparable when no scalar cross-tier primary metric exists", () => { + const alpha = report({ provider: "alpha", passAccuracy: 0.9 }) + const beta = report({ provider: "beta", passAccuracy: 0.8 }) + const comparison = comparePrimaryMetrics([ + { provider: alpha.provider, report: alpha }, + { provider: beta.provider, report: beta }, + ]) + + expect(comparison.comparable).toBe(false) + expect(comparison.rows[0].primaryMetric).toBeUndefined() + expect(comparison.mismatchReasons.join(" ")).toContain("no scalar primary metric") + expect(comparison.winners).toEqual([]) + }) + + test("compares legacy providers with derived input identity and rejects input drift", () => { + const legacyReport = (provider: string, accuracy: number): BenchmarkResult => ({ + ...report({ + provider, + primaryKey: "accuracy", + primaryValue: accuracy, + passAccuracy: accuracy, + }), + benchmark: "locomo", + benchmarkScope: { displayName: "LoCoMo", includedTiers: [], coverage: "full" }, + datasetIdentity: undefined, + benchmarkInputFingerprint: "legacy-selected-input-a", + protocolIdentity: { id: "memorybench.legacy", version: "1.0.0" }, + providerPromptFingerprint: "shared-legacy-prompt", + }) + const alpha = legacyReport("alpha", 0.7) + const beta = legacyReport("beta", 0.8) + + const comparable = comparePrimaryMetrics([ + { provider: alpha.provider, report: alpha }, + { provider: beta.provider, report: beta }, + ]) + expect(comparable.comparable).toBe(true) + expect(comparable.winners).toEqual(["beta"]) + + const drifted = { + ...beta, + benchmarkInputFingerprint: "legacy-selected-input-with-haystack-drift", + } + const mismatch = comparePrimaryMetrics([ + { provider: alpha.provider, report: alpha }, + { provider: drifted.provider, report: drifted }, + ]) + expect(mismatch.comparable).toBe(false) + expect(mismatch.winners).toEqual([]) + }) + + test("does not rank a partial provider set", () => { + const alpha = report({ provider: "alpha", primaryValue: 0.8, passAccuracy: 0.8 }) + const comparison = comparePrimaryMetrics([{ provider: alpha.provider, report: alpha }], 2) + + expect(comparison.comparable).toBe(false) + expect(comparison.mismatchReasons).toContain("only 1 of 2 provider reports are complete") + expect(comparison.winners).toEqual([]) + }) +}) + +describe("batch comparison preflight barrier", () => { + test("preflights every run without starting provider execution", async () => { + const alpha = report({ provider: "alpha", primaryValue: 0.7, passAccuracy: 0.7 }) + const beta = report({ provider: "beta", primaryValue: 0.6, passAccuracy: 0.6 }) + const calls: Array<{ runId: string; preflightOnly?: boolean }> = [] + const manager = new BatchManager({ + async run(options) { + calls.push({ runId: options.runId, preflightOnly: options.preflightOnly }) + }, + }) + + await manager.preflightRuns(manifest([alpha, beta])) + + expect(calls).toEqual([ + { runId: alpha.runId, preflightOnly: true }, + { runId: beta.runId, preflightOnly: true }, + ]) + }) + + test("reports every provider preflight failure before execution starts", async () => { + const alpha = report({ provider: "alpha", primaryValue: 0.7, passAccuracy: 0.7 }) + const beta = report({ provider: "beta", primaryValue: 0.6, passAccuracy: 0.6 }) + const manager = new BatchManager({ + async run(options) { + if (options.runId === beta.runId) throw new Error("resume identity changed") + }, + }) + + await expect(manager.preflightRuns(manifest([alpha, beta]))).rejects.toThrow( + "beta: resume identity changed" + ) + }) +}) diff --git a/test/beam-dataset.test.ts b/test/beam-dataset.test.ts new file mode 100644 index 0000000..13d5f8b --- /dev/null +++ b/test/beam-dataset.test.ts @@ -0,0 +1,940 @@ +import { afterEach, describe, expect, test } from "bun:test" +import { mkdtemp, readFile, readdir, rm, unlink, writeFile } from "node:fs/promises" +import { tmpdir } from "node:os" +import { join } from "node:path" +import { + BEAM_EXPECTED_COUNTS, + canonicalizeBeamRows, + computeDatasetFingerprint, + computeManifestHash, + describeBeamTemporalCoverage, + getUnverifiedBeamDerivationTiers, + loadPreparedBeamDataset, + loadPreparedBeamTestFixture, + parseBeamTimeAnchorStrict, + serializeBeamJsonl, + sha256Bytes, + stableBeamStringify, + validateCanonicalBeamTier, + verifyPreparedBeamSourceDerivation, +} from "../src/benchmarks/beam/dataset" +import { prepareBeamDataset } from "../src/benchmarks/beam/prepare" +import { createBeamDatasetIdentity } from "../src/benchmarks/beam" +import { + BEAM_QUESTION_TYPE_IDS, + type BeamDatasetManifest, + type BeamScale, +} from "../src/benchmarks/beam/types" + +const tempRoots: string[] = [] + +afterEach(async () => { + await Promise.all(tempRoots.splice(0).map((path) => rm(path, { recursive: true, force: true }))) +}) + +async function tempRoot(): Promise { + const path = await mkdtemp(join(tmpdir(), "memorybench-beam-test-")) + tempRoots.push(path) + return path +} + +function probingQuestions(chatId: string): Record { + return Object.fromEntries( + BEAM_QUESTION_TYPE_IDS.map((questionType) => [ + questionType, + [0, 1].map((ordinal) => ({ + question: `Question ${chatId} ${questionType} ${ordinal}`, + rubric: [ + `Nugget ${chatId} ${questionType} ${ordinal} A`, + `Nugget ${chatId} ${questionType} ${ordinal} B`, + ], + answer: `Answer ${chatId} ${questionType} ${ordinal}`, + difficulty: ordinal === 0 ? "easy" : "hard", + })), + ]) + ) +} + +function sourceRows( + scale: BeamScale, + options?: { invalidBatchAnchorWithMessageFallback?: boolean } +): Record[] { + const count = BEAM_EXPECTED_COUNTS[scale].chats + return Array.from({ length: count }, (_, index) => { + const chatId = String(index + 1) + const invalidWithFallback = index === 0 && options?.invalidBatchAnchorWithMessageFallback + return { + conversation_id: chatId, + chat: [ + { + batch_number: 1, + time_anchor: invalidWithFallback ? "February-30-2024" : "March-01-2024", + turns: [ + [ + { + role: "user", + content: `User ${scale}/${chatId}`, + ...(invalidWithFallback ? { time_anchor: "March-02-2024" } : {}), + }, + { role: "assistant", content: `Assistant ${scale}/${chatId}` }, + ], + ], + }, + ], + probing_questions: probingQuestions(chatId), + } + }) +} + +function parquetResponse(payload = "fixture"): Response { + const bytes = Buffer.concat([Buffer.from("PAR1"), Buffer.from(payload), Buffer.from("PAR1")]) + return new Response(bytes, { + status: 200, + headers: { "content-length": String(bytes.byteLength) }, + }) +} + +function flipOneByte(bytes: Uint8Array): Buffer { + const changed = Buffer.from(bytes) + changed[Math.min(4, changed.length - 1)]! ^= 1 + return changed +} + +async function rewriteManifest( + snapshotPath: string, + mutate: (manifest: BeamDatasetManifest) => void +): Promise { + const current = JSON.parse( + await readFile(join(snapshotPath, "manifest.json"), "utf8") + ) as BeamDatasetManifest + mutate(current) + const { + datasetFingerprint: _oldDatasetFingerprint, + manifestHash: _oldManifestHash, + ...core + } = current + const datasetFingerprint = computeDatasetFingerprint(core) + const withoutHash = { ...core, datasetFingerprint } + const manifest: BeamDatasetManifest = { + ...withoutHash, + manifestHash: computeManifestHash(withoutHash), + } + await writeFile(join(snapshotPath, "manifest.json"), stableBeamStringify(manifest) + "\n") + await writeFile( + join(snapshotPath, ".complete"), + stableBeamStringify({ + datasetFingerprint: manifest.datasetFingerprint, + manifestHash: manifest.manifestHash, + }) + "\n" + ) + return manifest +} + +function fixtureDecoder(rows: Record[]) { + return async (filePath: string): Promise => { + if (filePath.includes("0-10M-")) return rows.slice(0, 5) + if (filePath.includes("1-10M-")) return rows.slice(5) + return rows + } +} + +function toPythonLiteral(value: unknown): string { + return JSON.stringify(value) + .replaceAll("'", "\\'") + .replaceAll('"', "'") + .replace(/\btrue\b/g, "True") + .replace(/\bfalse\b/g, "False") + .replace(/\bnull\b/g, "None") +} + +describe("BEAM canonical dataset", () => { + test("tracks source derivation independently for every tier in a combined snapshot", () => { + const fingerprint = "f".repeat(64) + const verified = new Set([`${fingerprint}:1M`]) + + expect(getUnverifiedBeamDerivationTiers(fingerprint, ["1M", "10M"], verified)).toEqual(["10M"]) + }) + + test("validates exact 1M counts, abilities, stable IDs, and literal transcript content", () => { + const canonical = canonicalizeBeamRows("1M", sourceRows("1M")) + + expect(canonical.chats).toHaveLength(35) + expect(canonical.questions).toHaveLength(700) + expect(canonical.counts.byQuestionType.abstention).toBe(70) + expect(canonical.chats[0].sessions[0].messages[0].content).toContain("") + expect(canonical.chats[0].sessions[0].documentDate).toBe("2024-03-01") + expect(canonical.questions[0].questionId).toMatch( + /^beam:1M:[a-zA-Z0-9_-]+:[a-z_]+:[a-f0-9]{64}$/ + ) + }) + + test("validates exact 10M counts", () => { + const canonical = canonicalizeBeamRows("10M", sourceRows("10M")) + expect(canonical.chats).toHaveLength(10) + expect(canonical.questions).toHaveLength(200) + for (const questionType of BEAM_QUESTION_TYPE_IDS) { + expect(canonical.counts.byQuestionType[questionType]).toBe(20) + } + }) + + test("parses the published 1M flat batches and Python-style probing questions", () => { + const rows = sourceRows("1M") + const questions = probingQuestions("1") + ;(questions.abstention[0] as Record).publishedFlag = true + ;(questions.abstention[0] as Record).publishedOptional = null + rows[0] = { + conversation_id: "1", + chat: [ + [ + { + id: 0, + role: "user", + content: "First official user message", + index: "1,1", + question_type: "main_question", + time_anchor: "March-01-2024", + }, + { id: 1, role: "assistant", content: "First official answer" }, + { id: 2, role: "user", content: "Follow-up" }, + { id: 3, role: "assistant", content: "Follow-up answer" }, + ], + [ + { + id: 4, + role: "user", + content: "Second batch", + time_anchor: "March-12-2024", + }, + { id: 5, role: "assistant", content: "Second batch answer" }, + ], + ], + probing_questions: toPythonLiteral(questions), + } + + const canonical = canonicalizeBeamRows("1M", rows) + const chat = canonical.chats.find((item) => item.chatId === "1")! + expect(chat.sessions).toHaveLength(3) + expect(chat.sessions.map((session) => session.sessionId)).toEqual([ + "beam-1M-1-batch-1-turn-1", + "beam-1M-1-batch-1-turn-2", + "beam-1M-1-batch-2-turn-1", + ]) + expect(chat.sessions.map((session) => session.documentDate)).toEqual([ + "2024-03-01", + "2024-03-01", + "2024-03-12", + ]) + expect(chat.sessions[1].messages.map((message) => message.role)).toEqual(["user", "assistant"]) + }) + + test("selects only the reviewed published transcript column", () => { + const rows = sourceRows("1M") + rows[0].chat_truncated = rows[0].chat + delete rows[0].chat + + expect(() => canonicalizeBeamRows("1M", rows)).toThrow("published transcript is missing") + }) + + test("parses nullable 10M plan fields without colliding repeated batch numbers", () => { + const rows = sourceRows("10M") + rows[0] = { + conversation_id: "1", + chat: [ + { + "plan-1": [ + { + batch_number: 1, + time_anchor: null, + turns: [ + [ + { role: "user", content: "Plan one", time_anchor: "July-01-2024" }, + { role: "assistant", content: "Plan one answer" }, + ], + ], + }, + ], + "plan-2": null, + }, + { + "plan-1": null, + "plan-2": [ + { + batch_number: 1, + time_anchor: null, + turns: [ + [ + { role: "user", content: "Plan two", time_anchor: "July-16-2024" }, + { role: "assistant", content: "Plan two answer" }, + ], + ], + }, + ], + }, + ], + probing_questions: toPythonLiteral(probingQuestions("1")), + } + + const canonical = canonicalizeBeamRows("10M", rows) + const chat = canonical.chats.find((item) => item.chatId === "1")! + expect(chat.sessions.map((session) => session.sessionId)).toEqual([ + "beam-10M-1-plan-1-batch-1-turn-1", + "beam-10M-1-plan-2-batch-1-turn-1", + ]) + expect(chat.sessions.map((session) => session.planNumber)).toEqual([1, 2]) + expect(chat.sessions.map((session) => session.documentDate)).toEqual([ + "2024-07-01", + "2024-07-16", + ]) + }) + + test("splits pinned 10M complete variable-length blocks into strict pairs", () => { + const rows = sourceRows("10M") + const fivePairBlock = Array.from({ length: 5 }, (_, pairIndex) => [ + { + id: pairIndex * 2, + index: `1,${pairIndex + 1}`, + question_type: pairIndex === 0 ? "main_question" : "followup_question", + role: "user", + content: `Official 10M user ${pairIndex + 1}`, + time_anchor: pairIndex === 0 ? "July-01-2024" : null, + }, + { + id: pairIndex * 2 + 1, + index: null, + question_type: null, + role: "assistant", + content: `Official 10M assistant ${pairIndex + 1}`, + time_anchor: null, + }, + ]).flat() + rows[0] = { + conversation_id: "1", + chat: [ + { + "plan-1": [ + { + batch_number: 1, + time_anchor: null, + turns: [fivePairBlock], + }, + ], + }, + ], + probing_questions: toPythonLiteral(probingQuestions("1")), + } + + const canonical = canonicalizeBeamRows("10M", rows) + const chat = canonical.chats.find((item) => item.chatId === "1")! + + expect(chat.sessions).toHaveLength(5) + expect(chat.sessions.map((session) => session.sessionId)).toEqual( + Array.from({ length: 5 }, (_, index) => `beam-10M-1-plan-1-batch-1-turn-${index + 1}`) + ) + expect( + chat.sessions.every( + (session) => + session.messages.length === 2 && + session.messages[0]?.role === "user" && + session.messages[1]?.role === "assistant" + ) + ).toBe(true) + expect(chat.sessions.map((session) => session.documentDate)).toEqual( + Array.from({ length: 5 }, () => "2024-07-01") + ) + expect(canonical.counts.sessionsWithPaddedAssistant).toBe(0) + }) + + test("pads only the two audited pinned 10M missing-assistant source identities", () => { + const rows = sourceRows("10M") + const pair = (label: string) => [ + { role: "user", content: `User ${label}` }, + { role: "assistant", content: `Assistant ${label}` }, + ] + const incomplete = (literalUserContent: string) => [ + { role: "user", content: `Complete user before ${literalUserContent}` }, + { role: "assistant", content: `Complete assistant before ${literalUserContent}` }, + { + role: "user", + content: literalUserContent, + question_type: "followup_question", + }, + ] + + rows[0] = { + conversation_id: "1", + chat: [ + { + "plan-7": [ + { + batch_number: 10, + time_anchor: null, + turns: [ + ...Array.from({ length: 18 }, (_, index) => pair(`one-${index + 1}`)), + incomplete("Literal "), + ], + }, + ], + }, + ], + probing_questions: toPythonLiteral(probingQuestions("1")), + } + rows[1] = { + conversation_id: "2", + chat: [ + { + "plan-7": [ + { + batch_number: 8, + time_anchor: null, + turns: [ + ...Array.from({ length: 50 }, (_, index) => pair(`two-${index + 1}`)), + incomplete("Literal "), + ], + }, + ], + }, + ], + probing_questions: toPythonLiteral(probingQuestions("2")), + } + + const canonical = canonicalizeBeamRows("10M", rows) + const padded = canonical.chats.flatMap((chat) => + chat.sessions.filter((session) => session.hasPaddedAssistant) + ) + + expect(canonical.counts.sessionsWithPaddedAssistant).toBe(2) + expect(padded.map((session) => session.sessionId)).toEqual([ + "beam-10M-1-plan-7-batch-10-turn-20", + "beam-10M-2-plan-7-batch-8-turn-52", + ]) + expect(padded.map((session) => session.messages)).toEqual([ + [ + { role: "user", content: "Literal " }, + { role: "assistant", content: "N/A" }, + ], + [ + { role: "user", content: "Literal " }, + { role: "assistant", content: "N/A" }, + ], + ]) + expect(describeBeamTemporalCoverage({ "10M": canonical.counts }, ["10M"])).toContain( + `2/${canonical.counts.sessions} use an audited N/A assistant padding` + ) + }) + + test("fails closed when a published flat batch does not alternate user and assistant", () => { + const rows = sourceRows("1M") + rows[0].chat = [ + [ + { role: "user", content: "one" }, + { role: "user", content: "two" }, + ], + ] + expect(() => canonicalizeBeamRows("1M", rows)).toThrow("must alternate user then assistant") + }) + + test("fails closed for malformed structured 10M message blocks", () => { + const invalidCases: Array<{ turn: unknown[]; error: string }> = [ + { + turn: [{ role: "user", content: "missing response" }], + error: "at least one user message followed by one assistant message", + }, + { + turn: [ + { role: "assistant", content: "wrong first role" }, + { role: "user", content: "wrong second role" }, + ], + error: "must alternate user then assistant", + }, + { + turn: [ + { role: "user", content: "first" }, + { role: "assistant", content: "second" }, + { role: "assistant", content: "extra" }, + ], + error: "must alternate user then assistant", + }, + { + turn: [ + { role: "user", content: "first" }, + { role: "assistant", content: "second" }, + { role: "user", content: "unclassified dangling user" }, + ], + error: "must contain complete user/assistant pairs", + }, + { + turn: [ + { role: "user", content: "first" }, + { role: "assistant", content: "second" }, + { + role: "user", + content: "classified but identity-mismatched dangling user", + question_type: "followup_question", + }, + ], + error: "10M:1:plan-0:batch-1:source-turn-1", + }, + { + turn: [ + { role: "user", content: "first" }, + { role: "assistant", content: "second" }, + { role: "user", content: "third" }, + { role: "assistant", content: "fourth" }, + { + role: "user", + content: "unexpected larger odd block", + question_type: "followup_question", + }, + ], + error: "must contain complete user/assistant pairs", + }, + ] + + for (const invalid of invalidCases) { + const rows = sourceRows("10M") + rows[0].chat = [ + { + batch_number: 1, + turns: [invalid.turn], + }, + ] + expect(() => canonicalizeBeamRows("10M", rows)).toThrow(invalid.error) + } + }) + + test("rejects structural tampering at an audited padded-assistant identity", () => { + const rows = sourceRows("10M") + rows[0].chat = [ + { + "plan-7": [ + { + batch_number: 10, + turns: [ + ...Array.from({ length: 18 }, (_, index) => [ + { role: "user", content: `Prior user ${index}` }, + { role: "assistant", content: `Prior assistant ${index}` }, + ]), + [ + { role: "user", content: "Complete user" }, + { role: "assistant", content: "Complete assistant" }, + { + role: "user", + content: "Tampered classification", + question_type: "main_question", + }, + ], + ], + }, + ], + }, + ] + + expect(() => canonicalizeBeamRows("10M", rows)).toThrow( + "must contain complete user/assistant pairs" + ) + }) + + test("canonical output and IDs are independent of source-row order", () => { + const rows = sourceRows("1M") + const forward = canonicalizeBeamRows("1M", rows) + const reverse = canonicalizeBeamRows("1M", [...rows].reverse()) + + expect(serializeBeamJsonl(reverse.chats)).toBe(serializeBeamJsonl(forward.chats)) + expect(serializeBeamJsonl(reverse.questions)).toBe(serializeBeamJsonl(forward.questions)) + }) + + test("rejects incomplete tiers before selection", () => { + expect(() => canonicalizeBeamRows("1M", sourceRows("1M").slice(0, 34))).toThrow( + "expected 35 chats" + ) + }) + + test("rejects empty rubrics and unknown abilities", () => { + const emptyRubric = sourceRows("10M") + ;( + emptyRubric[0].probing_questions as Record>> + ).abstention[0].rubric = [] + expect(() => canonicalizeBeamRows("10M", emptyRubric)).toThrow("rubric must be a non-empty") + + const unknown = sourceRows("10M") + ;(unknown[0].probing_questions as Record).made_up_ability = [ + { question: "Bad", rubric: ["Bad"] }, + ] + expect(() => canonicalizeBeamRows("10M", unknown)).toThrow("unknown BEAM question type") + }) + + test("uses the first valid message anchor when a batch anchor is invalid", () => { + const canonical = canonicalizeBeamRows( + "10M", + sourceRows("10M", { invalidBatchAnchorWithMessageFallback: true }) + ) + const session = canonical.chats[0].sessions[0] + expect(session.documentDate).toBe("2024-03-02") + expect(session.hadInvalidTimeAnchor).toBe(true) + expect(canonical.counts.sessionsWithInvalidTimeAnchor).toBe(1) + expect(describeBeamTemporalCoverage({ "10M": canonical.counts }, ["10M"])).toContain( + "10M: 0/10 sessions without a valid date; 1/10 encountered an invalid source time anchor" + ) + expect(describeBeamTemporalCoverage({ "10M": canonical.counts }, ["10M"])).toContain( + "separate, not additive" + ) + }) + + test("strict date parsing rejects impossible dates", () => { + expect(parseBeamTimeAnchorStrict("March-01-2024")).toBe("2024-03-01") + expect(parseBeamTimeAnchorStrict("2024-02-29")).toBe("2024-02-29") + expect(parseBeamTimeAnchorStrict("February-30-2024")).toBeUndefined() + expect(parseBeamTimeAnchorStrict("2023-02-29")).toBeUndefined() + expect(parseBeamTimeAnchorStrict(null)).toBeUndefined() + }) + + test("rejects malformed canonical session ordinals, identity, and message pairing", () => { + const valid = canonicalizeBeamRows("1M", sourceRows("1M")) + const cases = [ + (chats: typeof valid.chats) => { + chats[0]!.sessions[0]!.sessionId = "" + }, + (chats: typeof valid.chats) => { + chats[0]!.sessions[0]!.sessionId = "arbitrary-session" + }, + (chats: typeof valid.chats) => { + chats[0]!.sessions[0]!.batchNumber = -1 + }, + (chats: typeof valid.chats) => { + chats[0]!.sessions[0]!.messages = [{ role: "assistant", content: "orphan" }] + }, + (chats: typeof valid.chats) => { + chats[0]!.sessions[0]!.messages = [ + { role: "assistant", content: "wrong first" }, + { role: "user", content: "wrong second" }, + ] + }, + (chats: typeof valid.chats) => { + chats[0]!.sessions[0]!.hasPaddedAssistant = true + }, + ] + + for (const mutate of cases) { + const chats = structuredClone(valid.chats) + mutate(chats) + expect(() => validateCanonicalBeamTier("1M", chats, valid.questions)).toThrow() + } + }) +}) + +describe("BEAM preparation and loading", () => { + test("publishes a validated snapshot with manifest and completion marker", async () => { + const outputRoot = await tempRoot() + const result = await prepareBeamDataset({ + tiers: ["1M"], + outputRoot, + fetchImpl: async () => parquetResponse(), + parquetDecoder: async (_path, tier) => sourceRows(tier), + unsafeSkipPublishedHashCheckForTests: true, + }) + + const loaded = await loadPreparedBeamTestFixture({ + snapshotPath: result.snapshotPath, + tiers: ["1M"], + expectedDatasetFingerprint: result.manifest.datasetFingerprint, + }) + expect(loaded.manifest.sources[0].sourceIdentity).toBe("injected-test-fixture") + expect(loaded.manifest.sources[0].revision).toHaveLength(40) + expect(loaded.chatsByTier["1M"]).toHaveLength(35) + expect(loaded.questionsByTier["1M"]).toHaveLength(700) + expect(loaded.manifest.counts["1M"]?.byChat["1"]).toMatchObject({ + sessions: 1, + questions: 20, + byQuestionType: expect.objectContaining({ abstention: 2, temporal_reasoning: 2 }), + }) + expect(loaded.manifest.orderedQuestionIds["1M"]).toHaveLength(700) + expect(loaded.manifest.orderedChatIdsDigest["1M"]).toMatch(/^[a-f0-9]{64}$/) + expect(loaded.manifest.orderedQuestionIdsDigest["1M"]).toMatch(/^[a-f0-9]{64}$/) + }) + + test("uses selected-tier identity while retaining enclosing snapshot provenance", async () => { + const outputRoot = await tempRoot() + const oneTier = await prepareBeamDataset({ + tiers: ["1M"], + outputRoot, + fetchImpl: async () => parquetResponse(), + parquetDecoder: async (_path, tier) => sourceRows(tier), + unsafeSkipPublishedHashCheckForTests: true, + }) + const combined = structuredClone(oneTier.manifest) + combined.includedTiers = ["1M", "10M"] + combined.datasetFingerprint = "a".repeat(64) + combined.manifestHash = "b".repeat(64) + combined.sources.push({ + ...structuredClone(combined.sources[0]!), + tier: "10M", + split: "10M", + }) + combined.canonicalFiles.push({ + ...structuredClone(combined.canonicalFiles[0]!), + path: "canonical/10M/chats.jsonl", + }) + combined.counts["10M"] = structuredClone(combined.counts["1M"]!) + combined.orderedChatIds["10M"] = ["unused"] + combined.orderedChatIdsDigest["10M"] = "c".repeat(64) + combined.orderedQuestionIds["10M"] = ["unused"] + combined.orderedQuestionIdsDigest["10M"] = "d".repeat(64) + + const standaloneIdentity = createBeamDatasetIdentity(oneTier.manifest, ["1M"]) + const combinedIdentity = createBeamDatasetIdentity(combined, ["1M"]) + expect(combinedIdentity.datasetFingerprint).toBe(standaloneIdentity.datasetFingerprint) + expect(combinedIdentity.manifestHash).toBe(standaloneIdentity.manifestHash) + expect(combinedIdentity.sourceFiles).toEqual(standaloneIdentity.sourceFiles) + expect(combinedIdentity.canonicalFiles).toEqual(standaloneIdentity.canonicalFiles) + expect(combinedIdentity.snapshotFingerprint).toBe(combined.datasetFingerprint) + expect(combinedIdentity.snapshotFingerprint).not.toBe( + standaloneIdentity.snapshotFingerprint + ) + }) + + test("validates every manifest count and ordered identity field", async () => { + const mutations: Array<{ + message: string + mutate: (manifest: BeamDatasetManifest) => void + }> = [ + { + message: "validated counts do not match manifest", + mutate(manifest) { + manifest.counts["1M"]!.byChat["1"]!.questions = 19 + }, + }, + { + message: "validated counts do not match manifest", + mutate(manifest) { + manifest.counts["1M"]!.sessionsWithPaddedAssistant = 1 + }, + }, + { + message: "chat identity digest does not match manifest", + mutate(manifest) { + manifest.orderedChatIdsDigest["1M"] = "0".repeat(64) + }, + }, + { + message: "ordered question identity does not match manifest", + mutate(manifest) { + manifest.orderedQuestionIds["1M"]![0] = "tampered-question" + }, + }, + ] + + for (const mutation of mutations) { + const outputRoot = await tempRoot() + const result = await prepareBeamDataset({ + tiers: ["1M"], + outputRoot, + fetchImpl: async () => parquetResponse(), + parquetDecoder: async (_path, tier) => sourceRows(tier), + unsafeSkipPublishedHashCheckForTests: true, + }) + await rewriteManifest(result.snapshotPath, mutation.mutate) + await expect( + loadPreparedBeamTestFixture({ snapshotPath: result.snapshotPath, tiers: ["1M"] }) + ).rejects.toThrow(mutation.message) + } + }) + + test("ordinary scored-run loading rejects an injected test-source snapshot", async () => { + const outputRoot = await tempRoot() + const result = await prepareBeamDataset({ + tiers: ["1M"], + outputRoot, + fetchImpl: async () => parquetResponse(), + parquetDecoder: async (_path, tier) => sourceRows(tier), + unsafeSkipPublishedHashCheckForTests: true, + }) + + await expect( + loadPreparedBeamDataset({ snapshotPath: result.snapshotPath, tiers: ["1M"] }) + ).rejects.toThrow("injected test-source identity") + }) + + test("ordinary loading authenticates source bytes against the reviewed published pin", async () => { + const outputRoot = await tempRoot() + const result = await prepareBeamDataset({ + tiers: ["1M"], + outputRoot, + fetchImpl: async () => parquetResponse(), + parquetDecoder: async (_path, tier) => sourceRows(tier), + unsafeSkipPublishedHashCheckForTests: true, + }) + await rewriteManifest(result.snapshotPath, (manifest) => { + manifest.sources[0]!.sourceIdentity = "reviewed-published" + }) + + await expect( + loadPreparedBeamDataset({ snapshotPath: result.snapshotPath, tiers: ["1M"] }) + ).rejects.toThrow("SHA-256 does not match the reviewed published pin") + }) + + test("source-byte identity changes the dataset fingerprint even when canonical rows match", async () => { + const firstRoot = await tempRoot() + const secondRoot = await tempRoot() + const first = await prepareBeamDataset({ + tiers: ["1M"], + outputRoot: firstRoot, + fetchImpl: async () => parquetResponse("fixture-a"), + parquetDecoder: async (_path, tier) => sourceRows(tier), + unsafeSkipPublishedHashCheckForTests: true, + }) + const second = await prepareBeamDataset({ + tiers: ["1M"], + outputRoot: secondRoot, + fetchImpl: async () => parquetResponse("fixture-b"), + parquetDecoder: async (_path, tier) => sourceRows(tier), + unsafeSkipPublishedHashCheckForTests: true, + }) + + expect(second.manifest.canonicalFiles).toEqual(first.manifest.canonicalFiles) + expect(second.manifest.datasetFingerprint).not.toBe(first.manifest.datasetFingerprint) + }) + + test("detects source and canonical file tampering", async () => { + for (const target of ["source", "canonical"] as const) { + const outputRoot = await tempRoot() + const result = await prepareBeamDataset({ + tiers: ["1M"], + outputRoot, + fetchImpl: async () => parquetResponse(), + parquetDecoder: async (_path, tier) => sourceRows(tier), + unsafeSkipPublishedHashCheckForTests: true, + }) + const relativePath = + target === "source" + ? result.manifest.sources[0]!.files[0]!.snapshotPath + : result.manifest.canonicalFiles[0]!.path + const path = join(result.snapshotPath, relativePath) + await writeFile(path, flipOneByte(await readFile(path))) + + await expect( + loadPreparedBeamTestFixture({ + snapshotPath: result.snapshotPath, + tiers: ["1M"], + }) + ).rejects.toThrow(`${target} file hash mismatch`) + } + }) + + test("produces the same fingerprint for shuffled source rows", async () => { + const firstRoot = await tempRoot() + const secondRoot = await tempRoot() + const rows = sourceRows("10M") + const first = await prepareBeamDataset({ + tiers: ["10M"], + outputRoot: firstRoot, + fetchImpl: async () => parquetResponse(), + parquetDecoder: fixtureDecoder(rows), + unsafeSkipPublishedHashCheckForTests: true, + }) + const second = await prepareBeamDataset({ + tiers: ["10M"], + outputRoot: secondRoot, + fetchImpl: async () => parquetResponse(), + parquetDecoder: fixtureDecoder([...rows].reverse()), + unsafeSkipPublishedHashCheckForTests: true, + }) + + expect(second.manifest.datasetFingerprint).toBe(first.manifest.datasetFingerprint) + expect(second.manifest.manifestHash).toBe(first.manifest.manifestHash) + }) + + test("re-derives canonical bytes from source instead of trusting the manifest", async () => { + const outputRoot = await tempRoot() + const rows = sourceRows("1M") + const result = await prepareBeamDataset({ + tiers: ["1M"], + outputRoot, + fetchImpl: async () => parquetResponse(), + parquetDecoder: async () => rows, + unsafeSkipPublishedHashCheckForTests: true, + }) + const prepared = await loadPreparedBeamTestFixture({ + snapshotPath: result.snapshotPath, + tiers: ["1M"], + }) + await expect( + verifyPreparedBeamSourceDerivation(prepared, ["1M"], async () => rows) + ).resolves.toBeUndefined() + + const chatsPath = join(result.snapshotPath, "canonical/1M/chats.jsonl") + const chats = (await readFile(chatsPath, "utf8")) + .trim() + .split("\n") + .map((line) => JSON.parse(line)) + chats[0].sessions[0].messages[0].content = "forged but schema-valid transcript" + const forgedContent = serializeBeamJsonl(chats) + await writeFile(chatsPath, forgedContent) + await rewriteManifest(result.snapshotPath, (manifest) => { + const entry = manifest.canonicalFiles.find( + (file) => file.path === "canonical/1M/chats.jsonl" + )! + const bytes = Buffer.from(forgedContent) + entry.byteSize = bytes.byteLength + entry.sha256 = sha256Bytes(bytes) + }) + const forged = await loadPreparedBeamTestFixture({ + snapshotPath: result.snapshotPath, + tiers: ["1M"], + }) + await expect( + verifyPreparedBeamSourceDerivation(forged, ["1M"], async () => rows) + ).rejects.toThrow("source-to-canonical derivation mismatch") + }) + + test("fails closed and removes staging output after malformed conversion", async () => { + const outputRoot = await tempRoot() + await expect( + prepareBeamDataset({ + tiers: ["1M"], + outputRoot, + fetchImpl: async () => parquetResponse(), + parquetDecoder: async () => sourceRows("1M").slice(0, 1), + unsafeSkipPublishedHashCheckForTests: true, + }) + ).rejects.toThrow("expected 35 chats") + + expect(await readdir(outputRoot)).toEqual([]) + }) + + test("validates all bytes before publishing a completion marker", async () => { + const outputRoot = await tempRoot() + await expect( + prepareBeamDataset({ + tiers: ["1M"], + outputRoot, + fetchImpl: async () => parquetResponse(), + parquetDecoder: async (filePath, tier) => { + await writeFile(filePath, flipOneByte(await readFile(filePath))) + return sourceRows(tier) + }, + unsafeSkipPublishedHashCheckForTests: true, + }) + ).rejects.toThrow("source file hash mismatch") + + expect(await readdir(outputRoot)).toEqual([]) + }) + + test("refuses a snapshot without its atomic completion marker", async () => { + const outputRoot = await tempRoot() + const result = await prepareBeamDataset({ + tiers: ["10M"], + outputRoot, + fetchImpl: async () => parquetResponse(), + parquetDecoder: fixtureDecoder(sourceRows("10M")), + unsafeSkipPublishedHashCheckForTests: true, + }) + await unlink(join(result.snapshotPath, ".complete")) + + await expect( + loadPreparedBeamTestFixture({ + snapshotPath: result.snapshotPath, + tiers: ["10M"], + }) + ).rejects.toThrow("snapshot is incomplete") + }) +}) diff --git a/test/beam-evaluator.test.ts b/test/beam-evaluator.test.ts new file mode 100644 index 0000000..556db00 --- /dev/null +++ b/test/beam-evaluator.test.ts @@ -0,0 +1,811 @@ +import { describe, expect, test } from "bun:test" +import type { + EvaluationRuntime, + QuestionEvaluation, + StructuredModelRequest, +} from "../src/types/protocol" +import type { JudgeInput, JudgeResult } from "../src/types/judge" +import type { UnifiedQuestion, UnifiedSearchResult } from "../src/types/unified" +import { + BEAM_ABILITY_IDS, + BEAM_EVENT_EQUIVALENCE_SCHEMA, + BEAM_EVENT_ORDERING_SCORING_VERSION, + BEAM_EVALUATOR_IDENTITY, + BEAM_AGGREGATION_IMPLEMENTATION_SHA256, + BEAM_STRUCTURED_OUTPUT_SCHEMA_SHA256, + BEAM_NUGGET_JUDGMENT_SCHEMA, + BEAM_PASS_THRESHOLD, + BEAM_PAPER_PROTOCOL_VERSION, + BEAM_RETRIEVAL_TOP_K_VALUES, + BeamPaperProtocol, + alignBeamEvents, + computeKendallTauB, + extractBeamPredictedEvents, + scoreAlignedBeamEvents, +} from "../src/protocols/beam-paper" + +class ScriptedRuntime implements EvaluationRuntime { + readonly requests: StructuredModelRequest[] = [] + private readonly outputs: unknown[] + + constructor(outputs: unknown[]) { + this.outputs = [...outputs] + } + + async evaluateLegacy(_input: JudgeInput): Promise { + throw new Error("Legacy evaluation is not expected in BEAM tests") + } + + async generateStructured(request: StructuredModelRequest): Promise { + this.requests.push(request as StructuredModelRequest) + if (this.outputs.length === 0) throw new Error("No scripted structured output remains") + return request.schema.parse(this.outputs.shift()) + } +} + +class ExactEventRuntime implements EvaluationRuntime { + readonly requests: StructuredModelRequest[] = [] + + async evaluateLegacy(_input: JudgeInput): Promise { + throw new Error("Legacy evaluation is not expected in BEAM tests") + } + + async generateStructured(request: StructuredModelRequest): Promise { + this.requests.push(request as StructuredModelRequest) + const match = request.prompt.match(/^First snippet: ([\s\S]*) \n Second snippet: ([\s\S]*)$/u) + if (!match) throw new Error(`Unexpected event prompt: ${request.prompt}`) + return request.schema.parse({ answer: match[1] === match[2] ? "YES" : "NO" }) + } +} + +function makeQuestion( + input: { + id?: string + type?: string + question?: string + rubric?: unknown + scale?: "1M" | "10M" + } = {} +): UnifiedQuestion { + return { + questionId: input.id ?? "beam-1m-1-abstention-0", + question: input.question ?? "What happened?", + questionType: input.type ?? "abstention", + groundTruth: "ground truth", + haystackSessionIds: ["session-1"], + metadata: { + rubric: input.rubric ?? ["First nugget"], + ...(input.scale ? { scale: input.scale } : {}), + }, + } +} + +function evaluateInput(question: UnifiedQuestion, hypothesis: string) { + const protocol = new BeamPaperProtocol() + return { + protocol, + input: { + question, + hypothesis, + results: [], + retrieval: protocol.createRetrievalPlan({ question }), + }, + } +} + +function makeEvaluation(question: UnifiedQuestion, score: number): QuestionEvaluation { + return { + questionId: question.questionId, + questionType: question.questionType, + primaryScore: score, + passed: score >= BEAM_PASS_THRESHOLD, + label: score >= BEAM_PASS_THRESHOLD ? "pass" : "fail", + explanation: "fixture", + } +} + +describe("BEAM paper nugget evaluation", () => { + test("preserves 0, 0.5, and 1 and averages without rounding", async () => { + const question = makeQuestion({ + question: "Which requirements were met?", + rubric: ["Nugget alpha", "Nugget beta", "Nugget gamma"], + }) + const runtime = new ScriptedRuntime([ + { score: 0, reason: "absent" }, + { score: 0.5, reason: "partial" }, + { score: 1, reason: "complete" }, + ]) + const { protocol, input } = evaluateInput(question, "Model answer") + + const result = await protocol.evaluateQuestion(input, runtime) + + expect(result.primaryScore).toBe(0.5) + expect(result.passed).toBe(true) + expect(result.label).toBe("pass") + expect(runtime.requests).toHaveLength(3) + expect(runtime.requests.map((request) => request.schemaName)).toEqual([ + "beam_nugget_judgment", + "beam_nugget_judgment", + "beam_nugget_judgment", + ]) + + const prompts = runtime.requests.map((request) => request.prompt) + expect(prompts[0]).toContain("Which requirements were met?") + expect(prompts[0]).toContain("Nugget alpha") + expect(prompts[0]).not.toContain("Nugget beta") + expect(prompts[1]).toContain("Nugget beta") + expect(prompts[1]).not.toContain("Nugget alpha") + expect(prompts[2]).toContain("Model answer") + + const details = result.details as { nuggetJudgments: unknown[] } + expect(details.nuggetJudgments).toEqual([ + { nugget: "Nugget alpha", score: 0, reason: "absent" }, + { nugget: "Nugget beta", score: 0.5, reason: "partial" }, + { nugget: "Nugget gamma", score: 1, reason: "complete" }, + ]) + }) + + test("checkpoints every nugget and resumes without repeating paid judgments", async () => { + const question = makeQuestion({ rubric: ["one", "two", "three"] }) + const first = evaluateInput(question, "answer") + const firstRuntime = new ScriptedRuntime([{ score: 1, reason: "done one" }]) + let progress: Record | undefined + + await expect( + first.protocol.evaluateQuestion( + { + ...first.input, + onProtocolProgress: async (next) => { + progress = structuredClone(next) + }, + }, + firstRuntime + ) + ).rejects.toThrow("No scripted structured output remains") + expect((progress?.judgments as unknown[]).length).toBe(1) + + const resumedRuntime = new ScriptedRuntime([ + { score: 0.5, reason: "done two" }, + { score: 0, reason: "done three" }, + ]) + const resumed = await first.protocol.evaluateQuestion( + { ...first.input, protocolProgress: progress }, + resumedRuntime + ) + + expect(resumedRuntime.requests).toHaveLength(2) + expect(resumed.primaryScore).toBe(0.5) + }) + + test("strict schemas reject clamped, missing, malformed, and extra output", () => { + expect(BEAM_NUGGET_JUDGMENT_SCHEMA.safeParse({ score: 0.75, reason: "close" }).success).toBe( + false + ) + expect(BEAM_NUGGET_JUDGMENT_SCHEMA.safeParse({ score: 1 }).success).toBe(false) + expect(BEAM_NUGGET_JUDGMENT_SCHEMA.safeParse({ score: "1", reason: "string" }).success).toBe( + false + ) + expect( + BEAM_NUGGET_JUDGMENT_SCHEMA.safeParse({ score: 1, reason: "ok", extra: true }).success + ).toBe(false) + expect(BEAM_EVENT_EQUIVALENCE_SCHEMA.safeParse({ answer: "yes" }).success).toBe(false) + expect(BEAM_EVENT_EQUIVALENCE_SCHEMA.safeParse({ answer: "YES" }).success).toBe(true) + }) + + test("fails closed on missing or malformed rubric metadata", async () => { + const protocol = new BeamPaperProtocol() + for (const rubric of [undefined, [], [""], ["valid", 1]]) { + const question = makeQuestion({ rubric }) + if (rubric === undefined) question.metadata = {} + expect(() => protocol.validateQuestion(question)).toThrow("non-empty string rubric") + } + }) + + test("rejects a degenerate event-ordering rubric before any model call", () => { + const protocol = new BeamPaperProtocol() + const question = makeQuestion({ + id: "singleton-event-rubric", + type: "event_ordering", + rubric: ["Only one event"], + }) + + expect(() => protocol.validateQuestion(question)).toThrow( + "at least two reference events for Kendall tau-b" + ) + }) +}) + +describe("BEAM paper event ordering", () => { + test("scores perfect and reversed event sequences as 1 and 0", async () => { + const question = makeQuestion({ + id: "event-question", + type: "event_ordering", + rubric: ["A", "B", "C"], + }) + + const perfect = evaluateInput(question, "A\nB\nC") + const perfectResult = await perfect.protocol.evaluateQuestion( + perfect.input, + new ExactEventRuntime() + ) + expect(perfectResult.primaryScore).toBe(1) + expect(perfectResult.passed).toBe(true) + + const reversed = evaluateInput(question, "C\nB\nA") + const reversedResult = await reversed.protocol.evaluateQuestion( + reversed.input, + new ExactEventRuntime() + ) + expect(reversedResult.primaryScore).toBe(0) + expect(reversedResult.passed).toBe(false) + }) + + test("checkpoints event-pair equivalence and reuses it after resume", async () => { + const question = makeQuestion({ + id: "event-resume", + type: "event_ordering", + rubric: ["A", "B", "C"], + }) + const first = evaluateInput(question, "A\nB\nC") + let progress: Record | undefined + + await expect( + first.protocol.evaluateQuestion( + { + ...first.input, + onProtocolProgress: async (next) => { + progress = structuredClone(next) + }, + }, + new ScriptedRuntime([{ answer: "YES" }]) + ) + ).rejects.toThrow("No scripted structured output remains") + expect((progress?.judgments as unknown[]).length).toBe(1) + + const resumedRuntime = new ScriptedRuntime([{ answer: "YES" }, { answer: "YES" }]) + const resumed = await first.protocol.evaluateQuestion( + { ...first.input, protocolProgress: progress }, + resumedRuntime + ) + + expect(resumedRuntime.requests).toHaveLength(2) + expect(resumed.primaryScore).toBe(1) + }) + + test("uses the union rank vectors for partial and missing sequences", async () => { + const question = makeQuestion({ + id: "partial-event-question", + type: "event_ordering", + rubric: ["A", "B", "C"], + }) + const partial = evaluateInput(question, "A\nC") + const result = await partial.protocol.evaluateQuestion(partial.input, new ExactEventRuntime()) + const event = (result.details as { eventOrdering: { rankVectors: unknown } }).eventOrdering as { + rankVectors: { referenceRanks: number[]; predictedRanks: number[] } + missingReferenceEvents: unknown[] + } + + expect(result.primaryScore).toBeCloseTo(2 / 3, 12) + expect(result.passed).toBe(true) + expect(result.metrics?.normalizedKendallTauB).toBeCloseTo(2 / 3, 12) + expect(result.metrics?.eventF1).toBeCloseTo(0.8, 12) + expect(result.metrics?.authorsHelperFinalScore).toBeCloseTo(8 / 15, 12) + expect(event.rankVectors.referenceRanks).toEqual([1, 2, 3]) + expect(event.rankVectors.predictedRanks).toEqual([1, 4, 2]) + expect(event.missingReferenceEvents).toEqual([{ referenceIndex: 1, event: "B" }]) + + const oneMatched = scoreAlignedBeamEvents({ + referenceEvents: ["A", "B", "C"], + predictedEvents: ["A"], + alignments: [ + { predictedIndex: 0, predictedEvent: "A", referenceIndex: 0, referenceEvent: "A" }, + ], + }) + expect(oneMatched.normalizedKendallTauB).toBeCloseTo(0.9082482904638631, 12) + expect(oneMatched.f1).toBe(0.5) + expect(oneMatched.finalScore).toBeCloseTo(0.45412414523193156, 12) + expect(oneMatched.rankVectors.predictedRanks).toEqual([1, 4, 4]) + }) + + test("keeps the authors helper F1 product diagnostic while Table 1 uses tau", async () => { + const question = makeQuestion({ + id: "partial-precision-recall-event-question", + type: "event_ordering", + rubric: ["A", "B", "C"], + }) + const value = evaluateInput(question, "A\nC\nX") + const result = await value.protocol.evaluateQuestion(value.input, new ExactEventRuntime()) + + expect(result.metrics?.normalizedKendallTauB).toBeCloseTo(2 / 3, 12) + expect(result.metrics?.eventPrecision).toBeCloseTo(2 / 3, 12) + expect(result.metrics?.eventRecall).toBeCloseTo(2 / 3, 12) + expect(result.metrics?.eventF1).toBeCloseTo(2 / 3, 12) + expect(result.metrics?.authorsHelperFinalScore).toBeCloseTo(4 / 9, 12) + expect(result.primaryScore).toBeCloseTo(2 / 3, 12) + expect(result.passed).toBe(true) + + const eventOrdering = (result.details as { eventOrdering: { finalScore: number } }) + .eventOrdering + expect(eventOrdering.finalScore).toBeCloseTo(4 / 9, 12) + }) + + test("matches the authors' duplicate-line union and last-rank semantics", async () => { + const { attempts, alignments } = await alignBeamEvents( + ["A", "B"], + ["A", "A", "B"], + async ({ referenceEvent, predictedEvent }) => referenceEvent === predictedEvent + ) + + expect(attempts).toHaveLength(3) + expect(alignments).toEqual([ + { + predictedIndex: 0, + predictedEvent: "A", + referenceIndex: 0, + referenceEvent: "A", + }, + { predictedIndex: 1, predictedEvent: "A" }, + { + predictedIndex: 2, + predictedEvent: "B", + referenceIndex: 1, + referenceEvent: "B", + }, + ]) + + const score = scoreAlignedBeamEvents({ + referenceEvents: ["A", "B"], + predictedEvents: ["A", "A", "B"], + attempts, + alignments, + }) + expect(score.unmatchedPredictedEvents).toEqual([{ predictedIndex: 1, event: "A" }]) + expect(score.canonicalPredictedEvents).toEqual(["A", "A", "B"]) + expect(score.rankVectors.union.map((item) => item.id)).toEqual(["reference:0", "reference:1"]) + expect(score.rankVectors.predictedRanks).toEqual([2, 3]) + expect(score.normalizedKendallTauB).toBe(1) + expect(score.f1).toBe(1) + expect(score.finalScore).toBe(1) + }) + + test("matches the authors' literal newline split, including blank events", async () => { + expect(extractBeamPredictedEvents("A\n\nB")).toEqual(["A", "", "B"]) + expect(extractBeamPredictedEvents("")).toEqual([""]) + + const question = makeQuestion({ + id: "blank-event-question", + type: "event_ordering", + rubric: ["A", "B"], + }) + const value = evaluateInput(question, "A\n\nB") + const result = await value.protocol.evaluateQuestion(value.input, new ExactEventRuntime()) + expect(result.primaryScore).toBeCloseTo(2 / 3, 12) + expect(result.metrics?.authorsHelperFinalScore).toBeCloseTo(8 / 15, 12) + + let progress: Record | undefined + await expect( + value.protocol.evaluateQuestion( + { + ...value.input, + onProtocolProgress: async (next) => { + progress = structuredClone(next) + }, + }, + new ScriptedRuntime([{ answer: "YES" }, { answer: "NO" }]) + ) + ).rejects.toThrow("No scripted structured output remains") + expect( + (progress!.judgments as Array<{ predictedEvent: string }>).some( + (judgment) => judgment.predictedEvent === "" + ) + ).toBe(true) + + const resumed = await value.protocol.evaluateQuestion( + { ...value.input, protocolProgress: progress }, + new ScriptedRuntime([{ answer: "YES" }]) + ) + expect(resumed.primaryScore).toBeCloseTo(2 / 3, 12) + expect(resumed.metrics?.authorsHelperFinalScore).toBeCloseTo(8 / 15, 12) + }) + + test("matches the authors' undefined tau-b edge and fails closed instead of inventing a score", () => { + const identicalSingleton = computeKendallTauB([1], [1]) + expect(identicalSingleton.degenerate).toBe(true) + expect(Number.isNaN(identicalSingleton.tauB)).toBe(true) + + expect(() => + scoreAlignedBeamEvents({ + referenceEvents: ["A", "B"], + predictedEvents: [], + alignments: [], + }) + ).toThrow("undefined for degenerate rank vectors") + }) +}) + +describe("BEAM paper aggregation and identity", () => { + test("reports fractional average and PASS-at-0.5 separately", () => { + const protocol = new BeamPaperProtocol() + const questions = [0, 1, 2].map((index) => + makeQuestion({ id: `q-${index}`, type: "abstention" }) + ) + const evaluations = [0, 0.5, 1].map((score, index) => makeEvaluation(questions[index]!, score)) + + const report = protocol.aggregateQuality({ questions, evaluations }) + + expect(report.primaryMetric).toEqual({ + key: "beamScorePartial", + value: 0.5, + higherIsBetter: true, + }) + expect(report.metrics.passAccuracy).toBeCloseTo(2 / 3, 12) + expect(report.bySlice?.abstention.averageScore).toBe(0.5) + expect(report.bySlice?.abstention.passAccuracy).toBeCloseTo(2 / 3, 12) + + const boundaryQuestions = [0.499, 0.5].map((score) => + makeQuestion({ id: `boundary-${score}`, type: "abstention" }) + ) + const boundary = protocol.aggregateQuality({ + questions: boundaryQuestions, + evaluations: boundaryQuestions.map((question, index) => + makeEvaluation(question, [0.499, 0.5][index]!) + ), + }) + expect(boundary.metrics.passAccuracy).toBe(0.5) + }) + + test("reports each tier independently and labels cross-tier macro aggregation", () => { + const protocol = new BeamPaperProtocol() + const questions: UnifiedQuestion[] = [] + const evaluations: QuestionEvaluation[] = [] + for (const ability of BEAM_ABILITY_IDS) { + for (let index = 0; index < 2; index++) { + const question = makeQuestion({ + id: `1m-${ability}-${index}`, + type: ability, + scale: "1M", + }) + questions.push(question) + evaluations.push(makeEvaluation(question, 1)) + } + const question = makeQuestion({ id: `10m-${ability}`, type: ability, scale: "10M" }) + questions.push(question) + evaluations.push(makeEvaluation(question, 0)) + } + + const report = protocol.aggregateQuality({ questions, evaluations }) + + expect(report.primaryMetric).toBeUndefined() + expect(report.metrics.beamScore1MPartial).toBe(1) + expect(report.metrics.beamScore10MPartial).toBe(0) + expect(report.metrics.beamTierMacroAverageSecondaryPartial).toBe(0.5) + expect(report.metrics.beamAbilityPooledAverage).toBeCloseTo(2 / 3, 12) + expect(report.metrics.beamQuestionMicroAverage).toBeCloseTo(2 / 3, 12) + expect(report.bySlice?.["tier:1M"].averageScore).toBe(1) + expect(report.bySlice?.["tier:10M"].averageScore).toBe(0) + expect(report.bySlice?.["tier:1M/ability:abstention"].questionCount).toBe(2) + }) + + test("macro-averages the ten abilities instead of question-weighting", () => { + const protocol = new BeamPaperProtocol() + const questions: UnifiedQuestion[] = [] + const evaluations: QuestionEvaluation[] = [] + + for (const ability of BEAM_ABILITY_IDS) { + const count = ability === "abstention" ? 9 : 1 + const score = ability === "abstention" ? 1 : 0 + for (let index = 0; index < count; index++) { + const question = makeQuestion({ id: `${ability}-${index}`, type: ability }) + questions.push(question) + evaluations.push(makeEvaluation(question, score)) + } + } + + const report = protocol.aggregateQuality({ questions, evaluations }) + + expect(report.primaryMetric?.key).toBe("beamScorePartial") + expect(report.primaryMetric?.value).toBeCloseTo(0.1, 12) + expect(report.metrics.passAccuracy).toBe(0.5) + expect(report.bySlice?.abstention.averageScore).toBe(1) + expect(report.bySlice?.event_ordering.averageScore).toBe(0) + }) + + test("reserves the official beamScore key for a complete public tier", () => { + const protocol = new BeamPaperProtocol() + const questions: UnifiedQuestion[] = [] + const evaluations: QuestionEvaluation[] = [] + + for (const ability of BEAM_ABILITY_IDS) { + for (let index = 0; index < 70; index++) { + const question = makeQuestion({ + id: `official-1m-${ability}-${index}`, + type: ability, + scale: "1M", + }) + questions.push(question) + evaluations.push(makeEvaluation(question, 0.75)) + } + } + + const report = protocol.aggregateQuality({ questions, evaluations }) + expect(report.primaryMetric).toEqual({ + key: "beamScore", + value: 0.75, + higherIsBetter: true, + }) + expect(report.metrics.beamScore1M).toBe(0.75) + expect(report.metrics.officialQuestionSet1M).toBe(1) + + const partial = protocol.aggregateQuality({ + questions: questions.slice(0, 20), + evaluations: evaluations.slice(0, 20), + }) + expect(partial.primaryMetric?.key).toBe("beamScorePartial") + expect(partial.metrics.beamScore1MPartial).toBe(0.75) + expect(partial.metrics.officialQuestionSet1M).toBe(0) + }) + + test("defaults to Top-K 5, allows paper ablations, and fingerprints config", () => { + const defaultProtocol = new BeamPaperProtocol() + expect(defaultProtocol.retrievalTopK).toBe(5) + expect(defaultProtocol.auxiliaryRetrievalEvaluation).toBe("disabled") + expect(defaultProtocol.identity.version).toBe(BEAM_PAPER_PROTOCOL_VERSION) + expect(BEAM_PAPER_PROTOCOL_VERSION).toBe("1.5.0") + expect(BEAM_EVALUATOR_IDENTITY.eventOrderingScoringVersion).toBe( + BEAM_EVENT_ORDERING_SCORING_VERSION + ) + + for (const topK of BEAM_RETRIEVAL_TOP_K_VALUES) { + const protocol = new BeamPaperProtocol({ retrievalTopK: topK }) + const question = makeQuestion() + const retrieval = protocol.createRetrievalPlan({ question }) + expect(retrieval).toMatchObject({ + requestedTopK: topK, + answerCutoff: topK, + }) + const results: UnifiedSearchResult[] = Array.from({ length: 25 }, (_, index) => ({ + id: `evidence-${index}`, + rank: index + 1, + text: `Evidence ${index}`, + provider: "filesystem", + resultType: "chunk", + })) + expect( + protocol.createAnswerPlan({ question, sessions: [], results, retrieval }) + .answerEvidenceCount + ).toBe(topK) + expect(protocol.identity.evaluatorHash).toHaveLength(64) + expect(protocol.identity.details?.evaluatorIdentity).toEqual(BEAM_EVALUATOR_IDENTITY) + } + + expect(new BeamPaperProtocol({ retrievalTopK: 5 }).identity.configFingerprint).not.toBe( + new BeamPaperProtocol({ retrievalTopK: 10 }).identity.configFingerprint + ) + expect(BEAM_EVALUATOR_IDENTITY.nuggetPromptSha256).toBe( + "5318eec79b7c650bbd532009636e67af066116b9a9d13dbad794c0850e1189dd" + ) + expect(BEAM_EVALUATOR_IDENTITY.eventEquivalencePromptSha256).toBe( + "3460ec68c19c6974d998b610af5143559748b02cd4a5273590a236646cd247ad" + ) + expect(BEAM_AGGREGATION_IMPLEMENTATION_SHA256).toHaveLength(64) + expect(BEAM_EVALUATOR_IDENTITY.structuredOutputSchemaSha256).toBe( + BEAM_STRUCTURED_OUTPUT_SCHEMA_SHA256 + ) + expect( + ( + new BeamPaperProtocol().identity.details?.aggregation as { + implementationSha256: string + } + ).implementationSha256 + ).toBe(BEAM_AGGREGATION_IMPLEMENTATION_SHA256) + expect(() => new BeamPaperProtocol({ retrievalTopK: 100 })).toThrow("must be one of") + }) + + test("fingerprints method-only protocol orchestration drift", () => { + class ChangedAnswerPlanProtocol extends BeamPaperProtocol { + override createAnswerPlan(input: Parameters[0]) { + return super.createAnswerPlan(input) + } + } + + const base = new BeamPaperProtocol() + const changed = new ChangedAnswerPlanProtocol() + expect(changed.identity.answerPromptHash).not.toBe(base.identity.answerPromptHash) + expect(changed.identity.implementationFingerprint).not.toBe( + base.identity.implementationFingerprint + ) + }) + + test("renders the approved BEAM session document without inventing a date", () => { + const protocol = new BeamPaperProtocol() + const question = makeQuestion() + const withDate = protocol.createIngestionPlan({ + question, + sessions: [ + { + sessionId: "session-1", + messages: [ + { role: "user", content: "Hello" }, + { role: "assistant", content: "Hi" }, + ], + metadata: { date: "2024-03-01" }, + }, + ], + }) + expect(withDate[0]).toEqual({ + customId: "session-1", + content: "DOCUMENT_DATE: 2024-03-01\n\n[USER]\nHello\n\n[ASSISTANT]\nHi", + metadata: { sessionId: "session-1", documentDate: "2024-03-01" }, + messages: [ + { role: "user", content: "Hello" }, + { role: "assistant", content: "Hi" }, + ], + }) + + const withoutDate = protocol.createIngestionPlan({ + question, + sessions: [ + { + sessionId: "session-2", + messages: [ + { role: "user", content: "No date here" }, + { role: "assistant", content: "Acknowledged" }, + ], + }, + ], + }) + expect(withoutDate[0]?.content).toBe("[USER]\nNo date here\n\n[ASSISTANT]\nAcknowledged") + expect(withoutDate[0]?.metadata).toEqual({ sessionId: "session-2" }) + + expect(() => + protocol.createIngestionPlan({ + question, + sessions: [ + { + sessionId: "session-invalid-date", + messages: [ + { role: "user", content: "Impossible date" }, + { role: "assistant", content: "Acknowledged" }, + ], + metadata: { documentDate: "2024-02-30" }, + }, + ], + }) + ).toThrow("invalid document date") + + expect(() => + protocol.createIngestionPlan({ + question, + sessions: [ + { + sessionId: "session-conflicting-dates", + messages: [ + { role: "user", content: "Conflicting date" }, + { role: "assistant", content: "Acknowledged" }, + ], + metadata: { documentDate: "2024-03-01", date: "2024-03-02" }, + }, + ], + }) + ).toThrow("conflicting document dates") + + expect(() => + protocol.createIngestionPlan({ + question, + sessions: [ + { + sessionId: "session-malformed-date", + messages: [ + { role: "user", content: "Malformed date" }, + { role: "assistant", content: "Acknowledged" }, + ], + metadata: { documentDate: "March 1, 2024" }, + }, + ], + }) + ).toThrow("invalid document date") + + expect(() => + protocol.createIngestionPlan({ + question, + sessions: [ + { + sessionId: "session-orphan-user", + messages: [{ role: "user", content: "Missing assistant" }], + }, + ], + }) + ).toThrow("exactly one non-empty user message followed by one non-empty assistant message") + + for (const messages of [ + [ + { role: "assistant", content: "Wrong role" }, + { role: "user", content: "Wrong role" }, + ], + [ + { role: "user", content: "One" }, + { role: "assistant", content: "Two" }, + { role: "assistant", content: "Extra" }, + ], + [ + { role: "user", content: " " }, + { role: "assistant", content: "Two" }, + ], + [ + { role: "user", content: 42 }, + { role: "assistant", content: "Two" }, + ], + ]) { + expect(() => + protocol.createIngestionPlan({ + question, + sessions: [{ sessionId: "session-malformed-turns", messages } as never], + }) + ).toThrow("exactly one non-empty user message followed by one non-empty assistant message") + } + }) + + test("formats only normalized evidence text and preserves result-level dates", () => { + const protocol = new BeamPaperProtocol() + const question = makeQuestion() + const retrieval = protocol.createRetrievalPlan({ question }) + const plan = protocol.createAnswerPlan({ + question, + sessions: [], + retrieval, + results: [ + { + id: "result-1", + rank: 1, + text: "Normalized evidence text", + provider: "fake", + resultType: "chunk", + documentDate: "2024-05-06", + rawArtifactRef: "raw/result-1.json", + }, + ], + }) + + expect(plan.answerEvidenceCount).toBe(1) + expect(plan.request.prompt).toContain("[2024-05-06] Normalized evidence text") + expect(plan.request.prompt).not.toContain("raw/result-1.json") + expect(plan.request.prompt).not.toContain('"resultType"') + }) + + test("makes the authors' newline event scorer contract explicit only for event answers", () => { + const protocol = new BeamPaperProtocol() + const eventQuestion = makeQuestion({ + id: "event-answer-format", + type: "event_ordering", + rubric: ["First event", "Second event"], + }) + const eventRetrieval = protocol.createRetrievalPlan({ question: eventQuestion }) + const eventPlan = protocol.createAnswerPlan({ + question: eventQuestion, + sessions: [], + results: [], + retrieval: eventRetrieval, + }) + expect(eventPlan.request.prompt).toContain("output exactly one event per line") + expect(eventPlan.request.prompt).toContain("Do not use bullets, numbering, headings") + expect(eventPlan.baseRequest.prompt).toContain("output exactly one event per line") + + const nuggetQuestion = makeQuestion({ id: "ordinary-answer-format", type: "abstention" }) + const nuggetPlan = protocol.createAnswerPlan({ + question: nuggetQuestion, + sessions: [], + results: [], + retrieval: protocol.createRetrievalPlan({ question: nuggetQuestion }), + }) + expect(nuggetPlan.request.prompt).not.toContain("output exactly one event per line") + + const answerIdentity = protocol.identity.details?.answerPrompt as { + eventOrderingAnswerFormatVersion: string + eventOrderingPromptSha256: string + } + expect(answerIdentity.eventOrderingAnswerFormatVersion).toBe( + "authors-newline-scorer-compatible-v1" + ) + expect(answerIdentity.eventOrderingPromptSha256).toMatch(/^[a-f0-9]{64}$/) + }) +}) diff --git a/test/beam-mem0-profile.test.ts b/test/beam-mem0-profile.test.ts new file mode 100644 index 0000000..7ba8f3f --- /dev/null +++ b/test/beam-mem0-profile.test.ts @@ -0,0 +1,283 @@ +import { describe, expect, test } from "bun:test" +import type { JudgeInput, JudgeResult } from "../src/types/judge" +import type { + EvaluationRuntime, + QuestionEvaluation, + StructuredModelRequest, +} from "../src/types/protocol" +import type { UnifiedQuestion, UnifiedSearchResult, UnifiedSession } from "../src/types/unified" +import { OpenAIJudge } from "../src/judges/openai" +import { BeamPaperProtocol } from "../src/protocols/beam-paper" +import { parseRunArgs } from "../src/cli/commands/run" +import { + BEAM_MEM0_JUDGE_SYSTEM_PROMPT, + BEAM_MEM0_NUGGET_PROFILE, + BeamMem0NuggetProtocol, + clampMem0NuggetScore, +} from "../src/protocols/beam-mem0" + +class ScriptedRuntime implements EvaluationRuntime { + readonly requests: StructuredModelRequest[] = [] + + constructor(private readonly outputs: unknown[]) {} + + async evaluateLegacy(_input: JudgeInput): Promise { + throw new Error("Legacy evaluation is not expected") + } + + async generateStructured(request: StructuredModelRequest): Promise { + this.requests.push(request as StructuredModelRequest) + const output = this.outputs.shift() + if (output === undefined) throw new Error("No scripted output remains") + return request.schema.parse(output) + } +} + +function eventQuestion(id = "event-question"): UnifiedQuestion { + return { + questionId: id, + question: "In what order did the events happen?", + questionType: "event_ordering", + groundTruth: "Alpha then beta", + haystackSessionIds: ["s1"], + metadata: { rubric: ["Alpha happened", "Beta happened"], scale: "1M" }, + } +} + +const sessions: UnifiedSession[] = [ + { + sessionId: "s1", + messages: [ + { role: "user", content: "Alpha happened." }, + { role: "assistant", content: "Then beta happened." }, + ], + metadata: { documentDate: "2024-01-01" }, + }, +] + +describe("BEAM mem0 nugget comparison profile", () => { + test("parses the explicit comparison and source-reuse CLI identity", () => { + const parsed = parseRunArgs([ + "-p", + "supermemory", + "-b", + "beam-1m", + "-r", + "target", + "--source-run", + "source", + "--from-phase", + "search", + "--evaluation-profile", + "mem0-nugget", + "--retrieval-top-k", + "50", + "--answer-cutoff", + "50", + "-m", + "gpt-5", + "-j", + "gpt-5", + ]) + + expect(parsed).toMatchObject({ + provider: "supermemory", + benchmark: "beam-1m", + runId: "target", + sourceRunId: "source", + fromPhase: "search", + evaluationProfile: "mem0-nugget", + retrievalTopK: 50, + answerCutoff: 50, + answeringModel: "gpt-5", + judgeModel: "gpt-5", + }) + }) + + test("is isolated from the paper protocol while sharing its exact ingestion contract", () => { + const paper = new BeamPaperProtocol() + const comparison = new BeamMem0NuggetProtocol({ retrievalTopK: 50, answerCutoff: 50 }) + const question = eventQuestion() + + expect(comparison.identity.id).toBe("beam-mem0-nugget") + expect(comparison.identity.version).toBe("1.2.0") + expect(comparison.identity.details?.comparisonProfile).toBe(BEAM_MEM0_NUGGET_PROFILE) + expect(comparison.identity.details?.evaluatorIdentity).toMatchObject({ + sourceCommit: "4b61c5d31b9c668a12b4f5e78064248a02c82d2b", + judgeModel: "gpt-5", + eventOrderingPolicy: "ordinary-nugget-average-primary", + temperature: null, + maxOutputTokens: 4096, + maxAttempts: 5, + innerMaxRetries: 2, + timeoutMs: 120_000, + retryBackoffMs: 2_000, + transport: "openai-chat-completions", + runtimeExecutionVersion: "chat-transport-outer-retry-v1", + parseFallback: "none-fail-closed-deviation-from-mem0-raw-text-marker-fallback", + }) + expect(comparison.identity.details?.answerPrompt).toMatchObject({ + innerMaxRetries: 2, + terminalEmptyOutputPolicy: "accept-and-evaluate", + runtimeExecutionVersion: "chat-transport-durable-outer-retry-v1", + }) + expect(comparison.identity.ingestionPolicyHash).toBe(paper.identity.ingestionPolicyHash) + expect(comparison.createIngestionPlan({ question, sessions })).toEqual( + paper.createIngestionPlan({ question, sessions }) + ) + expect(comparison.requiredJudge).toEqual({ + provider: "openai", + modelId: "gpt-5", + modelAlias: "gpt-5", + }) + expect(comparison.createRetrievalPlan({ question })).toMatchObject({ + requestedTopK: 50, + answerCutoff: 50, + threshold: 0, + }) + expect(() => new BeamPaperProtocol({ retrievalTopK: 50 })).toThrow("must be one of") + }) + + test("uses the mem0 answer format without the paper event-line rule", () => { + const protocol = new BeamMem0NuggetProtocol({ retrievalTopK: 50, answerCutoff: 1 }) + const question = eventQuestion() + const results: UnifiedSearchResult[] = [ + { + id: "m1", + rank: 1, + text: "Alpha happened and beta followed.", + sessionId: "s1", + documentDate: "2024-01-01", + provider: "supermemory", + resultType: "memory", + }, + ] + const plan = protocol.createAnswerPlan({ + question, + sessions, + results, + retrieval: protocol.createRetrievalPlan({ question }), + }) + + expect(plan.answerEvidenceCount).toBe(1) + expect(plan.request).toMatchObject({ + maxOutputTokens: 4096, + transport: "openai-chat-completions", + maxAttempts: 5, + innerMaxRetries: 2, + timeoutMs: 120_000, + retryBackoffMs: 2_000, + terminalEmptyOutputPolicy: "accept-and-evaluate", + }) + expect(plan.request.prompt).toContain("[2024-01-01] Alpha happened and beta followed.") + expect(plan.request.prompt).not.toContain("output exactly one event per line") + + const paperPlan = new BeamPaperProtocol().createAnswerPlan({ + question, + sessions, + results, + retrieval: new BeamPaperProtocol().createRetrievalPlan({ question }), + }) + expect(paperPlan.request.terminalEmptyOutputPolicy).toBeUndefined() + }) + + test("scores event ordering only through mem0 nugget averages and clamps arbitrary numbers", async () => { + const protocol = new BeamMem0NuggetProtocol({ retrievalTopK: 50, answerCutoff: 50 }) + const question = eventQuestion() + const runtime = new ScriptedRuntime([ + { score: 0.74, reason: "partial" }, + { score: 0.75, reason: "complete" }, + ]) + + const evaluation = await protocol.evaluateQuestion( + { + question, + hypothesis: "Beta happened before alpha.", + results: [], + retrieval: protocol.createRetrievalPlan({ question }), + }, + runtime + ) + + expect(evaluation.primaryScore).toBe(0.75) + expect(evaluation.metrics).toEqual({ nuggetAverage: 0.75, nuggetCount: 2 }) + expect(evaluation.details).toMatchObject({ eventOrderingScoreUsed: 0 }) + expect(runtime.requests).toHaveLength(2) + expect( + runtime.requests.every((request) => request.system === BEAM_MEM0_JUDGE_SYSTEM_PROMPT) + ).toBe(true) + expect(runtime.requests.every((request) => request.temperature === undefined)).toBe(true) + expect(runtime.requests.every((request) => request.maxOutputTokens === 4096)).toBe(true) + expect(runtime.requests.every((request) => request.maxAttempts === 5)).toBe(true) + expect(runtime.requests.every((request) => request.innerMaxRetries === 2)).toBe(true) + expect(runtime.requests.every((request) => request.timeoutMs === 120_000)).toBe(true) + expect(runtime.requests.every((request) => request.retryBackoffMs === 2_000)).toBe(true) + expect( + runtime.requests.every((request) => request.transport === "openai-chat-completions") + ).toBe(true) + expect( + runtime.requests.every((request) => request.schemaName === "beam_mem0_nugget_judgment") + ).toBe(true) + }) + + test("judges an explicitly accepted empty hypothesis instead of skipping it", async () => { + const protocol = new BeamMem0NuggetProtocol() + const question = eventQuestion("empty-answer") + const runtime = new ScriptedRuntime([ + { score: 0, reason: "missing" }, + { score: 0, reason: "missing" }, + ]) + + const evaluation = await protocol.evaluateQuestion( + { + question, + hypothesis: "", + results: [], + retrieval: protocol.createRetrievalPlan({ question }), + }, + runtime + ) + + expect(evaluation.primaryScore).toBe(0) + expect(runtime.requests).toHaveLength(2) + expect(runtime.requests.every((request) => request.prompt.includes("LLM RESPONSE:\n\n"))).toBe( + true + ) + }) + + test("reports the question-micro nugget average as its distinct primary metric", () => { + const protocol = new BeamMem0NuggetProtocol() + const questions = [ + eventQuestion("event"), + { ...eventQuestion("abstain"), questionType: "abstention" }, + ] + const evaluations: QuestionEvaluation[] = questions.map((question, index) => ({ + questionId: question.questionId, + questionType: question.questionType, + primaryScore: index === 0 ? 0.5 : 1, + passed: true, + explanation: "fixture", + })) + + const quality = protocol.aggregateQuality({ questions, evaluations }) + expect(quality.primaryMetric).toEqual({ + key: "mem0NuggetAverage", + value: 0.75, + higherIsBetter: true, + }) + expect(quality.metrics.mem0NuggetAverage).toBe(0.75) + }) + + test("matches mem0 clamp thresholds", () => { + expect([0.1, 0.25, 0.74, 0.75, 3].map(clampMem0NuggetScore)).toEqual([0, 0.5, 0.5, 1, 1]) + }) + + test("selects OpenAI Chat Completions for the pinned mem0 judge transport", async () => { + const judge = new OpenAIJudge() + await judge.initialize({ apiKey: "test-key", model: "gpt-5" }) + const model = judge.getModel("openai-chat-completions") + + expect(model.provider).toBe("openai.chat") + expect(model.modelId).toBe("gpt-5") + }) +}) diff --git a/test/beam-preflight.test.ts b/test/beam-preflight.test.ts new file mode 100644 index 0000000..13cf8db --- /dev/null +++ b/test/beam-preflight.test.ts @@ -0,0 +1,46 @@ +import { afterEach, describe, expect, spyOn, test } from "bun:test" +import { mkdtemp, rm } from "node:fs/promises" +import { tmpdir } from "node:os" +import { join } from "node:path" +import { CheckpointManager } from "../src/orchestrator/checkpoint" +import { Orchestrator } from "../src/orchestrator" +import { SupermemoryProvider } from "../src/providers/supermemory" + +const tempRoots: string[] = [] + +afterEach(async () => { + await Promise.all(tempRoots.splice(0).map((path) => rm(path, { recursive: true, force: true }))) +}) + +async function tempRoot(): Promise { + const path = await mkdtemp(join(tmpdir(), "memorybench-beam-preflight-test-")) + tempRoots.push(path) + return path +} + +describe("BEAM orchestration preflight", () => { + test("rejects an invalid dataset before provider initialization or checkpoint creation", async () => { + const root = await tempRoot() + const manager = new CheckpointManager(join(root, "runs")) + const initialize = spyOn(SupermemoryProvider.prototype, "initialize") + + try { + await expect( + new Orchestrator(manager).run({ + provider: "supermemory", + benchmark: "beam-1m", + judgeModel: "gpt-4.1-mini", + runId: "invalid-beam-dataset", + dataPath: join(root, "missing-dataset"), + datasetRevision: "a".repeat(64), + limit: 1, + }) + ).rejects.toThrow("run the BEAM prepare command") + + expect(initialize).not.toHaveBeenCalled() + expect(manager.exists("invalid-beam-dataset")).toBe(false) + } finally { + initialize.mockRestore() + } + }) +}) diff --git a/test/benchmark-api.test.ts b/test/benchmark-api.test.ts new file mode 100644 index 0000000..8b7daa9 --- /dev/null +++ b/test/benchmark-api.test.ts @@ -0,0 +1,19 @@ +import { describe, expect, test } from "bun:test" +import { handleBenchmarksRoutes } from "../src/server/routes/benchmarks" + +describe("benchmark API dataset failures", () => { + test("keeps an unknown benchmark distinct from a known benchmark with invalid data", async () => { + const unknownUrl = new URL("http://localhost/api/benchmarks/not-a-benchmark/questions") + const unknown = await handleBenchmarksRoutes(new Request(unknownUrl), unknownUrl) + expect(unknown?.status).toBe(404) + + const beamUrl = new URL("http://localhost/api/benchmarks/beam-1m/questions") + beamUrl.searchParams.set("dataPath", "/tmp/memorybench-definitely-missing-beam-api") + beamUrl.searchParams.set("datasetRevision", "a".repeat(64)) + const beam = await handleBenchmarksRoutes(new Request(beamUrl), beamUrl) + const body = (await beam?.json()) as { error?: string } + + expect(beam?.status).toBe(400) + expect(body.error).toContain("run the BEAM prepare command") + }) +}) diff --git a/test/benchmark-scope-label.test.ts b/test/benchmark-scope-label.test.ts new file mode 100644 index 0000000..73fe2fe --- /dev/null +++ b/test/benchmark-scope-label.test.ts @@ -0,0 +1,19 @@ +import { describe, expect, test } from "bun:test" +import { getBenchmarkDisplayName } from "../ui/lib/utils" + +describe("benchmark scope labels", () => { + test("qualifies every supported BEAM tier when checkpoint scope is unavailable", () => { + expect(getBenchmarkDisplayName("beam-1m")).toBe("BEAM 1M") + expect(getBenchmarkDisplayName("beam-10m")).toBe("BEAM 10M") + expect(getBenchmarkDisplayName("beam-1m-10m")).toBe("BEAM 1M/10M") + }) + + test("uses the recorded benchmark scope as the authoritative display name", () => { + expect( + getBenchmarkDisplayName("beam-1m", { + displayName: "BEAM 1M reviewed snapshot", + includedTiers: ["1M"], + }) + ).toBe("BEAM 1M reviewed snapshot") + }) +}) diff --git a/test/cli-exit.test.ts b/test/cli-exit.test.ts new file mode 100644 index 0000000..613547a --- /dev/null +++ b/test/cli-exit.test.ts @@ -0,0 +1,62 @@ +import { describe, expect, test } from "bun:test" +import { resolve } from "path" + +const PROJECT_ROOT = resolve(import.meta.dir, "..") + +async function runCli(args: string[]) { + const child = Bun.spawn([process.execPath, "run", "src/index.ts", ...args], { + cwd: PROJECT_ROOT, + stdout: "pipe", + stderr: "pipe", + }) + const [exitCode, stdout, stderr] = await Promise.all([ + child.exited, + new Response(child.stdout).text(), + new Response(child.stderr).text(), + ]) + return { exitCode, stdout, stderr } +} + +describe("CLI process status", () => { + test("returns a nonzero exit status when BEAM preparation arguments are invalid", async () => { + const result = await runCli(["beam", "prepare", "--tiers", "invalid"]) + + expect(result.exitCode).not.toBe(0) + expect(result.stderr).toContain("Invalid BEAM tier") + }) + + test("returns a nonzero exit status for an invalid comparison provider", async () => { + const result = await runCli([ + "compare", + "--providers", + "not-a-provider", + "--benchmark", + "locomo", + ]) + + expect(result.exitCode).not.toBe(0) + expect(result.stderr).toContain("Invalid provider") + }) + + test("returns a nonzero exit status when list-questions cannot validate BEAM data", async () => { + const result = await runCli([ + "list-questions", + "--benchmark", + "beam-1m", + "--data-path", + "/tmp/memorybench-definitely-missing-beam", + "--dataset-revision", + "a".repeat(64), + ]) + + expect(result.exitCode).not.toBe(0) + expect(result.stderr).toContain("run the BEAM prepare command") + }) + + test("returns a nonzero exit status for an unknown list-questions benchmark", async () => { + const result = await runCli(["list-questions", "--benchmark", "not-a-benchmark"]) + + expect(result.exitCode).not.toBe(0) + expect(result.stderr).toContain("Invalid benchmark") + }) +}) diff --git a/test/evaluation-runtime.test.ts b/test/evaluation-runtime.test.ts new file mode 100644 index 0000000..cf41a86 --- /dev/null +++ b/test/evaluation-runtime.test.ts @@ -0,0 +1,124 @@ +import { describe, expect, test } from "bun:test" +import { z } from "zod" +import type { Judge } from "../src/types/judge" +import { + JudgeEvaluationRuntime, + executeStructuredWithRetries, + type StructuredGenerationExecutor, +} from "../src/orchestrator/evaluation-runtime" + +const schema = z.object({ score: z.union([z.literal(0), z.literal(0.5), z.literal(1)]) }).strict() + +const fakeJudge = { + name: "fake", + async initialize() {}, + async evaluate() { + return { score: 1, label: "correct" as const, explanation: "fixture" } + }, + getPromptForQuestionType() { + return "fixture" + }, + getModel() { + throw new Error("The injected structured executor should not request a model") + }, +} satisfies Judge + +describe("structured evaluation runtime", () => { + test("retries malformed structured output and preserves a later valid half score", async () => { + const outputs: unknown[] = [{ score: "0.5" }, { score: 0.75 }, { score: 0.5 }] + let calls = 0 + const result = await executeStructuredWithRetries( + { + schema, + schemaName: "fixture_score", + prompt: "score this", + maxAttempts: 3, + }, + async () => { + calls += 1 + return outputs.shift() + } + ) + + expect(result).toEqual({ score: 0.5 }) + expect(calls).toBe(3) + }) + + test("fails explicitly after exhausting schema-invalid output", async () => { + let calls = 0 + await expect( + executeStructuredWithRetries( + { + schema, + schemaName: "fixture_score", + prompt: "score this", + maxAttempts: 2, + }, + async () => { + calls += 1 + return { score: "1.0 contains 0.5" } + } + ) + ).rejects.toThrow("failed after 2 attempts") + expect(calls).toBe(2) + }) + + test("counts every retry and records complete, partial, and unknown token coverage", async () => { + const paidFailure = Object.assign(new Error("provider rejected structured output"), { + usage: { inputTokens: 10, outputTokens: 2, totalTokens: 12 }, + }) + const attempts: Array<{ object: unknown; usage?: unknown } | { error: Error }> = [ + { error: paidFailure }, + { object: { score: "invalid" }, usage: { inputTokens: 3 } }, + { error: new Error("transport failed without usage") }, + { + object: { score: 0.5 }, + usage: { inputTokens: 7, outputTokens: 1, totalTokens: 8 }, + }, + ] + const executor: StructuredGenerationExecutor = async () => { + const attempt = attempts.shift()! + if ("error" in attempt) throw attempt.error + return attempt + } + const runtime = new JudgeEvaluationRuntime(fakeJudge, executor) + + const result = await runtime.generateStructured({ + schema, + schemaName: "usage_retry_fixture", + prompt: "score this", + maxAttempts: 4, + }) + + expect(result).toEqual({ score: 0.5 }) + expect(runtime.getUsage()).toEqual({ + requestCount: 4, + tokenUsageCompleteRequestCount: 2, + tokenUsagePartialRequestCount: 1, + tokenUsageUnknownRequestCount: 1, + inputTokens: 20, + outputTokens: 3, + totalTokens: 20, + }) + }) + + test("retains unknown paid-attempt coverage after terminal failure", async () => { + const executor: StructuredGenerationExecutor = async () => { + throw new Error("timeout without usage") + } + const runtime = new JudgeEvaluationRuntime(fakeJudge, executor) + + await expect( + runtime.generateStructured({ + schema, + schemaName: "terminal_failure_fixture", + prompt: "score this", + maxAttempts: 2, + }) + ).rejects.toThrow("failed after 2 attempts") + expect(runtime.getUsage()).toEqual({ + requestCount: 2, + tokenUsageUnknownRequestCount: 2, + }) + }) +}) diff --git a/test/ingest-cli.test.ts b/test/ingest-cli.test.ts new file mode 100644 index 0000000..7af706c --- /dev/null +++ b/test/ingest-cli.test.ts @@ -0,0 +1,89 @@ +import { describe, expect, test } from "bun:test" +import { parseIngestArgs } from "../src/cli/commands/ingest" +import { mergeResumeConcurrency } from "../src/orchestrator" + +describe("ingest CLI concurrency", () => { + test("passes default build concurrency to ingest runs", () => { + expect( + parseIngestArgs([ + "--provider", + "supermemory", + "--benchmark", + "beam-1m", + "--run-id", + "beam-ingest-c5", + "--concurrency", + "5", + ]) + ).toMatchObject({ + concurrency: { default: 5 }, + }) + }) + + test("supports separate ingest and indexing limits", () => { + expect( + parseIngestArgs([ + "--provider", + "supermemory", + "--benchmark", + "beam-1m", + "--run-id", + "beam-ingest-split-concurrency", + "--concurrency-ingest", + "7", + "--concurrency-indexing", + "3", + ]) + ).toMatchObject({ + concurrency: { ingest: 7, indexing: 3 }, + }) + }) + + test("overrides persisted concurrency when resuming a run", () => { + expect(mergeResumeConcurrency({ default: 5, indexing: 8 }, { default: 20 })).toEqual({ + default: 20, + indexing: 8, + }) + }) + + test("parses an ordered provider ingest batch size", () => { + expect( + parseIngestArgs([ + "--provider", + "supermemory", + "--benchmark", + "beam-1m", + "--run-id", + "beam-ingest-c20-b5", + "--ingest-batch-size", + "5", + ]) + ).toMatchObject({ ingestBatchSize: 5 }) + }) + + test("rejects an invalid ingest batch size", () => { + expect(() => + parseIngestArgs([ + "--provider", + "supermemory", + "--benchmark", + "beam-1m", + "--ingest-batch-size", + "0", + ]) + ).toThrow("--ingest-batch-size must be an integer between 1 and 600") + }) + + test("parses the per-readiness timeout in seconds", () => { + expect( + parseIngestArgs([ + "--provider", + "supermemory", + "--benchmark", + "beam-1m", + "--ingest-timeout-seconds", + "300", + ]) + ).toMatchObject({ ingestReadinessTimeoutMs: 300_000 }) + }) +}) diff --git a/test/input-identity.test.ts b/test/input-identity.test.ts new file mode 100644 index 0000000..9310ecb --- /dev/null +++ b/test/input-identity.test.ts @@ -0,0 +1,177 @@ +import { describe, expect, test } from "bun:test" +import type { Benchmark } from "../src/types/benchmark" +import type { UnifiedQuestion, UnifiedSession } from "../src/types/unified" +import { legacyBenchmarkProtocol } from "../src/protocols/legacy" +import { + canonicalizeSelectedQuestionIds, + fingerprintSelectedBenchmarkInput, + resolveEffectiveDatasetRevision, +} from "../src/orchestrator/input-identity" +import { stableSha256 } from "../src/utils/stable" + +function question(questionId: string, text = `Question ${questionId}`): UnifiedQuestion { + return { + questionId, + question: text, + questionType: "test", + groundTruth: `Answer ${questionId}`, + haystackSessionIds: ["session-1"], + metadata: { fixture: true }, + } +} + +function session(content = "Shared user content", metadata: Record = {}) { + return [ + { + sessionId: "session-1", + messages: [ + { role: "user" as const, content }, + { role: "assistant" as const, content: "Shared assistant content" }, + ], + metadata, + }, + ] +} + +function benchmark(input: { + questions: UnifiedQuestion[] + sessionsByQuestion: Record + groups?: Record +}): Benchmark { + return { + name: "legacy-fixture", + scope: { displayName: "Legacy fixture", includedTiers: [], coverage: "full" }, + protocol: legacyBenchmarkProtocol, + async load() {}, + getQuestions() { + return input.questions + }, + getHaystackSessions(questionId) { + return input.sessionsByQuestion[questionId] ?? [] + }, + getGroundTruth(questionId) { + return input.questions.find((value) => value.questionId === questionId)?.groundTruth ?? "" + }, + getQuestionTypes() { + return { test: { id: "test", alias: "test", description: "Test" } } + }, + getIngestionGroupId: input.groups + ? (questionId) => input.groups![questionId] ?? questionId + : undefined, + } +} + +describe("selected benchmark input identity", () => { + test("persists the enclosing snapshot fingerprint as the effective revision", () => { + expect(resolveEffectiveDatasetRevision(undefined, undefined)).toBeUndefined() + expect(resolveEffectiveDatasetRevision("configured", undefined)).toBe("configured") + expect( + resolveEffectiveDatasetRevision("configured", { + datasetFingerprint: "resolved-fingerprint", + snapshotFingerprint: "snapshot-fingerprint", + manifestHash: "manifest", + manifestSchemaVersion: 1, + canonicalSchemaVersion: 1, + converterVersion: "1", + converterImplementationHash: "converter", + includedTiers: ["1M"], + counts: {}, + orderedQuestionIdsDigest: {}, + sourceFiles: [], + canonicalFiles: [], + sources: [], + }) + ).toBe("snapshot-fingerprint") + }) + + test("canonicalizes identical ID sets independently of caller order and detects set drift", () => { + const questions = [question("q1"), question("q2"), question("q3")] + const forward = canonicalizeSelectedQuestionIds(questions, ["q1", "q3"]) + const reverse = canonicalizeSelectedQuestionIds(questions, ["q3", "q1"]) + + expect(forward).toEqual(["q1", "q3"]) + expect(reverse).toEqual(forward) + expect(stableSha256(reverse)).toBe(stableSha256(forward)) + expect(stableSha256(canonicalizeSelectedQuestionIds(questions, ["q1", "q2"]))).not.toBe( + stableSha256(forward) + ) + const sharedSessions = session() + const fixture = benchmark({ + questions, + sessionsByQuestion: { q1: sharedSessions, q2: sharedSessions, q3: sharedSessions }, + }) + expect(fingerprintSelectedBenchmarkInput(fixture, [questions[2], questions[0]])).toBe( + fingerprintSelectedBenchmarkInput(fixture, [questions[0], questions[2]]) + ) + expect(() => canonicalizeSelectedQuestionIds(questions, ["q1", "q1"])).toThrow("duplicates") + }) + + test("changes for question or raw haystack drift and rejects declared group collisions", () => { + const questions = [question("q1"), question("q2")] + const sharedSessions = session() + const base = benchmark({ + questions, + sessionsByQuestion: { q1: sharedSessions, q2: sharedSessions }, + groups: { q1: "chat-1", q2: "chat-1" }, + }) + const baseFingerprint = fingerprintSelectedBenchmarkInput(base, questions) + + const changedQuestion = [question("q1", "Changed question"), question("q2")] + expect( + fingerprintSelectedBenchmarkInput( + benchmark({ + questions: changedQuestion, + sessionsByQuestion: { q1: sharedSessions, q2: sharedSessions }, + groups: { q1: "chat-1", q2: "chat-1" }, + }), + changedQuestion + ) + ).not.toBe(baseFingerprint) + + const changedSessions = session("Changed user content") + expect( + fingerprintSelectedBenchmarkInput( + benchmark({ + questions, + sessionsByQuestion: { q1: changedSessions, q2: changedSessions }, + groups: { q1: "chat-1", q2: "chat-1" }, + }), + questions + ) + ).not.toBe(baseFingerprint) + + expect(() => + fingerprintSelectedBenchmarkInput( + benchmark({ + questions, + sessionsByQuestion: { q1: sharedSessions, q2: changedSessions }, + groups: { q1: "chat-1", q2: "chat-1" }, + }), + questions + ) + ).toThrow("different raw benchmark haystacks") + }) + + test("serializes one shared raw haystack once rather than once per question", () => { + let metadataReads = 0 + const metadata: Record = {} + Object.defineProperty(metadata, "probe", { + enumerable: true, + get() { + metadataReads++ + return "value" + }, + }) + const sharedSessions = session("Shared", metadata) + const questions = Array.from({ length: 20 }, (_, index) => question(`q${index + 1}`)) + const fixture = benchmark({ + questions, + sessionsByQuestion: Object.fromEntries( + questions.map((value) => [value.questionId, sharedSessions]) + ), + }) + + fingerprintSelectedBenchmarkInput(fixture, questions) + expect(metadataReads).toBe(2) + }) +}) diff --git a/test/leaderboard-identity.test.ts b/test/leaderboard-identity.test.ts new file mode 100644 index 0000000..acc7a51 --- /dev/null +++ b/test/leaderboard-identity.test.ts @@ -0,0 +1,640 @@ +import { describe, expect, test } from "bun:test" +import { + createLeaderboardComparisonIdentity, + rankLeaderboardEntries, + validateLeaderboardReportForPublication, + type LeaderboardAggregationContext, +} from "../src/server/leaderboard-identity" +import type { RunCheckpoint } from "../src/types/checkpoint" +import type { BenchmarkResult, LatencyStats } from "../src/types/unified" +import { sha256Text, stableSha256 } from "../src/utils/stable" +import { resolveAnsweringRuntimeIdentity } from "../src/utils/models" +import { BEAM_ABILITY_IDS, BeamPaperProtocol } from "../src/protocols/beam-paper" + +function source(overrides: Record = {}) { + return { + benchmark: "beam-1m", + benchmarkScope: { + displayName: "BEAM 1M", + includedTiers: ["1M"], + coverage: "full", + }, + datasetIdentity: { + datasetFingerprint: "dataset-a", + revision: "pinned", + }, + selectedQuestionIdsDigest: "questions-a", + benchmarkInputFingerprint: "benchmark-input-a", + protocolIdentity: { + id: "beam-paper", + version: "1.1.0", + configFingerprint: "config-a", + }, + retrievalTopK: 5, + judgeModel: "gpt-4.1-mini", + answeringModel: "gpt-4.1-mini", + answeringRuntimeIdentity: resolveAnsweringRuntimeIdentity("gpt-4.1-mini"), + primaryMetric: { + key: "beamScore", + value: 0.5, + higherIsBetter: true, + }, + accuracy: 0.6, + ...overrides, + } +} + +describe("leaderboard comparison identity", () => { + test("score values vary within one cohort", () => { + const first = createLeaderboardComparisonIdentity(source()) + const second = createLeaderboardComparisonIdentity( + source({ primaryMetric: { key: "beamScore", value: 0.8, higherIsBetter: true } }) + ) + + expect(first.cohortKey).toBe(second.cohortKey) + expect(first.primaryMetric.value).toBe(0.5) + expect(second.primaryMetric.value).toBe(0.8) + }) + + test("dataset, question set, protocol, Top-K, models, and metric semantics split cohorts", () => { + const baseline = createLeaderboardComparisonIdentity(source()).cohortKey + const variants = [ + source({ datasetIdentity: { datasetFingerprint: "dataset-b" } }), + source({ selectedQuestionIdsDigest: "questions-b" }), + source({ benchmarkInputFingerprint: "benchmark-input-b" }), + source({ + protocolIdentity: { + id: "beam-paper", + version: "1.1.0", + configFingerprint: "config-b", + }, + }), + source({ retrievalTopK: 10 }), + source({ judgeModel: "different-judge" }), + source({ answeringModel: "different-answering-model" }), + source({ + answeringRuntimeIdentity: { + ...resolveAnsweringRuntimeIdentity("gpt-4.1-mini"), + modelId: "different-effective-model", + }, + }), + source({ primaryMetric: { key: "passAccuracy", value: 0.5, higherIsBetter: true } }), + source({ primaryMetric: { key: "beamScore", value: 0.5, higherIsBetter: false } }), + ] + + for (const variant of variants) { + expect(createLeaderboardComparisonIdentity(variant).cohortKey).not.toBe(baseline) + } + }) + + test("legacy provider-prompt drift splits cohorts", () => { + const legacyProtocol = { id: "memorybench.legacy", version: "1.0.0" } + const first = createLeaderboardComparisonIdentity( + source({ protocolIdentity: legacyProtocol, providerPromptFingerprint: "prompt-a" }) + ) + const second = createLeaderboardComparisonIdentity( + source({ protocolIdentity: legacyProtocol, providerPromptFingerprint: "prompt-b" }) + ) + + expect(first.cohortKey).not.toBe(second.cohortKey) + }) + + test("rejects a run that reports mixed retrieval Top-K values", () => { + expect(() => + createLeaderboardComparisonIdentity( + source({ questionMetrics: [{ configuredTopK: 5 }, { configuredTopK: 10 }] }) + ) + ).toThrow("mixed retrieval Top-K") + }) + + test("rejects disagreement between configured and recorded retrieval Top-K", () => { + expect(() => + createLeaderboardComparisonIdentity( + source({ retrievalTopK: 10, questionMetrics: [{ configuredTopK: 5 }] }) + ) + ).toThrow("differs from its recorded question Top-K") + }) +}) + +describe("leaderboard cohort ranking", () => { + test("ranks by primary metric only within the same cohort", () => { + const low = createLeaderboardComparisonIdentity(source()) + const high = createLeaderboardComparisonIdentity( + source({ primaryMetric: { key: "beamScore", value: 0.8, higherIsBetter: true } }) + ) + const otherDataset = createLeaderboardComparisonIdentity( + source({ + datasetIdentity: { datasetFingerprint: "dataset-b" }, + primaryMetric: { key: "beamScore", value: 0.99, higherIsBetter: true }, + }) + ) + + const ranked = rankLeaderboardEntries([ + { + id: 1, + benchmark: "beam-1m", + accuracy: 0.9, + comparisonIdentity: low, + }, + { + id: 2, + benchmark: "beam-1m", + accuracy: 0.1, + comparisonIdentity: high, + }, + { + id: 3, + benchmark: "beam-1m", + accuracy: 1, + comparisonIdentity: otherDataset, + }, + ]) + + expect(ranked.find((entry) => entry.id === 2)).toMatchObject({ cohortRank: 1, cohortSize: 2 }) + expect(ranked.find((entry) => entry.id === 1)).toMatchObject({ cohortRank: 2, cohortSize: 2 }) + expect(ranked.find((entry) => entry.id === 3)).toMatchObject({ cohortRank: 1, cohortSize: 1 }) + }) + + test("honors lower-is-better metrics and gives exact ties the same rank", () => { + const identity = createLeaderboardComparisonIdentity( + source({ primaryMetric: { key: "latency", value: 20, higherIsBetter: false } }) + ) + const faster = { + ...identity, + primaryMetric: { ...identity.primaryMetric, value: 10 }, + } + const ranked = rankLeaderboardEntries([ + { id: 1, benchmark: "beam-1m", accuracy: 0, comparisonIdentity: identity }, + { id: 2, benchmark: "beam-1m", accuracy: 1, comparisonIdentity: faster }, + { id: 3, benchmark: "beam-1m", accuracy: 0.5, comparisonIdentity: faster }, + ]) + + expect(ranked.find((entry) => entry.id === 2)?.cohortRank).toBe(1) + expect(ranked.find((entry) => entry.id === 3)?.cohortRank).toBe(1) + expect(ranked.find((entry) => entry.id === 1)?.cohortRank).toBe(3) + }) + + test("legacy rows fall back to benchmark and accuracy semantics", () => { + const ranked = rankLeaderboardEntries([ + { id: 1, benchmark: "locomo", accuracy: 0.4 }, + { id: 2, benchmark: "locomo", accuracy: 0.7 }, + { id: 3, benchmark: "convomem", accuracy: 0.9 }, + ]) + + expect(ranked.find((entry) => entry.id === 2)).toMatchObject({ cohortRank: 1, cohortSize: 2 }) + expect(ranked.find((entry) => entry.id === 1)).toMatchObject({ cohortRank: 2, cohortSize: 2 }) + expect(ranked.find((entry) => entry.id === 3)).toMatchObject({ cohortRank: 1, cohortSize: 1 }) + }) +}) + +const ZERO_LATENCY: LatencyStats = { + min: 0, + max: 0, + mean: 0, + median: 0, + p95: 0, + p99: 0, + stdDev: 0, + count: 1, +} + +function publishableFixture(): { + checkpoint: RunCheckpoint + report: BenchmarkResult + aggregation: LeaderboardAggregationContext +} { + const questionIds = ["q1"] + const digest = stableSha256(questionIds) + const benchmarkScope = { + displayName: "BEAM 1M", + includedTiers: ["1M"], + coverage: "subset" as const, + } + const datasetIdentity = { datasetFingerprint: "dataset-a" } + const protocolIdentity = { id: "test-paper-protocol", version: "1.0.0" } + const checkpoint = { + runId: "run-a", + provider: "supermemory", + providerPromptFingerprint: "provider-prompt-a", + benchmark: "beam-1m", + benchmarkScope, + datasetIdentity, + selectedQuestionIdsDigest: digest, + benchmarkInputFingerprint: "benchmark-input-a", + protocolIdentity, + judge: "gpt-4.1-mini", + answeringModel: "gpt-4.1-mini", + answeringRuntimeIdentity: resolveAnsweringRuntimeIdentity("gpt-4.1-mini"), + retrievalTopK: 5, + targetQuestionIds: questionIds, + builds: {}, + questions: { + q1: { + questionId: "q1", + question: "When?", + questionType: "temporal", + groundTruth: "Tuesday", + phases: { + evaluate: { + status: "completed", + evaluation: { + questionId: "q1", + questionType: "temporal", + primaryScore: 0.5, + passed: true, + explanation: "ok", + }, + }, + }, + }, + }, + } as unknown as RunCheckpoint + const report: BenchmarkResult = { + provider: "supermemory", + providerPromptFingerprint: "provider-prompt-a", + benchmark: "beam-1m", + runId: "run-a", + dataSourceRunId: "run-a", + judge: "gpt-4.1-mini", + answeringModel: "gpt-4.1-mini", + answeringRuntimeIdentity: resolveAnsweringRuntimeIdentity("gpt-4.1-mini"), + timestamp: "2026-08-03T00:00:00.000Z", + selectedQuestionIdsDigest: digest, + benchmarkInputFingerprint: "benchmark-input-a", + retrievalTopK: 5, + benchmarkScope, + datasetIdentity, + protocolIdentity, + quality: { + primaryMetric: { key: "beamScore", value: 0.5, higherIsBetter: true }, + metrics: { passAccuracy: 1 }, + }, + summary: { totalQuestions: 1, correctCount: 1, accuracy: 1, averageScore: 0.5 }, + builds: { + uniqueBuildCount: 0, + sumContainerBuildWorkMs: 0, + buildPhaseWallClockMs: 0, + totalBuildCostUsd: null, + knownCostBuildCount: 0, + totalCostBuildCount: 0, + items: [], + }, + questionMetrics: [ + { + questionId: "q1", + buildId: "build-1", + searchLatencyMs: 0, + answerLatencyMs: 0, + onlineQueryLatencyMs: 0, + evaluationLatencyMs: 0, + queryCostUsd: null, + evaluationCostUsd: null, + configuredTopK: 5, + providerRequestLimit: 5, + rawReturnedCount: 5, + returnedCount: 5, + normalizedCount: 5, + droppedCount: 0, + answerCutoff: 5, + answerEvidenceCount: 5, + contextTokens: 0, + providerRequests: [], + }, + ], + latency: { + ingest: ZERO_LATENCY, + indexing: ZERO_LATENCY, + search: ZERO_LATENCY, + answer: ZERO_LATENCY, + evaluate: ZERO_LATENCY, + total: ZERO_LATENCY, + }, + byQuestionType: {}, + evaluations: [ + { + questionId: "q1", + questionType: "temporal", + question: "When?", + score: 0.5, + primaryScore: 0.5, + passed: true, + label: "correct", + explanation: "ok", + hypothesis: "Tuesday", + groundTruth: "Tuesday", + searchResults: [], + }, + ], + } + const aggregation: LeaderboardAggregationContext = { + protocol: { + identity: protocolIdentity as RunCheckpoint["protocolIdentity"], + aggregateQuality({ evaluations }) { + const score = + evaluations.reduce((sum, evaluation) => sum + evaluation.primaryScore, 0) / + evaluations.length + return { + primaryMetric: { key: "beamScore", value: score, higherIsBetter: true }, + metrics: { + passAccuracy: + evaluations.filter((evaluation) => evaluation.passed).length / evaluations.length, + }, + } + }, + }, + questions: [ + { + questionId: "q1", + question: "When?", + questionType: "temporal", + groundTruth: "Tuesday", + haystackSessionIds: [], + }, + ], + } + return { checkpoint, report, aggregation } +} + +describe("leaderboard publication report gate", () => { + test("accepts a complete report whose full identity matches its checkpoint", () => { + const { checkpoint, report, aggregation } = publishableFixture() + const identity = validateLeaderboardReportForPublication(checkpoint, report, 1, aggregation) + + expect(identity.judgeModel).toBe("gpt-4.1-mini") + expect(identity.answeringModel).toBe("gpt-4.1-mini") + expect(identity.benchmarkInputFingerprint).toBe("benchmark-input-a") + expect(identity.retrievalTopK).toBe(5) + }) + + test("rejects reports without a scalar official primary metric", () => { + const { checkpoint, report, aggregation } = publishableFixture() + report.quality.primaryMetric = undefined + + expect(() => + validateLeaderboardReportForPublication(checkpoint, report, 1, aggregation) + ).toThrow("protocol quality aggregation") + }) + + test("rejects a sampled single-tier BEAM run from ranked publication", () => { + const { checkpoint, report, aggregation } = publishableFixture() + const protocol = new BeamPaperProtocol({ retrievalTopK: 5 }) + const question = aggregation.questions[0]! + question.questionType = "temporal_reasoning" + question.metadata = { scale: "1M", rubric: ["expected detail"] } + const evaluation = { + questionId: question.questionId, + questionType: question.questionType, + primaryScore: 0.5, + passed: true, + explanation: "ok", + } + + checkpoint.protocolIdentity = protocol.identity + checkpoint.questions.q1!.questionType = question.questionType + checkpoint.questions.q1!.phases.evaluate = { status: "completed", evaluation } + report.protocolIdentity = protocol.identity + report.evaluations[0]!.questionType = question.questionType + report.quality = protocol.aggregateQuality({ questions: [question], evaluations: [evaluation] }) + + expect(report.quality.primaryMetric?.key).toBe("beamScorePartial") + expect(() => + validateLeaderboardReportForPublication(checkpoint, report, 1, { + protocol, + questions: [question], + }) + ).toThrow("official complete-tier BEAM") + }) + + test("accepts the exact full 1M question set using the manifest digest scheme", () => { + const { checkpoint, report } = publishableFixture() + const protocol = new BeamPaperProtocol({ retrievalTopK: 5 }) + const questions = BEAM_ABILITY_IDS.flatMap((ability) => + Array.from({ length: 70 }, (_, index) => ({ + questionId: `official-1m-${ability}-${index}`, + question: `Question ${ability} ${index}`, + questionType: ability, + groundTruth: "Expected answer", + haystackSessionIds: [], + metadata: { scale: "1M", rubric: ["Expected detail", "Second detail"] }, + })) + ) + const evaluations = questions.map((question) => ({ + questionId: question.questionId, + questionType: question.questionType, + primaryScore: 0.5, + passed: true, + explanation: "ok", + })) + const questionIds = questions.map((question) => question.questionId) + const selectedQuestionIdsDigest = stableSha256(questionIds) + const datasetIdentity = { + datasetFingerprint: "official-dataset", + orderedQuestionIdsDigest: { "1M": sha256Text(questionIds.join("\n")) }, + } as RunCheckpoint["datasetIdentity"] + + checkpoint.protocolIdentity = protocol.identity + checkpoint.datasetIdentity = datasetIdentity + checkpoint.targetQuestionIds = questionIds + checkpoint.selectedQuestionIdsDigest = selectedQuestionIdsDigest + checkpoint.questions = Object.fromEntries( + questions.map((question, index) => [ + question.questionId, + { + questionId: question.questionId, + question: question.question, + questionType: question.questionType, + groundTruth: question.groundTruth, + phases: { evaluate: { status: "completed", evaluation: evaluations[index] } }, + }, + ]) + ) as RunCheckpoint["questions"] + + report.protocolIdentity = protocol.identity + report.datasetIdentity = datasetIdentity + report.selectedQuestionIdsDigest = selectedQuestionIdsDigest + report.quality = protocol.aggregateQuality({ questions, evaluations }) + report.summary = { + totalQuestions: questions.length, + correctCount: questions.length, + accuracy: 1, + averageScore: 0.5, + } + report.evaluations = questions.map((question) => ({ + questionId: question.questionId, + questionType: question.questionType, + question: question.question, + score: 0.5, + primaryScore: 0.5, + passed: true, + label: "correct", + explanation: "ok", + hypothesis: "answer", + groundTruth: question.groundTruth, + searchResults: [], + })) + report.questionMetrics = questions.map((question, index) => ({ + ...report.questionMetrics[0]!, + questionId: question.questionId, + buildId: `build-${index}`, + })) + + const identity = validateLeaderboardReportForPublication( + checkpoint, + report, + questions.length, + { protocol, questions } + ) + expect(identity.primaryMetric).toMatchObject({ key: "beamScore", value: 0.5 }) + }) + + test("rejects report identity or aggregation drift", () => { + const { checkpoint, report, aggregation } = publishableFixture() + report.answeringModel = "different-model" + report.summary.correctCount = 0 + + expect(() => + validateLeaderboardReportForPublication(checkpoint, report, 1, aggregation) + ).toThrow(/answering model.*correct-count aggregation/) + }) + + test("rejects effective answering-runtime or benchmark-input drift", () => { + const { checkpoint, report, aggregation } = publishableFixture() + report.answeringRuntimeIdentity = { + ...report.answeringRuntimeIdentity, + modelId: "same-alias-different-effective-model", + } + report.benchmarkInputFingerprint = "benchmark-input-b" + + expect(() => + validateLeaderboardReportForPublication(checkpoint, report, 1, aggregation) + ).toThrow(/answering runtime.*benchmark-input fingerprint/) + }) + + test("rejects deletion of an official report dataset identity", () => { + const { checkpoint, report, aggregation } = publishableFixture() + report.datasetIdentity = undefined + + expect(() => + validateLeaderboardReportForPublication(checkpoint, report, 1, aggregation) + ).toThrow("missing dataset identity") + }) + + test("rejects tampered protocol-owned primary and secondary quality metrics", () => { + const { checkpoint, report, aggregation } = publishableFixture() + report.quality.primaryMetric!.value = 0.99 + report.quality.metrics.passAccuracy = 0 + + expect(() => + validateLeaderboardReportForPublication(checkpoint, report, 1, aggregation) + ).toThrow("protocol quality aggregation") + }) + + test("rejects a coherently rewritten report evaluation against checkpoint state", () => { + const { checkpoint, report, aggregation } = publishableFixture() + report.evaluations[0]!.score = 0.99 + report.evaluations[0]!.primaryScore = 0.99 + report.summary.averageScore = 0.99 + report.quality.primaryMetric!.value = 0.99 + + expect(() => + validateLeaderboardReportForPublication(checkpoint, report, 1, aggregation) + ).toThrow(/report\/checkpoint evaluations.*protocol quality aggregation/) + }) + + test("keeps combined BEAM without an official scalar out of the leaderboard", () => { + const { checkpoint, report } = publishableFixture() + const protocol = new BeamPaperProtocol({ retrievalTopK: 5 }) + const questions = [ + { + questionId: "q1", + question: "Question one", + questionType: "temporal_reasoning", + groundTruth: "Answer one", + haystackSessionIds: [], + metadata: { scale: "1M" }, + }, + { + questionId: "q2", + question: "Question two", + questionType: "temporal_reasoning", + groundTruth: "Answer two", + haystackSessionIds: [], + metadata: { scale: "10M" }, + }, + ] + const evaluations = questions.map((question) => ({ + questionId: question.questionId, + questionType: question.questionType, + primaryScore: 0.5, + passed: true, + explanation: "ok", + })) + const digest = stableSha256(questions.map((question) => question.questionId)) + + checkpoint.benchmark = "beam-1m-10m" + checkpoint.benchmarkScope = { + displayName: "BEAM 1M/10M", + includedTiers: ["1M", "10M"], + coverage: "subset", + } + checkpoint.protocolIdentity = protocol.identity + checkpoint.selectedQuestionIdsDigest = digest + checkpoint.targetQuestionIds = ["q1", "q2"] + checkpoint.questions = Object.fromEntries( + questions.map((question, index) => [ + question.questionId, + { + questionId: question.questionId, + question: question.question, + questionType: question.questionType, + groundTruth: question.groundTruth, + phases: { + evaluate: { status: "completed", evaluation: evaluations[index] }, + }, + }, + ]) + ) as RunCheckpoint["questions"] + + report.benchmark = checkpoint.benchmark + report.benchmarkScope = checkpoint.benchmarkScope + report.protocolIdentity = protocol.identity + report.selectedQuestionIdsDigest = digest + report.summary = { totalQuestions: 2, correctCount: 2, accuracy: 1, averageScore: 0.5 } + report.quality = protocol.aggregateQuality({ questions, evaluations }) + report.evaluations = questions.map((question) => ({ + questionId: question.questionId, + questionType: question.questionType, + question: question.question, + score: 0.5, + primaryScore: 0.5, + passed: true, + label: "correct", + explanation: "ok", + hypothesis: "answer", + groundTruth: question.groundTruth, + searchResults: [], + searchDurationMs: 0, + answerDurationMs: 0, + totalDurationMs: 0, + })) + report.questionMetrics = questions.map((question, index) => ({ + ...report.questionMetrics[0]!, + questionId: question.questionId, + buildId: `build-${index + 1}`, + })) + const aggregation = { protocol, questions } + + expect(() => + validateLeaderboardReportForPublication(checkpoint, report, 2, aggregation) + ).toThrow("finite scalar primary metric") + + report.quality.primaryMetric = { + key: "beamScore", + value: 0.5, + higherIsBetter: true, + } + expect(() => + validateLeaderboardReportForPublication(checkpoint, report, 2, aggregation) + ).toThrow("protocol quality aggregation") + }) +}) diff --git a/test/normalization-contract.test.ts b/test/normalization-contract.test.ts new file mode 100644 index 0000000..70ea2eb --- /dev/null +++ b/test/normalization-contract.test.ts @@ -0,0 +1,103 @@ +import { describe, expect, test } from "bun:test" +import { validateProviderSearchResponse } from "../src/orchestrator/phases/search" +import { renderFilesystemMemoryFile } from "../src/providers/filesystem" +import { createProviderSearchResponse, resolveDocumentDate } from "../src/providers/normalization" +import { normalizeSupermemorySearchResults } from "../src/providers/supermemory" +import { configureMem0Project } from "../src/providers/mem0" +import type { ProviderResultDropDiagnostic } from "../src/types/unified" + +describe("normalized result contract", () => { + test("Mem0 project configuration fails closed", async () => { + const failure = new Error("project update rejected") + await expect( + configureMem0Project({ + updateProject: async () => { + throw failure + }, + }) + ).rejects.toBe(failure) + }) + + test("records a reason for every malformed or empty provider result", () => { + const droppedResults: ProviderResultDropDiagnostic[] = [] + const results = normalizeSupermemorySearchResults( + [null, { id: "missing-text", memory: " " }, { chunk: "missing id" }], + 3, + droppedResults + ) + + expect(results).toEqual([]) + expect(droppedResults).toEqual([ + { index: 0, reason: "malformed-result" }, + { index: 1, reason: "empty-text" }, + { index: 2, reason: "missing-id" }, + ]) + expect( + createProviderSearchResponse({ + results, + requestedLimit: 3, + rawReturnedCount: 3, + droppedResults, + providerRequests: [{ operation: "search.hybrid", limit: 3 }], + }).diagnostics + ).toMatchObject({ droppedCount: 3, droppedResults }) + }) + + test("rejects unrecorded drops and unsupported normalized result types", () => { + expect(() => + createProviderSearchResponse({ + results: [], + requestedLimit: 1, + rawReturnedCount: 1, + providerRequests: [{ operation: "search", limit: 1 }], + }) + ).toThrow("recorded 0/1 drop reasons") + + expect(() => + validateProviderSearchResponse( + { + results: [ + { + id: "bad-type", + rank: 1, + text: "evidence", + provider: "supermemory", + resultType: "raw-json" as never, + }, + ], + diagnostics: { + requestedLimit: 1, + rawReturnedCount: 1, + normalizedCount: 1, + droppedCount: 0, + droppedResults: [], + providerRequests: [{ operation: "search", limit: 1 }], + }, + }, + { name: "supermemory", searchRequestStructure: { kind: "single" } }, + 1 + ) + ).toThrow("unsupported result type") + }) + + test("never promotes provider timestamps or unknown sentinels to source document dates", () => { + expect( + resolveDocumentDate({ createdAt: "2026-01-01", updatedAt: "2026-01-02" }) + ).toBeUndefined() + expect(resolveDocumentDate({ documentDate: "Unknown date" })).toBeUndefined() + expect(resolveDocumentDate({ temporalContext: { documentDate: "unknown" } })).toBeUndefined() + expect(resolveDocumentDate({ date: "not specified" })).toBeUndefined() + expect(resolveDocumentDate({ documentDate: "2024-03-01" })).toBe("2024-03-01") + }) + + test("filesystem undated documents omit the date header and sentinel", () => { + const undated = renderFilesystemMemoryFile("session-1", undefined, "Stored fact") + expect(undated).toBe("# Memory: session-1\n\nStored fact") + expect(undated).not.toContain("**Date:**") + expect(undated).not.toContain("Unknown date") + + expect(renderFilesystemMemoryFile("session-1", "2024-03-01", "Stored fact")).toContain( + "**Date:** 2024-03-01" + ) + }) +}) diff --git a/test/protocol-contract.test.ts b/test/protocol-contract.test.ts new file mode 100644 index 0000000..c1b639e --- /dev/null +++ b/test/protocol-contract.test.ts @@ -0,0 +1,473 @@ +import { describe, expect, test } from "bun:test" +import { createBenchmark, getAvailableBenchmarks } from "../src/benchmarks" +import type { Benchmark, BenchmarkName } from "../src/types/benchmark" +import { CHECKPOINT_SCHEMA_VERSION, type RunCheckpoint } from "../src/types/checkpoint" +import type { BenchmarkProtocol, ProtocolIdentity } from "../src/types/protocol" +import type { UnifiedQuestion, UnifiedSession } from "../src/types/unified" +import { BeamPaperProtocol } from "../src/protocols/beam-paper" +import { legacyBenchmarkProtocol } from "../src/protocols/legacy" +import { prepareValidatedBuildPlans } from "../src/orchestrator/builds" +import { assertResumeIdentity, resolveEffectiveRetrievalTopK } from "../src/orchestrator" +import { fingerprintProviderPrompts } from "../src/providers/prompt-identity" +import { resolveAnsweringRuntimeIdentity } from "../src/utils/models" + +function beamQuestion(): UnifiedQuestion { + return { + questionId: "beam:test:chat:abstention:question", + question: "What should be remembered?", + questionType: "abstention", + groundTruth: "Nothing", + haystackSessionIds: ["s1", "s2"], + metadata: { rubric: ["The answer should abstain when evidence is absent"] }, + } +} + +function legacyQuestion(): UnifiedQuestion { + return { + questionId: "legacy-question", + question: "What happened?", + questionType: "legacy-type", + groundTruth: "An event", + haystackSessionIds: ["s1", "s2"], + metadata: { rubric: ["Array-valued rubric metadata must not select BEAM"] }, + } +} + +function sessions(): UnifiedSession[] { + return [ + { + sessionId: "s1", + messages: [ + { role: "user", content: "First message" }, + { role: "assistant", content: "First response" }, + ], + metadata: { documentDate: "2024-01-01", date: "2024-01-01" }, + }, + { + sessionId: "s2", + messages: [ + { role: "user", content: "Second message" }, + { role: "assistant", content: "Second response" }, + ], + metadata: { documentDate: "2024-01-02", date: "2024-01-02" }, + }, + ] +} + +function fakeBenchmark(input: { + name: string + protocol: BenchmarkProtocol + question: UnifiedQuestion +}): Benchmark { + return { + name: input.name, + scope: { displayName: "Fake benchmark", includedTiers: [], coverage: "full" }, + protocol: input.protocol, + async load() {}, + getQuestions() { + return [input.question] + }, + getHaystackSessions() { + return sessions() + }, + getGroundTruth() { + return input.question.groundTruth + }, + getQuestionTypes() { + return { + [input.question.questionType]: { + id: input.question.questionType, + alias: input.question.questionType, + description: "Test", + }, + } + }, + getIngestionGroupId() { + return "shared-chat" + }, + } +} + +function cloneProtocolIdentity( + identity: ProtocolIdentity, + changes: Partial +): ProtocolIdentity { + return { ...identity, ...changes } +} + +function protocolWithIdentity( + base: BenchmarkProtocol, + identity: ProtocolIdentity +): BenchmarkProtocol { + return { + identity, + auxiliaryRetrievalEvaluation: base.auxiliaryRetrievalEvaluation, + ingestionExecutionPolicy: base.ingestionExecutionPolicy, + requiredJudge: base.requiredJudge, + validateQuestion: base.validateQuestion.bind(base), + createIngestionPlan: base.createIngestionPlan.bind(base), + createRetrievalPlan: base.createRetrievalPlan.bind(base), + createAnswerPlan: base.createAnswerPlan.bind(base), + evaluateQuestion: base.evaluateQuestion.bind(base), + aggregateQuality: base.aggregateQuality.bind(base), + } +} + +class SameBytesChangedBeamIngestionProtocol extends BeamPaperProtocol { + override createIngestionPlan(input: Parameters[0]) { + return super.createIngestionPlan(input).map((document) => ({ + ...document, + metadata: { ...document.metadata }, + messages: document.messages ? [...document.messages] : undefined, + })) + } +} + +function buildPlan( + benchmark: Benchmark, + question: UnifiedQuestion, + providerPromptFingerprint = "prompt-v1" +) { + return prepareValidatedBuildPlans({ + benchmark, + questions: [question], + provider: "fake-provider", + providerAdapterVersion: "adapter-v1", + providerPromptFingerprint, + providerIngestionConfigFingerprint: "ingestion-config-v1", + dataSourceRunId: "source-run", + })[0] +} + +function checkpoint(protocolIdentity: ProtocolIdentity): RunCheckpoint { + return { + schemaVersion: CHECKPOINT_SCHEMA_VERSION, + runId: "resume-run", + dataSourceRunId: "resume-run", + status: "running", + provider: "fake-provider", + providerAdapterVersion: "adapter-v1", + providerPromptFingerprint: "prompt-v1", + benchmark: "renamed-benchmark", + benchmarkScope: { + displayName: "Renamed benchmark", + includedTiers: ["1M"], + coverage: "subset", + }, + selectedQuestionIdsDigest: "selected-question-digest", + benchmarkInputFingerprint: "benchmark-input-digest", + protocolIdentity, + judge: "gpt-4.1-mini", + answeringModel: "answer-model", + answeringRuntimeIdentity: resolveAnsweringRuntimeIdentity("answer-model"), + retrievalTopK: 5, + createdAt: "2024-01-01T00:00:00.000Z", + updatedAt: "2024-01-01T00:00:00.000Z", + buildPhaseAttempts: [], + builds: {}, + questions: {}, + } +} + +function resumeInput(protocolIdentity: ProtocolIdentity) { + return { + provider: "fake-provider", + providerAdapterVersion: "adapter-v1", + providerPromptFingerprint: "prompt-v1", + benchmark: "renamed-benchmark", + benchmarkScope: { + displayName: "Renamed benchmark", + includedTiers: ["1M"], + coverage: "subset" as const, + }, + datasetIdentity: undefined, + selectedQuestionIdsDigest: "selected-question-digest", + benchmarkInputFingerprint: "benchmark-input-digest", + protocolIdentity, + retrievalTopK: 5, + judge: "gpt-4.1-mini", + answeringModel: "answer-model", + answeringRuntimeIdentity: resolveAnsweringRuntimeIdentity("answer-model"), + } +} + +describe("benchmark registry scope", () => { + test("exposes only the explicit supported BEAM scope identifiers", () => { + expect(getAvailableBenchmarks()).toContain("beam-1m") + expect(getAvailableBenchmarks()).toContain("beam-10m") + expect(getAvailableBenchmarks()).toContain("beam-1m-10m") + expect(getAvailableBenchmarks()).not.toContain("beam" as BenchmarkName) + + expect(createBenchmark("beam-1m").name).toBe("beam-1m") + expect(createBenchmark("beam-10m").name).toBe("beam-10m") + expect(createBenchmark("beam-1m-10m").name).toBe("beam-1m-10m") + expect(() => createBenchmark("beam" as BenchmarkName)).toThrow("Unknown benchmark: beam") + }) +}) + +describe("explicit protocol ownership", () => { + test("a renamed benchmark remains BEAM when it explicitly owns BeamPaperProtocol", () => { + const protocol = new BeamPaperProtocol({ retrievalTopK: 5 }) + const value = beamQuestion() + const benchmark = fakeBenchmark({ + name: "completely-renamed-benchmark", + protocol, + question: value, + }) + + expect(benchmark.name.startsWith("beam")).toBe(false) + expect(benchmark.protocol.identity.id).toBe("beam-paper") + expect(benchmark.protocol.createRetrievalPlan({ question: value }).requestedTopK).toBe(5) + expect(buildPlan(benchmark, value).documents[0].content).toStartWith( + "DOCUMENT_DATE: 2024-01-01" + ) + }) + + test("BEAM ingestion projects provider-visible messages to role and content only", () => { + const protocol = new BeamPaperProtocol() + const value = beamQuestion() + const source = sessions() + source[0]!.messages[0] = { + ...source[0]!.messages[0]!, + speaker: "source-user", + timestamp: "2024-01-01T12:34:56Z", + } + + const [document] = protocol.createIngestionPlan({ question: value, sessions: source }) + expect(document!.messages).toEqual([ + { role: "user", content: "First message" }, + { role: "assistant", content: "First response" }, + ]) + expect(document!.content).toBe( + "DOCUMENT_DATE: 2024-01-01\n\n[USER]\nFirst message\n\n[ASSISTANT]\nFirst response" + ) + }) + + test("a BEAM-looking legacy benchmark with array rubric remains Legacy", () => { + const value = legacyQuestion() + const benchmark = fakeBenchmark({ + name: "beam-lookalike", + protocol: legacyBenchmarkProtocol, + question: value, + }) + + expect(Array.isArray(value.metadata?.rubric)).toBe(true) + expect(benchmark.protocol.identity.id).toBe("memorybench.legacy") + expect(benchmark.protocol.createRetrievalPlan({ question: value }).requestedTopK).toBe(10) + expect(buildPlan(benchmark, value).documents[0].content).toContain( + "session as a stringified JSON" + ) + }) +}) + +describe("resume protocol identity", () => { + test("accepts a byte-identical runtime identity", () => { + const identity = new BeamPaperProtocol({ retrievalTopK: 5 }).identity + expect(() => assertResumeIdentity(checkpoint(identity), resumeInput(identity))).not.toThrow() + }) + + test("rejects protocol, protocol-config, retrieval-config, and adapter drift", () => { + const identity = new BeamPaperProtocol({ retrievalTopK: 5 }).identity + const cases: Array<{ + expected: string + mutate: (input: ReturnType) => void + }> = [ + { + expected: "benchmark protocol", + mutate(input) { + input.protocolIdentity = cloneProtocolIdentity(identity, { id: "other-protocol" }) + }, + }, + { + expected: "benchmark protocol", + mutate(input) { + input.protocolIdentity = cloneProtocolIdentity(identity, { + configFingerprint: "changed-config", + }) + }, + }, + { + expected: "retrieval Top-K", + mutate(input) { + input.retrievalTopK = 10 + }, + }, + { + expected: "provider adapter version", + mutate(input) { + input.providerAdapterVersion = "adapter-v2" + }, + }, + { + expected: "answering runtime", + mutate(input) { + input.answeringRuntimeIdentity = { + ...input.answeringRuntimeIdentity, + modelId: "changed-effective-model-id", + } + }, + }, + ] + + for (const item of cases) { + const input = resumeInput(identity) + item.mutate(input) + expect(() => assertResumeIdentity(checkpoint(identity), input)).toThrow(item.expected) + } + }) + + test("rejects legacy provider-prompt drift while BEAM ignores unused provider prompts", () => { + const legacyCheckpoint = checkpoint(legacyBenchmarkProtocol.identity) + const legacyInput = resumeInput(legacyBenchmarkProtocol.identity) + legacyInput.providerPromptFingerprint = "prompt-v2" + expect(() => assertResumeIdentity(legacyCheckpoint, legacyInput)).toThrow("provider prompt") + + const beamIdentity = new BeamPaperProtocol({ retrievalTopK: 5 }).identity + const beamInput = resumeInput(beamIdentity) + beamInput.providerPromptFingerprint = "prompt-v2" + expect(() => assertResumeIdentity(checkpoint(beamIdentity), beamInput)).not.toThrow() + }) + + test("rejects schema-v3 checkpoints with missing effective input or answering identity", () => { + const identity = new BeamPaperProtocol({ retrievalTopK: 5 }).identity + const missingRuntime = checkpoint(identity) as Partial + delete missingRuntime.answeringRuntimeIdentity + expect(() => + assertResumeIdentity(missingRuntime as RunCheckpoint, resumeInput(identity)) + ).toThrow("answering runtime") + + const missingInput = checkpoint(identity) as Partial + delete missingInput.benchmarkInputFingerprint + expect(() => + assertResumeIdentity(missingInput as RunCheckpoint, resumeInput(identity)) + ).toThrow("benchmark input") + }) +}) + +describe("effective retrieval identity", () => { + test("resolves the BEAM default to 5 when the CLI option is omitted", () => { + const protocol = new BeamPaperProtocol() + expect(resolveEffectiveRetrievalTopK(protocol, [beamQuestion()])).toBe(5) + }) + + test("records a supported override and rejects configuration/plan disagreement", () => { + const protocol = new BeamPaperProtocol({ retrievalTopK: 20 }) + expect(resolveEffectiveRetrievalTopK(protocol, [beamQuestion()], 20)).toBe(20) + expect(() => resolveEffectiveRetrievalTopK(protocol, [beamQuestion()], 10)).toThrow( + "differs from protocol plan" + ) + }) +}) + +describe("provider prompt identity", () => { + test("is deterministic and changes for prompt strings or function source", () => { + const first = fingerprintProviderPrompts({ + answerPrompt: "Answer {{question}}", + judgePrompt: (question) => ({ default: `Judge ${question}` }), + }) + const same = fingerprintProviderPrompts({ + answerPrompt: "Answer {{question}}", + judgePrompt: (question) => ({ default: `Judge ${question}` }), + }) + const changedString = fingerprintProviderPrompts({ + answerPrompt: "Changed {{question}}", + judgePrompt: (question) => ({ default: `Judge ${question}` }), + }) + const changedFunction = fingerprintProviderPrompts({ + answerPrompt: "Answer {{question}}", + judgePrompt: (question) => ({ default: `Strictly judge ${question}` }), + }) + + expect(same).toBe(first) + expect(changedString).not.toBe(first) + expect(changedFunction).not.toBe(first) + }) +}) + +describe("build identity includes only protocol ingestion identity", () => { + test("answer and judge prompt drift does not change legacy or BEAM builds", () => { + const legacyValue = legacyQuestion() + const legacy = fakeBenchmark({ + name: "legacy", + protocol: legacyBenchmarkProtocol, + question: legacyValue, + }) + expect(buildPlan(legacy, legacyValue, "prompt-v2").buildFingerprint).toBe( + buildPlan(legacy, legacyValue, "prompt-v1").buildFingerprint + ) + + const beamValue = beamQuestion() + const beam = fakeBenchmark({ + name: "beam", + protocol: new BeamPaperProtocol(), + question: beamValue, + }) + expect(buildPlan(beam, beamValue, "prompt-v2").buildFingerprint).toBe( + buildPlan(beam, beamValue, "prompt-v1").buildFingerprint + ) + }) + + test("BEAM Top-K changes full protocol identity but not build or haystack identity", () => { + const value = beamQuestion() + const topK5 = new BeamPaperProtocol({ retrievalTopK: 5 }) + const topK10 = new BeamPaperProtocol({ retrievalTopK: 10 }) + const first = buildPlan( + fakeBenchmark({ name: "renamed", protocol: topK5, question: value }), + value + ) + const second = buildPlan( + fakeBenchmark({ name: "renamed", protocol: topK10, question: value }), + value + ) + + expect(topK10.identity.configFingerprint).not.toBe(topK5.identity.configFingerprint) + expect(topK10.identity.retrievalPolicyHash).not.toBe(topK5.identity.retrievalPolicyHash) + expect(topK10.identity.ingestionPolicyHash).toBe(topK5.identity.ingestionPolicyHash) + expect(second.haystack.fingerprint).toBe(first.haystack.fingerprint) + expect(second.buildFingerprint).toBe(first.buildFingerprint) + expect(second.buildId).toBe(first.buildId) + }) + + test("non-ingestion implementation drift does not change the build", () => { + const value = beamQuestion() + const base = new BeamPaperProtocol({ retrievalTopK: 5 }) + const changedImplementation = protocolWithIdentity( + base, + cloneProtocolIdentity(base.identity, { + implementationFingerprint: "changed-implementation", + }) + ) + const first = buildPlan( + fakeBenchmark({ name: "renamed", protocol: base, question: value }), + value + ) + const second = buildPlan( + fakeBenchmark({ name: "renamed", protocol: changedImplementation, question: value }), + value + ) + + expect(second.haystack.fingerprint).toBe(first.haystack.fingerprint) + expect(second.buildFingerprint).toBe(first.buildFingerprint) + }) + + test("an ingestion method change changes the build even when emitted bytes match", () => { + const value = beamQuestion() + const base = new BeamPaperProtocol({ retrievalTopK: 5 }) + const changedIngestion = new SameBytesChangedBeamIngestionProtocol({ retrievalTopK: 5 }) + const first = buildPlan( + fakeBenchmark({ name: "renamed", protocol: base, question: value }), + value + ) + const second = buildPlan( + fakeBenchmark({ name: "renamed", protocol: changedIngestion, question: value }), + value + ) + + expect(changedIngestion.identity.ingestionPolicyHash).not.toBe( + base.identity.ingestionPolicyHash + ) + expect(second.haystack.fingerprint).toBe(first.haystack.fingerprint) + expect(second.buildFingerprint).not.toBe(first.buildFingerprint) + expect(second.buildId).not.toBe(first.buildId) + }) +}) diff --git a/test/provider-build-lifecycle.test.ts b/test/provider-build-lifecycle.test.ts new file mode 100644 index 0000000..953201f --- /dev/null +++ b/test/provider-build-lifecycle.test.ts @@ -0,0 +1,158 @@ +import { describe, expect, test } from "bun:test" +import { SupermemoryProvider, classifySupermemoryReadiness } from "../src/providers/supermemory" +import { ZepProvider } from "../src/providers/zep" +import { HybridSearchEngine, type Chunk } from "../src/providers/rag/search" + +function chunk(id: string, sessionId: string, chunkIndex: number, content: string): Chunk { + return { + id, + sessionId, + chunkIndex, + content, + embedding: [1, 0], + } +} + +describe("provider build lifecycle safety", () => { + test("RAG retry replaces a complete session and removes stale trailing chunks", () => { + const engine = new HybridSearchEngine() + engine.addChunks("container", [ + chunk("session-1-0", "session-1", 0, "old zero"), + chunk("session-1-1", "session-1", 1, "stale trailing chunk"), + chunk("session-2-0", "session-2", 0, "unrelated session"), + ]) + + engine.replaceSessionChunks( + "container", + ["session-1"], + [chunk("session-1-0", "session-1", 0, "new only chunk")] + ) + + expect(engine.getChunks("container").map(({ id, content }) => ({ id, content }))).toEqual([ + { id: "session-2-0", content: "unrelated session" }, + { id: "session-1-0", content: "new only chunk" }, + ]) + }) + + test("Supermemory indexing has a bounded timeout", async () => { + const provider = new SupermemoryProvider(0) + ;(provider as unknown as { client: object }).client = {} + + await expect( + provider.awaitIndexing({ documentIds: ["document-1"] }, "container") + ).rejects.toThrow("Supermemory indexing timed out") + }) + + test("Supermemory indexing times out when an SDK polling request hangs", async () => { + const provider = new SupermemoryProvider(20) + ;( + provider as unknown as { + client: { + documents: { get: () => Promise } + } + } + ).client = { + documents: { get: () => new Promise(() => {}) }, + } + + await expect( + provider.awaitIndexing({ documentIds: ["document-1"] }, "container") + ).rejects.toThrow("Supermemory indexing timed out") + }) + + test("Supermemory readiness requires both document processing and dreaming", () => { + expect(classifySupermemoryReadiness({ status: "queued" })).toBe("pending") + expect(classifySupermemoryReadiness({ status: "done" })).toBe("pending") + expect(classifySupermemoryReadiness({ status: "done", dreamingStatus: "dreaming" })).toBe( + "pending" + ) + expect(classifySupermemoryReadiness({ status: "done", dreamingStatus: "done" })).toBe( + "completed" + ) + expect(classifySupermemoryReadiness({ status: "failed", dreamingStatus: "done" })).toBe( + "failed" + ) + }) + + test("Supermemory polls one document endpoint until dreaming is complete", async () => { + const responses = [ + { status: "done", dreamingStatus: "dreaming" }, + { status: "done", dreamingStatus: "done" }, + ] + const documentGetCalls: string[] = [] + const provider = new SupermemoryProvider(100, 0) + ;( + provider as unknown as { + client: { documents: { get: (id: string) => Promise<(typeof responses)[number]> } } + } + ).client = { + documents: { + get: async (id: string) => { + documentGetCalls.push(id) + return responses.shift() ?? { status: "done", dreamingStatus: "done" } + }, + }, + } + + let finalProgress: unknown + await provider.awaitIndexing({ documentIds: ["document-1"] }, "container", (progress) => { + finalProgress = progress + }) + + expect(documentGetCalls).toEqual(["document-1", "document-1"]) + expect(finalProgress).toEqual({ + completedIds: ["document-1"], + failedIds: [], + total: 1, + }) + }) + + test("Zep indexing has a bounded timeout", async () => { + const provider = new ZepProvider(0) + ;(provider as unknown as { client: object }).client = {} + + await expect( + provider.awaitIndexing({ documentIds: ["episode-1"] }, "container") + ).rejects.toThrow("Zep indexing timed out") + }) + + test("Zep indexing times out when an SDK polling request hangs", async () => { + const provider = new ZepProvider(20) + ;( + provider as unknown as { + client: { + task: { get: () => Promise } + graph: { episode: { get: () => Promise } } + } + } + ).client = { + task: { get: () => new Promise(() => {}) }, + graph: { episode: { get: () => new Promise(() => {}) } }, + } + + await expect( + provider.awaitIndexing({ documentIds: ["episode-1"] }, "container") + ).rejects.toThrow("Zep indexing timed out") + }) + + test("provider build fingerprints are deterministic and include non-secret configuration", () => { + const defaultSupermemory = new SupermemoryProvider().getIngestionConfigFingerprint({ + apiKey: "ignored-secret-a", + baseUrl: "https://api.supermemory.ai", + }) + const samePublicConfig = new SupermemoryProvider().getIngestionConfigFingerprint({ + apiKey: "ignored-secret-b", + baseUrl: "https://api.supermemory.ai", + }) + const differentEndpoint = new SupermemoryProvider().getIngestionConfigFingerprint({ + apiKey: "ignored-secret-a", + baseUrl: "https://staging.example.test", + }) + + expect(samePublicConfig).toBe(defaultSupermemory) + expect(differentEndpoint).not.toBe(defaultSupermemory) + expect(new ZepProvider().getIngestionConfigFingerprint({ apiKey: "secret" })).toMatch( + /^[a-f0-9]{64}$/ + ) + }) +}) diff --git a/test/provider-normalization.test.ts b/test/provider-normalization.test.ts new file mode 100644 index 0000000..7d4bcfb --- /dev/null +++ b/test/provider-normalization.test.ts @@ -0,0 +1,788 @@ +import { describe, expect, test } from "bun:test" +import { mkdtemp, rm } from "node:fs/promises" +import { tmpdir } from "node:os" +import { join } from "node:path" +import type { CanonicalIngestionDocument } from "../src/types/unified" +import { + SupermemoryProvider, + normalizeSupermemorySearchResults, +} from "../src/providers/supermemory" +import { buildSupermemoryAnswerPrompt } from "../src/providers/supermemory/prompts" +import { Mem0Provider, normalizeMem0SearchResults } from "../src/providers/mem0" +import { buildMem0AnswerPrompt } from "../src/providers/mem0/prompts" +import { + ZepProvider, + allocateZepSearchBudget, + normalizeZepSearchResults, +} from "../src/providers/zep" +import { buildZepAnswerPrompt } from "../src/providers/zep/prompts" +import { normalizeFilesystemSearchResults } from "../src/providers/filesystem" +import { buildFilesystemAnswerPrompt } from "../src/providers/filesystem/prompts" +import { + loadPersistedRagChunks, + normalizeRagSearchResults, + persistRagChunks, +} from "../src/providers/rag" +import { HybridSearchEngine, type Chunk } from "../src/providers/rag/search" +import { buildRAGAnswerPrompt } from "../src/providers/rag/prompts" + +const DOCUMENT: CanonicalIngestionDocument = { + customId: "session-1", + content: "DOCUMENT_DATE: 2025-03-14T10:00:00Z\n\nUSER: I moved to Pune.\nASSISTANT: Got it.", + metadata: { + sessionId: "session-1", + documentDate: "2025-03-14T10:00:00Z", + }, + messages: [ + { role: "user", content: "I moved to Pune." }, + { role: "assistant", content: "Got it." }, + ], +} + +describe("Supermemory provider boundary", () => { + test("normalizes both memory and singular chunk results", () => { + const results = normalizeSupermemorySearchResults( + [ + { + id: "memory-1", + memory: "Vedant moved to Pune.", + similarity: 0.91, + metadata: { + sessionId: "session-1", + temporalContext: { documentDate: "2025-03-14T10:00:00Z" }, + }, + }, + { + id: "chunk-1", + chunk: "USER: I moved to Pune.", + similarity: 0.82, + metadata: null, + documents: [ + { + metadata: { + sessionId: "session-2", + documentDate: "2025-03-15T11:00:00Z", + }, + }, + ], + }, + ], + 2 + ) + + expect(results).toEqual([ + { + id: "memory-1", + rank: 1, + text: "Vedant moved to Pune.", + score: 0.91, + sessionId: "session-1", + documentDate: "2025-03-14T10:00:00Z", + provider: "supermemory", + resultType: "memory", + }, + { + id: "chunk-1", + rank: 2, + text: "USER: I moved to Pune.", + score: 0.82, + sessionId: "session-2", + documentDate: "2025-03-15T11:00:00Z", + provider: "supermemory", + resultType: "chunk", + }, + ]) + }) + + test("resolves each source field across all associated document metadata", () => { + const [result] = normalizeSupermemorySearchResults( + [ + { + id: "chunk-with-later-source-metadata", + chunk: "Normalized chunk text", + similarity: 0.77, + metadata: { providerLabel: "hybrid" }, + documents: [ + { metadata: { unrelated: "first metadata is non-empty" } }, + { metadata: { sessionId: "session-later" } }, + { metadata: { documentDate: "2025-04-02" } }, + ], + }, + ], + 1 + ) + + expect(result).toMatchObject({ + sessionId: "session-later", + documentDate: "2025-04-02", + }) + }) + + test("prefers result metadata and fails closed on ambiguous document fallback metadata", () => { + const [authoritative] = normalizeSupermemorySearchResults( + [ + { + id: "result-authoritative", + memory: "Authoritative metadata", + metadata: { sessionId: "result-session", documentDate: "2025-05-01" }, + documents: [ + { metadata: { sessionId: "document-a", documentDate: "2025-05-02" } }, + { metadata: { sessionId: "document-b", documentDate: "2025-05-03" } }, + ], + }, + ], + 1 + ) + expect(authoritative).toMatchObject({ + sessionId: "result-session", + documentDate: "2025-05-01", + }) + + expect(() => + normalizeSupermemorySearchResults( + [ + { + id: "ambiguous-documents", + chunk: "Ambiguous source", + documents: [ + { metadata: { sessionId: "session-a", documentDate: "2025-05-01" } }, + { metadata: { sessionId: "session-b", documentDate: "2025-05-02" } }, + ], + }, + ], + 1 + ) + ).toThrow("conflicting document sessionId") + }) + + test("sends the canonical document unchanged and honors zero threshold plus Top-K", async () => { + const addCalls: unknown[] = [] + const addRequestOptions: unknown[] = [] + const searchCalls: unknown[] = [] + const provider = new SupermemoryProvider() + const fakeClient = { + add: async (request: unknown, requestOptions: unknown) => { + addCalls.push(request) + addRequestOptions.push(requestOptions) + return { id: "document-1", status: "queued" } + }, + search: { + memories: async (request: unknown) => { + searchCalls.push(request) + return { + results: [ + { + id: "memory-1", + memory: "Vedant moved to Pune.", + similarity: 0.9, + metadata: DOCUMENT.metadata, + }, + ], + } + }, + }, + } + ;(provider as unknown as { client: typeof fakeClient }).client = fakeClient + + await provider.ingest([DOCUMENT], { containerTag: "beam_run" }) + const response = await provider.search("Where did Vedant move?", { + containerTag: "beam_run", + limit: 5, + threshold: 0, + searchMode: "hybrid", + }) + const results = response.results + + expect(addCalls).toEqual([ + { + content: DOCUMENT.content, + containerTag: "beam_run", + customId: DOCUMENT.customId, + metadata: DOCUMENT.metadata, + }, + ]) + expect(addRequestOptions).toEqual([{ idempotencyKey: expect.stringMatching(/^[a-f0-9]{64}$/) }]) + expect(searchCalls).toEqual([ + { + q: "Where did Vedant move?", + containerTag: "beam_run", + limit: 5, + threshold: 0, + searchMode: "hybrid", + include: { documents: true }, + rerank: false, + rewriteQuery: false, + }, + ]) + expect(searchCalls[0]).not.toHaveProperty("include.chunks") + expect(results).toHaveLength(1) + expect(response.diagnostics).toEqual({ + requestedLimit: 5, + providerRequests: [ + { + operation: "search.hybrid", + limit: 5, + parameters: { + searchMode: "hybrid", + threshold: 0, + includeDocuments: true, + includeChunks: false, + rerank: false, + rewriteQuery: false, + }, + }, + ], + rawReturnedCount: 1, + normalizedCount: 1, + droppedCount: 0, + droppedResults: [], + }) + }) + + test("requests immediate dreaming for a causal ingestion barrier", async () => { + const addCalls: unknown[] = [] + const provider = new SupermemoryProvider() + const fakeClient = { + add: async (request: unknown) => { + addCalls.push(request) + return { id: "document-1", status: "queued" } + }, + } + ;(provider as unknown as { client: typeof fakeClient }).client = fakeClient + + await provider.ingest([DOCUMENT], { + containerTag: "beam_run", + processingMode: "instant", + }) + + expect(addCalls).toEqual([ + { + content: DOCUMENT.content, + containerTag: "beam_run", + customId: DOCUMENT.customId, + metadata: DOCUMENT.metadata, + dreaming: "instant", + }, + ]) + }) + + test("uses the V3 batch endpoint for multiple ordered sessions", async () => { + const calls: Array<{ body: unknown; options: unknown }> = [] + const provider = new SupermemoryProvider() + const fakeClient = { + documents: { + batchAdd: async (body: unknown, options: unknown) => { + calls.push({ body, options }) + return { + results: [ + { id: "document-1", status: "queued" }, + { id: "document-2", status: "queued" }, + ], + failed: 0, + success: 2, + } + }, + }, + } + ;(provider as unknown as { client: typeof fakeClient }).client = fakeClient + const second: CanonicalIngestionDocument = { + ...DOCUMENT, + customId: "session-2", + content: "DOCUMENT_DATE: 2025-03-15T10:00:00Z\n\nUSER: I moved again.\nASSISTANT: Got it.", + metadata: { + sessionId: "session-2", + documentDate: "2025-03-15T10:00:00Z", + }, + } + + const result = await provider.ingest([DOCUMENT, second], { + containerTag: "beam_run", + processingMode: "instant", + }) + + expect(result.documentIds).toEqual(["document-1", "document-2"]) + expect(calls).toHaveLength(1) + expect(calls[0]?.body).toEqual({ + documents: [ + { + content: DOCUMENT.content, + customId: DOCUMENT.customId, + metadata: DOCUMENT.metadata, + }, + { + content: second.content, + customId: second.customId, + metadata: second.metadata, + }, + ], + containerTag: "beam_run", + dreaming: "instant", + }) + expect(calls[0]?.options).toMatchObject({ idempotencyKey: expect.any(String) }) + }) + + test("preserves successful batch items and attributes validation failures by custom ID", async () => { + const provider = new SupermemoryProvider() + const fakeClient = { + documents: { + batchAdd: async () => ({ + results: [ + { id: "document-1", status: "queued" }, + { id: "document-3", status: "queued" }, + { id: "session-2", status: "error", error: "invalid document" }, + ], + failed: 1, + success: 2, + }), + }, + } + ;(provider as unknown as { client: typeof fakeClient }).client = fakeClient + + const result = await provider.ingest( + [DOCUMENT, { ...DOCUMENT, customId: "session-2" }, { ...DOCUMENT, customId: "session-3" }], + { + containerTag: "beam_run", + processingMode: "instant", + } + ) + + expect(result).toEqual({ + documentIds: ["document-1", "document-3"], + items: [ + { customId: "session-1", documentIds: ["document-1"] }, + { customId: "session-2", documentIds: [], error: "invalid document" }, + { customId: "session-3", documentIds: ["document-3"] }, + ], + }) + }) + + test("fails closed when the backend exceeds the requested evidence budget", () => { + expect(() => + normalizeSupermemorySearchResults( + [ + { id: "one", memory: "one", similarity: 1 }, + { id: "two", memory: "two", similarity: 0.9 }, + ], + 1 + ) + ).toThrow("returned 2 results for requested Top-K 1") + }) + + test("propagates every supported BEAM ablation Top-K unchanged", async () => { + const requestLimits: number[] = [] + const provider = new SupermemoryProvider() + const fakeClient = { + search: { + memories: async (request: { limit: number }) => { + requestLimits.push(request.limit) + return { results: [] } + }, + }, + } + ;(provider as unknown as { client: typeof fakeClient }).client = fakeClient + + for (const limit of [5, 10, 15, 20]) { + const response = await provider.search("query", { + containerTag: "beam_run", + limit, + threshold: 0, + }) + expect(response.diagnostics.requestedLimit).toBe(limit) + expect(response.diagnostics.providerRequests.map((request) => request.limit)).toEqual([limit]) + } + + expect(requestLimits).toEqual([5, 10, 15, 20]) + }) +}) + +describe("Mem0 normalization", () => { + test("supports direct and nested v1.1 result shapes", () => { + const results = normalizeMem0SearchResults( + [ + { + id: "memory-1", + memory: "Direct memory", + score: 0.8, + metadata: { sessionId: "session-1", documentDate: "2025-01-01" }, + }, + { + data: { + id: "memory-2", + memory: "Nested memory", + score: 0.7, + metadata: { sessionId: "session-2", documentDate: "2025-01-02" }, + }, + }, + ], + 2 + ) + + expect( + results.map(({ id, rank, text, score, sessionId, documentDate }) => ({ + id, + rank, + text, + score, + sessionId, + documentDate, + })) + ).toEqual([ + { + id: "memory-1", + rank: 1, + text: "Direct memory", + score: 0.8, + sessionId: "session-1", + documentDate: "2025-01-01", + }, + { + id: "memory-2", + rank: 2, + text: "Nested memory", + score: 0.7, + sessionId: "session-2", + documentDate: "2025-01-02", + }, + ]) + }) + + test("reconciles by deterministic run_id before adding and uses synchronous ingestion", async () => { + const addCalls: Array<{ messages: unknown; options: Record }> = [] + let stored = false + const provider = new Mem0Provider() + const fakeClient = { + getAll: async (options: Record) => + stored ? [{ id: "memory-1", metadata: DOCUMENT.metadata, run_id: options.run_id }] : [], + add: async (messages: unknown, options: Record) => { + addCalls.push({ messages, options }) + stored = true + return [{ id: "memory-1" }] + }, + } + ;(provider as unknown as { client: typeof fakeClient }).client = fakeClient + + const first = await provider.ingest([DOCUMENT], { containerTag: "beam_run" }) + const resumed = await provider.ingest([DOCUMENT], { containerTag: "beam_run" }) + + expect(first.documentIds).toEqual(["memory-1"]) + expect(resumed.documentIds).toEqual(["memory-1"]) + expect(addCalls).toHaveLength(1) + expect(addCalls[0].options).toMatchObject({ + user_id: "beam_run", + run_id: DOCUMENT.customId, + async_mode: false, + enable_graph: false, + }) + expect(addCalls[0].messages).toEqual(DOCUMENT.messages) + + let progress: unknown + await provider.awaitIndexing(first, "beam_run", (value) => { + progress = value + }) + expect(progress).toEqual({ + completedIds: ["memory-1"], + failedIds: [], + total: 1, + }) + }) +}) + +describe("Zep normalization and shared retrieval budget", () => { + test("splits one total Top-K budget between edges and nodes", () => { + for (const limit of [1, 5, 10, 20]) { + const budget = allocateZepSearchBudget(limit) + expect(budget.edgeLimit + budget.nodeLimit).toBe(limit) + } + expect(allocateZepSearchBudget(5)).toEqual({ edgeLimit: 3, nodeLimit: 2 }) + }) + + test("fails closed when ontology setup fails and retries setup on the next ingest", async () => { + let ontologyAttempts = 0 + let addAttempts = 0 + const provider = new ZepProvider() + const fakeClient = { + graph: { + create: async () => ({}), + setOntology: async () => { + ontologyAttempts++ + if (ontologyAttempts === 1) throw new Error("ontology unavailable") + return {} + }, + episode: { + getByGraphId: async () => ({ episodes: [] }), + }, + add: async () => { + addAttempts++ + return { uuid: "episode-1" } + }, + }, + } + ;(provider as unknown as { client: typeof fakeClient }).client = fakeClient + + await expect(provider.ingest([DOCUMENT], { containerTag: "beam_run" })).rejects.toThrow( + "ontology unavailable" + ) + expect(ontologyAttempts).toBe(1) + expect(addAttempts).toBe(0) + + const retried = await provider.ingest([DOCUMENT], { containerTag: "beam_run" }) + + expect(retried.documentIds).toEqual(["episode-1"]) + expect(ontologyAttempts).toBe(2) + expect(addAttempts).toBe(1) + }) + + test("reconciles deterministic episodes and preserves canonical transcript order", async () => { + const episodes: Array<{ + uuid: string + sourceDescription: string + content: string + processed: boolean + createdAt: string + }> = [] + const addCalls: unknown[] = [] + const provider = new ZepProvider() + const fakeClient = { + graph: { + create: async () => ({}), + setOntology: async () => ({}), + episode: { + getByGraphId: async () => ({ episodes }), + get: async (uuid: string) => episodes.find((episode) => episode.uuid === uuid)!, + }, + add: async (request: { data: string; sourceDescription: string; createdAt: string }) => { + addCalls.push(request) + const episode = { + uuid: `episode-${episodes.length + 1}`, + sourceDescription: request.sourceDescription, + content: request.data, + processed: true, + createdAt: request.createdAt, + } + episodes.push(episode) + return episode + }, + }, + task: { get: async () => ({ status: "completed" }) }, + } + ;(provider as unknown as { client: typeof fakeClient }).client = fakeClient + + const first = await provider.ingest([DOCUMENT], { containerTag: "beam_run" }) + const resumed = await provider.ingest([DOCUMENT], { containerTag: "beam_run" }) + + expect(first.documentIds).toEqual(["episode-1"]) + expect(resumed.documentIds).toEqual(["episode-1"]) + expect(addCalls).toHaveLength(1) + expect(addCalls[0]).toMatchObject({ + data: DOCUMENT.content, + sourceDescription: `memorybench:${DOCUMENT.customId}:1/1`, + createdAt: "2025-03-14T10:00:00Z", + }) + + let progress: unknown + await provider.awaitIndexing(first, "beam_run", (value) => { + progress = value + }) + expect(progress).toEqual({ + completedIds: ["episode-1"], + failedIds: [], + total: 1, + }) + }) + + test("requests at most K combined graph results and normalizes both types", async () => { + const searchCalls: Array<{ limit: number; scope: string }> = [] + const provider = new ZepProvider() + const fakeClient = { + graph: { + search: async (request: { limit: number; scope: string }) => { + searchCalls.push(request) + if (request.scope === "edges") { + return { + edges: Array.from({ length: request.limit }, (_, index) => ({ + uuid: `edge-${index}`, + fact: `Fact ${index}`, + relevance: 0.9 - index * 0.1, + })), + } + } + return { + nodes: Array.from({ length: request.limit }, (_, index) => ({ + uuid: `node-${index}`, + name: `Entity ${index}`, + summary: `Summary ${index}`, + relevance: 0.6 - index * 0.1, + })), + } + }, + }, + } + const internals = provider as unknown as { + client: typeof fakeClient + graphIds: Map + } + internals.client = fakeClient + internals.graphIds.set("beam_run", "memorybench_beam_run") + + const response = await provider.search("query", { containerTag: "beam_run", limit: 5 }) + const results = response.results + + expect(searchCalls.map(({ limit, scope }) => ({ limit, scope }))).toEqual([ + { limit: 3, scope: "edges" }, + { limit: 2, scope: "nodes" }, + ]) + expect(results).toHaveLength(5) + expect(results.map((result) => result.rank)).toEqual([1, 2, 3, 4, 5]) + expect(new Set(results.map((result) => result.resultType))).toEqual( + new Set(["graph-edge", "graph-node"]) + ) + expect(response.diagnostics.providerRequests).toEqual([ + { + operation: "graph.edges", + limit: 3, + parameters: { scope: "edges", reranker: "cross_encoder" }, + }, + { + operation: "graph.nodes", + limit: 2, + parameters: { scope: "nodes", reranker: "cross_encoder" }, + }, + ]) + }) + + test("orders mixed graph evidence by the common score", () => { + const results = normalizeZepSearchResults( + [ + { _type: "node", uuid: "node-1", name: "Vedant", summary: "Lives in Pune", relevance: 0.7 }, + { _type: "edge", uuid: "edge-1", fact: "Vedant moved to Pune", relevance: 0.9 }, + ], + 2 + ) + expect(results.map((result) => result.id)).toEqual(["edge-1", "node-1"]) + expect(results.map((result) => result.score)).toEqual([0.9, 0.7]) + }) +}) + +describe("local provider normalization", () => { + test("filesystem preserves canonical session and document date", () => { + expect( + normalizeFilesystemSearchResults( + [ + { + id: "document-1", + sessionId: "session-1", + content: "Stored memory", + score: 0.75, + documentDate: "2025-02-01T12:00:00Z", + }, + ], + 1 + ) + ).toEqual([ + { + id: "document-1", + rank: 1, + text: "Stored memory", + score: 0.75, + sessionId: "session-1", + documentDate: "2025-02-01T12:00:00Z", + provider: "filesystem", + resultType: "document", + }, + ]) + }) + + test("RAG preserves exact dates and omits fake unknown dates", () => { + const base = { + vectorScore: 0.8, + bm25Score: 0.7, + chunkIndex: 0, + } + const results = normalizeRagSearchResults( + [ + { + ...base, + id: "chunk-1", + content: "Dated chunk", + score: 0.9, + sessionId: "session-1", + date: "2025-03-14T10:00:00Z", + }, + { + ...base, + id: "chunk-2", + content: "Undated chunk", + score: 0.8, + sessionId: "session-2", + date: "unknown", + }, + ], + 2 + ) + + expect(results[0].documentDate).toBe("2025-03-14T10:00:00Z") + expect(results[1]).not.toHaveProperty("documentDate") + expect(results).toHaveLength(2) + }) + + test("RAG persists reusable indexes and deterministic chunk retries are idempotent", async () => { + const root = await mkdtemp(join(tmpdir(), "memorybench-rag-index-")) + const containerTag = "mb:durable-rag" + const initial: Chunk = { + id: "chunk-1", + content: "Vedant moved to Pune", + sessionId: "session-1", + chunkIndex: 0, + embedding: [1, 0], + date: "2025-03-14", + metadata: { documentDate: "2025-03-14" }, + } + try { + const firstProcess = new HybridSearchEngine() + firstProcess.addChunks(containerTag, [initial]) + await persistRagChunks(root, containerTag, firstProcess.getChunks(containerTag)) + + const persisted = await loadPersistedRagChunks(root, containerTag) + expect(persisted).toEqual([initial]) + const resumedProcess = new HybridSearchEngine() + resumedProcess.replaceChunks(containerTag, persisted!) + resumedProcess.addChunks(containerTag, [{ ...initial, content: "Vedant moved to Mumbai" }]) + + expect(resumedProcess.getChunkCount(containerTag)).toBe(1) + expect(resumedProcess.search(containerTag, [1, 0], "Mumbai", 5)).toHaveLength(1) + expect(resumedProcess.getChunks(containerTag)[0].content).toBe("Vedant moved to Mumbai") + } finally { + await rm(root, { recursive: true, force: true }) + } + }) +}) + +describe("provider prompt safety", () => { + test("renders only normalized evidence fields and never stringifies raw payloads", () => { + const evidence = [ + { + id: "result-1", + rank: 1, + text: "Safe normalized evidence", + score: 0.9, + sessionId: "session-1", + documentDate: "2025-03-14", + provider: "supermemory", + resultType: "memory", + rawPayload: { secret: "RAW_DEBUG_PAYLOAD_MUST_NOT_LEAK" }, + }, + ] + const prompts = [ + buildSupermemoryAnswerPrompt("Question?", evidence), + buildMem0AnswerPrompt("Question?", evidence), + buildZepAnswerPrompt("Question?", evidence), + buildFilesystemAnswerPrompt("Question?", evidence), + buildRAGAnswerPrompt("Question?", evidence), + ] + + for (const prompt of prompts) { + expect(prompt).toContain("Safe normalized evidence") + expect(prompt).not.toContain("RAW_DEBUG_PAYLOAD_MUST_NOT_LEAK") + } + }) +}) diff --git a/test/report.test.ts b/test/report.test.ts new file mode 100644 index 0000000..af4f86e --- /dev/null +++ b/test/report.test.ts @@ -0,0 +1,567 @@ +import { describe, expect, test } from "bun:test" +import type { Benchmark } from "../src/types/benchmark" +import { CHECKPOINT_SCHEMA_VERSION, type RunCheckpoint } from "../src/types/checkpoint" +import type { BenchmarkProtocol, ProtocolIdentity } from "../src/types/protocol" +import type { UnifiedQuestion } from "../src/types/unified" +import { generateReport } from "../src/orchestrator/phases/report" +import { stableSha256 } from "../src/utils/stable" + +const identity: ProtocolIdentity = { + id: "test.continuous-score", + version: "1.0.0", + configFingerprint: "config", + implementationFingerprint: "implementation", + ingestionPolicyHash: "ingestion", + retrievalPolicyHash: "retrieval", + answerPromptHash: "answer", + evaluatorHash: "evaluator", + aggregationHash: "aggregation", +} + +function questions(): UnifiedQuestion[] { + return [0, 1, 2].map((index) => ({ + questionId: `q${index}`, + question: `Question ${index}`, + questionType: "ability", + groundTruth: `Ground truth ${index}`, + haystackSessionIds: ["s1"], + })) +} + +function reportProtocol(): BenchmarkProtocol { + return { + identity, + auxiliaryRetrievalEvaluation: "disabled", + ingestionExecutionPolicy: { + readinessBarrier: "after-build", + processingMode: "provider-default", + }, + validateQuestion() {}, + createIngestionPlan() { + return [] + }, + createRetrievalPlan({ question }) { + return { query: question.question, requestedTopK: 5, answerCutoff: 5 } + }, + createAnswerPlan() { + return { + request: { prompt: "answer" }, + baseRequest: { prompt: "answer" }, + answerEvidenceCount: 0, + } + }, + async evaluateQuestion({ question }) { + return { + questionId: question.questionId, + questionType: question.questionType, + primaryScore: 0, + passed: false, + explanation: "unused", + } + }, + aggregateQuality({ evaluations }) { + const averageScore = + evaluations.length === 0 + ? 0 + : evaluations.reduce((sum, evaluation) => sum + evaluation.primaryScore, 0) / + evaluations.length + const passed = evaluations.filter((evaluation) => evaluation.passed).length + const passAccuracy = evaluations.length === 0 ? 0 : passed / evaluations.length + return { + primaryMetric: { + key: "continuous_score", + value: averageScore, + higherIsBetter: true, + }, + metrics: { averageScore, passAccuracy, passed, total: evaluations.length }, + } + }, + } +} + +function benchmark(values: UnifiedQuestion[]): Benchmark { + return { + name: "test-benchmark", + scope: { displayName: "Test benchmark", includedTiers: ["test"], coverage: "full" }, + protocol: reportProtocol(), + async load() {}, + getQuestions() { + return values + }, + getHaystackSessions() { + return [] + }, + getGroundTruth(questionId) { + return values.find((question) => question.questionId === questionId)?.groundTruth ?? "" + }, + getQuestionTypes() { + return { + ability: { id: "ability", alias: "ability", description: "Test ability" }, + } + }, + } +} + +function checkpoint(values: UnifiedQuestion[]): RunCheckpoint { + const scores = [0, 0.5, 1] + const searchDurations = [10, 20, 0] + const answerDurations = [20, 20, 30] + const checkpointQuestions: RunCheckpoint["questions"] = {} + for (let index = 0; index < values.length; index++) { + const question = values[index] + const score = scores[index] + checkpointQuestions[question.questionId] = { + questionId: question.questionId, + buildId: "build-1", + question: question.question, + groundTruth: question.groundTruth, + questionType: question.questionType, + phases: { + search: { + status: "completed", + results: [], + retrievalPlan: { + query: question.question, + requestedTopK: 5, + answerCutoff: 5, + threshold: 0, + }, + requestedCount: 5, + returnedCount: 0, + normalizedCount: 0, + droppedCount: 0, + providerRequests: [{ operation: "fake.search", limit: 5 }], + answerCutoff: 5, + answerEvidenceCount: 0, + durationMs: searchDurations[index], + usage: { requestCount: 1, totalTokens: index + 1 }, + costUsd: 0.1, + }, + answer: { + status: "completed", + hypothesis: `Answer ${index}`, + durationMs: answerDurations[index], + contextTokens: index, + evidenceCount: 0, + usage: { requestCount: 1, totalTokens: index + 2 }, + costUsd: 0.2, + }, + evaluate: { + status: "completed", + durationMs: 1000, + usage: { + requestCount: 1, + ...(index === 0 + ? { tokenUsageUnknownRequestCount: 1 } + : { tokenUsageCompleteRequestCount: 1 }), + totalTokens: index + 3, + }, + costUsd: 0.4, + evaluation: { + questionId: question.questionId, + questionType: question.questionType, + primaryScore: score, + passed: score >= 0.5, + explanation: `Score ${score}`, + }, + }, + }, + } + } + + return { + schemaVersion: CHECKPOINT_SCHEMA_VERSION, + runId: "report-run", + dataSourceRunId: "report-run", + status: "completed", + provider: "fake", + providerAdapterVersion: "1", + providerPromptFingerprint: "fake-prompts", + benchmark: "test-benchmark", + benchmarkScope: { + displayName: "Test benchmark", + includedTiers: ["test"], + coverage: "full", + }, + selectedQuestionIdsDigest: stableSha256(values.map((question) => question.questionId)), + benchmarkInputFingerprint: "benchmark-input", + protocolIdentity: identity, + judge: "fake-judge", + answeringModel: "fake-answer", + answeringRuntimeIdentity: { + schemaVersion: 1, + transport: "ai-sdk-generate-text-v1", + modelAlias: "fake-answer", + provider: "openai", + modelId: "fake-answer", + supportsTemperature: true, + effectiveDefaultTemperature: 0, + effectiveDefaultMaxOutputTokens: 1000, + }, + retrievalTopK: 5, + createdAt: "2024-01-01T00:00:00.000Z", + updatedAt: "2024-01-01T00:00:02.000Z", + targetQuestionIds: values.map((question) => question.questionId), + buildPhaseAttempts: [ + { + startedAt: "2024-01-01T00:00:00.000Z", + completedAt: "2024-01-01T00:00:00.120Z", + durationMs: 120, + status: "completed", + }, + ], + builds: { + "build-1": { + buildId: "build-1", + ingestionGroupId: "chat-1", + memberQuestionIds: values.map((question) => question.questionId), + containerTag: "container-1", + haystack: { + schemaVersion: 2, + algorithm: "sha256", + fingerprint: "haystack", + orderedSessionIds: ["s1"], + sessionFingerprints: ["session"], + }, + buildFingerprint: "build-fingerprint", + providerIngestionConfigFingerprint: "provider-ingestion-config", + sessions: [{ sessionId: "s1", documentDate: "2024-01-01", messageCount: 2 }], + missingDocumentDateCount: 0, + reused: false, + ingest: { + status: "completed", + completedSessionIds: ["s1"], + documentIds: ["document-1"], + taskIds: [], + startedAt: "2024-01-01T00:00:00.000Z", + completedAt: "2024-01-01T00:00:00.100Z", + durationMs: 100, + attempts: [ + { + phase: "ingest", + attempt: 1, + startedAt: "2024-01-01T00:00:00.000Z", + completedAt: "2024-01-01T00:00:00.100Z", + durationMs: 100, + status: "completed", + usage: { requestCount: 1, totalTokens: 10 }, + costUsd: 1, + }, + ], + }, + indexing: { + status: "completed", + completedIds: ["document-1"], + failedIds: [], + startedAt: "2024-01-01T00:00:00.100Z", + completedAt: "2024-01-01T00:00:00.150Z", + durationMs: 50, + attempts: [ + { + phase: "indexing", + attempt: 1, + startedAt: "2024-01-01T00:00:00.100Z", + completedAt: "2024-01-01T00:00:00.150Z", + durationMs: 50, + status: "completed", + usage: { requestCount: 2, totalTokens: 20 }, + costUsd: 2, + }, + ], + }, + }, + }, + questions: checkpointQuestions, + } +} + +describe("continuous quality reporting", () => { + test("reports the average score and >=0.5 pass accuracy independently", () => { + const values = questions() + const report = generateReport(benchmark(values), checkpoint(values)) + + expect(report.summary.totalQuestions).toBe(3) + expect(report.summary.averageScore).toBe(0.5) + expect(report.summary.correctCount).toBe(2) + expect(report.summary.accuracy).toBeCloseTo(2 / 3) + expect(report.quality.primaryMetric).toEqual({ + key: "continuous_score", + value: 0.5, + higherIsBetter: true, + }) + expect(report.quality.metrics.averageScore).toBe(0.5) + expect(report.quality.metrics.passAccuracy).toBeCloseTo(2 / 3) + expect(report.evaluations.map((evaluation) => evaluation.score)).toEqual([0, 0.5, 1]) + expect(report.evaluations.map((evaluation) => evaluation.primaryScore)).toEqual([0, 0.5, 1]) + expect(report.evaluations.map((evaluation) => evaluation.passed)).toEqual([false, true, true]) + expect(report.evaluations.map((evaluation) => evaluation.label)).toEqual([ + "incorrect", + "correct", + "correct", + ]) + }) + + test("refuses to score a selected question set with incomplete evaluation state", () => { + const values = questions() + const state = checkpoint(values) + state.questions[values[0]!.questionId]!.phases.evaluate = { status: "pending" } + + expect(() => generateReport(benchmark(values), state)).toThrow( + "Cannot generate a scored report with incomplete evaluations" + ) + }) + + test("never turns a protocol-declared failure into a pass because its score is 1", () => { + const values = questions() + const state = checkpoint(values) + state.questions.q2.phases.evaluate.evaluation!.passed = false + + const report = generateReport(benchmark(values), state) + expect(report.evaluations.at(-1)).toMatchObject({ + primaryScore: 1, + passed: false, + label: "incorrect", + }) + expect(report.summary.correctCount).toBe(1) + expect(report.quality.metrics.passAccuracy).toBeCloseTo(1 / 3) + }) +}) + +describe("build and query metrics", () => { + test("charges one shared build once and keeps evaluation outside online latency", () => { + const values = questions() + const report = generateReport(benchmark(values), checkpoint(values)) + + expect(report.builds.uniqueBuildCount).toBe(1) + expect(report.builds.sumContainerBuildWorkMs).toBe(150) + expect(report.builds.buildPhaseWallClockMs).toBe(120) + expect(report.builds.totalBuildCostUsd).toBe(3) + expect(report.costs.query).toMatchObject({ knownCostCount: 3, totalCostCount: 3 }) + expect(report.costs.query.totalCostUsd).toBeCloseTo(0.9) + expect(report.costs.evaluation).toMatchObject({ knownCostCount: 3, totalCostCount: 3 }) + expect(report.costs.evaluation.totalCostUsd).toBeCloseTo(1.2) + expect(report.builds.items[0]).toMatchObject({ + ingestLatencyMs: 100, + indexingLatencyMs: 50, + buildWorkMs: 150, + attemptCount: 2, + attempts: [ + { phase: "ingest", attempt: 1, durationMs: 100 }, + { phase: "indexing", attempt: 1, durationMs: 50 }, + ], + usage: { requestCount: 3, totalTokens: 30 }, + costUsd: 3, + sessionCount: 1, + documentCount: 1, + taskCount: 0, + completedIndexingCount: 1, + failedIndexingCount: 0, + }) + + expect(report.latency.ingest).toMatchObject({ count: 1, mean: 100 }) + expect(report.latency.indexing).toMatchObject({ count: 1, mean: 50 }) + expect(report.latency.total.count).toBe(3) + expect(report.latency.total.mean).toBeCloseTo((30 + 40 + 30) / 3) + expect(report.latency.evaluate).toMatchObject({ count: 3, mean: 1000 }) + expect(report.evaluations.map((evaluation) => evaluation.totalDurationMs)).toEqual([30, 40, 30]) + + expect(report.questionMetrics).toHaveLength(3) + expect(report.questionMetrics[0]).toMatchObject({ + configuredTopK: 5, + providerRequestLimit: 5, + rawReturnedCount: 0, + returnedCount: 0, + normalizedCount: 0, + droppedCount: 0, + answerCutoff: 5, + answerEvidenceCount: 0, + contextTokens: 0, + threshold: 0, + providerRequests: [{ operation: "fake.search", limit: 5 }], + }) + for (const metrics of report.questionMetrics) { + expect(metrics.evaluationLatencyMs).toBe(1000) + expect(metrics.buildAllocationQuestionCount).toBe(3) + expect(metrics.allocatedBuildWorkMs).toBe(50) + expect(metrics.amortizedOnlinePlusBuildWorkMs).toBe(metrics.onlineQueryLatencyMs + 50) + expect(metrics.queryCostUsd).toBeCloseTo(0.3) + expect(metrics.evaluationCostUsd).toBe(0.4) + expect(metrics.queryUsage?.requestCount).toBe(2) + expect(metrics.evaluationUsage?.requestCount).toBe(1) + } + expect(report.questionMetrics[0]?.evaluationUsage?.tokenUsageUnknownRequestCount).toBe(1) + expect(report.questionMetrics[1]?.evaluationUsage?.tokenUsageCompleteRequestCount).toBe(1) + }) + + test("reports query and evaluation totals only with complete cost coverage", () => { + const values = questions() + const state = checkpoint(values) + state.questions.q1.phases.search.costUsd = null + state.questions.q2.phases.evaluate.costUsd = null + + const report = generateReport(benchmark(values), state) + + expect( + report.questionMetrics.find((metrics) => metrics.questionId === "q1")?.queryCostUsd + ).toBeNull() + expect( + report.questionMetrics.find((metrics) => metrics.questionId === "q2")?.evaluationCostUsd + ).toBeNull() + expect(report.costs.query).toEqual({ + totalCostUsd: null, + knownCostCount: 2, + totalCostCount: 3, + }) + expect(report.costs.evaluation).toEqual({ + totalCostUsd: null, + knownCostCount: 2, + totalCostCount: 3, + }) + }) + + test("does not recharge a reused build", () => { + const values = questions() + const reusedCheckpoint = checkpoint(values) + reusedCheckpoint.builds["build-1"].reused = true + reusedCheckpoint.builds["build-1"].sourceRunId = "original-run" + const report = generateReport(benchmark(values), reusedCheckpoint) + + expect(report.builds.uniqueBuildCount).toBe(1) + expect(report.builds.sumContainerBuildWorkMs).toBe(0) + expect(report.builds.totalBuildCostUsd).toBeNull() + expect(report.builds.items[0]).toMatchObject({ + reused: true, + sourceRunId: "original-run", + ingestLatencyMs: 0, + indexingLatencyMs: 0, + buildWorkMs: 0, + attemptCount: 0, + costUsd: null, + }) + }) + + test("charges only indexing when a copied run reuses ingestion", () => { + const values = questions() + const partialCheckpoint = checkpoint(values) + const build = partialCheckpoint.builds["build-1"] + build.sourceRunId = "original-run" + build.reused = false + build.reusedPhases = { ingest: true, indexing: false } + + const report = generateReport(benchmark(values), partialCheckpoint) + expect(report.builds.sumContainerBuildWorkMs).toBe(50) + expect(report.builds.totalBuildCostUsd).toBe(2) + expect(report.builds.items[0]).toMatchObject({ + reused: false, + reusedPhases: { ingest: true, indexing: false }, + ingestLatencyMs: 0, + indexingLatencyMs: 50, + buildWallClockMs: 50, + buildWorkMs: 50, + attemptCount: 1, + attempts: [{ phase: "indexing", attempt: 1 }], + costUsd: 2, + }) + expect(report.latency.ingest.count).toBe(0) + expect(report.latency.indexing).toMatchObject({ count: 1, mean: 50 }) + }) + + test("retains zero build durations and reports unknown incurred cost with coverage", () => { + const values = questions() + const zeroState = checkpoint(values) + const zeroBuild = zeroState.builds["build-1"] + zeroBuild.ingest.durationMs = 0 + zeroBuild.ingest.completedAt = zeroBuild.ingest.startedAt + zeroBuild.ingest.attempts[0].durationMs = 0 + zeroBuild.ingest.attempts[0].completedAt = zeroBuild.ingest.attempts[0].startedAt + zeroBuild.indexing.durationMs = 0 + zeroBuild.indexing.startedAt = zeroBuild.ingest.startedAt + zeroBuild.indexing.completedAt = zeroBuild.ingest.startedAt + zeroBuild.indexing.attempts[0].startedAt = zeroBuild.ingest.startedAt! + zeroBuild.indexing.attempts[0].durationMs = 0 + zeroBuild.indexing.attempts[0].completedAt = zeroBuild.indexing.attempts[0].startedAt + zeroState.buildPhaseAttempts[0].durationMs = 0 + + const zeroReport = generateReport(benchmark(values), zeroState) + expect(zeroReport.builds.items[0]).toMatchObject({ + ingestLatencyMs: 0, + indexingLatencyMs: 0, + buildWorkMs: 0, + buildWallClockMs: 0, + }) + expect(zeroReport.latency.ingest).toMatchObject({ count: 1, mean: 0 }) + expect(zeroReport.latency.indexing).toMatchObject({ count: 1, mean: 0 }) + expect(zeroReport.builds.buildPhaseWallClockMs).toBe(0) + + const unknownCostState = checkpoint(values) + unknownCostState.builds["build-1"].ingest.attempts[0].costUsd = null + const unknownCostReport = generateReport(benchmark(values), unknownCostState) + expect(unknownCostReport.builds.items[0].costUsd).toBeNull() + expect(unknownCostReport.builds.totalBuildCostUsd).toBeNull() + expect(unknownCostReport.builds.knownCostBuildCount).toBe(0) + expect(unknownCostReport.builds.totalCostBuildCount).toBe(1) + }) + + test("includes failed retry attempts in build work and cost", () => { + const values = questions() + const state = checkpoint(values) + const build = state.builds["build-1"] + build.ingest.attempts.unshift({ + phase: "ingest", + attempt: 1, + startedAt: "2023-12-31T23:59:59.900Z", + completedAt: "2023-12-31T23:59:59.925Z", + durationMs: 25, + status: "failed", + costUsd: 0.5, + error: "transient failure", + }) + build.ingest.attempts[1].attempt = 2 + build.ingest.durationMs = 125 + + const report = generateReport(benchmark(values), state) + expect(report.builds.items[0]).toMatchObject({ + ingestLatencyMs: 125, + buildWorkMs: 175, + attemptCount: 3, + costUsd: 3.5, + attempts: [ + { phase: "ingest", attempt: 1, status: "failed", durationMs: 25 }, + { phase: "ingest", attempt: 2, status: "completed", durationMs: 100 }, + { phase: "indexing", attempt: 1, status: "completed", durationMs: 50 }, + ], + }) + }) + + test("reports concurrent container work separately from build-phase wall clock", () => { + const values = questions() + const state = checkpoint(values) + const firstBuild = state.builds["build-1"] + firstBuild.memberQuestionIds = values.slice(0, 2).map((value) => value.questionId) + const secondBuild = structuredClone(firstBuild) + secondBuild.buildId = "build-2" + secondBuild.ingestionGroupId = "chat-2" + secondBuild.containerTag = "container-2" + secondBuild.memberQuestionIds = [values[2].questionId] + state.builds[secondBuild.buildId] = secondBuild + state.questions[values[2].questionId].buildId = secondBuild.buildId + + const report = generateReport(benchmark(values), state) + expect(report.builds.uniqueBuildCount).toBe(2) + expect(report.builds.sumContainerBuildWorkMs).toBe(300) + expect(report.builds.buildPhaseWallClockMs).toBe(120) + expect(report.latency.ingest).toMatchObject({ count: 2, mean: 100 }) + expect(report.latency.indexing).toMatchObject({ count: 2, mean: 50 }) + expect( + report.questionMetrics.map((metrics) => ({ + questionId: metrics.questionId, + denominator: metrics.buildAllocationQuestionCount, + allocatedBuildWorkMs: metrics.allocatedBuildWorkMs, + })) + ).toEqual([ + { questionId: "q0", denominator: 2, allocatedBuildWorkMs: 75 }, + { questionId: "q1", denominator: 2, allocatedBuildWorkMs: 75 }, + { questionId: "q2", denominator: 1, allocatedBuildWorkMs: 150 }, + ]) + }) +}) diff --git a/test/retrieval-eval.test.ts b/test/retrieval-eval.test.ts new file mode 100644 index 0000000..978c48b --- /dev/null +++ b/test/retrieval-eval.test.ts @@ -0,0 +1,113 @@ +import { describe, expect, test } from "bun:test" +import type { JudgeInput, JudgeResult } from "../src/types/judge" +import type { EvaluationRuntime, StructuredModelRequest } from "../src/types/protocol" +import type { UnifiedSearchResult } from "../src/types/unified" +import { + calculateProtocolRetrievalMetrics, + calculateRetrievalMetrics, +} from "../src/orchestrator/phases/retrieval-eval" + +class RetrievalRuntime implements EvaluationRuntime { + readonly requests: StructuredModelRequest[] = [] + + constructor( + private readonly output: unknown, + private readonly failure?: Error + ) {} + + async evaluateLegacy(_input: JudgeInput): Promise { + throw new Error("Legacy answer evaluation is not expected") + } + + async generateStructured(request: StructuredModelRequest): Promise { + this.requests.push(request as StructuredModelRequest) + if (this.failure) throw this.failure + return request.schema.parse(this.output) + } +} + +function results(): UnifiedSearchResult[] { + return [ + { + id: "one", + rank: 1, + text: "Unrelated normalized text", + provider: "supermemory", + resultType: "memory", + rawArtifactRef: "debug/secret-raw-result.json", + }, + { + id: "two", + rank: 2, + text: "Vedant moved to Pune.", + score: 0.9, + sessionId: "session-2", + documentDate: "2025-01-02", + provider: "supermemory", + resultType: "chunk", + }, + ] +} + +describe("protocol-owned retrieval relevance diagnostics", () => { + test("does not run the non-paper auxiliary judge for BEAM", async () => { + const runtime = new RetrievalRuntime(undefined, new Error("must not be called")) + const metrics = await calculateProtocolRetrievalMetrics( + "disabled", + runtime, + "Where did Vedant move?", + "Pune", + results(), + 2 + ) + + expect(metrics).toBeUndefined() + expect(runtime.requests).toHaveLength(0) + }) + + test("uses structured runtime output and only normalized prompt fields for legacy diagnostics", async () => { + const runtime = new RetrievalRuntime({ + results: [ + { id: "result_1", relevant: 0 }, + { id: "result_2", relevant: 1 }, + ], + }) + const metrics = await calculateRetrievalMetrics( + runtime, + "Where did Vedant move?", + "Pune", + results(), + 2 + ) + + expect(metrics).toMatchObject({ + hitAtK: 1, + precisionAtK: 0.5, + recallAtK: 1, + mrr: 0.5, + k: 2, + relevantRetrieved: 1, + }) + expect(runtime.requests).toHaveLength(1) + expect(runtime.requests[0]?.schemaName).toBe("legacy_retrieval_relevance") + expect(runtime.requests[0]?.prompt).toContain("Vedant moved to Pune.") + expect(runtime.requests[0]?.prompt).not.toContain("secret-raw-result") + }) + + test("fails on mismatched IDs and propagates judge errors instead of fabricating zeros", async () => { + const mismatched = new RetrievalRuntime({ + results: [ + { id: "result_2", relevant: 0 }, + { id: "result_1", relevant: 1 }, + ], + }) + await expect( + calculateRetrievalMetrics(mismatched, "question", "answer", results(), 2) + ).rejects.toThrow("output ID mismatch") + + const failed = new RetrievalRuntime(undefined, new Error("judge unavailable")) + await expect( + calculateRetrievalMetrics(failed, "question", "answer", results(), 2) + ).rejects.toThrow("judge unavailable") + }) +}) diff --git a/test/run-list-identity.test.ts b/test/run-list-identity.test.ts new file mode 100644 index 0000000..ea428ec --- /dev/null +++ b/test/run-list-identity.test.ts @@ -0,0 +1,25 @@ +import { describe, expect, test } from "bun:test" +import type { RunCheckpoint } from "../src/types/checkpoint" +import { getRunListIdentity } from "../src/server/run-identity" + +describe("run-list identity", () => { + test("exposes benchmark scope and immutable dataset/protocol identity", () => { + const checkpoint = { + benchmarkScope: { displayName: "BEAM 1M", includedTiers: ["1M"], coverage: "subset" }, + datasetIdentity: { datasetFingerprint: "dataset-fingerprint" }, + benchmarkInputFingerprint: "benchmark-input-digest", + selectedQuestionIdsDigest: "question-set-digest", + protocolIdentity: { id: "beam-paper", version: "1.1.0" }, + providerPromptFingerprint: "provider-prompt-digest", + } as RunCheckpoint + + expect(getRunListIdentity(checkpoint)).toEqual({ + benchmarkScope: checkpoint.benchmarkScope, + datasetIdentity: checkpoint.datasetIdentity, + benchmarkInputFingerprint: "benchmark-input-digest", + selectedQuestionIdsDigest: "question-set-digest", + protocolIdentity: checkpoint.protocolIdentity, + providerPromptFingerprint: "provider-prompt-digest", + }) + }) +}) diff --git a/test/search-contract.test.ts b/test/search-contract.test.ts new file mode 100644 index 0000000..3fd47ed --- /dev/null +++ b/test/search-contract.test.ts @@ -0,0 +1,87 @@ +import { describe, expect, test } from "bun:test" +import { validateProviderSearchResponse } from "../src/orchestrator/phases/search" +import type { ProviderSearchRequestStructure, ProviderSearchResponse } from "../src/types/provider" + +function provider(name: string, searchRequestStructure: ProviderSearchRequestStructure) { + return { name, searchRequestStructure } +} + +function response(input?: { + provider?: string + rawReturnedCount?: number + normalizedCount?: number + droppedCount?: number + requestLimits?: number[] +}): ProviderSearchResponse { + const providerName = input?.provider ?? "fake" + const normalizedCount = input?.normalizedCount ?? 2 + const rawReturnedCount = input?.rawReturnedCount ?? 3 + return { + results: Array.from({ length: normalizedCount }, (_, index) => ({ + id: `result-${index + 1}`, + rank: index + 1, + text: `Evidence ${index + 1}`, + provider: providerName, + resultType: "memory", + })), + diagnostics: { + requestedLimit: 5, + providerRequests: (input?.requestLimits ?? [5]).map((limit, index) => ({ + operation: `search.${index + 1}`, + limit, + })), + rawReturnedCount, + normalizedCount, + droppedCount: input?.droppedCount ?? rawReturnedCount - normalizedCount, + droppedResults: Array.from( + { length: input?.droppedCount ?? rawReturnedCount - normalizedCount }, + (_, index) => ({ index: normalizedCount + index, reason: "empty-text" as const }) + ), + }, + } +} + +describe("search response contract", () => { + test("accepts valid normalization drops and rejects malformed drop bookkeeping", () => { + const adapter = provider("fake", { kind: "single" }) + + expect(() => validateProviderSearchResponse(response(), adapter, 5)).not.toThrow() + expect(() => validateProviderSearchResponse(response({ droppedCount: 0 }), adapter, 5)).toThrow( + "droppedCount 0 does not equal raw minus normalized count" + ) + }) + + test("rejects a single-call adapter that under-requests benchmark Top-K", () => { + expect(() => + validateProviderSearchResponse( + response({ requestLimits: [4] }), + provider("fake", { kind: "single" }), + 5 + ) + ).toThrow("single provider request limit 4 does not equal benchmark Top-K 5") + }) + + test("rejects a single-call adapter that over-requests benchmark Top-K", () => { + expect(() => + validateProviderSearchResponse( + response({ requestLimits: [6] }), + provider("fake", { kind: "single" }), + 5 + ) + ).toThrow("single provider request limit 6 does not equal benchmark Top-K 5") + }) + + test("accepts Zep's edge/node split only when the shared request budget equals Top-K", () => { + const zep = provider("zep", { kind: "split", budget: "shared-total" }) + + expect(() => + validateProviderSearchResponse(response({ provider: "zep", requestLimits: [3, 2] }), zep, 5) + ).not.toThrow() + expect(() => + validateProviderSearchResponse(response({ provider: "zep", requestLimits: [2, 2] }), zep, 5) + ).toThrow("split provider request limits total 4, expected benchmark Top-K 5") + expect(() => + validateProviderSearchResponse(response({ provider: "zep", requestLimits: [3, 3] }), zep, 5) + ).toThrow("split provider request limits total 6, expected benchmark Top-K 5") + }) +}) diff --git a/test/shared-build.test.ts b/test/shared-build.test.ts new file mode 100644 index 0000000..aa620b4 --- /dev/null +++ b/test/shared-build.test.ts @@ -0,0 +1,1508 @@ +import { afterEach, describe, expect, test } from "bun:test" +import { mkdir, mkdtemp, readFile, rm, writeFile } from "node:fs/promises" +import { tmpdir } from "node:os" +import { basename, join } from "node:path" +import type { Benchmark } from "../src/types/benchmark" +import type { BuildCheckpoint, RunCheckpoint } from "../src/types/checkpoint" +import type { BenchmarkProtocol, ProtocolIdentity } from "../src/types/protocol" +import type { + AwaitIndexingOptions, + IngestOptions, + Provider, + SearchOptions, +} from "../src/types/provider" +import type { + CanonicalIngestionDocument, + UnifiedQuestion, + UnifiedSearchResult, + UnifiedSession, +} from "../src/types/unified" +import { + assertCompletedSessionsAreOrderedPrefix, + assertBuildCheckpointConsistency, + assertResumeBuilds, + cloneCompletedBuildsForReuse, + createBuildCheckpoint, + prepareValidatedBuildPlans, +} from "../src/orchestrator/builds" +import { CheckpointManager } from "../src/orchestrator/checkpoint" +import { runIngestPhase } from "../src/orchestrator/phases/ingest" +import { runIndexingPhase } from "../src/orchestrator/phases/indexing" +import { runSearchPhase } from "../src/orchestrator/phases/search" + +const tempRoots: string[] = [] + +afterEach(async () => { + await Promise.all(tempRoots.splice(0).map((path) => rm(path, { recursive: true, force: true }))) +}) + +async function createTempRoot(): Promise { + const path = await mkdtemp(join(tmpdir(), "memorybench-shared-build-")) + tempRoots.push(path) + return path +} + +const protocolIdentity: ProtocolIdentity = { + id: "test.protocol", + version: "1.0.0", + configFingerprint: "config", + implementationFingerprint: "implementation", + ingestionPolicyHash: "ingestion", + retrievalPolicyHash: "retrieval", + answerPromptHash: "answer", + evaluatorHash: "evaluator", + aggregationHash: "aggregation", +} + +function question(questionId: string, haystackSessionIds = ["s1", "s2", "s3"]): UnifiedQuestion { + return { + questionId, + question: `Question ${questionId}`, + questionType: "test", + groundTruth: `Ground truth ${questionId}`, + haystackSessionIds, + } +} + +function sessions(): UnifiedSession[] { + return [ + { + sessionId: "s1", + messages: [ + { role: "user", content: "First user message" }, + { role: "assistant", content: "First assistant message" }, + ], + metadata: { documentDate: "2024-01-01" }, + }, + { + sessionId: "s2", + messages: [ + { role: "user", content: "Second user message" }, + { role: "assistant", content: "Second assistant message" }, + ], + metadata: { documentDate: "2024-01-02" }, + }, + { + sessionId: "s3", + messages: [ + { role: "user", content: "Third user message" }, + { role: "assistant", content: "Third assistant message" }, + ], + }, + ] +} + +function testProtocol(input?: { + planningCalls?: string[] + requestedTopK?: number + answerCutoff?: number + readinessBarrier?: "after-build" | "after-each-document" + processingMode?: "provider-default" | "instant" +}): BenchmarkProtocol { + const requestedTopK = input?.requestedTopK ?? 5 + const answerCutoff = input?.answerCutoff ?? requestedTopK + return { + identity: protocolIdentity, + auxiliaryRetrievalEvaluation: "disabled", + ingestionExecutionPolicy: { + readinessBarrier: input?.readinessBarrier ?? "after-build", + processingMode: input?.processingMode ?? "provider-default", + }, + validateQuestion(value) { + if (!value.questionId) throw new Error("missing question ID") + }, + createIngestionPlan({ question: value, sessions: sourceSessions }) { + input?.planningCalls?.push(value.questionId) + return sourceSessions.map((session): CanonicalIngestionDocument => { + const documentDate = + typeof session.metadata?.documentDate === "string" + ? session.metadata.documentDate + : undefined + const transcript = session.messages + .map((message) => `[${message.role.toUpperCase()}]\n${message.content}`) + .join("\n\n") + return { + customId: session.sessionId, + content: documentDate ? `DOCUMENT_DATE: ${documentDate}\n\n${transcript}` : transcript, + metadata: { + sessionId: session.sessionId, + ...(documentDate ? { documentDate } : {}), + }, + messages: session.messages, + } + }) + }, + createRetrievalPlan({ question: value }) { + return { + query: value.question, + requestedTopK, + answerCutoff, + threshold: 0.2, + searchMode: "hybrid", + } + }, + createAnswerPlan({ results }) { + return { + request: { prompt: "answer" }, + baseRequest: { prompt: "answer" }, + answerEvidenceCount: results.length, + } + }, + async evaluateQuestion({ question: value }) { + return { + questionId: value.questionId, + questionType: value.questionType, + primaryScore: 1, + passed: true, + explanation: "test", + } + }, + aggregateQuality({ evaluations }) { + const average = + evaluations.length === 0 + ? 0 + : evaluations.reduce((sum, evaluation) => sum + evaluation.primaryScore, 0) / + evaluations.length + return { + primaryMetric: { key: "score", value: average, higherIsBetter: true }, + metrics: { average }, + } + }, + } +} + +function testBenchmark(input: { + questions: UnifiedQuestion[] + sessionsByQuestion: Record + protocol?: BenchmarkProtocol + groupId?: string +}): Benchmark { + return { + name: "test-benchmark", + scope: { displayName: "Test", includedTiers: ["test"], coverage: "full" }, + protocol: input.protocol ?? testProtocol(), + async load() {}, + getQuestions() { + return input.questions + }, + getHaystackSessions(questionId) { + return input.sessionsByQuestion[questionId] ?? [] + }, + getGroundTruth(questionId) { + return input.questions.find((item) => item.questionId === questionId)?.groundTruth ?? "" + }, + getQuestionTypes() { + return { test: { id: "test", alias: "test", description: "Test" } } + }, + getIngestionGroupId() { + return input.groupId ?? "shared-chat" + }, + } +} + +function preparePlans(benchmark: Benchmark, questions: UnifiedQuestion[], ingestBatchSize = 1) { + return prepareValidatedBuildPlans({ + benchmark, + questions, + provider: "fake", + providerAdapterVersion: "1", + providerPromptFingerprint: "fake-prompts", + providerIngestionConfigFingerprint: "fake-ingestion-config", + dataSourceRunId: "source-run", + ingestBatchSize, + }) +} + +async function initializeCheckpoint(input: { + manager: CheckpointManager + plans: ReturnType + questions: UnifiedQuestion[] +}): Promise { + const checkpoint = input.manager.create( + "test-run", + "fake", + "test-benchmark", + "fake-judge", + "fake-answer", + { + providerAdapterVersion: "1", + providerPromptFingerprint: "fake-prompts", + benchmarkScope: { displayName: "Test", includedTiers: ["test"], coverage: "full" }, + protocolIdentity, + selectedQuestionIdsDigest: "selected", + benchmarkInputFingerprint: "benchmark-input", + retrievalTopK: 5, + concurrency: { default: 1 }, + } + ) + for (const plan of input.plans) input.manager.initBuild(checkpoint, createBuildCheckpoint(plan)) + for (const value of input.questions) { + const plan = input.plans.find((candidate) => + candidate.memberQuestionIds.includes(value.questionId) + )! + input.manager.initQuestion(checkpoint, value.questionId, plan.buildId, { + question: value.question, + groundTruth: value.groundTruth, + questionType: value.questionType, + }) + } + input.manager.save(checkpoint) + await input.manager.flush(checkpoint.runId) + return checkpoint +} + +class FakeProvider implements Provider { + name = "fake" + adapterVersion = "1" + searchRequestStructure = { kind: "single" } as const + concurrency = { default: 1 } + ingestAttempts: string[] = [] + successfulIngests: string[] = [] + indexingCalls = 0 + searchCalls: SearchOptions[] = [] + failOnceForSession?: string + failAlwaysForSession?: string + omitIndexingProgress = false + searchResults: UnifiedSearchResult[] = [] + searchRawReturnedCount?: number + + async initialize(): Promise {} + + getIngestionConfigFingerprint(): string { + return "fake-ingestion-config" + } + + async ingest(documents: CanonicalIngestionDocument[]) { + const sessionIds = documents.map((document) => document.metadata.sessionId) + this.ingestAttempts.push(...sessionIds) + const failedSession = sessionIds.find( + (sessionId) => + this.failAlwaysForSession === sessionId || this.failOnceForSession === sessionId + ) + if (failedSession) { + if (this.failOnceForSession === failedSession) this.failOnceForSession = undefined + throw new Error(`injected failure for ${failedSession}`) + } + this.successfulIngests.push(...sessionIds) + return { documentIds: sessionIds.map((sessionId) => `doc-${sessionId}`) } + } + + async awaitIndexing( + result: { documentIds: string[]; taskIds?: string[] }, + _containerTag: string, + onProgress?: (progress: { completedIds: string[]; failedIds: string[]; total: number }) => void + ) { + this.indexingCalls++ + if (this.omitIndexingProgress) return + const completedIds = [...result.documentIds, ...(result.taskIds ?? [])] + onProgress?.({ completedIds, failedIds: [], total: completedIds.length }) + } + + async search(_query: string, options: SearchOptions) { + this.searchCalls.push(options) + const rawReturnedCount = this.searchRawReturnedCount ?? this.searchResults.length + return { + results: this.searchResults, + diagnostics: { + requestedLimit: options.limit, + providerRequests: [{ operation: "search", limit: options.limit }], + rawReturnedCount, + normalizedCount: this.searchResults.length, + droppedCount: rawReturnedCount - this.searchResults.length, + droppedResults: Array.from( + { length: rawReturnedCount - this.searchResults.length }, + (_, offset) => ({ + index: this.searchResults.length + offset, + reason: "malformed-result" as const, + }) + ), + }, + } + } + + async clear(): Promise {} +} + +class CountingCheckpointManager extends CheckpointManager { + saveCalls = 0 + + override save(checkpoint: RunCheckpoint): void { + this.saveCalls++ + super.save(checkpoint) + } +} + +class CausalProvider extends FakeProvider { + events: string[] = [] + processingModes: Array = [] + omitReadinessOnceForSession?: string + readinessTimeouts: Array = [] + + override async ingest(documents: CanonicalIngestionDocument[], options: IngestOptions) { + const sessionIds = documents.map((document) => document.metadata.sessionId) + this.events.push(`add:${sessionIds.join(",")}`) + this.processingModes.push(options.processingMode) + return super.ingest(documents) + } + + override async awaitIndexing( + result: { documentIds: string[]; taskIds?: string[] }, + containerTag: string, + onProgress?: (progress: { completedIds: string[]; failedIds: string[]; total: number }) => void, + options?: AwaitIndexingOptions + ) { + this.readinessTimeouts.push(options?.timeoutMs) + const sessionIds = result.documentIds.map((documentId) => documentId.replace(/^doc-/, "")) + this.events.push(`wait:${sessionIds.join(",")}`) + if (this.omitReadinessOnceForSession && sessionIds.includes(this.omitReadinessOnceForSession)) { + this.omitReadinessOnceForSession = undefined + this.indexingCalls++ + onProgress?.({ completedIds: [], failedIds: [], total: result.documentIds.length }) + return + } + await super.awaitIndexing(result, containerTag, onProgress) + this.events.push(`ready:${sessionIds.join(",")}`) + } +} + +class PartialBatchCausalProvider extends CausalProvider { + private returnedPartialFailure = false + + override async ingest(documents: CanonicalIngestionDocument[], options: IngestOptions) { + if (!this.returnedPartialFailure && documents.some((document) => document.customId === "s2")) { + this.returnedPartialFailure = true + const sessionIds = documents.map((document) => document.metadata.sessionId) + this.events.push(`add:${sessionIds.join(",")}`) + this.processingModes.push(options.processingMode) + this.ingestAttempts.push(...sessionIds) + this.successfulIngests.push(...sessionIds.filter((sessionId) => sessionId !== "s2")) + return { + documentIds: sessionIds + .filter((sessionId) => sessionId !== "s2") + .map((sessionId) => `doc-${sessionId}`), + items: documents.map((document) => + document.customId === "s2" + ? { customId: document.customId, documentIds: [], error: "invalid document" } + : { customId: document.customId, documentIds: [`doc-${document.customId}`] } + ), + } + } + return super.ingest(documents, options) + } +} + +class EventCheckpointManager extends CheckpointManager { + constructor( + root: string, + private readonly events: string[] + ) { + super(root) + } + + override recordIngestProgress( + checkpoint: RunCheckpoint, + buildId: string, + input: Parameters[2] + ): void { + this.events.push(`checkpoint:${input.sessionId}`) + super.recordIngestProgress(checkpoint, buildId, input) + } +} + +class ConcurrentCausalProvider extends CausalProvider { + firstSessionContainers = new Set() + private resolveBothFirstSessions!: () => void + private readonly bothFirstSessions = new Promise((resolve) => { + this.resolveBothFirstSessions = resolve + }) + + override async ingest(documents: CanonicalIngestionDocument[], options: IngestOptions) { + const result = await super.ingest(documents, options) + if (documents[0].metadata.sessionId === "s1") { + this.firstSessionContainers.add(options.containerTag) + if (this.firstSessionContainers.size === 2) this.resolveBothFirstSessions() + } + return result + } + + override async awaitIndexing( + result: { documentIds: string[]; taskIds?: string[] }, + containerTag: string, + onProgress?: (progress: { completedIds: string[]; failedIds: string[]; total: number }) => void + ) { + if (result.documentIds[0] === "doc-s1") { + await Promise.race([ + this.bothFirstSessions, + new Promise((_, reject) => + setTimeout(() => reject(new Error("independent builds did not run concurrently")), 250) + ), + ]) + } + return super.awaitIndexing(result, containerTag, onProgress) + } +} + +class CrashWindowProvider extends FakeProvider { + remoteDocuments = new Map() + remoteCreates = 0 + crashAfterFirstRemoteSuccess = true + + override async ingest(documents: CanonicalIngestionDocument[]) { + const sessionId = documents[0].metadata.sessionId + this.ingestAttempts.push(sessionId) + const existing = this.remoteDocuments.get(sessionId) + if (existing) return { documentIds: [existing] } + + const documentId = `doc-${sessionId}` + this.remoteDocuments.set(sessionId, documentId) + this.remoteCreates++ + if (this.crashAfterFirstRemoteSuccess) { + this.crashAfterFirstRemoteSuccess = false + throw new Error("simulated crash after remote success") + } + this.successfulIngests.push(sessionId) + return { documentIds: [documentId] } + } +} + +function markBuildIngested(build: BuildCheckpoint): void { + build.ingest.status = "completed" + build.ingest.completedSessionIds = [...build.haystack.orderedSessionIds] + build.ingest.documentIds = build.haystack.orderedSessionIds.map((id) => `doc-${id}`) +} + +function markBuildIndexed(build: BuildCheckpoint): void { + markBuildIngested(build) + build.indexing.status = "completed" + build.indexing.completedIds = [...build.ingest.documentIds, ...build.ingest.taskIds] + build.indexing.failedIds = [] +} + +describe("shared-build planning and identity", () => { + test("twenty independently planned questions create one build and twenty references", async () => { + const questions = Array.from({ length: 20 }, (_, index) => question(`q${index + 1}`)) + const planningCalls: string[] = [] + const sessionsByQuestion = Object.fromEntries( + questions.map((value) => [value.questionId, sessions()]) + ) + const benchmark = testBenchmark({ + questions, + sessionsByQuestion, + protocol: testProtocol({ planningCalls }), + }) + + const plans = preparePlans(benchmark, questions) + expect(planningCalls).toEqual(questions.map((value) => value.questionId)) + expect(plans).toHaveLength(1) + expect(plans[0].memberQuestionIds).toHaveLength(20) + expect(plans[0].documents).toHaveLength(3) + + const manager = new CheckpointManager(await createTempRoot()) + const checkpoint = await initializeCheckpoint({ manager, plans, questions }) + expect(Object.keys(checkpoint.builds)).toHaveLength(1) + expect(new Set(Object.values(checkpoint.questions).map((value) => value.buildId)).size).toBe(1) + }) + + test("distinct groups with identical haystacks receive distinct containers", () => { + const questions = [question("q1"), question("q2")] + const benchmark = testBenchmark({ + questions, + sessionsByQuestion: { q1: sessions(), q2: sessions() }, + }) + benchmark.getIngestionGroupId = (questionId) => questionId + + const plans = preparePlans(benchmark, questions) + expect(plans).toHaveLength(2) + expect(plans[0].buildId).not.toBe(plans[1].buildId) + expect(plans[0].containerTag).not.toBe(plans[1].containerTag) + }) + + test("validates declared IDs against raw sessions before protocol planning", () => { + const value = question("q1") + const base = testProtocol() + const droppingProtocol: BenchmarkProtocol = { + ...base, + createIngestionPlan(input) { + return base.createIngestionPlan(input).slice(0, -1) + }, + } + const benchmark = testBenchmark({ + questions: [value], + sessionsByQuestion: { q1: sessions() }, + protocol: droppingProtocol, + }) + expect(() => preparePlans(benchmark, [value])).toThrow("ordered ingestion plan") + + const undeclared = question("undeclared", ["s1", "s2"]) + expect(() => + preparePlans( + testBenchmark({ + questions: [undeclared], + sessionsByQuestion: { undeclared: sessions() }, + }), + [undeclared] + ) + ).toThrow("ordered sessions returned by getHaystackSessions()") + }) + + const mismatches: Array<{ + name: string + mutate: (value: UnifiedSession[]) => UnifiedSession[] + }> = [ + { + name: "content", + mutate(value) { + value[0].messages[0].content = "changed content" + return value + }, + }, + { + name: "role", + mutate(value) { + value[0].messages[0].role = "assistant" + return value + }, + }, + { + name: "date", + mutate(value) { + value[0].metadata = { documentDate: "2025-01-01" } + return value + }, + }, + { + name: "message order", + mutate(value) { + value[0].messages.reverse() + return value + }, + }, + { + name: "message speaker", + mutate(value) { + value[0].messages[0]!.speaker = "Different speaker" + return value + }, + }, + { + name: "message timestamp", + mutate(value) { + value[0].messages[0]!.timestamp = "2025-01-01T12:34:56Z" + return value + }, + }, + ] + + for (const mismatch of mismatches) { + test(`rejects grouped questions with a ${mismatch.name} mismatch`, () => { + const questions = [question("q1"), question("q2")] + const benchmark = testBenchmark({ + questions, + sessionsByQuestion: { q1: sessions(), q2: mismatch.mutate(sessions()) }, + }) + expect(() => preparePlans(benchmark, questions)).toThrow("different haystacks") + }) + } + + test("fingerprints structured messages even when their rendered transcripts collide", () => { + const questions = [question("q1"), question("q2")] + const first = sessions() + const second = sessions() + first[0].messages = [{ role: "user", content: "First\n\n[ASSISTANT]\nSecond" }] + second[0].messages = [ + { role: "user", content: "First" }, + { role: "assistant", content: "Second" }, + ] + + const benchmark = testBenchmark({ + questions, + sessionsByQuestion: { q1: first, q2: second }, + }) + expect(() => preparePlans(benchmark, questions)).toThrow("different haystacks") + }) + + test("rejects grouped questions with session order or session identity drift", () => { + const reorderedQuestions = [question("q1"), question("q2", ["s2", "s1", "s3"])] + const reordered = sessions() + reordered.splice(0, 2, reordered[1], reordered[0]) + expect(() => + preparePlans( + testBenchmark({ + questions: reorderedQuestions, + sessionsByQuestion: { q1: sessions(), q2: reordered }, + }), + reorderedQuestions + ) + ).toThrow("different haystacks") + + const changedId = sessions() + changedId[1].sessionId = "different-session" + const changedIdQuestions = [question("q1"), question("q2", ["s1", "different-session", "s3"])] + expect(() => + preparePlans( + testBenchmark({ + questions: changedIdQuestions, + sessionsByQuestion: { q1: sessions(), q2: changedId }, + }), + changedIdQuestions + ) + ).toThrow("different haystacks") + }) + + test("rejects duplicate and reordered declared haystack IDs", () => { + const duplicate = question("duplicate", ["s1", "s1", "s3"]) + expect(() => + preparePlans( + testBenchmark({ questions: [duplicate], sessionsByQuestion: { duplicate: sessions() } }), + [duplicate] + ) + ).toThrow("duplicate haystackSessionIds") + + const reordered = question("reordered", ["s2", "s1", "s3"]) + expect(() => + preparePlans( + testBenchmark({ questions: [reordered], sessionsByQuestion: { reordered: sessions() } }), + [reordered] + ) + ).toThrow("do not exactly match") + }) + + test("accepts only an ordered completed-session prefix", () => { + const value = question("q1") + const benchmark = testBenchmark({ questions: [value], sessionsByQuestion: { q1: sessions() } }) + const build = createBuildCheckpoint(preparePlans(benchmark, [value])[0]) + + for (const valid of [[], ["s1"], ["s1", "s2"], ["s1", "s2", "s3"]]) { + build.ingest.completedSessionIds = valid + expect(() => assertCompletedSessionsAreOrderedPrefix(build)).not.toThrow() + } + for (const invalid of [ + ["s2"], + ["s1", "s1"], + ["s1", "s3"], + ["s2", "s1"], + ["s1", "s2", "unknown"], + ]) { + build.ingest.completedSessionIds = invalid + expect(() => assertCompletedSessionsAreOrderedPrefix(build)).toThrow() + } + }) + + test("rejects nested haystack schema and algorithm drift", () => { + const value = question("q1") + const benchmark = testBenchmark({ questions: [value], sessionsByQuestion: { q1: sessions() } }) + const build = createBuildCheckpoint(preparePlans(benchmark, [value])[0]) + + const schemaDrift = structuredClone(build) + ;(schemaDrift.haystack as { schemaVersion: number }).schemaVersion = 1 + expect(() => assertBuildCheckpointConsistency(schemaDrift)).toThrow( + "unsupported haystack schema" + ) + + const algorithmDrift = structuredClone(build) + ;(algorithmDrift.haystack as { algorithm: string }).algorithm = "md5" + expect(() => assertBuildCheckpointConsistency(algorithmDrift)).toThrow( + "unsupported haystack algorithm" + ) + }) + + test("provider ingestion configuration changes the build but not the haystack", () => { + const value = question("q1") + const benchmark = testBenchmark({ questions: [value], sessionsByQuestion: { q1: sessions() } }) + const makePlan = (providerIngestionConfigFingerprint: string) => + prepareValidatedBuildPlans({ + benchmark, + questions: [value], + provider: "fake", + providerAdapterVersion: "1", + providerPromptFingerprint: "fake-prompts", + providerIngestionConfigFingerprint, + dataSourceRunId: "source-run", + })[0] + + const first = makePlan("config-a") + const second = makePlan("config-b") + expect(second.haystack.fingerprint).toBe(first.haystack.fingerprint) + expect(second.buildFingerprint).not.toBe(first.buildFingerprint) + expect(second.buildId).not.toBe(first.buildId) + }) + + test("non-ingested question fields do not change the haystack identity", () => { + const firstQuestion = question("q1") + firstQuestion.metadata = { rubric: ["first rubric"], difficulty: "easy" } + const secondQuestion = { + ...firstQuestion, + question: "Completely different probe", + groundTruth: "Different answer", + metadata: { rubric: ["different rubric"], difficulty: "hard" }, + } + const first = preparePlans( + testBenchmark({ questions: [firstQuestion], sessionsByQuestion: { q1: sessions() } }), + [firstQuestion] + )[0] + const second = preparePlans( + testBenchmark({ questions: [secondQuestion], sessionsByQuestion: { q1: sessions() } }), + [secondQuestion] + )[0] + + expect(second.haystack.fingerprint).toBe(first.haystack.fingerprint) + }) + + test("rejects tampered persisted haystack identity and question references on resume", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const questions = [question("q1"), question("q2")] + const benchmark = testBenchmark({ + questions, + sessionsByQuestion: { q1: sessions(), q2: sessions() }, + }) + const plans = preparePlans(benchmark, questions) + const checkpoint = await initializeCheckpoint({ manager, plans, questions }) + expect(() => assertResumeBuilds(checkpoint, plans)).not.toThrow() + + const tamperedHaystack = structuredClone(checkpoint) + Object.values(tamperedHaystack.builds)[0].haystack.orderedSessionIds[0] = "tampered" + expect(() => assertResumeBuilds(tamperedHaystack, plans)).toThrow( + "tampered haystack fingerprint" + ) + + const rewiredQuestion = structuredClone(checkpoint) + rewiredQuestion.questions.q1.buildId = "missing-build" + expect(() => assertResumeBuilds(rewiredQuestion, plans)).toThrow( + "invalid member question reference" + ) + + const ghostMember = structuredClone(checkpoint) + Object.values(ghostMember.builds)[0].memberQuestionIds.push("ghost-question") + expect(() => assertResumeBuilds(ghostMember, plans)).toThrow( + "invalid member question reference" + ) + + const wrongTargets = structuredClone(checkpoint) + wrongTargets.targetQuestionIds = ["q1"] + expect(() => assertResumeBuilds(wrongTargets, plans)).toThrow("targetQuestionIds do not match") + + const missingQuestion = structuredClone(checkpoint) + delete missingQuestion.questions.q2 + expect(() => assertResumeBuilds(missingQuestion, plans)).toThrow( + "invalid member question reference" + ) + }) + + test("rejects checkpoint phases marked complete before their underlying work", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const value = question("q1") + const benchmark = testBenchmark({ questions: [value], sessionsByQuestion: { q1: sessions() } }) + const plans = preparePlans(benchmark, [value]) + const checkpoint = await initializeCheckpoint({ manager, plans, questions: [value] }) + const build = Object.values(checkpoint.builds)[0] + build.ingest.status = "completed" + build.ingest.completedSessionIds = ["s1"] + expect(() => assertResumeBuilds(checkpoint, plans)).toThrow( + "marked ingested before every session completed" + ) + + markBuildIndexed(build) + checkpoint.questions.q1.phases.search.status = "completed" + expect(() => assertResumeBuilds(checkpoint, plans)).toThrow("incomplete completed-search state") + }) + + test("rejects a causal build that claims ingestion completed without readiness", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const value = question("q1") + const benchmark = testBenchmark({ + questions: [value], + sessionsByQuestion: { q1: sessions() }, + protocol: testProtocol({ + readinessBarrier: "after-each-document", + processingMode: "instant", + }), + }) + const plans = preparePlans(benchmark, [value]) + const checkpoint = await initializeCheckpoint({ manager, plans, questions: [value] }) + const build = Object.values(checkpoint.builds)[0] + markBuildIngested(build) + + expect(() => assertResumeBuilds(checkpoint, plans)).toThrow( + "completed causal ingestion without completing its per-document indexing barriers" + ) + }) +}) + +describe("shared-build lifecycle", () => { + test("clones only completed ingest/index builds for a clean query-time run", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const value = question("q1") + const benchmark = testBenchmark({ questions: [value], sessionsByQuestion: { q1: sessions() } }) + const plans = preparePlans(benchmark, [value]) + const source = await initializeCheckpoint({ manager, plans, questions: [value] }) + markBuildIndexed(Object.values(source.builds)[0]) + + const [reused] = cloneCompletedBuildsForReuse(source, plans) + + expect(reused.sourceRunId).toBe(source.runId) + expect(reused.reused).toBe(true) + expect(reused.reusedPhases).toEqual({ ingest: true, indexing: true }) + expect(reused.containerTag).toBe(Object.values(source.builds)[0].containerTag) + expect(reused).not.toBe(Object.values(source.builds)[0]) + }) + + test("refuses source-build reuse before indexing completes", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const value = question("q1") + const benchmark = testBenchmark({ questions: [value], sessionsByQuestion: { q1: sessions() } }) + const plans = preparePlans(benchmark, [value]) + const source = await initializeCheckpoint({ manager, plans, questions: [value] }) + + expect(() => cloneCompletedBuildsForReuse(source, plans)).toThrow("ingestion is incomplete") + }) + + test("stores build state once and questions own only query phases", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const questions = [question("q1"), question("q2")] + const benchmark = testBenchmark({ + questions, + sessionsByQuestion: { q1: sessions(), q2: sessions() }, + }) + const plans = preparePlans(benchmark, questions) + const checkpoint = await initializeCheckpoint({ manager, plans, questions }) + + expect(Object.keys(checkpoint.builds)).toHaveLength(1) + for (const value of Object.values(checkpoint.questions)) { + expect(Object.keys(value.phases).sort()).toEqual(["answer", "evaluate", "search"]) + expect("ingest" in value.phases).toBe(false) + expect("indexing" in value.phases).toBe(false) + expect("sessions" in value).toBe(false) + expect("containerTag" in value).toBe(false) + } + }) + + test("copied runs distinguish fully reused builds from reused ingestion", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const value = question("q1") + const benchmark = testBenchmark({ + questions: [value], + sessionsByQuestion: { q1: sessions() }, + }) + const plans = preparePlans(benchmark, [value]) + const source = await initializeCheckpoint({ manager, plans, questions: [value] }) + const sourceBuild = Object.values(source.builds)[0] + markBuildIndexed(sourceBuild) + manager.save(source) + await manager.flush(source.runId) + + const indexingCopy = manager.copyCheckpoint(source.runId, "copy-indexing", "indexing") + const indexingBuild = Object.values(indexingCopy.builds)[0] + expect(indexingBuild.reused).toBe(false) + expect(indexingBuild.reusedPhases).toEqual({ ingest: true, indexing: false }) + expect(indexingBuild.ingest.status).toBe("completed") + expect(indexingBuild.indexing.status).toBe("pending") + + const searchCopy = manager.copyCheckpoint(source.runId, "copy-search", "search") + const searchBuild = Object.values(searchCopy.builds)[0] + expect(searchBuild.reused).toBe(true) + expect(searchBuild.reusedPhases).toEqual({ ingest: true, indexing: true }) + expect(searchBuild.indexing.status).toBe("completed") + }) + + test("copied runs own their reused search artifact after the source is deleted", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const value = question("q1") + const benchmark = testBenchmark({ questions: [value], sessionsByQuestion: { q1: sessions() } }) + const plans = preparePlans(benchmark, [value]) + const source = await initializeCheckpoint({ manager, plans, questions: [value] }) + markBuildIndexed(Object.values(source.builds)[0]) + const sourceResultPath = manager.getQuestionResultsPath(source.runId, value.questionId) + await writeFile(sourceResultPath, JSON.stringify({ questionId: value.questionId })) + source.questions.q1.phases.search = { + status: "completed", + retrievalPlan: { + query: value.question, + requestedTopK: 5, + answerCutoff: 5, + threshold: 0.2, + searchMode: "hybrid", + }, + resultFile: sourceResultPath, + results: [], + requestedCount: 5, + rawReturnedCount: 0, + returnedCount: 0, + normalizedCount: 0, + droppedCount: 0, + droppedResults: [], + providerRequests: [{ operation: "search", limit: 5 }], + answerCutoff: 5, + } + manager.save(source) + await manager.flush(source.runId) + + const copy = manager.copyCheckpoint(source.runId, "answer-copy-with-results", "answer") + await manager.flush(copy.runId) + const copiedResultPath = manager.getQuestionResultsPath(copy.runId, value.questionId) + expect(copy.questions.q1.phases.search.resultFile).toBe(copiedResultPath) + manager.delete(source.runId) + expect(JSON.parse(await readFile(copiedResultPath, "utf8"))).toEqual({ questionId: "q1" }) + }) + + test("copy overrides require rerunning the phase that owns the changed model", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const value = question("q1") + const benchmark = testBenchmark({ questions: [value], sessionsByQuestion: { q1: sessions() } }) + const plans = preparePlans(benchmark, [value]) + const source = await initializeCheckpoint({ manager, plans, questions: [value] }) + markBuildIndexed(Object.values(source.builds)[0]) + source.answeringRuntimeIdentity = { + ...source.answeringRuntimeIdentity!, + modelId: "answer-runtime-used-by-source", + } + manager.save(source) + await manager.flush(source.runId) + + expect(() => + manager.copyCheckpoint(source.runId, "bad-report-judge", "report", { + judge: "different-judge", + }) + ).toThrow("rerun evaluate or an earlier phase") + expect(() => + manager.copyCheckpoint(source.runId, "bad-evaluate-answer", "evaluate", { + answeringModel: "different-answer-model", + }) + ).toThrow("rerun answer or an earlier phase") + + const evaluateCopy = manager.copyCheckpoint(source.runId, "evaluate-copy", "evaluate", { + judge: "different-judge", + }) + expect(evaluateCopy.judge).toBe("different-judge") + expect(evaluateCopy.answeringRuntimeIdentity).toEqual(source.answeringRuntimeIdentity) + expect(evaluateCopy.questions.q1.phases.evaluate.status).toBe("pending") + + const answerCopy = manager.copyCheckpoint(source.runId, "answer-copy", "answer", { + judge: "different-judge", + answeringModel: "different-answer-model", + }) + expect(answerCopy.judge).toBe("different-judge") + expect(answerCopy.answeringModel).toBe("different-answer-model") + expect(answerCopy.answeringRuntimeIdentity?.modelAlias).toBe("different-answer-model") + expect(answerCopy.answeringRuntimeIdentity).not.toEqual(source.answeringRuntimeIdentity) + expect(answerCopy.questions.q1.phases.answer.status).toBe("pending") + }) + + test("a chained copy records its immediate source while retaining data-source identity", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const value = question("q1") + const benchmark = testBenchmark({ questions: [value], sessionsByQuestion: { q1: sessions() } }) + const plans = preparePlans(benchmark, [value]) + const source = await initializeCheckpoint({ manager, plans, questions: [value] }) + markBuildIndexed(Object.values(source.builds)[0]) + manager.save(source) + await manager.flush(source.runId) + + const firstCopy = manager.copyCheckpoint(source.runId, "first-copy", "search") + await manager.flush(firstCopy.runId) + const secondCopy = manager.copyCheckpoint(firstCopy.runId, "second-copy", "search") + const build = Object.values(secondCopy.builds)[0] + expect(build.sourceRunId).toBe("first-copy") + expect(secondCopy.dataSourceRunId).toBe(source.dataSourceRunId) + }) + + test("refuses to reuse incomplete ingestion or indexing", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const value = question("q1") + const benchmark = testBenchmark({ questions: [value], sessionsByQuestion: { q1: sessions() } }) + const plans = preparePlans(benchmark, [value]) + const source = await initializeCheckpoint({ manager, plans, questions: [value] }) + const build = Object.values(source.builds)[0] + + expect(() => manager.copyCheckpoint(source.runId, "bad-index-copy", "indexing")).toThrow( + "incomplete ingestion" + ) + + markBuildIngested(build) + manager.save(source) + await manager.flush(source.runId) + expect(() => manager.copyCheckpoint(source.runId, "bad-search-copy", "search")).toThrow( + "not fully indexed" + ) + }) + + test("continues after a failed session and retries it at the end of the build", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const questions = [question("q1"), question("q2")] + const benchmark = testBenchmark({ + questions, + sessionsByQuestion: { q1: sessions(), q2: sessions() }, + }) + const plans = preparePlans(benchmark, questions) + const checkpoint = await initializeCheckpoint({ manager, plans, questions }) + const provider = new FakeProvider() + provider.failOnceForSession = "s2" + + await runIngestPhase(provider, checkpoint, manager, plans) + await manager.flush(checkpoint.runId) + expect(provider.ingestAttempts).toEqual(["s1", "s2", "s3", "s2"]) + expect(provider.successfulIngests).toEqual(["s1", "s3", "s2"]) + expect(provider.ingestAttempts.filter((sessionId) => sessionId === "s1")).toHaveLength(1) + const build = Object.values(checkpoint.builds)[0] + expect(build.ingest.completedSessionIds).toEqual(["s1", "s2", "s3"]) + expect(build.ingest.deferredSessions).toEqual([]) + + await runIndexingPhase(provider, checkpoint, manager) + await runIndexingPhase(provider, checkpoint, manager) + expect(provider.indexingCalls).toBe(1) + expect(build.indexing.status).toBe("completed") + }) + + test("persists an unresolved session, finishes the first pass, and resumes only its retry", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const value = question("q1") + const benchmark = testBenchmark({ questions: [value], sessionsByQuestion: { q1: sessions() } }) + const plans = preparePlans(benchmark, [value]) + const checkpoint = await initializeCheckpoint({ manager, plans, questions: [value] }) + const provider = new FakeProvider() + provider.failAlwaysForSession = "s2" + + await expect(runIngestPhase(provider, checkpoint, manager, plans)).rejects.toThrow( + "sessions across 1 builds still need retry" + ) + await manager.flush(checkpoint.runId) + const persisted = manager.load(checkpoint.runId)! + const failedBuild = Object.values(persisted.builds)[0] + expect(provider.ingestAttempts).toEqual(["s1", "s2", "s3", "s2"]) + expect(failedBuild.ingest.completedSessionIds).toEqual(["s1", "s2", "s3"]) + expect(failedBuild.ingest.deferredSessions).toEqual([ + expect.objectContaining({ + sequence: 1, + sessionId: "s2", + customId: "s2", + stage: "submission", + attempts: 2, + lastError: "injected failure for s2", + }), + ]) + expect(failedBuild.containerTag).toBe(plans[0]!.containerTag) + + provider.failAlwaysForSession = undefined + await runIngestPhase(provider, persisted, manager, plans) + expect(provider.ingestAttempts).toEqual(["s1", "s2", "s3", "s2", "s2"]) + expect(failedBuild.ingest.status).toBe("completed") + expect(failedBuild.ingest.deferredSessions).toEqual([]) + }) + + test("causal builds add, await readiness, and checkpoint every session in order", async () => { + const root = await createTempRoot() + const events: string[] = [] + const manager = new EventCheckpointManager(root, events) + const value = question("q1") + const benchmark = testBenchmark({ + questions: [value], + sessionsByQuestion: { q1: sessions() }, + protocol: testProtocol({ + readinessBarrier: "after-each-document", + processingMode: "instant", + }), + }) + const plans = preparePlans(benchmark, [value]) + const checkpoint = await initializeCheckpoint({ manager, plans, questions: [value] }) + const provider = new CausalProvider() + provider.events = events + + await runIngestPhase(provider, checkpoint, manager, plans) + await runIndexingPhase(provider, checkpoint, manager) + + expect(events).toEqual([ + "add:s1", + "wait:s1", + "ready:s1", + "checkpoint:s1", + "add:s2", + "wait:s2", + "ready:s2", + "checkpoint:s2", + "add:s3", + "wait:s3", + "ready:s3", + "checkpoint:s3", + ]) + expect(provider.processingModes).toEqual(["instant", "instant", "instant"]) + expect(provider.readinessTimeouts).toEqual([300_000, 300_000, 300_000]) + expect(provider.indexingCalls).toBe(3) + const build = Object.values(checkpoint.builds)[0] + expect(build.ingest.status).toBe("completed") + expect(build.indexing.status).toBe("completed") + expect(build.indexing.completedIds).toEqual(["doc-s1", "doc-s2", "doc-s3"]) + }) + + test("causal builds submit ordered session batches and wait between batches", async () => { + const root = await createTempRoot() + const events: string[] = [] + const manager = new EventCheckpointManager(root, events) + const value = question("q1") + const benchmark = testBenchmark({ + questions: [value], + sessionsByQuestion: { q1: sessions() }, + protocol: testProtocol({ + readinessBarrier: "after-each-document", + processingMode: "instant", + }), + }) + const plans = preparePlans(benchmark, [value], 2) + const checkpoint = await initializeCheckpoint({ manager, plans, questions: [value] }) + const provider = new CausalProvider() + provider.events = events + + await runIngestPhase(provider, checkpoint, manager, plans) + + expect(events).toEqual([ + "add:s1,s2", + "wait:s1,s2", + "ready:s1,s2", + "checkpoint:s1", + "checkpoint:s2", + "add:s3", + "wait:s3", + "ready:s3", + "checkpoint:s3", + ]) + expect(provider.processingModes).toEqual(["instant", "instant"]) + expect(provider.indexingCalls).toBe(2) + const build = Object.values(checkpoint.builds)[0] + expect(build.ingest.completedSessionIds).toEqual(["s1", "s2", "s3"]) + expect(build.indexing.completedIds).toEqual(["doc-s1", "doc-s2", "doc-s3"]) + }) + + test("a partial provider batch keeps successful IDs and retries only the failed session", async () => { + const root = await createTempRoot() + const events: string[] = [] + const manager = new EventCheckpointManager(root, events) + const value = question("q1") + const benchmark = testBenchmark({ + questions: [value], + sessionsByQuestion: { q1: sessions() }, + protocol: testProtocol({ + readinessBarrier: "after-each-document", + processingMode: "instant", + }), + }) + const plans = preparePlans(benchmark, [value], 3) + const checkpoint = await initializeCheckpoint({ manager, plans, questions: [value] }) + const provider = new PartialBatchCausalProvider() + provider.events = events + + await runIngestPhase(provider, checkpoint, manager, plans) + + expect(events).toEqual([ + "add:s1,s2,s3", + "wait:s1,s3", + "ready:s1,s3", + "checkpoint:s1", + "checkpoint:s2", + "checkpoint:s3", + "add:s2", + "wait:s2", + "ready:s2", + ]) + expect(provider.ingestAttempts).toEqual(["s1", "s2", "s3", "s2"]) + const build = Object.values(checkpoint.builds)[0] + expect(build.ingest.deferredSessions).toEqual([]) + expect(new Set(build.indexing.completedIds)).toEqual(new Set(["doc-s1", "doc-s2", "doc-s3"])) + }) + + test("ingest batch size changes build and container identity", () => { + const value = question("q1") + const benchmark = testBenchmark({ + questions: [value], + sessionsByQuestion: { q1: sessions() }, + }) + const single = preparePlans(benchmark, [value], 1)[0]! + const batched = preparePlans(benchmark, [value], 5)[0]! + + expect(single.buildFingerprint).not.toBe(batched.buildFingerprint) + expect(single.containerTag).not.toBe(batched.containerTag) + expect(batched.ingestBatchSize).toBe(5) + }) + + test("causal ingest defers failed readiness, continues, and retries it at build end", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const value = question("q1") + const benchmark = testBenchmark({ + questions: [value], + sessionsByQuestion: { q1: sessions() }, + protocol: testProtocol({ + readinessBarrier: "after-each-document", + processingMode: "instant", + }), + }) + const plans = preparePlans(benchmark, [value]) + const checkpoint = await initializeCheckpoint({ manager, plans, questions: [value] }) + const provider = new CausalProvider() + provider.omitReadinessOnceForSession = "s2" + + await runIngestPhase(provider, checkpoint, manager, plans) + + expect(provider.ingestAttempts).toEqual(["s1", "s2", "s3", "s2"]) + expect(Object.values(checkpoint.builds)[0].ingest.completedSessionIds).toEqual([ + "s1", + "s2", + "s3", + ]) + expect(Object.values(checkpoint.builds)[0].ingest.deferredSessions).toEqual([]) + expect(Object.values(checkpoint.builds)[0].indexing.status).toBe("completed") + }) + + test("independent causal conversation builds run concurrently", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const questions = [question("q1"), question("q2")] + const benchmark = testBenchmark({ + questions, + sessionsByQuestion: { q1: sessions(), q2: sessions() }, + protocol: testProtocol({ + readinessBarrier: "after-each-document", + processingMode: "instant", + }), + }) + benchmark.getIngestionGroupId = (questionId) => questionId + const plans = preparePlans(benchmark, questions) + const checkpoint = await initializeCheckpoint({ manager, plans, questions }) + checkpoint.concurrency = { default: 2 } + const provider = new ConcurrentCausalProvider() + + await runIngestPhase(provider, checkpoint, manager, plans) + + expect(provider.firstSessionContainers.size).toBe(2) + expect( + Object.values(checkpoint.builds).every((build) => build.indexing.status === "completed") + ).toBe(true) + }) + + test("flush surfaces a completed checkpoint write failure and a later full save can recover", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const value = question("q1") + const benchmark = testBenchmark({ questions: [value], sessionsByQuestion: { q1: sessions() } }) + const plans = preparePlans(benchmark, [value]) + const checkpoint = await initializeCheckpoint({ manager, plans, questions: [value] }) + const temporaryCheckpointPath = join(root, checkpoint.runId, "checkpoint.json.tmp") + + // A directory at the atomic-write temporary path deterministically makes + // writeFile fail after save() has returned. + await mkdir(temporaryCheckpointPath) + manager.updateStatus(checkpoint, "running") + await Promise.resolve() + await expect(manager.flush(checkpoint.runId)).rejects.toThrow() + + await rm(temporaryCheckpointPath, { recursive: true, force: true }) + manager.updateStatus(checkpoint, "completed") + await manager.flush(checkpoint.runId) + expect(manager.load(checkpoint.runId)?.status).toBe("completed") + }) + + test("uses portable hashed result filenames for colon-bearing BEAM question IDs", async () => { + const manager = new CheckpointManager(await createTempRoot()) + const resultPath = manager.getQuestionResultsPath("run-id", "beam:1M:10:event_ordering:abcdef") + expect(basename(resultPath)).toMatch(/^[a-f0-9]{64}\.json$/) + expect(basename(resultPath)).not.toContain(":") + }) + + test("remote-success/local-crash reconciliation does not create a duplicate document", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const value = question("q1") + const benchmark = testBenchmark({ questions: [value], sessionsByQuestion: { q1: sessions() } }) + const plans = preparePlans(benchmark, [value]) + const checkpoint = await initializeCheckpoint({ manager, plans, questions: [value] }) + const provider = new CrashWindowProvider() + + await runIngestPhase(provider, checkpoint, manager, plans) + expect(provider.ingestAttempts).toEqual(["s1", "s2", "s3", "s1"]) + expect(provider.remoteCreates).toBe(3) + expect(provider.remoteDocuments.size).toBe(3) + expect(new Set(Object.values(checkpoint.builds)[0].ingest.documentIds)).toEqual( + new Set(["doc-s1", "doc-s2", "doc-s3"]) + ) + }) + + test("one shared build persists one progress update per session, not per question", async () => { + const root = await createTempRoot() + const manager = new CountingCheckpointManager(root) + const questions = Array.from({ length: 20 }, (_, index) => question(`q${index + 1}`)) + const benchmark = testBenchmark({ + questions, + sessionsByQuestion: Object.fromEntries( + questions.map((value) => [value.questionId, sessions()]) + ), + }) + const plans = preparePlans(benchmark, questions) + const checkpoint = await initializeCheckpoint({ manager, plans, questions }) + manager.saveCalls = 0 + + await runIngestPhase(new FakeProvider(), checkpoint, manager, plans) + // The large checkpoint is saved only at attempt start and completion. + // Per-session remote success is durably recorded in the small append-only journal. + expect(manager.saveCalls).toBe(2) + expect(Object.values(checkpoint.builds)[0].ingest.completedSessionIds).toHaveLength(3) + await expect( + readFile( + manager.getIngestProgressJournalPath( + checkpoint.runId, + Object.values(checkpoint.builds)[0].buildId + ), + "utf8" + ) + ).rejects.toThrow() + }) + + test("replays fsynced per-session ingest progress without a full checkpoint rewrite", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const value = question("q1") + const benchmark = testBenchmark({ questions: [value], sessionsByQuestion: { q1: sessions() } }) + const plans = preparePlans(benchmark, [value]) + const checkpoint = await initializeCheckpoint({ manager, plans, questions: [value] }) + const build = Object.values(checkpoint.builds)[0] + + manager.recordIngestProgress(checkpoint, build.buildId, { + sequence: 0, + sessionId: "s1", + documentIds: ["doc-s1"], + taskIds: ["task-s1"], + readyForNextSession: false, + }) + + const rawSnapshot = JSON.parse( + await readFile(manager.getCheckpointPath(checkpoint.runId), "utf8") + ) as RunCheckpoint + expect(Object.values(rawSnapshot.builds)[0].ingest.completedSessionIds).toEqual([]) + + const resumed = new CheckpointManager(root).load(checkpoint.runId)! + expect(Object.values(resumed.builds)[0].ingest.completedSessionIds).toEqual(["s1"]) + expect(Object.values(resumed.builds)[0].ingest.documentIds).toEqual(["doc-s1"]) + expect(Object.values(resumed.builds)[0].ingest.taskIds).toEqual(["task-s1"]) + }) + + test("does not mark indexing complete when the provider omits final completion", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const value = question("q1") + const benchmark = testBenchmark({ questions: [value], sessionsByQuestion: { q1: sessions() } }) + const plans = preparePlans(benchmark, [value]) + const checkpoint = await initializeCheckpoint({ manager, plans, questions: [value] }) + markBuildIngested(Object.values(checkpoint.builds)[0]) + const provider = new FakeProvider() + provider.omitIndexingProgress = true + + await expect(runIndexingPhase(provider, checkpoint, manager)).rejects.toThrow( + "returned before 3 IDs completed" + ) + expect(Object.values(checkpoint.builds)[0].indexing.status).toBe("failed") + }) + + test("search fails closed for missing, failed, or incomplete builds", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const value = question("q1") + const benchmark = testBenchmark({ questions: [value], sessionsByQuestion: { q1: sessions() } }) + const plans = preparePlans(benchmark, [value]) + const checkpoint = await initializeCheckpoint({ manager, plans, questions: [value] }) + const provider = new FakeProvider() + const build = Object.values(checkpoint.builds)[0] + + await expect( + runSearchPhase(provider, benchmark, checkpoint, manager, [value.questionId]) + ).rejects.toThrow("ingestion is pending") + + markBuildIngested(build) + build.indexing.status = "failed" + await expect( + runSearchPhase(provider, benchmark, checkpoint, manager, [value.questionId]) + ).rejects.toThrow("indexing is failed") + + delete checkpoint.builds[build.buildId] + await expect( + runSearchPhase(provider, benchmark, checkpoint, manager, [value.questionId]) + ).rejects.toThrow("references missing build") + expect(provider.searchCalls).toHaveLength(0) + }) + + test("passes benchmark Top-K unchanged and records requested/returned counts", async () => { + const root = await createTempRoot() + const manager = new CheckpointManager(root) + const value = question("q1") + const protocol = testProtocol({ requestedTopK: 5, answerCutoff: 5 }) + const benchmark = testBenchmark({ + questions: [value], + sessionsByQuestion: { q1: sessions() }, + protocol, + }) + const plans = preparePlans(benchmark, [value]) + const checkpoint = await initializeCheckpoint({ manager, plans, questions: [value] }) + checkpoint.datasetIdentity = { + datasetFingerprint: "dataset-fixture", + } as RunCheckpoint["datasetIdentity"] + checkpoint.benchmarkInputFingerprint = "benchmark-input-fixture" + const build = Object.values(checkpoint.builds)[0] + build.ingest.status = "completed" + build.indexing.status = "completed" + const provider = new FakeProvider() + provider.searchResults = [1, 2, 3].map((rank) => ({ + id: `result-${rank}`, + rank, + text: `Result ${rank}`, + provider: "fake", + resultType: "memory", + })) + provider.searchRawReturnedCount = 4 + + await runSearchPhase(provider, benchmark, checkpoint, manager, [value.questionId]) + expect(provider.searchCalls).toHaveLength(1) + expect(provider.searchCalls[0]).toMatchObject({ + containerTag: build.containerTag, + limit: 5, + threshold: 0.2, + searchMode: "hybrid", + }) + expect(checkpoint.questions.q1.phases.search.requestedCount).toBe(5) + expect(checkpoint.questions.q1.phases.search.rawReturnedCount).toBe(4) + expect(checkpoint.questions.q1.phases.search.returnedCount).toBe(3) + expect(checkpoint.questions.q1.phases.search.normalizedCount).toBe(3) + expect(checkpoint.questions.q1.phases.search.droppedCount).toBe(1) + expect(checkpoint.questions.q1.phases.search.answerCutoff).toBe(5) + const resultFile = checkpoint.questions.q1.phases.search.resultFile! + const artifact = JSON.parse(await readFile(resultFile, "utf8")) + expect(artifact).toMatchObject({ + benchmark: "test-benchmark", + benchmarkScope: { displayName: "Test", includedTiers: ["test"], coverage: "full" }, + datasetIdentity: { datasetFingerprint: "dataset-fixture" }, + benchmarkInputFingerprint: "benchmark-input-fixture", + selectedQuestionIdsDigest: "selected", + protocolIdentity, + }) + + checkpoint.questions.q1.phases.search = { status: "pending" } + provider.searchRawReturnedCount = undefined + provider.searchResults = Array.from({ length: 6 }, (_, index) => ({ + id: `too-many-${index + 1}`, + rank: index + 1, + text: `Too many ${index + 1}`, + provider: "fake", + resultType: "memory", + })) + await expect( + runSearchPhase(provider, benchmark, checkpoint, manager, [value.questionId]) + ).rejects.toThrow("inconsistent retrieval diagnostics") + }) +}) diff --git a/ui/app/compare/[compareId]/page.tsx b/ui/app/compare/[compareId]/page.tsx index bf5e0f8..ac715ae 100644 --- a/ui/app/compare/[compareId]/page.tsx +++ b/ui/app/compare/[compareId]/page.tsx @@ -11,8 +11,15 @@ import { type CompareDetail, type CompareReport, type CompareRunInfo, + type BenchmarkResult, } from "@/lib/api" -import { formatDate, getStatusColor, cn } from "@/lib/utils" +import { + formatDate, + getBenchmarkDisplayName, + getPipelineProgress, + getStatusColor, + cn, +} from "@/lib/utils" import { AccuracyBarChart } from "@/components/accuracy-bar-chart" import { DataTable, type Column } from "@/components/data-table" import { CircularProgress } from "@/components/circular-progress" @@ -21,6 +28,25 @@ import { Tooltip } from "@/components/tooltip" const POLL_INTERVAL = 2000 // 2 seconds +function getQuestionTypeQuality(report: BenchmarkResult, type: string) { + const slice = report.quality?.bySlice?.[type] + if (typeof slice?.averageScore === "number") { + return { + value: slice.averageScore, + metricKey: "averageScore", + passAccuracy: + typeof slice.passAccuracy === "number" + ? slice.passAccuracy + : report.byQuestionType?.[type]?.accuracy, + } + } + return { + value: report.byQuestionType?.[type]?.accuracy, + metricKey: "accuracy", + passAccuracy: report.byQuestionType?.[type]?.accuracy, + } +} + export default function CompareDetailPage() { const params = useParams() const router = useRouter() @@ -33,6 +59,9 @@ export default function CompareDetailPage() { const [error, setError] = useState(null) const [stopping, setStopping] = useState(false) const [continuing, setContinuing] = useState(false) + const comparablePrimaryMetric = useMemo(() => { + return report?.comparison.comparable ? (report.comparison.identity ?? null) : null + }, [report]) // Check if comparison is in progress const isRunning = compare?.status === "running" || compare?.status === "pending" @@ -66,25 +95,13 @@ export default function CompareDetailPage() { run.status === "stopping" const p = run.progress const total = p?.total || 0 - const phasesCompleted = - (p?.ingested || 0) + - (p?.indexed || 0) + - (p?.searched || 0) + - (p?.answered || 0) + - (p?.evaluated || 0) - const totalPhases = 5 * total - const progress = totalPhases > 0 ? phasesCompleted / totalPhases : 0 + const { progress, phasesFullyComplete } = p + ? getPipelineProgress(p) + : { progress: 0, phasesFullyComplete: 0 } const episodes = p?.indexingEpisodes const hasEpisodeData = episodes && episodes.total > 0 - let phasesFullyComplete = 0 - if ((p?.ingested || 0) === total && total > 0) phasesFullyComplete++ - if ((p?.indexed || 0) === total && total > 0) phasesFullyComplete++ - if ((p?.searched || 0) === total && total > 0) phasesFullyComplete++ - if ((p?.answered || 0) === total && total > 0) phasesFullyComplete++ - if ((p?.evaluated || 0) === total && total > 0) phasesFullyComplete++ - const progressContent = (
{runIsActive && } @@ -130,7 +147,7 @@ export default function CompareDetailPage() { }, { key: "accuracy", - header: "Accuracy", + header: "Pass Accuracy", align: "right", render: (run) => { const accuracyPct = @@ -321,7 +338,7 @@ export default function CompareDetailPage() {
Benchmark:{" "} - {compare.benchmark} + {getBenchmarkDisplayName(compare.benchmark, compare.benchmarkScope)} Judge: {compare.judge} @@ -388,9 +405,17 @@ export default function CompareDetailPage() { {/* Overall Accuracy Table */} {/* Accuracy and Latency side by side */}
- {/* Overall Accuracy - 35% width */} -
-

Accuracy

+ {/* Protocol quality with legacy accuracy fallback */} +
+

Quality

+ {report.reports.length > 1 && !comparablePrimaryMetric && ( +

+ These reports are not like-for-like, so no cross-provider winner is highlighted. + {report.comparison.mismatchReasons.length > 0 && ( + <> {report.comparison.mismatchReasons.join("; ")} + )} +

+ )}
@@ -399,27 +424,44 @@ export default function CompareDetailPage() { Provider + {(() => { - const rows = report.reports.map((r) => ({ - provider: r.provider, - correct: r.report.summary?.correctCount ?? r.report.correctCount, - total: r.report.summary?.totalQuestions ?? r.report.totalQuestions, - accuracy: r.report.summary?.accuracy ?? r.report.accuracy, - })) - const validAccuracies = rows - .map((r) => r.accuracy) - .filter((a): a is number => a != null) - const bestAccuracy = - validAccuracies.length > 0 ? Math.max(...validAccuracies) : null + const rows = report.reports.map((r) => { + const passAccuracy = r.report.summary?.accuracy ?? r.report.accuracy + return { + provider: r.provider, + correct: r.report.summary?.correctCount ?? r.report.correctCount, + total: r.report.summary?.totalQuestions ?? r.report.totalQuestions, + passAccuracy, + primaryKey: r.report.quality?.primaryMetric?.key, + primaryValue: r.report.quality?.primaryMetric?.value, + tierScores: Object.entries(r.report.quality?.metrics ?? {}).filter( + ([key, value]) => + (key === "beamScore1M" || key === "beamScore10M") && + typeof value === "number" + ), + } + }) + const validPrimaryValues = rows + .map((r) => r.primaryValue) + .filter((value): value is number => value != null) + const bestPrimary = + validPrimaryValues.length > 0 && comparablePrimaryMetric + ? comparablePrimaryMetric.higherIsBetter + ? Math.max(...validPrimaryValues) + : Math.min(...validPrimaryValues) + : null // Only highlight the FIRST occurrence of the best value const firstBestIndex = - bestAccuracy != null - ? rows.findIndex((r) => r.accuracy === bestAccuracy) + bestPrimary != null + ? rows.findIndex((r) => r.primaryValue === bestPrimary) : -1 return rows.map((row, index) => { @@ -435,12 +477,31 @@ export default function CompareDetailPage() { isBest ? "text-status-success font-semibold" : "text-text-primary" } > - {row.accuracy != null ? `${(row.accuracy * 100).toFixed(1)}%` : "—"} + {row.primaryValue != null + ? `${(row.primaryValue * 100).toFixed(1)}%` + : "—"} + +
+ {row.primaryKey + ? row.primaryKey.replace(/([a-z])([A-Z])/g, "$1 $2") + : "no scalar primary"} +
+ {row.tierScores.map(([key, value]) => ( +
+ {key}: {((value as number) * 100).toFixed(1)}% +
+ ))} + + @@ -452,9 +513,9 @@ export default function CompareDetailPage() { - {/* Latency - 65% width */} + {/* Latency */} {report.reports.some((r) => r.report.latency || r.report.latencyStats) && ( -
+

Latency (median ms)

@@ -463,9 +524,26 @@ export default function CompareDetailPage() {
- Score + Primary + + Pass
+ + {row.passAccuracy != null + ? `${(row.passAccuracy * 100).toFixed(1)}%` + : "—"} {row.correct != null && row.total != null && ( - +
({row.correct}/{row.total}) - +
)}
- + + + + @@ -482,7 +560,7 @@ export default function CompareDetailPage() { Evaluate @@ -557,6 +635,91 @@ export default function CompareDetailPage() { )} + {report.reports.some( + ({ report: runReport }) => runReport.builds || runReport.questionMetrics?.length + ) && ( +
+

+ Build and Query Metrics +

+
+
+ Provider + Build / container + + Per question +
Ingest - Total + Online
+ + + {[ + "Provider", + "Builds", + "Build work", + "Build wall-clock", + "Build cost", + "Online mean", + "Eval mean", + "Amortized mean", + ].map((heading, index) => ( + + ))} + + + + {report.reports.map(({ provider, report: runReport }) => { + const metrics = runReport.questionMetrics || [] + const mean = (values: number[]) => + values.length > 0 + ? values.reduce((sum, value) => sum + value, 0) / values.length + : undefined + const amortized = metrics + .map((metric) => metric.amortizedOnlinePlusBuildWorkMs) + .filter((value): value is number => value != null) + const onlineMean = mean(metrics.map((metric) => metric.onlineQueryLatencyMs)) + const evaluationMean = mean( + metrics.map((metric) => metric.evaluationLatencyMs) + ) + const amortizedMean = mean(amortized) + return ( + + + + + + + + + + + ) + })} + +
+ {heading} +
{provider} + {runReport.builds?.uniqueBuildCount ?? "—"} + + {runReport.builds + ? `${runReport.builds.sumContainerBuildWorkMs}ms` + : "—"} + + {runReport.builds ? `${runReport.builds.buildPhaseWallClockMs}ms` : "—"} + + {runReport.builds?.totalBuildCostUsd == null + ? "—" + : `$${runReport.builds.totalBuildCostUsd.toFixed(4)}`} + + {onlineMean == null ? "—" : `${onlineMean.toFixed(1)}ms`} + + {evaluationMean == null ? "—" : `${evaluationMean.toFixed(1)}ms`} + + {amortizedMean == null ? "—" : `${amortizedMean.toFixed(1)}ms`} +
+
+
+ )} + {/* Retrieval Metrics Table */} {report.reports.some((r) => r.report.retrieval) && (
@@ -709,11 +872,13 @@ export default function CompareDetailPage() { {/* By Question Type - Table and Chart */} {report.reports.some( - (r) => r.report.byQuestionType && Object.keys(r.report.byQuestionType).length > 0 + (r) => + (r.report.byQuestionType && Object.keys(r.report.byQuestionType).length > 0) || + (r.report.quality?.bySlice && Object.keys(r.report.quality.bySlice).length > 0) ) && (

- Accuracy by Question Type + Quality by Question Type

{/* Left: Table (50%) */} @@ -737,14 +902,14 @@ export default function CompareDetailPage() { {(() => { - // Collect all question types const allTypes = new Set() report.reports.forEach((r) => { - if (r.report.byQuestionType) { - Object.keys(r.report.byQuestionType).forEach((type) => - allTypes.add(type) - ) - } + Object.keys(r.report.byQuestionType || {}).forEach((type) => + allTypes.add(type) + ) + Object.keys(r.report.quality?.bySlice || {}) + .filter((type) => !type.startsWith("tier:")) + .forEach((type) => allTypes.add(type)) }) const rows = Array.from(allTypes) @@ -752,19 +917,25 @@ export default function CompareDetailPage() { .map((type) => { const values = report.reports.map((r) => ({ provider: r.provider, - accuracy: r.report.byQuestionType?.[type]?.accuracy, + ...getQuestionTypeQuality(r.report, type), })) const validValues = values - .map((v) => v.accuracy) - .filter((a) => a !== undefined) as number[] - const bestAccuracy = - validValues.length > 0 ? Math.max(...validValues) : undefined - - // Find the index of the FIRST best value (for tie-breaking) + .map((entry) => entry.value) + .filter((value): value is number => value != null) + const comparable = + new Set( + values + .filter((entry) => entry.value != null) + .map((entry) => entry.metricKey) + ).size === 1 + const bestValue = + comparable && validValues.length > 0 + ? Math.max(...validValues) + : undefined const firstBestIndex = - bestAccuracy !== undefined - ? values.findIndex((v) => v.accuracy === bestAccuracy) + bestValue !== undefined + ? values.findIndex((entry) => entry.value === bestValue) : -1 return ( @@ -772,48 +943,63 @@ export default function CompareDetailPage() { {type.replace(/[-_]/g, "-")} - {values.map(({ provider, accuracy }, index) => { - // Only highlight the FIRST occurrence of the best value - const isBest = index === firstBestIndex - return ( - - {accuracy !== undefined ? ( - - {(accuracy * 100).toFixed(1)}% - - ) : ( - — - )} - - ) - })} + {values.map( + ({ provider, value, metricKey, passAccuracy }, index) => { + const isBest = index === firstBestIndex + return ( + + {value !== undefined ? ( + <> + + {(value * 100).toFixed(1)}% + +
+ {metricKey === "averageScore" + ? `avg · ${passAccuracy != null ? `${(passAccuracy * 100).toFixed(1)}% pass` : "pass unavailable"}` + : "legacy accuracy"} +
+ + ) : ( + — + )} + + ) + } + )} ) }) - // Calculate overall accuracy for each provider const overallValues = report.reports.map((r) => { - const accuracy = r.report.summary?.accuracy ?? r.report.accuracy return { provider: r.provider, - accuracy, + value: r.report.quality?.primaryMetric?.value, + metricKey: + r.report.quality?.primaryMetric?.key ?? "no scalar primary", } }) const validOverall = overallValues - .map((v) => v.accuracy) - .filter((a) => a !== undefined) as number[] + .map((entry) => entry.value) + .filter((value): value is number => value != null) const bestOverall = - validOverall.length > 0 ? Math.max(...validOverall) : undefined + validOverall.length > 0 && comparablePrimaryMetric + ? comparablePrimaryMetric.higherIsBetter + ? Math.max(...validOverall) + : Math.min(...validOverall) + : undefined const firstBestOverallIndex = bestOverall !== undefined - ? overallValues.findIndex((v) => v.accuracy === bestOverall) + ? overallValues.findIndex((entry) => entry.value === bestOverall) : -1 return ( @@ -824,20 +1010,25 @@ export default function CompareDetailPage() { Overall - {overallValues.map(({ provider, accuracy }, index) => { + {overallValues.map(({ provider, value, metricKey }, index) => { const isBest = index === firstBestOverallIndex return ( - {accuracy !== undefined ? ( - - {(accuracy * 100).toFixed(1)}% - + {value !== undefined ? ( + <> + + {(value * 100).toFixed(1)}% + +
+ {metricKey.replace(/([a-z])([A-Z])/g, "$1 $2")} +
+ ) : ( — )} @@ -856,13 +1047,30 @@ export default function CompareDetailPage() { {/* Right: Bar Chart (50%) */}
{(() => { - // Prepare data for chart const allTypes = new Set() report.reports.forEach((r) => { - if (r.report.byQuestionType) { - Object.keys(r.report.byQuestionType).forEach((type) => allTypes.add(type)) - } + Object.keys(r.report.byQuestionType || {}).forEach((type) => + allTypes.add(type) + ) + Object.keys(r.report.quality?.bySlice || {}) + .filter((type) => !type.startsWith("tier:")) + .forEach((type) => allTypes.add(type)) + }) + + const hasMixedMetrics = Array.from(allTypes).some((type) => { + const metricKeys = report.reports + .map((entry) => getQuestionTypeQuality(entry.report, type)) + .filter((quality) => quality.value != null) + .map((quality) => quality.metricKey) + return new Set(metricKeys).size > 1 }) + if (hasMixedMetrics) { + return ( +
+ Chart hidden because these runs use different per-type metrics. +
+ ) + } const chartData = Array.from(allTypes) .sort() @@ -870,7 +1078,7 @@ export default function CompareDetailPage() { type, values: report.reports.map((r) => ({ provider: r.provider, - accuracy: r.report.byQuestionType?.[type]?.accuracy, + accuracy: getQuestionTypeQuality(r.report, type).value, })), })) diff --git a/ui/app/compare/new/page.tsx b/ui/app/compare/new/page.tsx index a9f3d86..283c0fe 100644 --- a/ui/app/compare/new/page.tsx +++ b/ui/app/compare/new/page.tsx @@ -11,6 +11,7 @@ import { type SelectionMode, type SampleType, type SamplingConfig, + type Benchmark, } from "@/lib/api" import { SingleSelect } from "@/components/single-select" import { MultiSelect } from "@/components/multi-select" @@ -22,7 +23,7 @@ export default function NewComparePage() { const [error, setError] = useState(null) const [providers, setProviders] = useState<{ name: string; displayName: string }[]>([]) - const [benchmarks, setBenchmarks] = useState<{ name: string; displayName: string }[]>([]) + const [benchmarks, setBenchmarks] = useState([]) const [models, setModels] = useState({}) const [form, setForm] = useState({ @@ -31,6 +32,9 @@ export default function NewComparePage() { compareId: "", judgeModel: "gpt-4o", answeringModel: "gpt-4o", + dataPath: "", + datasetRevision: "", + retrievalTopK: "5", selectionMode: "full" as SelectionMode, sampleType: "consecutive" as SampleType, perCategory: "2", @@ -80,6 +84,15 @@ export default function NewComparePage() { } const displayCompareId = form.compareId || generateCompareId() + const selectedBenchmark = benchmarks.find((benchmark) => benchmark.name === form.benchmark) + const isBeam = form.benchmark.startsWith("beam-") + + useEffect(() => { + const requiredJudge = selectedBenchmark?.requiredJudge?.modelAlias + if (requiredJudge && form.judgeModel !== requiredJudge) { + setForm((current) => ({ ...current, judgeModel: requiredJudge })) + } + }, [selectedBenchmark, form.judgeModel]) async function handleSubmit(e: React.FormEvent) { e.preventDefault() @@ -120,6 +133,9 @@ export default function NewComparePage() { judgeModel: form.judgeModel, answeringModel: form.answeringModel, sampling, + dataPath: isBeam ? form.dataPath || undefined : undefined, + datasetRevision: isBeam ? form.datasetRevision || undefined : undefined, + retrievalTopK: isBeam ? Number(form.retrievalTopK) : undefined, }) router.push(`/compare`) @@ -134,6 +150,9 @@ export default function NewComparePage() { const providerOptions = providers.map((p) => ({ value: p.name, label: p.displayName })) const benchmarkOptions = benchmarks.map((b) => ({ value: b.name, label: b.displayName })) const modelOptions = allModels.map((m) => ({ value: m.alias, label: m.displayName || m.alias })) + const judgeModelOptions = selectedBenchmark?.requiredJudge + ? modelOptions.filter((model) => model.value === selectedBenchmark.requiredJudge!.modelAlias) + : modelOptions if (loading) { return ( @@ -207,6 +226,49 @@ export default function NewComparePage() { )}
+ {isBeam && ( +
+
+ + setForm({ ...form, dataPath: event.target.value })} + placeholder="Default: data/benchmarks/beam" + className="w-full px-3 py-2.5 text-sm bg-[#222222] border border-[#333333] rounded text-text-primary placeholder-text-muted focus:outline-none focus:border-accent" + /> +
+
+ + setForm({ ...form, datasetRevision: event.target.value })} + placeholder="Optional pinned fingerprint" + className="w-full px-3 py-2.5 text-sm bg-[#222222] border border-[#333333] rounded text-text-primary placeholder-text-muted focus:outline-none focus:border-accent font-mono" + /> +
+
+ + ({ + value: String(value), + label: String(value), + }))} + selected={form.retrievalTopK} + onChange={(value) => setForm({ ...form, retrievalTopK: value })} + /> +
+
+ )} +
@@ -202,9 +220,7 @@ export default function RunDetailPage() {

Loading benchmark dataset...

-

- This may take a moment for first-time downloads -

+

Validating benchmark data and run state

) @@ -213,20 +229,48 @@ export default function RunDetailPage() { const allQuestions = Object.values(run.questions) // Only count questions that have been evaluated const evaluatedQuestions = allQuestions.filter((q) => q.phases.evaluate.status === "completed") - const failedQuestions = evaluatedQuestions.filter((q) => q.phases.evaluate.label === "incorrect") + const correctQuestionCount = evaluatedQuestions.filter((q) => + isEvaluationPassed(q.phases.evaluate) + ).length const accuracy = report?.summary?.accuracy ?? - (evaluatedQuestions.length > 0 - ? (evaluatedQuestions.filter((q) => q.phases.evaluate.score === 1).length / - evaluatedQuestions.length) * - 100 - : 0) + report?.accuracy ?? + (evaluatedQuestions.length > 0 ? correctQuestionCount / evaluatedQuestions.length : null) + const reportCorrectCount = report?.summary?.correctCount ?? report?.correctCount + const reportTotalQuestions = report?.summary?.totalQuestions ?? report?.totalQuestions + const primaryMetric = report?.quality?.primaryMetric + const averageScore = report?.summary?.averageScore + const combinedBeamScope = (report?.benchmarkScope?.includedTiers?.length ?? 0) > 1 + const datasetIdentity = report?.datasetIdentity ?? run.datasetIdentity + const datasetFingerprint = + typeof datasetIdentity?.datasetFingerprint === "string" + ? datasetIdentity.datasetFingerprint + : undefined + const protocolIdentity = report?.protocolIdentity ?? run.protocolIdentity + const protocolLabel = + typeof protocolIdentity?.id === "string" + ? `${protocolIdentity.id}${typeof protocolIdentity.version === "string" ? `@${protocolIdentity.version}` : ""}` + : undefined // Find error from failed phases const runError = (() => { + const builds = Object.values(run.builds || {}) + const usesSharedBuilds = builds.length > 0 || allQuestions.some((question) => question.buildId) + for (const build of builds) { + for (const phase of [build.ingest, build.indexing]) { + if (phase?.status === "failed" && phase.error) return phase.error + } + } + for (const q of allQuestions) { - const phases = q.phases as Record - for (const phase of ["ingest", "indexing", "search", "answer", "evaluate"]) { + if (q.buildId && !run.builds?.[q.buildId]) { + return `Question ${q.questionId} references missing build ${q.buildId}` + } + const phases = q.phases as unknown as Record + const questionOwnedPhases = usesSharedBuilds + ? ["search", "answer", "evaluate"] + : ["ingest", "indexing", "search", "answer", "evaluate"] + for (const phase of questionOwnedPhases) { if (phases[phase]?.status === "failed" && phases[phase]?.error) { return phases[phase].error } @@ -286,7 +330,7 @@ export default function RunDetailPage() { Benchmark:{" "} - {run.benchmark} + {getBenchmarkDisplayName(run.benchmark, run.benchmarkScope)} Judge: {run.judge} @@ -344,13 +388,62 @@ export default function RunDetailPage() { )} 0 + ? `${correctQuestionCount}/${evaluatedQuestions.length} correct` + : undefined, }, + ...(averageScore != null + ? [ + { + label: "average question score", + value: averageScore.toFixed(3), + }, + ] + : []), { label: "questions", value: run.summary.total, @@ -366,10 +459,50 @@ export default function RunDetailPage() { value: run.answeringModel || "—", mono: true, }, + ...(run.retrievalTopK != null + ? [ + { + label: "retrieval top-k", + value: run.retrievalTopK, + }, + ] + : []), + ...(protocolLabel + ? [ + { + label: "protocol", + value: protocolLabel, + mono: true, + }, + ] + : []), + ...(datasetFingerprint + ? [ + { + label: "dataset fingerprint", + value: datasetFingerprint, + mono: true, + }, + ] + : []), + ...(report?.builds + ? [ + { + label: "builds", + value: report.builds.uniqueBuildCount, + subtext: `${report.builds.sumContainerBuildWorkMs.toLocaleString()}ms total build work`, + }, + ] + : []), ]} /> - - + + + +
)} diff --git a/ui/app/runs/[runId]/questions/[questionId]/page.tsx b/ui/app/runs/[runId]/questions/[questionId]/page.tsx index 125a6ff..e1d0146 100644 --- a/ui/app/runs/[runId]/questions/[questionId]/page.tsx +++ b/ui/app/runs/[runId]/questions/[questionId]/page.tsx @@ -4,15 +4,17 @@ import { useState, useEffect } from "react" import Link from "next/link" import { useParams } from "next/navigation" import { Highlight, themes } from "prism-react-renderer" -import { getQuestion } from "@/lib/api" -import { cn } from "@/lib/utils" +import { getQuestion, type QuestionCheckpoint } from "@/lib/api" +import { cn, isEvaluationPassed } from "@/lib/utils" export default function QuestionDetailPage() { const params = useParams() const runId = decodeURIComponent(params.runId as string) const questionId = decodeURIComponent(params.questionId as string) - const [question, setQuestion] = useState(null) + const [question, setQuestion] = useState< + (QuestionCheckpoint & { searchResultsFile?: { results?: any[] } }) | null + >(null) const [loading, setLoading] = useState(true) const [error, setError] = useState(null) @@ -51,10 +53,25 @@ export default function QuestionDetailPage() { ) } - const isCorrect = question.phases?.evaluate?.label === "correct" + const isCorrect = isEvaluationPassed(question.phases?.evaluate) const searchResults = question.searchResultsFile?.results || question.phases?.search?.results || [] - const containerTag = question.containerTag || "" + const searchPhase = question.phases?.search + const answerPhase = question.phases?.answer + const retrievalPlan = searchPhase?.retrievalPlan + const requestedTopK = retrievalPlan?.requestedTopK ?? searchPhase?.requestedCount + const rawReturnedCount = searchPhase?.rawReturnedCount ?? searchPhase?.returnedCount + const normalizedCount = searchPhase?.normalizedCount ?? searchResults.length + const answerCutoff = retrievalPlan?.answerCutoff ?? searchPhase?.answerCutoff + const answerEvidenceCount = searchPhase?.answerEvidenceCount ?? answerPhase?.evidenceCount + const containerTag = question.build?.containerTag || question.containerTag || "" + const evaluatePhase = question.phases?.evaluate + const protocolEvaluation = evaluatePhase?.evaluation + const primaryScore = + protocolEvaluation?.primaryScore ?? evaluatePhase?.primaryScore ?? evaluatePhase?.score + const evaluationMetrics = protocolEvaluation?.metrics ?? evaluatePhase?.metrics + const evaluationDetails = protocolEvaluation?.details ?? evaluatePhase?.details + const evaluationExplanation = protocolEvaluation?.explanation ?? evaluatePhase?.explanation const copyContainerTag = () => { navigator.clipboard.writeText(containerTag) @@ -182,12 +199,109 @@ export default function QuestionDetailPage() {
{/* Evaluation */} - {question.phases?.evaluate?.explanation && ( + {(primaryScore != null || + evaluationExplanation || + evaluationMetrics || + evaluationDetails) && (
-

- Evaluation Explanation +
+

Evaluation

+ {primaryScore != null && ( +
+
{primaryScore.toFixed(4)}
+
primary score
+
+ )} +
+ + {evaluationMetrics && Object.keys(evaluationMetrics).length > 0 && ( +
+ {Object.entries(evaluationMetrics).map(([metric, value]) => ( +
+
{metric}
+
+ {Number.isFinite(value) ? value.toFixed(4) : String(value)} +
+
+ ))} +
+ )} + + {evaluationExplanation && ( +
+

Explanation

+

{evaluationExplanation}

+
+ )} + + {evaluationDetails && Object.keys(evaluationDetails).length > 0 && ( +
+

+ Protocol details +

+ + {({ style, tokens, getLineProps, getTokenProps }) => ( + 
+                    {tokens.map((line, i) => (
+                      

+                        {line.map((token, key) => (
+                          
+                        ))}
+                      

+                    ))}
+
+ )} + +
+ )} +

+ )} + + {/* Retrieval contract */} + {(requestedTopK != null || rawReturnedCount != null || searchResults.length > 0) && ( +
+

+ Retrieval contract

-

{question.phases.evaluate.explanation}

+
+ + sum + request.limit, 0) ?? + searchPhase?.requestedCount + } + /> + + + + + + +
+ {(retrievalPlan?.searchMode || retrievalPlan?.threshold != null) && ( +

+ {retrievalPlan.searchMode ? `mode ${retrievalPlan.searchMode}` : "default mode"} + {retrievalPlan.threshold != null ? ` · threshold ${retrievalPlan.threshold}` : ""} +

+ )} + {searchPhase?.providerRequests && searchPhase.providerRequests.length > 0 && ( +
+

+ Provider requests +

+ 
+                {JSON.stringify(searchPhase.providerRequests, null, 2)}
+
+
+ )}
)} @@ -241,3 +355,12 @@ export default function QuestionDetailPage() {
) } + +function RetrievalValue({ label, value }: { label: string; value?: number }) { + return ( +
+
{label}
+
{value ?? "—"}
+
+ ) +} diff --git a/ui/app/runs/new/page.tsx b/ui/app/runs/new/page.tsx index ae3b222..f677343 100644 --- a/ui/app/runs/new/page.tsx +++ b/ui/app/runs/new/page.tsx @@ -16,8 +16,10 @@ import { type SampleType, type SamplingConfig, type Provider, + type Benchmark, } from "@/lib/api" import { SingleSelect } from "@/components/single-select" +import { getBenchmarkDisplayName } from "@/lib/utils" type Tab = "new" | "advanced" @@ -29,7 +31,7 @@ export default function NewRunPage() { const [error, setError] = useState(null) const [providers, setProviders] = useState([]) - const [benchmarks, setBenchmarks] = useState<{ name: string; displayName: string }[]>([]) + const [benchmarks, setBenchmarks] = useState([]) const [models, setModels] = useState({}) const [completedRuns, setCompletedRuns] = useState([]) @@ -39,6 +41,9 @@ export default function NewRunPage() { runId: "", judgeModel: "gpt-4o", answeringModel: "gpt-4o", + dataPath: "", + datasetRevision: "", + retrievalTopK: "5", selectionMode: "full" as SelectionMode, sampleType: "consecutive" as SampleType, perCategory: "2", @@ -156,6 +161,15 @@ export default function NewRunPage() { const canChangeAnsweringModel = ["indexing", "search", "answer"].includes(advancedForm.fromPhase) const selectedProvider = providers.find((p) => p.name === form.provider) + const selectedBenchmark = benchmarks.find((benchmark) => benchmark.name === form.benchmark) + const isBeam = form.benchmark.startsWith("beam-") + + useEffect(() => { + const requiredJudge = selectedBenchmark?.requiredJudge?.modelAlias + if (requiredJudge && form.judgeModel !== requiredJudge) { + setForm((current) => ({ ...current, judgeModel: requiredJudge })) + } + }, [selectedBenchmark, form.judgeModel]) useEffect(() => { if (selectedProvider) { @@ -311,6 +325,24 @@ export default function NewRunPage() { force: activeTab === "new", fromPhase, sourceRunId, + dataPath: + activeTab === "advanced" + ? selectedSourceRun?.dataPath + : isBeam + ? form.dataPath || undefined + : undefined, + datasetRevision: + activeTab === "advanced" + ? selectedSourceRun?.datasetRevision + : isBeam + ? form.datasetRevision || undefined + : undefined, + retrievalTopK: + activeTab === "advanced" + ? selectedSourceRun?.retrievalTopK + : isBeam + ? Number(form.retrievalTopK) + : undefined, }) router.push(`/runs/${encodeURIComponent(runId)}`) @@ -325,11 +357,14 @@ export default function NewRunPage() { const providerOptions = providers.map((p) => ({ value: p.name, label: p.displayName })) const benchmarkOptions = benchmarks.map((b) => ({ value: b.name, label: b.displayName })) const modelOptions = allModels.map((m) => ({ value: m.alias, label: m.displayName || m.alias })) + const judgeModelOptions = selectedBenchmark?.requiredJudge + ? modelOptions.filter((model) => model.value === selectedBenchmark.requiredJudge!.modelAlias) + : modelOptions const runOptions = completedRuns.map((r) => ({ value: r.runId, label: r.runId, - sublabel: `${r.provider} · ${r.benchmark}${r.summary.total ? ` · ${r.summary.total}q` : ""}${r.accuracy !== null ? ` · ${(r.accuracy * 100).toFixed(0)}%` : ""}`, + sublabel: `${r.provider} · ${getBenchmarkDisplayName(r.benchmark, r.benchmarkScope)}${r.summary.total ? ` · ${r.summary.total}q` : ""}${r.accuracy !== null ? ` · ${(r.accuracy * 100).toFixed(0)}%` : ""}`, })) if (loading) { @@ -513,7 +548,12 @@ export default function NewRunPage() {
Benchmark:{" "} - {selectedSourceRun?.benchmark} + {selectedSourceRun + ? getBenchmarkDisplayName( + selectedSourceRun.benchmark, + selectedSourceRun.benchmarkScope + ) + : undefined}
@@ -814,6 +854,49 @@ export default function NewRunPage() { )}
+ {isBeam && ( +
+
+ + setForm({ ...form, dataPath: event.target.value })} + placeholder="Default: data/benchmarks/beam" + className="w-full px-3 py-2.5 text-sm bg-[#222222] border border-[#333333] rounded text-text-primary placeholder-text-muted focus:outline-none focus:border-accent" + /> +
+
+ + setForm({ ...form, datasetRevision: event.target.value })} + placeholder="Optional pinned fingerprint" + className="w-full px-3 py-2.5 text-sm bg-[#222222] border border-[#333333] rounded text-text-primary placeholder-text-muted focus:outline-none focus:border-accent font-mono" + /> +
+
+ + ({ + value: String(value), + label: String(value), + }))} + selected={form.retrievalTopK} + onChange={(value) => setForm({ ...form, retrievalTopK: value })} + /> +
+
+ )} +
@@ -846,7 +929,7 @@ export default function NewRunPage() { setForm({ ...form, judgeModel: value })} placeholder="Select model" diff --git a/ui/app/runs/page.tsx b/ui/app/runs/page.tsx index a92a7b4..613b7a2 100644 --- a/ui/app/runs/page.tsx +++ b/ui/app/runs/page.tsx @@ -4,7 +4,13 @@ import { useState, useEffect, useMemo, useRef, useCallback } from "react" import Link from "next/link" import { useRouter } from "next/navigation" import { getRuns, deleteRun, stopRun, startRun, addToLeaderboard, type RunSummary } from "@/lib/api" -import { formatDate, getStatusColor, cn } from "@/lib/utils" +import { + formatDate, + getBenchmarkDisplayName, + getPipelineProgress, + getStatusColor, + cn, +} from "@/lib/utils" import { FilterBar } from "@/components/filter-bar" import { DataTable, type Column } from "@/components/data-table" import { RunActionsMenu } from "@/components/run-actions-menu" @@ -113,6 +119,9 @@ export default function RunsPage() { runId: run.runId, judgeModel: run.judge, answeringModel: run.answeringModel, + dataPath: run.dataPath, + datasetRevision: run.datasetRevision, + retrievalTopK: run.retrievalTopK, }) await refreshRuns() } catch (e) { @@ -140,7 +149,10 @@ export default function RunsPage() { }) return Object.entries(counts).map(([value, count]) => ({ value, - label: value, + label: getBenchmarkDisplayName( + value, + runs.find((run) => run.benchmark === value)?.benchmarkScope + ), count, })) }, [runs]) @@ -212,7 +224,7 @@ export default function RunsPage() { { key: "benchmark", header: "Benchmark", - render: (run) => {run.benchmark}, + render: (run) => {getBenchmarkDisplayName(run.benchmark, run.benchmarkScope)}, }, { key: "status", @@ -224,16 +236,7 @@ export default function RunsPage() { run.status === "initializing" || run.status === "stopping" const s = run.summary - const phasesCompleted = s.ingested + s.indexed + s.searched + s.answered + s.evaluated - const totalPhases = 5 * s.total - const progress = totalPhases > 0 ? phasesCompleted / totalPhases : 0 - - let phasesFullyComplete = 0 - if (s.ingested === s.total) phasesFullyComplete++ - if (s.indexed === s.total) phasesFullyComplete++ - if (s.searched === s.total) phasesFullyComplete++ - if (s.answered === s.total) phasesFullyComplete++ - if (s.evaluated === s.total) phasesFullyComplete++ + const { progress, phasesFullyComplete } = getPipelineProgress(s) return (
@@ -248,7 +251,7 @@ export default function RunsPage() { }, { key: "accuracy", - header: "Accuracy", + header: "Pass Accuracy", align: "right", render: (run) => { const accuracyPct = @@ -279,6 +282,7 @@ export default function RunsPage() { runId={run.runId} provider={run.provider} benchmark={run.benchmark} + benchmarkScope={run.benchmarkScope} status={run.status} onAddToLeaderboard={(data) => handleAddToLeaderboard(run.runId, data)} onDelete={() => handleDelete(run.runId)} diff --git a/ui/components/benchmark-results.tsx b/ui/components/benchmark-results.tsx index 0fb6bb6..318e161 100644 --- a/ui/components/benchmark-results.tsx +++ b/ui/components/benchmark-results.tsx @@ -1,8 +1,9 @@ "use client" import { useState, useMemo } from "react" -import { cn } from "@/lib/utils" +import { cn, isEvaluationFailed, isEvaluationPassed } from "@/lib/utils" import { MultiSelect } from "@/components/multi-select" +import type { BuildReport, QuestionMetric, RunCostReport } from "@/lib/api" function Tooltip({ text, children }: { text: string; children: React.ReactNode }) { const [show, setShow] = useState(false) @@ -62,6 +63,240 @@ export function StatsGrid({ cards }: StatsGridProps) { ) } +export function BuildMetricsTable({ builds }: { builds?: BuildReport | null }) { + if (!builds) return null + + return ( +
+

Build Metrics

+

+ One-time cost per container; never multiplied by question count. +

+
+ + + + +
+ {builds.items.length > 0 && ( +
+ + + + + + + + + + + + + {builds.items.map((build) => ( + + + + + + + + + ))} + +
+ container + + ingest + + index + + wall-clock + + work + + cost +
+ {build.containerTag} + {build.reused && reused} + {!build.reused && build.reusedPhases?.ingest && ( + ingest reused + )} + + {build.ingestLatencyMs}ms + + {build.indexingLatencyMs}ms + + {build.buildWallClockMs}ms + + {build.buildWorkMs}ms + + {build.costUsd == null ? "—" : `$${build.costUsd.toFixed(4)}`} +
+
+ )} +
+ ) +} + +export function QuestionMetricsSummary({ + metrics, + costs, +}: { + metrics?: QuestionMetric[] | null + costs?: RunCostReport | null +}) { + if (!metrics?.length) return null + const mean = (values: number[]) => values.reduce((sum, value) => sum + value, 0) / values.length + const amortized = metrics + .map((metric) => metric.amortizedOnlinePlusBuildWorkMs) + .filter((value): value is number => value != null) + const allocationDenominators = Array.from( + new Set( + metrics + .map((metric) => metric.buildAllocationQuestionCount) + .filter((value): value is number => value != null) + ) + ).sort((left, right) => left - right) + const summarizeCosts = (values: Array): RunCostReport["query"] => { + const known = values.filter((value): value is number => value != null) + return { + totalCostUsd: + values.length > 0 && known.length === values.length + ? known.reduce((sum, value) => sum + value, 0) + : null, + knownCostCount: known.length, + totalCostCount: values.length, + } + } + const effectiveCosts = + costs ?? + ({ + query: summarizeCosts(metrics.map((metric) => metric.queryCostUsd)), + evaluation: summarizeCosts(metrics.map((metric) => metric.evaluationCostUsd)), + } satisfies RunCostReport) + const evaluationUsage = metrics + .map((metric) => metric.evaluationUsage) + .filter((usage): usage is NonNullable => usage != null) + const sumEvaluationUsage = (field: keyof (typeof evaluationUsage)[number]) => + evaluationUsage.reduce((sum, usage) => sum + (usage[field] ?? 0), 0) + const evaluationRequestCount = sumEvaluationUsage("requestCount") + const completeTokenUsageCount = sumEvaluationUsage("tokenUsageCompleteRequestCount") + const partialTokenUsageCount = sumEvaluationUsage("tokenUsagePartialRequestCount") + const unknownTokenUsageCount = sumEvaluationUsage("tokenUsageUnknownRequestCount") + const classifiedTokenUsageCount = + completeTokenUsageCount + partialTokenUsageCount + unknownTokenUsageCount + const formatCost = (value: number | null) => (value == null ? "—" : `$${value.toFixed(4)}`) + const amortizationSubtext = + allocationDenominators.length === 0 + ? "online + allocated build work; denominator unavailable" + : allocationDenominators.length === 1 + ? `online + build work ÷ ${allocationDenominators[0]} completed question${allocationDenominators[0] === 1 ? "" : "s"} for its build` + : `online + build work ÷ per-build completed-question counts (${allocationDenominators.join(", ")})` + const recordedRetrievalMetrics = metrics.filter((metric) => + [ + metric.configuredTopK, + metric.rawReturnedCount, + metric.normalizedCount, + metric.answerEvidenceCount, + ].some((value) => value != null) + ) + const onlyNumbers = (values: Array) => + values.filter((value): value is number => value != null) + const uniqueTopK = Array.from( + new Set(onlyNumbers(recordedRetrievalMetrics.map((metric) => metric.configuredTopK))) + ).sort((a, b) => a - b) + const uniqueProviderLimits = Array.from( + new Set(onlyNumbers(recordedRetrievalMetrics.map((metric) => metric.providerRequestLimit))) + ).sort((a, b) => a - b) + const formatContractValue = (values: number[]) => + values.length === 0 + ? "—" + : values.length === 1 + ? String(values[0]) + : `mixed (${values.join(", ")})` + const formatMean = (values: Array) => { + const recorded = onlyNumbers(values) + return recorded.length > 0 ? mean(recorded).toFixed(1) : "—" + } + const totalDropped = recordedRetrievalMetrics.reduce( + (sum, metric) => sum + (metric.droppedCount ?? 0), + 0 + ) + + return ( +
+

Per-question Metrics

+

+ Online query and offline evaluation are reported separately from build work. +

+
+ metric.onlineQueryLatencyMs)).toFixed(1)}ms`} + /> + metric.evaluationLatencyMs)).toFixed(1)}ms`} + /> + 0 ? `${mean(amortized).toFixed(1)}ms` : "—"} + subtext={amortizationSubtext} + /> + + {evaluationRequestCount > 0 && ( + 0 + ? `${completeTokenUsageCount} complete, ${partialTokenUsageCount} partial, ${unknownTokenUsageCount} unknown token usage` + : "token-usage coverage unavailable for this report" + } + /> + )} +
+ {recordedRetrievalMetrics.length > 0 && ( + <> +

+ Retrieval contract +

+
+ + metric.rawReturnedCount))} / ${formatMean(recordedRetrievalMetrics.map((metric) => metric.normalizedCount))}`} + subtext={`${totalDropped} result${totalDropped === 1 ? "" : "s"} dropped during normalization`} + /> + metric.answerEvidenceCount))} / ${formatMean(recordedRetrievalMetrics.map((metric) => metric.answerCutoff))}`} + subtext="results actually placed in the answer prompt" + /> + metric.contextTokens))} + subtext="retrieved-context prompt tokens" + /> +
+ + )} +
+ ) +} + export interface QuestionTypeStats { accuracy: number correct: number @@ -70,30 +305,45 @@ export interface QuestionTypeStats { export interface AccuracyByTypeProps { byQuestionType: Record + qualityBySlice?: Record> } -export function AccuracyByType({ byQuestionType }: AccuracyByTypeProps) { - if (!byQuestionType || Object.keys(byQuestionType).length === 0) { - return null - } +export function AccuracyByType({ byQuestionType, qualityBySlice }: AccuracyByTypeProps) { + const types = Array.from( + new Set([ + ...Object.keys(byQuestionType || {}), + ...Object.keys(qualityBySlice || {}).filter((type) => !type.startsWith("tier:")), + ]) + ).sort() + if (types.length === 0) return null return (
-

Accuracy by Question Type

+

Quality by Question Type

- {Object.entries(byQuestionType).map(([type, stats]) => ( -
-
- {type.replace(/[-_]/g, " ")} -
-
- {(stats.accuracy * 100).toFixed(0)}% -
-
- {stats.correct}/{stats.total} + {types.map((type) => { + const stats = byQuestionType?.[type] + const averageScore = qualityBySlice?.[type]?.averageScore + const passAccuracy = qualityBySlice?.[type]?.passAccuracy ?? stats?.accuracy + const displayValue = averageScore ?? stats?.accuracy + return ( +
+
+ {type.replace(/[-_]/g, " ")} +
+
+ {displayValue != null ? `${(displayValue * 100).toFixed(0)}%` : "—"} +
+
+ {averageScore != null + ? `avg score${passAccuracy != null ? ` · ${(passAccuracy * 100).toFixed(0)}% pass` : ""}` + : stats + ? `${stats.correct}/${stats.total}` + : "average score unavailable"} +
-
- ))} + ) + })}
) @@ -142,6 +392,9 @@ export function LatencyTable({ latency }: LatencyTableProps) { phase + + scope + min @@ -167,9 +420,17 @@ export function LatencyTable({ latency }: LatencyTableProps) { (phase) => { const stats = latency[phase] if (!stats) return null + const scope = + phase === "ingest" || phase === "indexing" + ? "build / container" + : phase === "evaluate" + ? "offline / question" + : "online / question" + const phaseLabel = phase === "total" ? "online total" : phase return ( - {phase} + {phaseLabel} + {scope} {stats.min} @@ -324,8 +585,20 @@ export interface EvaluationResult { groundTruth: string hypothesis?: string score?: number + primaryScore?: number + passed?: boolean label?: string explanation?: string + metrics?: Record + details?: Record + evaluation?: { + primaryScore?: number + passed?: boolean + label?: string + explanation?: string + metrics?: Record + details?: Record + } } export interface EvaluationListProps { @@ -353,12 +626,12 @@ export function EvaluationList({ evaluations, onViewDetails }: EvaluationListPro }, [evaluations]) const failureCount = useMemo(() => { - return evaluations.filter((e) => e.label === "incorrect" || e.score === 0).length + return evaluations.filter((e) => isEvaluationFailed(e)).length }, [evaluations]) const filtered = useMemo(() => { return evaluations.filter((e) => { - if (showFailuresOnly && e.label !== "incorrect" && e.score !== 0) { + if (showFailuresOnly && !isEvaluationFailed(e)) { return false } @@ -478,8 +751,13 @@ export function EvaluationList({ evaluations, onViewDetails }: EvaluationListPro
{filtered.map((evaluation, idx) => { const isExpanded = expandedId === evaluation.questionId - const isCorrect = evaluation.score === 1 || evaluation.label === "correct" + const isCorrect = isEvaluationPassed(evaluation) const isLast = idx === filtered.length - 1 + const primaryScore = + evaluation.evaluation?.primaryScore ?? evaluation.primaryScore ?? evaluation.score + const metrics = evaluation.evaluation?.metrics ?? evaluation.metrics + const details = evaluation.evaluation?.details ?? evaluation.details + const explanation = evaluation.evaluation?.explanation ?? evaluation.explanation return (
- {evaluation.label} + {evaluation.label || (isCorrect ? "correct" : "incorrect")} + {primaryScore != null && ( + + {primaryScore.toFixed(3)} + + )} + {onViewDetails && (
- {evaluation.explanation && ( + {metrics && Object.keys(metrics).length > 0 && ( +
+
+ Protocol metrics +
+
+ {Object.entries(metrics).map(([metric, value]) => ( +
+
+ {metric} +
+
+ {Number.isFinite(value) ? value.toFixed(4) : String(value)} +
+
+ ))} +
+
+ )} + + {explanation && (
Explanation
-
- {evaluation.explanation} +
{explanation}
+
+ )} + + {details && Object.keys(details).length > 0 && ( +
+
+ Protocol details
+ 
+                          {JSON.stringify(details, null, 2)}
+
)}
diff --git a/ui/components/phase-progress.tsx b/ui/components/phase-progress.tsx index 1b99904..ac1c2fa 100644 --- a/ui/components/phase-progress.tsx +++ b/ui/components/phase-progress.tsx @@ -1,17 +1,11 @@ "use client" import { useState } from "react" -import { cn } from "@/lib/utils" +import { cn, getPipelinePhaseTotal, type PipelinePhaseKey, type PipelineSummary } from "@/lib/utils" import { Tooltip } from "@/components/tooltip" interface PhaseProgressProps { - summary: { - total: number - ingested: number - indexed: number - searched: number - answered: number - evaluated: number + summary: PipelineSummary & { indexingEpisodes?: { total: number completed: number @@ -61,9 +55,10 @@ export function PhaseProgress({ summary }: PhaseProgressProps) {
{phases.map((phase) => { const count = summary[phase.key] - const progress = (count / summary.total) * 100 - const isComplete = count === summary.total - const isInProgress = count > 0 && count < summary.total + const phaseTotal = getPipelinePhaseTotal(summary, phase.key as PipelinePhaseKey) + const progress = phaseTotal > 0 ? (count / phaseTotal) * 100 : 0 + const isComplete = phaseTotal > 0 && count >= phaseTotal + const isInProgress = count > 0 && count < phaseTotal const isPending = count === 0 const episodes = summary.indexingEpisodes @@ -76,7 +71,7 @@ export function PhaseProgress({ summary }: PhaseProgressProps) { const displayLabel = isShowingEpisodes ? "Episodes Indexed" : phase.label const displayCount = isShowingEpisodes ? episodes.completed : count - const displayTotal = isShowingEpisodes ? episodes.total : summary.total + const displayTotal = isShowingEpisodes ? episodes.total : phaseTotal const displayProgress = isShowingEpisodes ? (episodes.completed / episodes.total) * 100 : progress @@ -131,7 +126,7 @@ export function PhaseProgress({ summary }: PhaseProgressProps) {
{phase.label} - {count}/{summary.total} + {count}/{phaseTotal}
diff --git a/ui/components/question-list.tsx b/ui/components/question-list.tsx index 697962f..3318583 100644 --- a/ui/components/question-list.tsx +++ b/ui/components/question-list.tsx @@ -2,7 +2,7 @@ import { useState, useMemo } from "react" import Link from "next/link" -import { cn } from "@/lib/utils" +import { cn, isEvaluationFailed, isEvaluationPassed } from "@/lib/utils" import { MultiSelect } from "./multi-select" import type { QuestionCheckpoint, QuestionTypeRegistry } from "@/lib/api" @@ -33,14 +33,14 @@ export function QuestionList({ runId, questions, questionTypeRegistry }: Questio // Count failures const failureCount = useMemo(() => { - return questions.filter((q) => q.phases.evaluate.label === "incorrect").length + return questions.filter((q) => isEvaluationFailed(q.phases.evaluate)).length }, [questions]) // Filter questions const filtered = useMemo(() => { return questions.filter((q) => { // Failures filter - if (showFailuresOnly && q.phases.evaluate.label !== "incorrect") { + if (showFailuresOnly && !isEvaluationFailed(q.phases.evaluate)) { return false } @@ -166,7 +166,7 @@ export function QuestionList({ runId, questions, questionTypeRegistry }: Questio ) : (
{filtered.map((q, idx) => { - const isCorrect = q.phases.evaluate.label === "correct" + const isCorrect = isEvaluationPassed(q.phases.evaluate) const isExpanded = expanded === q.questionId const isLast = idx === filtered.length - 1 diff --git a/ui/components/run-actions-menu.tsx b/ui/components/run-actions-menu.tsx index eba2857..389b1ea 100644 --- a/ui/components/run-actions-menu.tsx +++ b/ui/components/run-actions-menu.tsx @@ -3,12 +3,13 @@ import { useState, useRef, useEffect } from "react" import { createPortal } from "react-dom" import Link from "next/link" -import { cn } from "@/lib/utils" +import { cn, getBenchmarkDisplayName } from "@/lib/utils" interface RunActionsMenuProps { runId: string provider: string benchmark: string + benchmarkScope?: Record status: string onAddToLeaderboard: (data: { version?: string; notes?: string }) => Promise onDelete: () => void @@ -20,6 +21,7 @@ export function RunActionsMenu({ runId, provider, benchmark, + benchmarkScope, status, onAddToLeaderboard, onDelete, @@ -205,6 +207,7 @@ export function RunActionsMenu({ position={popoverPosition} provider={provider} benchmark={benchmark} + benchmarkScope={benchmarkScope} onSubmit={async (data) => { await onAddToLeaderboard(data) setShowLeaderboardPopover(false) @@ -224,12 +227,13 @@ interface LeaderboardPopoverProps { position: { top: number; left: number } provider: string benchmark: string + benchmarkScope?: Record onSubmit: (data: { version?: string; notes?: string }) => Promise onClose: () => void } const LeaderboardPopover = forwardRef( - ({ position, provider, benchmark, onSubmit, onClose }, ref) => { + ({ position, provider, benchmark, benchmarkScope, onSubmit, onClose }, ref) => { const [editingVersion, setEditingVersion] = useState(false) const [version, setVersion] = useState("") const [notes, setNotes] = useState("") @@ -291,7 +295,7 @@ const LeaderboardPopover = forwardRef(
{provider} / - {benchmark} + {getBenchmarkDisplayName(benchmark, benchmarkScope)}
{/* Version */} diff --git a/ui/lib/api.ts b/ui/lib/api.ts index 8a22190..74213bb 100644 --- a/ui/lib/api.ts +++ b/ui/lib/api.ts @@ -11,6 +11,7 @@ export interface RunSummary { status: "initializing" | "pending" | "running" | "stopping" | "completed" | "partial" | "failed" summary: { total: number + builds?: number ingested: number indexed: number searched: number @@ -23,24 +24,99 @@ export interface RunSummary { } } accuracy: number | null + dataPath?: string + datasetRevision?: string + retrievalTopK?: number + benchmarkScope?: Record + datasetIdentity?: Record + benchmarkInputFingerprint?: string + protocolIdentity?: Record + selectedQuestionIdsDigest?: string + providerPromptFingerprint?: string +} + +export interface ProtocolEvaluation { + primaryScore: number + passed?: boolean + label?: string + explanation?: string + metrics?: Record + details?: Record +} + +export interface BuildCheckpoint { + buildId: string + containerTag: string + ingest: { status: string; error?: string } + indexing: { status: string; error?: string; failedIds?: string[] } } export interface QuestionCheckpoint { questionId: string - containerTag: string + buildId?: string + /** Legacy mirror returned by the question-detail endpoint. */ + containerTag?: string + build?: BuildCheckpoint question: string groundTruth: string questionType: string phases: { - ingest: { status: string; completedSessions: string[] } - indexing: { status: string } - search: { status: string; results?: any[] } - answer: { status: string; hypothesis?: string } - evaluate: { status: string; score?: number; label?: string; explanation?: string } + /** Legacy schema only; shared-build checkpoints keep these phases on `build`. */ + ingest?: { status: string; completedSessions?: string[]; error?: string } + indexing?: { status: string; error?: string } + search: { + status: string + retrievalPlan?: { + query: string + requestedTopK: number + answerCutoff: number + threshold?: number + searchMode?: string + filters?: Record + } + results?: any[] + requestedCount?: number + rawReturnedCount?: number + returnedCount?: number + normalizedCount?: number + droppedCount?: number + providerRequests?: Array<{ + operation: string + limit: number + parameters?: Record + }> + answerCutoff?: number + answerEvidenceCount?: number + durationMs?: number + error?: string + } + answer: { + status: string + hypothesis?: string + promptTokens?: number + basePromptTokens?: number + contextTokens?: number + evidenceCount?: number + durationMs?: number + error?: string + } + evaluate: { + status: string + evaluation?: ProtocolEvaluation + score?: number + primaryScore?: number + passed?: boolean + label?: string + explanation?: string + metrics?: Record + details?: Record + error?: string + } } } export interface RunDetail extends RunSummary { + builds?: Record questions: Record } @@ -54,6 +130,16 @@ export interface Benchmark { name: string displayName: string description: string + scope?: { + displayName: string + includedTiers: string[] + coverage: "full" | "subset" + } + requiredJudge?: { + provider: string + modelId: string + modelAlias: string + } } export interface QuestionTypeInfo { @@ -103,7 +189,7 @@ export async function getRun(runId: string): Promise { return fetchApi(`/api/runs/${encodeURIComponent(runId)}`) } -export async function getRunReport(runId: string): Promise { +export async function getRunReport(runId: string): Promise { return fetchApi(`/api/runs/${encodeURIComponent(runId)}/report`) } @@ -180,6 +266,9 @@ export async function startRun(params: { force?: boolean fromPhase?: PhaseId sourceRunId?: string + dataPath?: string + datasetRevision?: string + retrievalTopK?: number }): Promise<{ message: string; runId: string }> { return fetchApi("/api/runs/start", { method: "POST", @@ -203,7 +292,14 @@ export async function getBenchmarks(): Promise<{ benchmarks: Benchmark[] }> { export async function getBenchmarkQuestions( benchmark: string, - params?: { page?: number; limit?: number; type?: string } + params?: { + page?: number + limit?: number + type?: string + dataPath?: string + datasetRevision?: string + retrievalTopK?: number + } ): Promise< PaginatedResponse<{ questionId: string @@ -216,6 +312,11 @@ export async function getBenchmarkQuestions( if (params?.page) searchParams.set("page", params.page.toString()) if (params?.limit) searchParams.set("limit", params.limit.toString()) if (params?.type) searchParams.set("type", params.type) + if (params?.dataPath) searchParams.set("dataPath", params.dataPath) + if (params?.datasetRevision) searchParams.set("datasetRevision", params.datasetRevision) + if (params?.retrievalTopK != null) { + searchParams.set("retrievalTopK", params.retrievalTopK.toString()) + } const query = searchParams.toString() return fetchApi(`/api/benchmarks/${benchmark}/questions${query ? `?${query}` : ""}`) @@ -248,6 +349,139 @@ export interface LatencyByPhase { total: LatencyStats } +export interface RetrievalSummary { + hitAtK: number + precisionAtK: number + recallAtK: number + f1AtK: number + mrr: number + ndcg: number + k: number +} + +export interface QuestionTypeSummary { + total: number + correct: number + accuracy: number + averageScore?: number + passAccuracy?: number + retrieval?: RetrievalSummary +} + +export interface BenchmarkQuality { + primaryMetric?: { key: string; value: number; higherIsBetter: boolean } + metrics: Record + bySlice?: Record> +} + +export interface LeaderboardComparisonIdentity { + schemaVersion: 3 + benchmark: string + benchmarkScope: Record + datasetIdentity: Record + datasetFingerprint: string + questionSetFingerprint: string + benchmarkInputFingerprint: string + protocolIdentity: Record + protocolFingerprint: string + retrievalTopK: number | null + judgeModel: string + answeringModel: string + answeringRuntimeFingerprint: string + providerPromptFingerprint: string | null + primaryMetric: { key: string; value: number; higherIsBetter: boolean } + cohortKey: string + legacy: boolean +} + +export interface BuildMetric { + buildId: string + containerTag: string + sourceRunId?: string + reused?: boolean + reusedPhases?: { ingest: boolean; indexing: boolean } + ingestLatencyMs: number + indexingLatencyMs: number + buildWallClockMs: number + buildWorkMs: number + attemptCount?: number + attempts?: Array<{ + phase: "ingest" | "indexing" + attempt: number + startedAt: string + completedAt?: string + durationMs?: number + status: "in_progress" | "completed" | "failed" + costUsd: number | null + error?: string + }> + usage?: UsageMetric + costUsd?: number | null +} + +export interface UsageMetric { + requestCount?: number + tokenUsageCompleteRequestCount?: number + tokenUsagePartialRequestCount?: number + tokenUsageUnknownRequestCount?: number + inputTokens?: number + outputTokens?: number + totalTokens?: number +} + +export interface BuildReport { + uniqueBuildCount: number + sumContainerBuildWorkMs: number + buildPhaseWallClockMs: number + totalBuildCostUsd: number | null + knownCostBuildCount: number + totalCostBuildCount: number + items: BuildMetric[] +} + +export interface CostCoverageReport { + totalCostUsd: number | null + knownCostCount: number + totalCostCount: number +} + +export interface RunCostReport { + query: CostCoverageReport + evaluation: CostCoverageReport +} + +export interface QuestionMetric { + questionId: string + buildId: string + searchLatencyMs: number + answerLatencyMs: number + onlineQueryLatencyMs: number + evaluationLatencyMs: number + queryUsage?: UsageMetric + evaluationUsage?: UsageMetric + queryCostUsd: number | null + evaluationCostUsd: number | null + configuredTopK?: number + providerRequestLimit?: number + rawReturnedCount?: number + returnedCount?: number + normalizedCount?: number + droppedCount?: number + answerCutoff?: number + answerEvidenceCount?: number + contextTokens?: number + searchMode?: string + threshold?: number + providerRequests?: Array<{ + operation: string + limit: number + parameters?: Record + }> + buildAllocationQuestionCount?: number + allocatedBuildWorkMs?: number + amortizedOnlinePlusBuildWorkMs?: number +} + // Evaluation result for individual questions export interface EvaluationResult { questionId: string @@ -256,8 +490,12 @@ export interface EvaluationResult { groundTruth: string hypothesis: string score: number + passed?: boolean + primaryScore?: number label: string explanation: string + metrics?: Record + details?: Record searchResults?: any[] searchDurationMs?: number answerDurationMs?: number @@ -274,9 +512,28 @@ export interface LeaderboardEntry { accuracy: number totalQuestions: number correctCount: number - byQuestionType: Record + averageScore?: number + benchmarkScope: Record + datasetIdentity: Record + datasetFingerprint: string + questionSetFingerprint: string + benchmarkInputFingerprint: string + protocolFingerprint: string + retrievalTopK: number | null + primaryMetric: { key: string; value: number; higherIsBetter: boolean } + comparisonIdentity: LeaderboardComparisonIdentity + cohortKey: string + cohortRank: number + cohortSize: number + quality?: BenchmarkQuality + builds?: BuildReport + costs?: RunCostReport + questionMetrics?: QuestionMetric[] + protocolIdentity?: Record + byQuestionType: Record questionTypeRegistry: QuestionTypeRegistry | null latencyStats: LatencyByPhase | null + retrieval?: RetrievalSummary evaluations: EvaluationResult[] providerCode: string promptsUsed: Record | null @@ -335,6 +592,7 @@ export interface CompareRunInfo { error?: string progress?: { total: number + builds?: number ingested: number indexed: number searched: number @@ -353,6 +611,7 @@ export interface CompareRunProgress { runId: string progress: { total: number + builds?: number ingested: number indexed: number searched: number @@ -366,6 +625,10 @@ export interface CompareSummary { compareId: string providers: string[] benchmark: string + benchmarkScope?: Record + datasetIdentity?: Record + benchmarkInputFingerprint?: string + selectedQuestionIdsDigest?: string judge: string answeringModel: string status: CompareStatus @@ -385,7 +648,18 @@ export interface BenchmarkResult { runId: string provider: string benchmark: string + selectedQuestionIdsDigest?: string + retrievalTopK?: number + datasetIdentity?: Record + providerPromptFingerprint?: string version?: string + memscore?: string + benchmarkScope?: { + displayName: string + includedTiers: string[] + coverage: "full" | "subset" + } + protocolIdentity?: Record // Fields can be at root level or nested in summary accuracy?: number totalQuestions?: number @@ -394,20 +668,17 @@ export interface BenchmarkResult { totalQuestions: number correctCount: number accuracy: number + averageScore?: number } - byQuestionType: Record + quality?: BenchmarkQuality + builds?: BuildReport + costs?: RunCostReport + questionMetrics?: QuestionMetric[] + byQuestionType: Record questionTypeRegistry: QuestionTypeRegistry | null latency?: LatencyByPhase latencyStats?: LatencyByPhase | null - retrieval?: { - hitAtK: number - precisionAtK: number - recallAtK: number - f1AtK: number - mrr: number - ndcg: number - k: number - } + retrieval?: RetrievalSummary evaluations?: EvaluationResult[] providerCode?: string promptsUsed?: Record | null @@ -420,6 +691,14 @@ export interface CompareReport { benchmark: string judge: string answeringModel: string + comparison: { + comparable: boolean + identity?: { key: string; higherIsBetter: boolean } + identities: string[] + mismatchReasons: string[] + bestValue?: number + winners: string[] + } reports: Array<{ provider: string report: BenchmarkResult @@ -445,6 +724,9 @@ export async function startCompare(params: { judgeModel: string answeringModel?: string sampling?: SamplingConfig + dataPath?: string + datasetRevision?: string + retrievalTopK?: number }): Promise<{ message: string; compareId: string }> { return fetchApi("/api/compare/start", { method: "POST", diff --git a/ui/lib/utils.ts b/ui/lib/utils.ts index 4a09f34..cebcf4d 100644 --- a/ui/lib/utils.ts +++ b/ui/lib/utils.ts @@ -51,6 +51,21 @@ export function formatDuration(ms: number): string { return `${(ms / 60000).toFixed(1)}m` } +export function getBenchmarkDisplayName( + benchmark: string, + scope?: Record | { displayName?: string } +): string { + if (typeof scope?.displayName === "string" && scope.displayName.trim()) { + return scope.displayName + } + const names: Record = { + "beam-1m": "BEAM 1M", + "beam-10m": "BEAM 10M", + "beam-1m-10m": "BEAM 1M/10M", + } + return names[benchmark] ?? benchmark +} + export function getStatusColor(status: string): string { switch (status) { case "completed": @@ -68,6 +83,85 @@ export function getStatusColor(status: string): string { } } +export interface EvaluationCompatibility { + passed?: boolean + label?: string + score?: number + primaryScore?: number + evaluation?: EvaluationCompatibility +} + +/** Resolve protocol-native judgments first, then compatibility mirrors and legacy binary scores. */ +export function getEvaluationPassState( + value?: EvaluationCompatibility | null +): boolean | undefined { + const protocolEvaluation = value?.evaluation + if (typeof protocolEvaluation?.passed === "boolean") return protocolEvaluation.passed + if (typeof value?.passed === "boolean") return value.passed + + for (const label of [protocolEvaluation?.label, value?.label]) { + if (typeof label !== "string") continue + const normalized = label.toLowerCase() + if (normalized === "pass" || normalized === "correct") return true + if (normalized === "fail" || normalized === "incorrect" || normalized === "wrong") { + return false + } + } + + const legacyScore = value?.score ?? protocolEvaluation?.primaryScore ?? value?.primaryScore + return typeof legacyScore === "number" ? legacyScore === 1 : undefined +} + +export function isEvaluationPassed(value?: EvaluationCompatibility | null): boolean { + return getEvaluationPassState(value) === true +} + +export function isEvaluationFailed(value?: EvaluationCompatibility | null): boolean { + return getEvaluationPassState(value) === false +} + +export type PipelinePhaseKey = "ingested" | "indexed" | "searched" | "answered" | "evaluated" + +export interface PipelineSummary { + total: number + builds?: number + ingested: number + indexed: number + searched: number + answered: number + evaluated: number +} + +/** Shared-build runs count ingest/index per build; legacy summaries counted every phase per question. */ +export function getPipelinePhaseTotal(summary: PipelineSummary, phase: PipelinePhaseKey): number { + return phase === "ingested" || phase === "indexed" + ? (summary.builds ?? summary.total) + : summary.total +} + +export function getPipelineProgress(summary: PipelineSummary): { + progress: number + phasesFullyComplete: number +} { + const phases: PipelinePhaseKey[] = ["ingested", "indexed", "searched", "answered", "evaluated"] + let completedWork = 0 + let totalWork = 0 + let phasesFullyComplete = 0 + + for (const phase of phases) { + const phaseTotal = getPipelinePhaseTotal(summary, phase) + const completed = summary[phase] + completedWork += Math.min(completed, phaseTotal) + totalWork += phaseTotal + if (phaseTotal > 0 && completed >= phaseTotal) phasesFullyComplete++ + } + + return { + progress: totalWork > 0 ? completedWork / totalWork : 0, + phasesFullyComplete, + } +} + export function calculateAccuracy( summary: { total: number; evaluated: number } & Record, questions?: Record @@ -79,6 +173,6 @@ export function calculateAccuracy( ) if (evaluated.length === 0) return null - const correct = evaluated.filter((q: any) => q.phases?.evaluate?.score === 1).length + const correct = evaluated.filter((q: any) => isEvaluationPassed(q.phases?.evaluate)).length return (correct / evaluated.length) * 100 }