From 3b3c4a9caa7e623bde2d8210a82a8a323957c4e8 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Tue, 7 Jul 2026 14:01:09 -0600 Subject: [PATCH] feat(retrieval): add RAG eval improvement loop --- README.md | 34 +- docs/architecture.md | 12 +- docs/eval/rag-eval-roadmap.md | 114 ++++++ src/index.ts | 1 + src/retrieval-eval.ts | 724 ++++++++++++++++++++++++++++++++++ tests/retrieval-eval.test.ts | 271 +++++++++++++ 6 files changed, 1150 insertions(+), 6 deletions(-) create mode 100644 docs/eval/rag-eval-roadmap.md create mode 100644 src/retrieval-eval.ts create mode 100644 tests/retrieval-eval.test.ts diff --git a/README.md b/README.md index 039c3ce..3713c6e 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ This package turns raw sources and generated markdown knowledge into a versionab - [Start here](#start-here) — pick CLI vs programmatic - [CLI](#cli) — `init` → `source-add` → `index` → `search` → `lint` - [Design](#design) — the invariants (immutable sources, cited claims, deterministic graph) -- [Agent-Eval integration](#agent-eval-integration) — readiness bundles + release reports +- [Agent-Eval integration](#agent-eval-integration) — retrieval eval + readiness bundles + release reports - [Memory adapters](#memory-adapters) — generic memory contract + Neo4j Agent Memory bridge - [Research loop](#research-loop) — `runKnowledgeResearchLoop` + control-loop adapter - [Researcher profile](#researcher-profile) — sandbox `AgentProfile` for `runLoop` @@ -31,6 +31,7 @@ Two ways in, depending on what you're doing: - *"Does the agent have enough context to run?"* → [`buildEvalKnowledgeBundle`](#agent-eval-integration) (block / ask / acquire before execution). - *"Grow the KB as a researcher"* → [`runKnowledgeResearchLoop`](#research-loop) (deterministic mechanics; your agent owns judgment), [`runTwoAgentResearchLoop`](#two-agent-research-loop) (researcher proposes, verifier checks + fills gaps, offline), or the sandbox [researcher profile](#researcher-profile) for `runLoop`. - *"Spawn one researcher per sub-topic and stop when the KB is ready"* → [`runResearchSupervisor`](#research-supervisor) (a supervisor brain sizes the topology over a `Scope`; LIVE, needs creds). + - *"Tune retrieval for a knowledge base"* → `runRetrievalImprovementLoop` in the [Agent-Eval integration](#agent-eval-integration) section. - *"Does this candidate KB actually improve task success?"* → run an [agent-eval improvement loop](#agent-eval-integration) over KB variants, then `knowledgeReleaseReport` for the promotion decision. - *"Keep live authorities fresh"* → [pluggable sources](#pluggable-knowledge-sources) + `detectChanges` → eval re-runs. @@ -101,6 +102,7 @@ from `@tangle-network/agent-knowledge`. when comparing against natural confidence thresholds. The normalization is within-set ranking, not a cross-query absolute confidence. - Release confidence uses `@tangle-network/agent-eval` release gates (`evaluateReleaseConfidence`) instead of reimplementing them. +- Retrieval eval turns retrieval/RAG configs into `agent-eval` surfaces, auto-searches candidate configs, and scores them against page, source, source-anchor, or source-span targets. - `buildEvalKnowledgeBundle()` maps wiki/search evidence into `agent-eval` `KnowledgeRequirement`, `KnowledgeBundle`, and `KnowledgeReadinessReport` contracts so control loops can block, ask, or @@ -114,6 +116,36 @@ readiness/eval machinery without making `agent-knowledge` own the database. ## Agent-Eval Integration +Use retrieval eval when the question is whether a retrieval/RAG config can find the right knowledge before an agent reasons over it. +The labels should name stable pages, source records, anchors, or source spans, not ephemeral chunk IDs. +The completion roadmap is in [`docs/eval/rag-eval-roadmap.md`](docs/eval/rag-eval-roadmap.md). + +```ts +import { runRetrievalImprovementLoop } from '@tangle-network/agent-knowledge' + +const result = await runRetrievalImprovementLoop({ + baseline: { k: 5, hybrid: false, reranker: null }, + scenarios: trainRetrievalScenarios, + holdoutScenarios: holdoutRetrievalScenarios, + index, + searchSpace: { + k: [5, 10, 20], + hybrid: [false, true], + reranker: [null, 'bge-reranker'], + }, + targetRecall: 0.9, + deltaThreshold: 0.02, + costCeiling: 15, + runDir: '.agent-knowledge/retrieval-runs', +}) + +console.log(result.winnerConfig) +``` + +Pass a custom `retrieve` function to `buildRetrievalEvalDispatch` when the config controls an external vector store, reranker, hybrid search service, or chunker. +The built-in fallback uses `searchKnowledge` over the local deterministic index. +Use `buildRetrievalEvalDispatch`, `retrievalRecallJudge`, and `retrievalParameterSweepProposer` directly only when you need custom `agent-eval` wiring. + To answer whether a candidate knowledge base actually improves agent task success, run an `@tangle-network/agent-eval` improvement loop (`runImprovementLoop`) over your KB variants on a real task corpus; each run is scored into a `RunRecord`. Use `knowledgeReleaseReport()` before promotion: pass the candidate and baseline `RunRecord[]` (plus optional `ReleaseTraceEvidence` and the gate decision) and it folds them into a `ReleaseConfidenceScorecard` and a `KnowledgeRelease` using `agent-eval`'s release gates and `RunRecord` validation. diff --git a/docs/architecture.md b/docs/architecture.md index 6800f7c..44feb7b 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -1,13 +1,14 @@ # Architecture -`@tangle-network/agent-knowledge` is a domain-agnostic knowledge growth layer for agents. +`@tangle-network/agent-knowledge` is a domain-agnostic knowledge-base construction layer for agents. -It does not try to be a vector database, a RAG framework, or a product-specific wiki. It owns the small set of primitives every serious agent knowledge system needs: +It owns the small set of primitives every serious agent knowledge system needs: - immutable source records - generated knowledge pages and units - claims with source references - deterministic indexing, graph construction, search, and lint +- retrieval/RAG candidate surfaces, gold-target scoring, and eval-loop adapters - safe LLM write proposals - eval-gated release confidence through `@tangle-network/agent-eval` - visualization DTOs under the `/viz` subpath @@ -20,9 +21,10 @@ It does not try to be a vector database, a RAG framework, or a product-specific `agent-eval` owns traces, ASI, improvement loops, run records, and promotion gates. -`agent-knowledge` owns sources, claims, pages, graph/search/lint, and knowledge base candidates. It calls `agent-eval` instead of reimplementing evaluation. +`agent-knowledge` owns sources, claims, pages, graph/search/lint, retrieval/RAG construction surfaces, and knowledge base candidates. +It calls `agent-eval` instead of reimplementing improvement loops or promotion math. -Product apps own domain policies, source adapters, task corpora, and promotion decisions. +Product apps own domain policies, provider accounts, vector stores, source adapters, task corpora, and promotion decisions. Core does not own a D1 schema or fleet dispatcher. Apps wire `KbStore` and `KnowledgeDiscoveryDispatcher` to their tenancy, queue, budget, auth, and sandbox systems. @@ -34,7 +36,7 @@ Core does not own a D1 schema or fleet dispatcher. Apps wire `KbStore` and `Know 4. Validate paths, citations, links, and schema. 5. Index generated knowledge pages. 6. Search and graph-lint the knowledge base. -7. Evaluate candidate KB variants with an `agent-eval` improvement loop, then fold the resulting run records into release confidence with `knowledgeReleaseReport`. +7. Evaluate candidate KB and retrieval variants with an `agent-eval` improvement loop, then fold the resulting run records into release confidence with `knowledgeReleaseReport`. 8. Promote only variants that pass downstream gates. ## CLI diff --git a/docs/eval/rag-eval-roadmap.md b/docs/eval/rag-eval-roadmap.md new file mode 100644 index 0000000..c5d0471 --- /dev/null +++ b/docs/eval/rag-eval-roadmap.md @@ -0,0 +1,114 @@ +# RAG Eval Completion Roadmap + +Verdict: `runRetrievalImprovementLoop()` is the right first loop, but it is only the retrieval layer. +SOTA RAG evaluation requires retrieval quality, context quality, generated-answer quality, abstention behavior, robustness, and operating budgets. + +## Research Basis + +| Source | What matters for us | +| --- | --- | +| [Ragas](https://docs.ragas.io/en/stable/concepts/metrics/available_metrics/) / [paper](https://arxiv.org/abs/2309.15217) | Standard RAG eval splits retrieval and generation into context precision, context recall, faithfulness, and answer relevance. | +| [ARES](https://arxiv.org/abs/2311.09476) | Strong RAG eval scores context relevance, answer faithfulness, and answer relevance, and uses small human-labeled sets to calibrate automated judges. | +| [TruLens RAG Triad](https://www.trulens.org/getting_started/core_concepts/rag_triad/) | The minimal end-to-end triad is context relevance, groundedness, and answer relevance. | +| [RAGChecker](https://papers.nips.cc/paper_files/paper/2024/hash/27245589131d17368cccdfa990cbf16e-Abstract-Datasets_and_Benchmarks_Track.html) | Fine-grained diagnosis should separate retrieval misses, noisy context, and unsupported generated claims. | +| [BEIR](https://arxiv.org/abs/2104.08663) / TREC-style retrieval | Retrieval still needs classical rank metrics: Recall@k, Precision@k, MRR, MAP, and nDCG. | +| [CRAG](https://arxiv.org/abs/2406.04744) | Real RAG evals must include long-tail, dynamic, multi-hop, and unanswerable questions, not only easy static facts. | +| [DeepEval RAG metrics](https://deepeval.com/docs/metrics-faithfulness) | Production tools converge on faithfulness, answer relevance, and context relevance as generator/retriever checks. | + +## Current Repo Status + +Done: + +- `runRetrievalImprovementLoop()` auto-searches retrieval configs through `agent-eval`. +- Retrieval scenarios can label pages, page paths, sources, source anchors, and source spans. +- The retrieval judge reports recall, MRR, nDCG, precision@k, cost, and held-out promotion. +- The loop is tested with a real `agent-eval` run where `{ k: 2 }` beats `{ k: 1 }`. + +Not done: + +- Generated-answer evaluation. +- Context relevance and context sufficiency judges. +- Citation support and claim-level groundedness. +- Abstention and unanswerable-question scoring. +- Slice-level reporting for freshness, distractors, multi-hop, and long-tail cases. + +## Completion Criteria + +### Phase 1: Retrieval Quality + +Build a retrieval eval pack with at least 100 labeled scenarios. +Use source-span labels wherever possible. + +Required slices: + +- 25 known-answer questions. +- 25 paraphrase questions. +- 20 distractor questions. +- 10 freshness/version questions. +- 10 multi-source questions. +- 10 unanswerable or forbidden-source questions. + +Ship criteria: + +- Holdout source-span Recall@5 is at least 0.90. +- Holdout nDCG@5 is at least 0.80. +- Train-to-holdout recall gap is at most 0.08. +- Stale or forbidden source hit rate is at most 0.02. +- p95 retrieval latency and cost do not regress by more than 10 percent versus baseline. + +### Phase 2: Answer Quality + +Add a generated-answer eval artifact that includes query, retrieved context, answer text, citations, cost, latency, and trace ids. +Score it with deterministic checks first and LLM judges only for semantic quality. + +Ship criteria: + +- Faithfulness or groundedness is at least 0.95 on holdout. +- Answer relevance is at least 0.90 on holdout. +- Answer correctness is at least 0.85 on human-labeled holdout. +- Citation support is at least 0.95 for claims that cite sources. +- Unsupported-answer rate on unanswerable questions is at most 0.05. + +### Phase 3: Diagnosis + +Add RAGChecker-style failure attribution. +Every failed case must classify as one primary cause. + +Required failure classes: + +- Retrieval miss. +- Retrieval noisy context. +- Stale retrieval. +- Missing multi-hop evidence. +- Generator ignored evidence. +- Generator hallucinated unsupported claim. +- Citation mismatch. +- Correct abstention. +- Incorrect abstention. + +Ship criteria: + +- Every failed eval has one primary failure class. +- At least 95 percent of generated claims can be mapped to supporting context, contradicted context, or no context. +- Reports show metrics by slice and by failure class, not only the aggregate score. + +### Phase 4: Production Loop + +Run the same eval pack on every retrieval or prompt change. +Keep train/dev/holdout isolated. +Never tune on holdout. + +Ship criteria: + +- `runRetrievalImprovementLoop()` gates retrieval config changes. +- Answer-quality eval gates prompt and synthesis changes. +- Reports persist run id, commit, config hash, dataset hash, metric versions, cost, latency, and traces. +- A promoted candidate must improve the target metric without violating faithfulness, abstention, cost, or latency limits. + +## Next Implementation Steps + +1. Add `RagAnswerEvalScenario`, `RagAnswerEvalArtifact`, and `ragAnswerQualityJudge()`. +2. Add context relevance and context sufficiency judges over retrieved hits. +3. Add forbidden/stale source targets to retrieval scenarios. +4. Add slice-level aggregation helpers for the required six eval slices. +5. Add a CLI command that runs the retrieval loop and writes a reproducible report under `.agent-knowledge/eval/`. diff --git a/src/index.ts b/src/index.ts index 2615d8d..7c55af8 100644 --- a/src/index.ts +++ b/src/index.ts @@ -25,6 +25,7 @@ export * from './release' export * from './research-driving-driver' export * from './research-loop' export * from './research-supervisor' +export * from './retrieval-eval' export * from './schemas' export * from './search' export * from './sources' diff --git a/src/retrieval-eval.ts b/src/retrieval-eval.ts new file mode 100644 index 0000000..babfa04 --- /dev/null +++ b/src/retrieval-eval.ts @@ -0,0 +1,724 @@ +import { + type DispatchContext, + type Gate, + type GenerationRecord, + heldOutGate, + type JsonValue, + type JudgeConfig, + type MutableSurface, + type ParameterCandidate, + parameterSweepProposer, + type RunImprovementLoopOptions, + type RunImprovementLoopResult, + runImprovementLoop, + type Scenario, + type SurfaceProposer, +} from '@tangle-network/agent-eval/campaign' +import { searchKnowledge } from './search' +import type { KnowledgeIndex, KnowledgeSearchResult } from './types' + +export type RetrievalConfig = Record +export type RetrievalParameterSearchSpace = Record + +export type RetrievalGoldTarget = + | { kind: 'page'; pageId: string } + | { kind: 'page-path'; path: string } + | { kind: 'source'; sourceId: string } + | { kind: 'source-anchor'; sourceId: string; anchorId: string } + | { kind: 'source-span'; sourceId: string; charStart: number; charEnd: number } + +export interface RetrievalEvalScenario extends Scenario { + kind: 'retrieval-eval' + query: string + expected: RetrievalGoldTarget | readonly RetrievalGoldTarget[] + k?: number +} + +export interface RetrievedSourceSpan { + sourceId: string + anchorId?: string + charStart?: number + charEnd?: number +} + +export interface RetrievedKnowledgeHit { + pageId: string + path: string + title?: string + rank: number + score?: number + normalizedScore?: number + sourceIds?: readonly string[] + sourceSpans?: readonly RetrievedSourceSpan[] + snippet?: string + metadata?: Record +} + +export interface RetrievalEvalArtifact { + config: RetrievalConfig + query: string + requestedK: number + hits: readonly RetrievedKnowledgeHit[] + durationMs: number + costUsd?: number + metadata?: Record +} + +export interface RetrievalMetricSummary { + recall: number + mrr: number + ndcg: number + precisionAtK: number + expectedCount: number + matchedCount: number + relevantHitCount: number + firstHitRank: number | null + matchedTargetIds: readonly string[] +} + +export interface RetrievalEvalRetrieverInput { + index?: KnowledgeIndex + config: RetrievalConfig + scenario: RetrievalEvalScenario + k: number + signal: AbortSignal + context: DispatchContext +} + +export interface RetrievalEvalRetrieverResult { + hits: readonly RetrievedKnowledgeHit[] + costUsd?: number + metadata?: Record +} + +export type RetrievalEvalRetriever = ( + input: RetrievalEvalRetrieverInput, +) => Promise + +export interface BuildRetrievalEvalDispatchOptions { + index?: KnowledgeIndex + defaultK?: number + retrieve?: RetrievalEvalRetriever +} + +export interface RetrievalMetricWeights { + recall?: number + mrr?: number + ndcg?: number + precisionAtK?: number +} + +export interface RetrievalRecallJudgeOptions { + name?: string + weights?: RetrievalMetricWeights +} + +export interface BuildRetrievalParameterCandidatesOptions { + baseline?: RetrievalConfig +} + +export interface RetrievalParameterSweepProposerOptions { + candidates?: readonly ParameterCandidate[] + searchSpace?: RetrievalParameterSearchSpace + baseline?: RetrievalConfig +} + +type RetrievalLoopBaseOptions = RunImprovementLoopOptions< + RetrievalEvalScenario, + RetrievalEvalArtifact +> + +export interface RunRetrievalImprovementLoopOptions { + baseline: RetrievalConfig + scenarios: readonly RetrievalEvalScenario[] + holdoutScenarios?: readonly RetrievalEvalScenario[] + index?: KnowledgeIndex + defaultK?: number + retrieve?: RetrievalEvalRetriever + candidates?: readonly ParameterCandidate[] + searchSpace?: RetrievalParameterSearchSpace + judges?: readonly JudgeConfig[] + gate?: Gate + metricWeights?: RetrievalMetricWeights + targetRecall?: number + holdoutFraction?: number + splitSeed?: number + deltaThreshold?: number + runDir?: RetrievalLoopBaseOptions['runDir'] + seed?: RetrievalLoopBaseOptions['seed'] + reps?: RetrievalLoopBaseOptions['reps'] + resumable?: RetrievalLoopBaseOptions['resumable'] + costCeiling?: RetrievalLoopBaseOptions['costCeiling'] + maxConcurrency?: RetrievalLoopBaseOptions['maxConcurrency'] + dispatchTimeoutMs?: RetrievalLoopBaseOptions['dispatchTimeoutMs'] + expectUsage?: RetrievalLoopBaseOptions['expectUsage'] + tracing?: RetrievalLoopBaseOptions['tracing'] + storage?: RetrievalLoopBaseOptions['storage'] + populationSize?: RetrievalLoopBaseOptions['populationSize'] + maxGenerations?: RetrievalLoopBaseOptions['maxGenerations'] + promoteTopK?: RetrievalLoopBaseOptions['promoteTopK'] + maxImprovementShots?: RetrievalLoopBaseOptions['maxImprovementShots'] + report?: RetrievalLoopBaseOptions['report'] + findings?: RetrievalLoopBaseOptions['findings'] + now?: RetrievalLoopBaseOptions['now'] +} + +export interface RunRetrievalImprovementLoopResult + extends RunImprovementLoopResult { + baselineConfig: RetrievalConfig + winnerConfig: RetrievalConfig + trainScenarios: readonly RetrievalEvalScenario[] + holdoutScenarios: readonly RetrievalEvalScenario[] + candidates: readonly ParameterCandidate[] + targetRecall?: number +} + +export function retrievalConfigSurface(config: RetrievalConfig): string { + return canonicalJson(config) +} + +export function retrievalConfigFromSurface(surface: MutableSurface): RetrievalConfig { + if (typeof surface !== 'string') { + throw new Error( + `retrievalConfigFromSurface expected a JSON string surface, got ${typeof surface}`, + ) + } + let parsed: unknown + try { + parsed = JSON.parse(surface) + } catch (error) { + throw new Error(`retrievalConfigFromSurface could not parse JSON: ${(error as Error).message}`) + } + if (!isJsonObject(parsed)) { + throw new Error('retrievalConfigFromSurface expected a JSON object') + } + return parsed +} + +export function buildRetrievalEvalDispatch(options: BuildRetrievalEvalDispatchOptions) { + if (!options.retrieve && !options.index) { + throw new Error('buildRetrievalEvalDispatch requires either index or retrieve') + } + const defaultK = options.defaultK ?? 5 + const retrieve = options.retrieve ?? defaultSearchRetriever + + return async ( + surface: MutableSurface, + scenario: RetrievalEvalScenario, + context: DispatchContext, + ): Promise => { + const config = retrievalConfigFromSurface(surface) + const k = retrievalK(config, scenario, defaultK) + const startedAt = Date.now() + const result = await retrieve({ + index: options.index, + config, + scenario, + k, + signal: context.signal, + context, + }) + const normalized = normalizeRetrieverResult(result) + observeRetrievalCost(context, normalized.costUsd) + return { + config, + query: scenario.query, + requestedK: k, + hits: normalized.hits, + durationMs: Date.now() - startedAt, + costUsd: normalized.costUsd, + metadata: normalized.metadata, + } + } +} + +export function retrievalRecallJudge( + options: RetrievalRecallJudgeOptions = {}, +): JudgeConfig { + const weights = normalizeWeights(options.weights ?? { recall: 1 }) + return { + name: options.name ?? 'retrieval-recall', + dimensions: [ + { key: 'recall', description: 'fraction of expected knowledge targets retrieved' }, + { key: 'mrr', description: 'reciprocal rank of the first matching hit' }, + { key: 'ndcg', description: 'rank-aware gain for newly matched targets' }, + { + key: 'precision_at_k', + description: 'share of returned hits that match at least one target', + }, + ], + appliesTo: (scenario) => scenario.kind === 'retrieval-eval', + async score({ artifact, scenario }) { + const metrics = scoreRetrievalArtifact(artifact, scenario) + const composite = + (metrics.recall * weights.recall + + metrics.mrr * weights.mrr + + metrics.ndcg * weights.ndcg + + metrics.precisionAtK * weights.precisionAtK) / + (weights.recall + weights.mrr + weights.ndcg + weights.precisionAtK) + return { + dimensions: { + recall: metrics.recall, + mrr: metrics.mrr, + ndcg: metrics.ndcg, + precision_at_k: metrics.precisionAtK, + }, + composite, + notes: `matched ${metrics.matchedCount}/${metrics.expectedCount}; first_hit_rank=${metrics.firstHitRank ?? 'none'}`, + } + }, + } +} + +export function scoreRetrievalArtifact( + artifact: RetrievalEvalArtifact, + scenario: RetrievalEvalScenario, +): RetrievalMetricSummary { + const targets = normalizeTargets(scenario.expected) + if (targets.length === 0) { + throw new Error(`retrieval eval scenario ${scenario.id} has no expected targets`) + } + + const matchedTargetIds = new Set() + let firstHitRank: number | null = null + let relevantHitCount = 0 + let dcg = 0 + + for (const hit of artifact.hits) { + const newlyMatched = targets.filter((target) => { + const targetId = retrievalTargetId(target) + return !matchedTargetIds.has(targetId) && hitMatchesTarget(hit, target) + }) + + if (newlyMatched.length === 0) { + continue + } + + relevantHitCount += 1 + if (firstHitRank === null) { + firstHitRank = hit.rank + } + dcg += 1 / Math.log2(hit.rank + 1) + for (const target of newlyMatched) { + matchedTargetIds.add(retrievalTargetId(target)) + } + } + + const expectedCount = targets.length + const matchedCount = matchedTargetIds.size + const idealRankCount = Math.min(expectedCount, Math.max(1, artifact.requestedK)) + const idcg = idealDcg(idealRankCount) + return { + recall: matchedCount / expectedCount, + mrr: firstHitRank === null ? 0 : 1 / firstHitRank, + ndcg: idcg === 0 ? 0 : dcg / idcg, + precisionAtK: artifact.hits.length === 0 ? 0 : relevantHitCount / artifact.hits.length, + expectedCount, + matchedCount, + relevantHitCount, + firstHitRank, + matchedTargetIds: [...matchedTargetIds].sort(), + } +} + +export function buildRetrievalParameterCandidates( + searchSpace: RetrievalParameterSearchSpace, + options: BuildRetrievalParameterCandidatesOptions = {}, +): ParameterCandidate[] { + const candidates: ParameterCandidate[] = [] + for (const [path, values] of Object.entries(searchSpace).sort(([a], [b]) => a.localeCompare(b))) { + for (const value of values) { + if ( + options.baseline && + canonicalJson(getConfigPath(options.baseline, path)) === canonicalJson(value) + ) { + continue + } + candidates.push({ + label: `${path}=${formatCandidateValue(value)}`, + rationale: `Set retrieval config ${path} to ${formatCandidateValue(value)}`, + changes: [{ path, value }], + }) + } + } + return candidates +} + +export function retrievalParameterSweepProposer( + options: RetrievalParameterSweepProposerOptions, +): SurfaceProposer { + const candidates = + options.candidates ?? + (options.searchSpace + ? buildRetrievalParameterCandidates(options.searchSpace, { baseline: options.baseline }) + : []) + + if (candidates.length === 0) { + throw new Error('retrievalParameterSweepProposer requires at least one candidate') + } + + return parameterSweepProposer({ candidates }) +} + +export async function runRetrievalImprovementLoop( + options: RunRetrievalImprovementLoopOptions, +): Promise { + const split = splitRetrievalScenarios(options) + const candidates = resolveRetrievalCandidates(options) + const populationSize = options.populationSize ?? Math.max(1, Math.min(4, candidates.length)) + const maxGenerations = + options.maxGenerations ?? Math.max(1, Math.ceil(candidates.length / populationSize)) + const proposer = withTargetRecallStop( + retrievalParameterSweepProposer({ candidates }), + options.targetRecall, + ) + const gate = + options.gate ?? + heldOutGate({ + scenarios: split.holdoutScenarios, + deltaThreshold: options.deltaThreshold ?? 0.02, + }) + const result = await runImprovementLoop({ + baselineSurface: retrievalConfigSurface(options.baseline), + scenarios: split.trainScenarios, + holdoutScenarios: split.holdoutScenarios, + dispatchWithSurface: buildRetrievalEvalDispatch({ + index: options.index, + defaultK: options.defaultK, + retrieve: options.retrieve, + }), + judges: [...(options.judges ?? [retrievalRecallJudge({ weights: options.metricWeights })])], + proposer, + gate, + autoOnPromote: 'none', + runDir: options.runDir ?? '.agent-knowledge/retrieval-improvement', + seed: options.seed, + reps: options.reps, + resumable: options.resumable, + costCeiling: options.costCeiling, + maxConcurrency: options.maxConcurrency, + dispatchTimeoutMs: options.dispatchTimeoutMs, + expectUsage: options.expectUsage ?? 'off', + tracing: options.tracing, + storage: options.storage, + populationSize, + maxGenerations, + promoteTopK: options.promoteTopK, + maxImprovementShots: options.maxImprovementShots, + report: options.report, + findings: options.findings, + now: options.now, + }) + + return { + ...result, + baselineConfig: options.baseline, + winnerConfig: retrievalConfigFromSurface(result.winnerSurface), + trainScenarios: split.trainScenarios, + holdoutScenarios: split.holdoutScenarios, + candidates, + targetRecall: options.targetRecall, + } +} + +function splitRetrievalScenarios(options: RunRetrievalImprovementLoopOptions): { + trainScenarios: RetrievalEvalScenario[] + holdoutScenarios: RetrievalEvalScenario[] +} { + const scenarios = [...options.scenarios] + if (scenarios.length === 0) { + throw new Error('runRetrievalImprovementLoop requires at least one training scenario') + } + if (options.holdoutScenarios) { + const holdoutScenarios = [...options.holdoutScenarios] + if (holdoutScenarios.length === 0) { + throw new Error('runRetrievalImprovementLoop holdoutScenarios must not be empty') + } + return { trainScenarios: scenarios, holdoutScenarios } + } + + if (scenarios.length < 2) { + throw new Error( + 'runRetrievalImprovementLoop requires at least 2 scenarios when holdoutScenarios are not provided', + ) + } + const holdoutFraction = options.holdoutFraction ?? 0.3 + if (!Number.isFinite(holdoutFraction) || holdoutFraction <= 0 || holdoutFraction >= 1) { + throw new Error( + `runRetrievalImprovementLoop holdoutFraction must be > 0 and < 1, got ${String(holdoutFraction)}`, + ) + } + const shuffled = seededShuffle(scenarios, options.splitSeed ?? options.seed ?? 42) + const holdoutCount = Math.min( + shuffled.length - 1, + Math.max(1, Math.round(shuffled.length * holdoutFraction)), + ) + const splitIndex = shuffled.length - holdoutCount + return { + trainScenarios: shuffled.slice(0, splitIndex), + holdoutScenarios: shuffled.slice(splitIndex), + } +} + +function resolveRetrievalCandidates( + options: RunRetrievalImprovementLoopOptions, +): ParameterCandidate[] { + const candidates = options.candidates + ? [...options.candidates] + : options.searchSpace + ? buildRetrievalParameterCandidates(options.searchSpace, { baseline: options.baseline }) + : [] + + if (candidates.length === 0) { + throw new Error('runRetrievalImprovementLoop requires candidates or searchSpace') + } + return candidates +} + +function withTargetRecallStop( + proposer: SurfaceProposer, + targetRecall: number | undefined, +): SurfaceProposer { + if (targetRecall === undefined) { + return proposer + } + if (!Number.isFinite(targetRecall) || targetRecall < 0 || targetRecall > 1) { + throw new Error(`targetRecall must be between 0 and 1, got ${String(targetRecall)}`) + } + return { + kind: `${proposer.kind}:target-recall`, + propose: (context) => proposer.propose(context), + decide(args) { + const baseDecision = proposer.decide?.(args) + if (baseDecision?.stop) { + return baseDecision + } + const bestRecall = bestObservedRecall(args.history) + if (bestRecall >= targetRecall) { + return { + stop: true, + reason: `target recall ${targetRecall} reached with train recall ${bestRecall}`, + } + } + return { stop: false } + }, + } +} + +function bestObservedRecall(history: GenerationRecord[]): number { + let best = Number.NEGATIVE_INFINITY + for (const generation of history) { + for (const candidate of generation.candidates) { + const recall = candidate.dimensions.recall + if (recall !== undefined && Number.isFinite(recall) && recall > best) { + best = recall + } + } + } + return best +} + +function retrievalK( + config: RetrievalConfig, + scenario: RetrievalEvalScenario, + defaultK: number, +): number { + const rawK = config.k ?? config.topK ?? scenario.k ?? defaultK + if (typeof rawK !== 'number' || !Number.isFinite(rawK) || rawK <= 0) { + throw new Error(`retrieval k must be a positive number, got ${String(rawK)}`) + } + return Math.floor(rawK) +} + +async function defaultSearchRetriever( + input: RetrievalEvalRetrieverInput, +): Promise { + if (!input.index) { + throw new Error('default retrieval eval search requires an index') + } + const results = searchKnowledge(input.index, input.scenario.query, input.k) + return { hits: results.map(hitFromSearchResult) } +} + +function observeRetrievalCost(context: DispatchContext, costUsd: number | undefined): void { + if (costUsd === undefined) { + return + } + if (!Number.isFinite(costUsd) || costUsd < 0) { + throw new Error( + `retrieval costUsd must be a non-negative finite number, got ${String(costUsd)}`, + ) + } + context.cost.observe(costUsd, 'agent-knowledge:retrieval') +} + +function hitFromSearchResult(result: KnowledgeSearchResult): RetrievedKnowledgeHit { + return { + pageId: result.page.id, + path: result.page.path, + title: result.page.title, + rank: result.rank, + score: result.score, + normalizedScore: result.normalizedScore, + sourceIds: result.page.sourceIds, + snippet: result.snippet, + } +} + +function normalizeRetrieverResult( + result: readonly RetrievedKnowledgeHit[] | RetrievalEvalRetrieverResult, +): RetrievalEvalRetrieverResult { + if (isHitArray(result)) { + return { hits: result } + } + return result +} + +function normalizeTargets( + expected: RetrievalGoldTarget | readonly RetrievalGoldTarget[], +): RetrievalGoldTarget[] { + return isTargetArray(expected) ? [...expected] : [expected] +} + +function isHitArray( + result: readonly RetrievedKnowledgeHit[] | RetrievalEvalRetrieverResult, +): result is readonly RetrievedKnowledgeHit[] { + return Array.isArray(result) +} + +function isTargetArray( + expected: RetrievalGoldTarget | readonly RetrievalGoldTarget[], +): expected is readonly RetrievalGoldTarget[] { + return Array.isArray(expected) +} + +function hitMatchesTarget(hit: RetrievedKnowledgeHit, target: RetrievalGoldTarget): boolean { + switch (target.kind) { + case 'page': + return hit.pageId === target.pageId + case 'page-path': + return hit.path === target.path + case 'source': + return Boolean(hit.sourceIds?.includes(target.sourceId)) + case 'source-anchor': + return Boolean( + hit.sourceSpans?.some( + (span) => span.sourceId === target.sourceId && span.anchorId === target.anchorId, + ), + ) + case 'source-span': + return Boolean( + hit.sourceSpans?.some( + (span) => span.sourceId === target.sourceId && spanOverlaps(span, target), + ), + ) + } +} + +function spanOverlaps( + hit: RetrievedSourceSpan, + target: Extract, +): boolean { + if (typeof hit.charStart !== 'number' || typeof hit.charEnd !== 'number') { + return false + } + return hit.charStart < target.charEnd && target.charStart < hit.charEnd +} + +function retrievalTargetId(target: RetrievalGoldTarget): string { + switch (target.kind) { + case 'page': + return `page:${target.pageId}` + case 'page-path': + return `page-path:${target.path}` + case 'source': + return `source:${target.sourceId}` + case 'source-anchor': + return `source-anchor:${target.sourceId}:${target.anchorId}` + case 'source-span': + return `source-span:${target.sourceId}:${target.charStart}:${target.charEnd}` + } +} + +function idealDcg(count: number): number { + let score = 0 + for (let i = 1; i <= count; i += 1) { + score += 1 / Math.log2(i + 1) + } + return score +} + +function normalizeWeights(weights: RetrievalMetricWeights): Required { + const normalized = { + recall: normalizeWeight(weights.recall), + mrr: normalizeWeight(weights.mrr), + ndcg: normalizeWeight(weights.ndcg), + precisionAtK: normalizeWeight(weights.precisionAtK), + } + const total = normalized.recall + normalized.mrr + normalized.ndcg + normalized.precisionAtK + if (total <= 0) { + throw new Error('retrievalRecallJudge requires at least one positive metric weight') + } + return normalized +} + +function normalizeWeight(value: number | undefined): number { + if (value === undefined) { + return 0 + } + if (!Number.isFinite(value) || value < 0) { + throw new Error( + `retrievalRecallJudge metric weights must be non-negative finite numbers, got ${String(value)}`, + ) + } + return value +} + +function getConfigPath(config: RetrievalConfig, path: string): JsonValue | undefined { + let current: JsonValue | undefined = config + for (const part of path.split('.')) { + if (!isJsonObject(current)) { + return undefined + } + current = current[part] + } + return current +} + +function formatCandidateValue(value: JsonValue): string { + if (typeof value === 'string') { + return value + } + return canonicalJson(value) +} + +function canonicalJson(value: JsonValue | undefined): string { + if (value === undefined) { + return 'undefined' + } + if (Array.isArray(value)) { + return `[${value.map(canonicalJson).join(',')}]` + } + if (isJsonObject(value)) { + const entries = Object.entries(value) + .sort(([a], [b]) => a.localeCompare(b)) + .map(([key, child]) => `${JSON.stringify(key)}:${canonicalJson(child)}`) + return `{${entries.join(',')}}` + } + return JSON.stringify(value) +} + +function isJsonObject(value: unknown): value is Record { + return Boolean(value) && typeof value === 'object' && !Array.isArray(value) +} + +function seededShuffle(items: readonly T[], seed: number): T[] { + const out = [...items] + let state = seed >>> 0 + for (let i = out.length - 1; i > 0; i -= 1) { + state = (state * 1664525 + 1013904223) >>> 0 + const j = state % (i + 1) + ;[out[i], out[j]] = [out[j]!, out[i]!] + } + return out +} diff --git a/tests/retrieval-eval.test.ts b/tests/retrieval-eval.test.ts new file mode 100644 index 0000000..cb0adc2 --- /dev/null +++ b/tests/retrieval-eval.test.ts @@ -0,0 +1,271 @@ +import { inMemoryCampaignStorage } from '@tangle-network/agent-eval/campaign' +import { describe, expect, it } from 'vitest' +import { + buildRetrievalEvalDispatch, + buildRetrievalParameterCandidates, + type KnowledgeIndex, + type RetrievalEvalArtifact, + type RetrievalEvalScenario, + retrievalConfigFromSurface, + retrievalConfigSurface, + retrievalParameterSweepProposer, + retrievalRecallJudge, + runRetrievalImprovementLoop, + scoreRetrievalArtifact, +} from '../src/index' + +const signal = new AbortController().signal + +function testContext() { + return { + cellId: 'cell-1', + rep: 0, + seed: 123, + signal, + } as Parameters>[2] +} + +function testContextWithCost(observed: Array<{ amountUsd: number; source: string }>) { + return { + ...testContext(), + cost: { + observe: (amountUsd: number, source: string) => observed.push({ amountUsd, source }), + }, + } as Parameters>[2] +} + +function fixtureIndex(): KnowledgeIndex { + return { + root: 'memory://retrieval-eval', + generatedAt: '2026-01-01T00:00:00.000Z', + sources: [ + { + id: 'src-flash', + uri: 'memory://flash', + title: 'Flash Attention Source', + contentHash: 'sha256:flash', + text: 'Flash Attention improves memory bandwidth usage.', + createdAt: '2026-01-01T00:00:00.000Z', + }, + { + id: 'src-cache', + uri: 'memory://cache', + title: 'Cache Source', + contentHash: 'sha256:cache', + text: 'Cache eviction policy notes.', + createdAt: '2026-01-01T00:00:00.000Z', + }, + ], + pages: [ + { + id: 'flash-attention', + path: 'knowledge/concepts/flash-attention.md', + title: 'Flash Attention', + text: 'IO aware attention improves memory bandwidth with tiled SRAM reads.', + frontmatter: {}, + sourceIds: ['src-flash'], + tags: ['attention'], + outLinks: [], + }, + { + id: 'cache-policy', + path: 'knowledge/concepts/cache-policy.md', + title: 'Cache Policy', + text: 'LRU and TTL cache eviction policies.', + frontmatter: {}, + sourceIds: ['src-cache'], + tags: ['runtime'], + outLinks: [], + }, + ], + graph: { nodes: [], edges: [] }, + } +} + +describe('retrieval eval', () => { + it('dispatches local search and scores page/source recall deterministically', async () => { + const scenario: RetrievalEvalScenario = { + id: 'q-memory-bandwidth', + kind: 'retrieval-eval', + query: 'memory bandwidth', + expected: { kind: 'page', pageId: 'flash-attention' }, + } + const dispatch = buildRetrievalEvalDispatch({ index: fixtureIndex() }) + const artifact = await dispatch(retrievalConfigSurface({ k: 2 }), scenario, testContext()) + + expect(artifact.requestedK).toBe(2) + expect(artifact.hits[0]?.pageId).toBe('flash-attention') + + const judge = retrievalRecallJudge() + const score = await judge.score({ artifact, scenario, signal }) + expect(score.dimensions.recall).toBe(1) + expect(score.dimensions.mrr).toBe(1) + expect(score.composite).toBe(1) + + const sourceScenario: RetrievalEvalScenario = { + id: 'q-source', + kind: 'retrieval-eval', + query: 'memory bandwidth', + expected: { kind: 'source', sourceId: 'src-flash' }, + } + const sourceScore = await judge.score({ artifact, scenario: sourceScenario, signal }) + expect(sourceScore.dimensions.recall).toBe(1) + }) + + it('requires span evidence for span targets instead of accepting broad source hits', () => { + const scenario: RetrievalEvalScenario = { + id: 'q-span', + kind: 'retrieval-eval', + query: 'specific source span', + expected: { kind: 'source-span', sourceId: 'src-1', charStart: 10, charEnd: 20 }, + } + const broadSourceHit: RetrievalEvalArtifact = { + config: { k: 1 }, + query: scenario.query, + requestedK: 1, + durationMs: 0, + hits: [ + { + pageId: 'page-1', + path: 'knowledge/page-1.md', + rank: 1, + sourceIds: ['src-1'], + }, + ], + } + expect(scoreRetrievalArtifact(broadSourceHit, scenario).recall).toBe(0) + + const spanHit: RetrievalEvalArtifact = { + ...broadSourceHit, + hits: [ + { + pageId: 'page-1', + path: 'knowledge/page-1.md', + rank: 1, + sourceIds: ['src-1'], + sourceSpans: [{ sourceId: 'src-1', charStart: 15, charEnd: 25 }], + }, + ], + } + expect(scoreRetrievalArtifact(spanHit, scenario).recall).toBe(1) + }) + + it('records retriever-reported cost in the agent-eval cost meter', async () => { + const observed: Array<{ amountUsd: number; source: string }> = [] + const scenario: RetrievalEvalScenario = { + id: 'q-cost', + kind: 'retrieval-eval', + query: 'costed retrieval', + expected: { kind: 'page', pageId: 'page-1' }, + } + const dispatch = buildRetrievalEvalDispatch({ + retrieve: async () => ({ + costUsd: 0.25, + hits: [{ pageId: 'page-1', path: 'knowledge/page-1.md', rank: 1 }], + }), + }) + const artifact = await dispatch( + retrievalConfigSurface({ k: 1 }), + scenario, + testContextWithCost(observed), + ) + + expect(artifact.costUsd).toBe(0.25) + expect(observed).toEqual([{ amountUsd: 0.25, source: 'agent-knowledge:retrieval' }]) + }) + + it('builds parameter candidates and delegates proposal to agent-eval', async () => { + const baseline = { k: 5, hybrid: false, reranker: null, chunk: { overlap: 100 } } + const candidates = buildRetrievalParameterCandidates( + { + 'chunk.overlap': [100, 200], + hybrid: [false, true], + k: [5, 10], + }, + { baseline }, + ) + + expect(candidates.map((candidate) => candidate.label)).toEqual([ + 'chunk.overlap=200', + 'hybrid=true', + 'k=10', + ]) + + const proposer = retrievalParameterSweepProposer({ candidates }) + const proposals = await proposer.propose({ + currentSurface: retrievalConfigSurface(baseline), + history: [], + findings: [], + populationSize: 2, + generation: 0, + signal, + }) + const surfaces = proposals.map((proposal) => + typeof proposal === 'string' ? proposal : 'surface' in proposal ? proposal.surface : proposal, + ) + + expect(surfaces).toHaveLength(2) + expect(JSON.parse(surfaces[0] as string)).toMatchObject({ chunk: { overlap: 200 } }) + expect(JSON.parse(surfaces[1] as string)).toMatchObject({ hybrid: true }) + }) + + it('runs an agent-eval loop that auto-selects the better retrieval config', async () => { + const trainScenario: RetrievalEvalScenario = { + id: 'q-train', + kind: 'retrieval-eval', + query: 'needs second result', + expected: { kind: 'page', pageId: 'gold' }, + } + const holdoutScenario: RetrievalEvalScenario = { + id: 'q-holdout', + kind: 'retrieval-eval', + query: 'held out needs second result', + expected: { kind: 'page', pageId: 'gold' }, + } + + const result = await runRetrievalImprovementLoop({ + baseline: { k: 1 }, + scenarios: [trainScenario], + holdoutScenarios: [holdoutScenario], + searchSpace: { k: [1, 2] }, + retrieve: async ({ k }) => ({ + hits: [ + { pageId: 'distractor', path: 'knowledge/distractor.md', rank: 1 }, + ...(k >= 2 ? [{ pageId: 'gold', path: 'knowledge/gold.md', rank: 2 }] : []), + ], + }), + targetRecall: 1, + deltaThreshold: 0.01, + populationSize: 1, + maxGenerations: 1, + runDir: 'memory://retrieval-loop-test', + storage: inMemoryCampaignStorage(), + expectUsage: 'off', + }) + + expect(result.winnerConfig).toMatchObject({ k: 2 }) + expect(result.trainScenarios).toHaveLength(1) + expect(result.holdoutScenarios).toHaveLength(1) + }) + + it('fails loudly on invalid config surfaces and empty expected labels', () => { + expect(() => retrievalConfigFromSurface('not-json')).toThrow(/could not parse JSON/) + expect(() => + scoreRetrievalArtifact( + { + config: { k: 1 }, + query: 'empty labels', + requestedK: 1, + durationMs: 0, + hits: [], + }, + { + id: 'q-empty', + kind: 'retrieval-eval', + query: 'empty labels', + expected: [], + }, + ), + ).toThrow(/has no expected targets/) + }) +})