feat: every optimization stage sees each example's full context (RAG documents, inputs, metadata, tags)

CLAUDE.md

@@ -242,6 +242,13 @@ register_vector_store_adapter("my_db", MyDbStore)  # upsert/query/delete
   after `fit()`. `generate_fit_report(history, baseline_eval=..., final_eval=...)`
   shows the prompt diff, every candidate's decision and reason, and per-example
   answers. Guide: `docs/optimization/`; runnable: `examples/prompt_optimization/`.
+- Dataset rows carry the whole case — `input.context.documents` (RAG), any other
+  `input.*` (plan, `messages` …), `metadata`, `tags`, `rubric` — and every stage
+  gets it: the agent receives `input` as written (a copy; datasets are never
+  mutated), metrics the whole row, LLM judges the documents (or what the agent
+  listed under `citations` / `retrieval_context` / `sources`) plus inputs,
+  metadata and tags, the rewrite model all of it per example, the augmenter the
+  seed's documents. See `docs/optimization/data.md` and `rag_context.py`.
 
 ## Project conventions (match these when editing the framework)
 

docs/optimization/data.md

@@ -2,49 +2,118 @@
 
 ## The dataset file
 
-Class-agent datasets are JSONL, one `AgentExample` per line (the format every
-template's `datasets/all.jsonl` uses):
+Class-agent datasets are JSONL, one `AgentExample` per line. Every
+template's `datasets/all.jsonl` uses this format. A row carries the whole
+case: the question, the documents the answer must come from, any other agent
+input, labels, tags and a rubric. All of it reaches the optimization, not
+just the question and the answer.
 
 ```json
-{"id": "refund_2",
- "split": "validation",
- "input": {"current_query": "I bought an annual plan last week. Can I still get my money back?"},
- "expected_output": {"response": "Yes. Annual plans can be refunded within 30 days of purchase."},
- "metadata": {"must_include": ["30 days"], "topic": "refund"},
- "rubric": {"correctness": "States the 30-day window for annual plans."},
- "tags": ["billing"]}
+{"id": "policy_05",
+ "split": "train",
+ "input": {"current_query": "How long are audit logs kept?",
+           "context": {"documents": [{"source": "audit.md",
+               "content": "Audit logs are kept for 400 days on Enterprise and 90 days on Team."}]},
+           "customer_plan": "team"},
+ "expected_output": {"response": "On the Team plan, audit logs are kept for 90 days (400 days on Enterprise)."},
+ "metadata": {"topic": "security", "difficulty": "medium", "must_include": ["90 days"]},
+ "tags": ["security", "compliance", "plan-specific"],
+ "rubric": {"groundedness": "Every claim comes from the row's context documents.",
+            "plan": "Applies the policy to the customer's plan when one is given."}}
 ```
 
 | Field | Required | Meaning |
 | --- | --- | --- |
 | `id` | recommended | Stable id; generated (`example_0003`) when missing. Reports join before/after results on it. |
 | `split` | recommended | `train`, `validation`, `holdout` or `test` (default `train`). See [data roles](generalization.md#data-roles). |
-| `input` | yes | What the agent receives — the same dict `input_to_state()` gets. The question is read from `question`, `query`, `current_query` or `request`. `input.context` becomes the judge's retrieval context. |
+| `input` | yes | What the agent receives, unchanged: the same dict `input_to_state()` gets. The question is read from `question`, `query`, `current_query` or `request`. |
+| `input.context` | no | Documents the answer must come from (RAG). The value can be `{"documents": [...]}`, a list, or one string. A document is a string or a dict with `content` (or `text` / `page_content`) and an optional `source`. |
+| other `input.*` keys | no | Any other agent input: `messages` (earlier turns), a plan, a locale, a user profile, and so on. |
 | `expected_output` | for scoring | The reference answer (`{"response": "..."}`, or your structured output). Judges receive it as the reference. |
-| `metadata` | no | Anything metrics need: required facts, topic, difficulty. Class-agent metrics read it (`example.metadata`). |
-| `rubric` | no | Per-dimension criteria, appended to the judge's reference. |
-| `tags` | no | Free-form labels. |
+| `metadata` | no | Labels metrics need: required facts (`must_include`), topic, difficulty. |
+| `tags` | no | Grouping labels. Filter with `data.filter_by_tags(...)`, and break results down by tag. |
+| `rubric` | no | Per-dimension criteria, added to the judge's reference. |
 
 ```python
 from agentomatic.optimize import load_data
 
 data = load_data("agents/support/datasets/all.jsonl")    # AgentDataset
 data.train, data.validation, data.holdout, data.test    # lists of AgentExample
+billing = data.filter_by_tags("billing")                # any-match
 ```
 
 `PromptFitter` works on the simpler `agentomatic.optimize.Dataset` of
-`DataPoint(query, expected_answer, context, metadata)`. Convert, or build one
-directly:
+`DataPoint(query, expected_answer, context, metadata, tags)`. Convert, or
+build one directly:
 
 ```python
 from agentomatic.optimize import Dataset
 
-valset = Dataset(points=[e.to_datapoint() for e in data.validation])   # keeps metadata
-Dataset.from_list([{"query": "Refund window?", "expected_answer": "30 days"}])
-Dataset.from_jsonl("eval.jsonl")        # {"query", "expected_answer", "context", "metadata"}
+valset = Dataset(points=[e.to_datapoint() for e in data.validation])   # keeps everything
+Dataset.from_list([{"query": "Refund window?", "expected_answer": "30 days",
+                    "context": ["Refunds: 30 days."], "tags": ["billing"],
+                    "metadata": {"invoke": {"customer_plan": "team"}}}])
+Dataset.from_jsonl("eval.jsonl")   # {"query", "expected_answer", "context", "metadata", "tags"}
 trainset, valset = dataset.split(ratio=0.8)
 ```
 
+`to_datapoint()` maps a row as follows:
+
+| Row field | `DataPoint` field |
+| --- | --- |
+| `input` (except the question) | `metadata.invoke`, sent to the agent exactly as written |
+| `input.context` | `context`, as document texts such as `"Audit logs are kept … (source: audit.md)"` |
+| `metadata` | `metadata` |
+| `tags` | `tags` |
+| `rubric` | the expected reference, ahead of the answer, for the judge |
+
+## What each stage sees
+
+| Stage | Receives |
+| --- | --- |
+| **Agent** | `input` exactly as written: the question, `context`, `customer_plan` and the rest. Evaluation hands it a copy, so an agent that edits its input never changes the dataset. |
+| **Metrics** | The whole row (`example.input`, `.metadata`, `.tags`, `.rubric`), the same inside `fit()`, in `evaluate()` and nested in a composite. |
+| **LLM judges** | Question, answer, expected answer plus rubric, and the context documents. If the row has no documents, they get what the agent retrieved instead. Inside `fit()` / `evaluate()` they also get the other inputs, the metadata and the tags. |
+| **Rewrite model** | Every failure, success and dataset sample in its [briefing](how-it-works.md) shows `Context documents`, `Other inputs`, `Example metadata`, `Tags` and what the agent retrieved. It is told to write instructions that *use* such context and never to copy one row's facts into the prompt. |
+| **Augmenter** | The seed's documents and other inputs (a "Seed context" section). Variations inherit a copy of them, plus the seed's tags (and `augmented`) and rubric. |
+| **Report** | A `context` column in the per-example table: tags, documents, inputs and metadata. |
+
+## RAG datasets
+
+There are two ways for a RAG agent to get its documents. Both work with every
+stage above.
+
+* **The row brings its documents** (`input.context.documents`). The agent
+  answers from exactly those. The judge checks groundedness against them, and
+  the rewrite model reads them next to each answer. Use this to optimize the
+  *answering* prompt independently of retrieval quality, or when the gold
+  passages are known. The `rag` template reads them in `input_to_state`
+  (`context.documents`, flattened to `documents` over REST) and falls back to
+  its own retrieval when there are none.
+* **The agent retrieves for itself.** List what it used in the output under
+  `retrieval_context`, `citations`, `sources` or `documents`. These can be
+  strings or dicts with `content` and `source`. Judges then score
+  groundedness against those, and the rewrite model sees them as
+  `Retrieved by the agent`.
+
+```python
+def input_to_state(self, data):
+    ctx = data.get("context") or {}
+    docs = ctx.get("documents") if isinstance(ctx, dict) else ctx
+    return MyState(question=data.get("current_query", ""), documents=list(docs or []))
+
+def state_to_output(self, state):
+    return {"response": state.answer, "citations": state.citations}   # what it used
+```
+
+Metrics can read the documents too. For example, a groundedness check that
+counts how many of the answer's words come from them:
+[`grounded_score`](examples.md).
+
+`examples/prompt_optimization/rag_context.py` runs the whole loop on rows
+that bring their own documents, plan, metadata, tags and rubric. It prints
+what each stage sees and reports test scores per tag.
+
 ## Writing good examples
 
 * **One fact per expectation you care about.** An expected answer is what the
@@ -122,6 +191,10 @@ before trusting scores computed on them.
 * Generated questions that duplicate an existing one, or are ≥ 90 % similar
   to any validation / holdout / test question, are dropped (a paraphrase of a
   test question in train would inflate every score measured on it).
+* The generator sees each seed's documents and other inputs (a "Seed
+  context" section) and is told to ask only what those documents answer.
+  Every variation inherits a copy of its seed's inputs, rubric and tags (plus
+  `augmented`), so a RAG row stays answerable.
 * Each new row records `metadata.source = "augment"`, its `strategy` and its
   `parent_id`.
 * `data.metadata["augment_stats"]` records `requested`, `calls`, `parsed`,

docs/optimization/examples.md

@@ -3,17 +3,25 @@
 [`examples/prompt_optimization`](https://github.com/UnicoLab/agentomatic/tree/main/examples/prompt_optimization)
 contains runnable scripts, from scoring a few answers to a full training run.
 Each one is commented step by step and lists every option of the methods it
-uses. All of them optimize the same small `SupportAgent` (it answers Acme
-Cloud customer questions from a policy knowledge base) on a 26-row dataset
-split into train / validation / holdout / test.
+uses. All of them optimize the same small RAG `SupportAgent`, which answers
+Acme Cloud customer questions from policy documents. Those are the documents a
+row sends (`context.documents`), or else the ones it retrieves from its own
+knowledge base, and it also takes the customer's plan into account. Two
+datasets are used, each split into train / validation / holdout / test:
+
+* `support.jsonl` has 26 rows, with tags and a `customer_plan` input where the
+  plan matters.
+* `policies.jsonl` has 20 rows. Each brings its own documents, plus a plan,
+  metadata, tags and a rubric.
 
 | Script | Level | Covers | Docs |
 | --- | --- | --- | --- |
 | `metrics_and_judges.py` | building blocks | Every metric family, LLM judges (single, per-dimension, panel), composite objectives, adapters, losses — scored on three fixed answers | [Metrics](metrics.md) |
 | `data_augmentation.py` | building blocks | `prepare_dataset` strategies and options, leak protection, `augment_stats` | [Datasets](data.md) |
 | `anti_overfitting.py` | mechanics | The acceptance rules on hand-made numbers; a fit in which a prompt memorising the validation answers is vetoed and a general one kept | [Generalization](generalization.md) |
 | `low_level_prompt_fitter.py` | low level | `PromptFitter` with every knob, inner-loop callbacks, before/after test evaluation, report | [PromptFitter](prompt-fitter.md) |
-| `train.py` | high level | `load_data` → `prepare_dataset` → metrics → loss → `PromptFitterBridge` → `compile` → `fit` → test evaluation → `generate_fit_report` | [Class agents](class-agents.md) |
+| `train.py` | high level | `load_data` → `prepare_dataset` → metrics → loss → `PromptFitterBridge` → `compile` → `fit` → test evaluation (overall and per tag) → `generate_fit_report` | [Class agents](class-agents.md) |
+| `rag_context.py` | RAG | Rows with their own documents, plan, metadata, tags and rubric. Prints what the agent, the judge and the rewrite model see; context-aware metrics (`facts`, `grounded`); per-tag test scores; `--augment` keeps each seed's documents | [Datasets](data.md#rag-datasets) |
 
 ## Running them
 
@@ -28,6 +36,7 @@ python examples/prompt_optimization/data_augmentation.py --n-examples 40
 python examples/prompt_optimization/anti_overfitting.py
 python examples/prompt_optimization/low_level_prompt_fitter.py --trials 8
 python examples/prompt_optimization/train.py --epochs 2 --trials 6 --augment
+python examples/prompt_optimization/rag_context.py --epochs 1 --trials 4
 ```
 
 Every script takes `--base-url`, `--api-key`, `--model`, `--judge-model`,
@@ -52,11 +61,48 @@ Step 10/10 — results
   report: examples/prompt_optimization/out/train_report.html
 ```
 
+What `rag_context.py` prints first, for one row (abridged):
+
+```text
+── What each stage sees — row policy_05
+agent (input_to_state receives exactly this):
+  {"current_query": "How long are audit logs kept?", "context": {"documents": [{"source": "audit.md", "content": "Audit logs are kept for 400 days on Enterprise and 90 days on Team."}]}, "customer_plan": "team"}
+LLM judge — context documents:
+  - Audit logs are kept for 400 days on Enterprise and 90 days on Team. (source: audit.md)
+LLM judge — other inputs, metadata and tags:
+  {"inputs": {"customer_plan": "team"}, "metadata": {"topic": "security", "difficulty": "medium", "must_include": ["90 days"]}, "tags": ["security", "compliance", "plan-specific"]}
+rewrite model (dataset sample in its briefing):
+1. Q: How long are audit logs kept?
+   - Context documents: Audit logs are kept for 400 days on Enterprise and 90 days on Team. (source: audit.md)
+   - Other inputs: {"customer_plan": "team"}
+   - Example metadata: {"topic": "security", "difficulty": "medium", "must_include": ["90 days"]}
+   - Tags: security, compliance, plan-specific
+   Expected:
+     ## Expected answer
+     On the Team plan, audit logs are kept for 90 days (400 days on Enterprise).
+```
+
+…and last, the test split per tag:
+
+```text
+── Test quality by tag (before → after)
+  backups         0.00 → 0.85
+  plan-specific   0.00 → 0.90
+  pricing         0.03 → 0.83
+  refusal         0.03 → 0.20
+```
+
 ## Tested
 
 `tests/test_prompt_optimization_examples.py` runs every script as a
 subprocess against a deterministic fake OpenAI-compatible server
 (`tests/fake_openai_server.py`) and checks the behaviour each one teaches:
 the prompt improves on validation *and* test, augmentation grows only train,
 the memorising prompt is rejected, and the report contains the prompts,
-candidates and per-example results.
+candidates and per-example results. For `rag_context.py` it also checks the
+requests the fake server received:
+
+* the agent answered from each row's documents and plan;
+* judges saw the documents, the plan, the metadata and the tags;
+* the rewrite model's briefing contained all of them;
+* augmented rows kept their seed's documents.

docs/optimization/how-it-works.md

@@ -63,6 +63,13 @@ flowchart TD
    never studies the examples it is then selected on (`reflect_on="train"`).
 3. **Proposal.** The optimizer (see the table below) proposes up to four
    candidate configurations per round from the incumbent and the reflection.
+   Its rewrite model reads a briefing: the current config, metrics, score
+   history, train samples, and the lowest- and highest-scoring examples. Each
+   example comes with everything it ran with: question, expected answer,
+   actual answer, judge feedback, context documents, other inputs, metadata,
+   tags, and what the agent retrieved. The briefing tells the model to write
+   instructions that use such context, never to copy one example's facts into
+   the prompt. See [what each stage sees](data.md#what-each-stage-sees).
    A candidate identical to one already scored is skipped for free: a weak
    rewrite model often re-proposes the current prompt.
 4. **Minibatch screen.** Each candidate runs on a slice of validation (30 %,

docs/optimization/metrics.md

@@ -29,7 +29,7 @@ automatically.
 | | `agentomatic.agents` (class agents) | `agentomatic.optimize` (the fitter) |
 | --- | --- | --- |
 | Signature | `score(example, prediction) -> float` (sync) | `await evaluate(query, response, expected, context) -> EvalResult` |
-| Sees | The whole `AgentExample` (input, expected output, **metadata**) and the output dict | The question, the answer text, the reference text, retrieval context |
+| Sees | The whole `AgentExample` (input incl. context documents and other inputs, expected output, **metadata**, **tags**, rubric) and the output dict, the same in `fit()` and `evaluate()` | The question, the answer text, the reference text, the context documents (or what the agent retrieved). Inside `fit()` / `evaluate()`, `current_scoring_run()` also exposes the example |
 | Returns | A float | `EvalResult(score, reason, metadata)` — judges add `dimensions`, rationale; failures set `metadata["evaluation_failed"]` |
 | Best for | Checks that need structured output or per-example metadata | Judges and text checks |
 
@@ -118,6 +118,12 @@ How judges stay honest:
 * **Reference.** Judges receive the example's expected answer (and, for
   `AgentExample`s, a rich reference: expected output, rubric, required facts).
   Write expected answers that state the facts a correct answer must contain.
+* **Context.** Judges receive the row's context documents. If the row has
+  none, they get what the agent says it retrieved (`retrieval_context` /
+  `citations` / `sources` in its output). Inside `fit()` and `evaluate()`
+  they also get the row's other inputs (a plan, a locale…), its metadata and
+  its tags, in an "Example inputs, metadata and tags" section, so an answer
+  is judged in that light. See [what each stage sees](data.md#what-each-stage-sees).
 * **Deterministic.** Keep `temperature=0.0`. A judge that scores the same
   answer differently each time makes every comparison noise; the
   [significance rule](generalization.md#the-acceptance-rules) exists because

docs/optimization/reports.md

@@ -50,7 +50,7 @@ additionally writes one per fit under `experiment_dir`.
 | **What changed in the prompt** | The initial → final diff (long lines wrapped so changed sentences stand out), both full prompts side by side, and every accepted change in order with why it was accepted. |
 | **Epochs** | Start / best / held-out / candidates / stop reason per epoch. |
 | **All candidates** | Every proposed candidate: epoch, round, minibatch / validation / held-out scores, confidence, **decision** and **reason**; the diff of each distinct candidate prompt. |
-| **Examples — before vs after** | Per example: question, expected answer, answer before and after, scores before and after, per-metric Δ, and the judge's rationale. Uses the test split when `baseline_eval`/`final_eval` are given, otherwise the validation examples the fitter scored. |
+| **Examples — before vs after** | Per example: question, its context (tags · documents · other inputs · metadata, when the data has any), expected answer, answer before and after, scores before and after, per-metric Δ, and the judge's rationale. Uses the test split when `baseline_eval`/`final_eval` are given, otherwise the validation examples the fitter scored. |
 | **Data & settings** | Split sizes, augmentation stats, and every fitter knob used. |
 | **Run configuration**, **Recommendations** | Model, optimizer, suggestions, parameter changes, metric deltas, a rollout recommendation. |
 | **Curves** | Best score / loss per round, Keras `loss` / `val_loss` and metrics per epoch, the held-out evaluation. |