[piotrlaczkowski]

Why

A dataset row carries more than a question and an answer:

• context documents (RAG);
• other agent inputs: a plan, conversation turns, …;
• metadata, tags and a rubric.

Prompt optimization dropped most of that:

• The rewrite model saw only question, expected answer, actual answer and judge feedback, never the row's documents, inputs, metadata or tags.
• LLM judges never saw the row's other inputs, metadata or tags. For a RAG agent that retrieves for itself, they also never saw what it retrieved.
• The augmenter saw only question and answer, so RAG variations could not stay answerable from the seed's documents.
• The fitter corrupted the dataset. The agent's context.documents was replaced by a JSON dump of the whole context dict. This was written back into the user's dataset in place and nested one level deeper every epoch. agent.evaluate() also passed example.input by reference.

What changed

Core (agentomatic.optimize, agentomatic.agents)

• DataPoint gains tags. Context is normalised: one string, a list, or {"documents": [...]}, with document dicts rendered as content (source: x).
• example_context() is the single view of a row's inputs, documents, metadata and tags.
• The runner deep-copies a point's inputs. It no longer re-sends an example's own context, and merges documents without duplicates.
• On the local and HTTP paths, the runner captures what an agent says it retrieved (retrieval_context, citations, sources or documents).
• Evaluation hands agents a copy of example.input; datasets are never mutated.
• Optimizers receive each example's view together with retrieval, tool calls and reasoning.
• The briefing shows Context documents, Other inputs, Example metadata and Tags for every failure, success and sample, with the expected answer before the rubric. It tells the rewrite model to write instructions that use context and never to copy one row's facts.
• During fit and evaluate, LLM judges also get the row's other inputs, metadata and tags. When a row has no documents, evaluate() judges against the agent's retrieved sources.
• ScoreMetricAdapter rebuilds the same example a metric sees in evaluate().
• The augmenter's prompt shows the seed's inputs and documents, and variations get a copy of them.
• The keyword metric reads metadata.must_include.
• The fit report shows a context column per example.

Templates

• The rag agent answers from documents sent with the request (context.documents), cites them as [n], and falls back to its own retrieval. It returns citations.
• The rag prompts.json now matches the agent's grounded prompt; before, fit() optimized the generic one.
• Seed datasets/all.jsonl for every family carries tags.
  • rag rows add documents, must_include facts, a rubric and a refusal row.
  • Multi-turn chatbot rows carry their earlier turns.
• The generated README documents the row format.

Examples (examples/prompt_optimization)

• New rag_context.py + datasets/policies.jsonl:
  • rows bring their own documents, plan, metadata, tags and rubric;
  • prints what the agent, the judge and the rewrite model see;
  • context-aware metrics, and test scores per tag.
• The agent uses supplied documents and the customer's plan.
• support.jsonl gets tags and plan inputs.
• train.py reports scores per tag.
• metrics_and_judges.py shows metrics that read documents and inputs.

Docs

• docs/optimization/data.md: the row format, a "What each stage sees" table, and a "RAG datasets" section.
• Updated: how-it-works, metrics, reports, examples and the agent primer.

Verification

Run locally:

• ruff check and format;
• mypy;
• mkdocs build --strict;
• the full suite with the release-gate selection (-m "not slow and not live"): 3023 passed. One test is deselected because it fails only as root.

New tests:

• tests/test_example_context_propagation.py (53 tests):
  • data model, local and HTTP runner, concurrency, briefing, judges, metric adapters and augmentation;
  • end-to-end fit() and evaluate() against the fake model server, asserting what each request actually contained;
  • the dataset is unchanged after two epochs with an agent that edits its input.
• tests/test_template_rag_context.py:
  • seed rows for every family, and the scaffolded RAG agent over REST;
  • the generated rag train.py run end to end in a subprocess: augmentation → fit → judge → report.
• tests/test_prompt_optimization_examples.py: rag_context.py with and without augmentation, plus per-tag results and the new metrics.

All examples run against tests/fake_openai_server.py; they have not been run against a real model.

🤖 Generated with Claude Code

https://claude.ai/code/session_011gEMypXBosTqaEh3QFVXCX

---

Generated by Claude Code