From 81f62279a7915dfa9dfc152246d3f35bbd5ccbf1 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 2 Oct 2026 10:43:13 +0000 Subject: [PATCH 1/5] feat(optimize): pass every example's full context to the optimizer, judges and augmenter A dataset row is more than a question and an answer, but the optimization only passed those along: * the rewrite model saw question / expected / actual / judge feedback; * judges saw no inputs, metadata or tags; * the augmenter saw only question and answer. The agent did get its inputs, but `fit()` also corrupted them. The example's own `context` dict was handed back with `documents` replaced by a JSON dump of itself, written into the user's dataset in place, and nested one level deeper every epoch. `agent.evaluate()` likewise passed `example.input` by reference, so an agent editing its input rewrote the dataset. Now: - `DataPoint` gains `tags`. Context is normalised: one string, a list, or `{"documents": [...]}`. Document dicts render as `content (source: x)`. `example_context()` is the one view of the inputs, documents, metadata and tags. - The runner deep-copies a point's inputs. It no longer re-sends an example's own context to the agent, and merges documents without duplicates. It also captures what an agent says it retrieved (`retrieval_context` / `citations` / `sources` / `documents`) on the local and HTTP paths. - Evaluation hands agents a copy of `example.input` (`copy_input`). - The fitter attaches the example view to every eval detail. Optimizers receive it together with retrieval, tool calls and reasoning. - The briefing shows `Context documents`, `Other inputs`, `Example metadata` and `Tags` for each failure, success and sample, and the expected answer before the rubric. It also tells the rewrite model to write instructions that use such context, never to copy a row's facts. - LLM judges (`LocalJudgeMetric`, `LLMJudgeMetric`) also get the example's other inputs, metadata and tags during fit/evaluate. With no documents in the row, `evaluate()` judges against the agent's retrieved sources. - `ScoreMetricAdapter` rebuilds the example a metric sees in `evaluate()`: the original inputs and context, plus tags. - The augmenter's prompt shows the seed's inputs and documents, and variations get a deep copy of the seed's inputs. - The keyword metric reads `must_include` from metadata too. - The fit report shows a context column per example. Persisted example views are clipped. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_011gEMypXBosTqaEh3QFVXCX --- src/agentomatic/agents/base.py | 5 +- src/agentomatic/agents/metrics.py | 30 +- src/agentomatic/agents/mixins/evaluation.py | 4 +- src/agentomatic/agents/optimizers.py | 4 +- src/agentomatic/agents/types.py | 44 +- src/agentomatic/optimize/briefing.py | 101 ++- src/agentomatic/optimize/dataset.py | 157 +++- src/agentomatic/optimize/fitter.py | 75 +- src/agentomatic/optimize/holysheet_reports.py | 39 +- src/agentomatic/optimize/judges.py | 9 + src/agentomatic/optimize/metrics.py | 60 +- src/agentomatic/optimize/runner.py | 96 ++- .../optimize/structured_metrics.py | 11 +- src/agentomatic/optimize/synthesizer.py | 42 + src/agentomatic/optimize/train_api.py | 5 +- tests/test_example_context_propagation.py | 745 ++++++++++++++++++ 16 files changed, 1376 insertions(+), 51 deletions(-) create mode 100644 tests/test_example_context_propagation.py diff --git a/src/agentomatic/agents/base.py b/src/agentomatic/agents/base.py index ec914d03..ef22f4ba 100644 --- a/src/agentomatic/agents/base.py +++ b/src/agentomatic/agents/base.py @@ -74,6 +74,7 @@ def state_to_output(self, state): Optimizer, StateT, TraceEvent, + copy_input, ) # Per-request system-prompt overrides, keyed by agent ``id()``. A ContextVar @@ -590,7 +591,7 @@ def evaluate( for example in examples: t0 = time.perf_counter() try: - prediction = self.transform(example.input) + prediction = self.transform(copy_input(example.input)) duration = (time.perf_counter() - t0) * 1000 scores: dict[str, float] = {} @@ -1076,7 +1077,7 @@ def _epoch_logs( for example in examples: try: - prediction = self.transform(example.input) + prediction = self.transform(copy_input(example.input)) except Exception as exc: # noqa: BLE001 logger.warning(f"fit: transform failed for '{example.id}': {exc}") if loss is not None: diff --git a/src/agentomatic/agents/metrics.py b/src/agentomatic/agents/metrics.py index beec5da9..30caa7c6 100644 --- a/src/agentomatic/agents/metrics.py +++ b/src/agentomatic/agents/metrics.py @@ -429,6 +429,17 @@ async def _evaluate_as_run( from agentomatic.optimize.metrics import scoring_run from agentomatic.optimize.runner import RunResult + # The agent inputs ride in ``metadata.invoke``, as they do in fit(). + raw_input = getattr(example, "input", None) + invoke = ( + { + k: v + for k, v in raw_input.items() + if k not in {"query", "request", "question", "current_query"} + } + if isinstance(raw_input, dict) + else {} + ) run = RunResult( query=query, response=response, @@ -436,10 +447,12 @@ async def _evaluate_as_run( context=list(context or []), metadata={ "example": { + **({"invoke": invoke} if invoke else {}), **(getattr(example, "metadata", None) or {}), "id": getattr(example, "id", ""), "split": getattr(example, "split", ""), }, + "example_tags": list(getattr(example, "tags", None) or []), "output": dict(prediction), }, ) @@ -508,14 +521,17 @@ def score( ) if context is None: + from agentomatic.optimize.dataset import normalize_context + inp = getattr(example, "input", None) or {} - raw_ctx = None - if hasattr(inp, "get"): - raw_ctx = inp.get("context") - if isinstance(raw_ctx, dict) and raw_ctx: - context = [_json.dumps(raw_ctx, ensure_ascii=False)] - elif isinstance(raw_ctx, list) and raw_ctx: - context = [str(c) for c in raw_ctx if c] + raw_ctx = inp.get("context") if hasattr(inp, "get") else None + context = normalize_context(raw_ctx) or None + if context is None: + # No reference documents in the example: judge groundedness + # against what the agent says it retrieved (as PromptFitter does). + from agentomatic.optimize.runner import retrieval_from_output + + context = retrieval_from_output(prediction) or None # --- the text judged: the same rule PromptFitter scores with --- # (structured ``output`` dict as JSON, else ``response``/``answer``), diff --git a/src/agentomatic/agents/mixins/evaluation.py b/src/agentomatic/agents/mixins/evaluation.py index acad543d..638e1aef 100644 --- a/src/agentomatic/agents/mixins/evaluation.py +++ b/src/agentomatic/agents/mixins/evaluation.py @@ -12,6 +12,8 @@ from loguru import logger +from ..types import copy_input + if TYPE_CHECKING: from ..types import ( AgentDataset, @@ -83,7 +85,7 @@ def evaluate( start = time.perf_counter() try: prediction: dict[str, Any] = self.transform( # type: ignore[attr-defined] - example.input, + copy_input(example.input), ) except Exception as exc: # noqa: BLE001 elapsed = (time.perf_counter() - start) * 1000 diff --git a/src/agentomatic/agents/optimizers.py b/src/agentomatic/agents/optimizers.py index dda73f77..edd752d2 100644 --- a/src/agentomatic/agents/optimizers.py +++ b/src/agentomatic/agents/optimizers.py @@ -16,7 +16,7 @@ from loguru import logger -from .types import AgentDataset, Metric +from .types import AgentDataset, Metric, copy_input if TYPE_CHECKING: from agentomatic.optimize.llm_types import LLMSpec @@ -124,7 +124,7 @@ def optimize( count = 0 for example in examples: try: - prediction = agent.transform(example.input) + prediction = agent.transform(copy_input(example.input)) for metric in metrics: total_score += metric.score( example, diff --git a/src/agentomatic/agents/types.py b/src/agentomatic/agents/types.py index f7fbfb7e..67e61846 100644 --- a/src/agentomatic/agents/types.py +++ b/src/agentomatic/agents/types.py @@ -212,6 +212,27 @@ def _field_quality_hint(field_name: str, input_data: dict[str, Any]) -> str: return base +def copy_input(value: Any) -> Any: + """A deep copy of an example's input, to hand to an agent. + + Evaluation must never change the dataset: an agent that edits the dict it + receives (appending retrieved documents to ``context``, say) would + otherwise rewrite the example for every later epoch and split. + + Args: + value: The example's ``input``. + + Returns: + A deep copy (a shallow one for uncopyable objects). + """ + import copy + + try: + return copy.deepcopy(value) + except Exception: # noqa: BLE001 - uncopyable user objects + return dict(value) if isinstance(value, dict) else value + + @dataclass class AgentExample: """Single evaluation example for an agent. @@ -278,9 +299,14 @@ def to_datapoint(self) -> Any: - ``input["request"]`` — REST-style alias - ``json.dumps(input)`` — last-resort serialisation - Context: ``input["context"]`` (dict or list) is serialised as a JSON - string and passed as the retrieval context so LLM-as-judge metrics can - evaluate groundedness against the project snapshot. + Context: ``input["context"]`` becomes the point's context documents so + LLM-as-judge metrics can evaluate groundedness and the optimizer sees + what the agent answered from — a list as is, one string as one + document, a dict's ``documents`` as the documents (any other keys as + one extra JSON document). The agent itself still receives the + original ``input`` (via ``metadata.invoke``), unchanged. + + Tags travel as ``DataPoint.tags``. Expected answer: boolean-flag outputs (``{"content": true, "name": true}``) are converted to a human-readable description so judge LLMs can compare @@ -307,13 +333,12 @@ def to_datapoint(self) -> Any: input_data=self.input or {}, ) - # ── context: extract project snapshot for judge groundedness eval ─ + # ── context: documents for judge groundedness eval (and the optimizer) ─ + # A list, one string, or a dict (``documents`` + any other keys). + from agentomatic.optimize.dataset import normalize_context + raw_ctx = self.input.get("context") or (self.metadata.get("invoke") or {}).get("context") - context_list: list[str] = [] - if isinstance(raw_ctx, dict) and raw_ctx: - context_list = [json.dumps(raw_ctx, ensure_ascii=False)] - elif isinstance(raw_ctx, list): - context_list = [str(c) for c in raw_ctx if c] + context_list = normalize_context(raw_ctx) # ── metadata / invoke ───────────────────────────────────────────── meta = dict(self.metadata or {}) @@ -334,6 +359,7 @@ def to_datapoint(self) -> Any: expected_answer=expected, context=context_list, metadata=meta, + tags=list(self.tags or []), ) diff --git a/src/agentomatic/optimize/briefing.py b/src/agentomatic/optimize/briefing.py index b9456833..58203494 100644 --- a/src/agentomatic/optimize/briefing.py +++ b/src/agentomatic/optimize/briefing.py @@ -129,6 +129,24 @@ def _clip(value: Any, limit: int = 400) -> str: return text[: limit - 1] + "…" +def _expected_first(value: Any) -> Any: + """Put an expected reference's answer before its rubric / guidance. + + ``AgentExample.to_datapoint`` renders ``## Judge guidance`` and + ``## Rubric`` ahead of ``## Expected answer`` — the order a judge reads + them in. A clipped briefing line would then lose the answer itself, the + one thing a rewrite model most needs, so it goes first here (and the JSON + ``## Expected structured output`` that repeats it goes last). + """ + if not isinstance(value, str) or "## Expected answer" not in value: + return value + sections = [sec.strip() for sec in re.split(r"\n(?=## )", value.strip())] + answer = [sec for sec in sections if sec.startswith("## Expected answer")] + structured = [sec for sec in sections if sec.startswith("## Expected structured output")] + rest = [sec for sec in sections if sec not in answer and sec not in structured] + return "\n\n".join(answer + rest + structured) + + def _labelled(label: str, value: Any, limit: int = 400, *, indent: str = "") -> str: """Render ``label: value``, keeping a multi-line value under its label. @@ -144,7 +162,7 @@ def _labelled(label: str, value: Any, limit: int = 400, *, indent: str = "") -> belongs to. That is the single field a rewrite model most needs to find, so give a multi-line value its own indented block instead. """ - text = _clip(value, limit) + text = _clip(_expected_first(value), limit) if "\n" not in text: return f"{indent}{label}: {text}" body = "\n".join(f"{indent} {line}" if line.strip() else "" for line in text.splitlines()) @@ -159,6 +177,51 @@ def _fmt_json(data: Any, limit: int = 800) -> str: return _clip(text, limit) +def _fmt_inline(data: Any, limit: int) -> str: + """Compact one-line JSON, clipped.""" + try: + text = json.dumps(data, ensure_ascii=False, default=str) + except TypeError: + text = str(data) + return _clip(text, limit) + + +def format_example_context(view: dict[str, Any] | None, limit: int = 400) -> list[str]: + """Bullet lines for an example's context documents, inputs, metadata and tags. + + Args: + view: An :func:`~agentomatic.optimize.dataset.example_context` view. + limit: Character budget per line (documents share it). + + Returns: + Lines such as ``- Context documents: …`` / ``- Other inputs: …`` / + ``- Example metadata: …`` / ``- Tags: …``; empty for a bare example. + """ + if not view: + return [] + from agentomatic.optimize.dataset import normalize_context + + lines: list[str] = [] + inputs = dict(view.get("inputs") or {}) + docs = list(view.get("context") or []) + if not docs and "context" in inputs: + # The example's own context, as the agent received it: show it as + # documents rather than as one JSON blob. + docs = normalize_context(inputs.pop("context")) + if docs: + shown = docs[:4] + per_doc = max(80, limit // len(shown)) + more = f" (+{len(docs) - len(shown)} more)" if len(docs) > len(shown) else "" + lines.append("- Context documents: " + " | ".join(_clip(d, per_doc) for d in shown) + more) + if inputs: + lines.append(f"- Other inputs: {_fmt_inline(inputs, limit)}") + if view.get("metadata"): + lines.append(f"- Example metadata: {_fmt_inline(view['metadata'], limit)}") + if view.get("tags"): + lines.append("- Tags: " + _clip(", ".join(str(t) for t in view["tags"]), limit)) + return lines + + def format_runtime_config(config: PromptRuntimeConfig) -> str: """Serialize the full runtime config surface for LLM consumption.""" lines = [ @@ -214,18 +277,25 @@ def format_search_space(space: PromptSearchSpace | None) -> str: def format_dataset_samples(samples: list[Any], *, max_items: int = 8) -> str: - """Format train/val sample queries + expected answers.""" + """Format train sample queries, expected answers and their example context.""" if not samples: return "No dataset samples provided." + from agentomatic.optimize.dataset import example_context + lines: list[str] = [] for i, raw in enumerate(samples[:max_items], 1): if isinstance(raw, dict): query = raw.get("query") or raw.get("input") or "" expected = raw.get("expected_answer") or raw.get("expected") or raw.get("output") or "" + meta, context, tags = raw.get("metadata"), raw.get("context"), raw.get("tags") else: query = getattr(raw, "query", "") expected = getattr(raw, "expected_answer", None) or getattr(raw, "expected", "") + meta = getattr(raw, "metadata", None) + context, tags = getattr(raw, "context", None), getattr(raw, "tags", None) lines.append(f"{i}. Q: {_clip(query, 220)}") + view = example_context(meta if isinstance(meta, dict) else None, context, tags) + lines.extend(f" {line}" for line in format_example_context(view, 260)) lines.append(_labelled("Expected", expected, 220, indent=" ")) return "\n".join(lines) @@ -253,6 +323,7 @@ def format_eval_io( score = float(fail.get("score", fail.get("avg_score", 0.0)) or 0.0) lines.append(f"\n**Failure {idx}** (score={score:.3f})") lines.append(f"- Input/query: {_clip(fail.get('query'), clip)}") + lines.extend(format_example_context(fail.get("example"), clip)) lines.append(_labelled("- Expected", fail.get("expected"), clip)) lines.append(f"- Actual output: {_clip(fail.get('response'), clip)}") issues = fail.get("feedback") or fail.get("reason") or fail.get("details") @@ -265,7 +336,9 @@ def format_eval_io( ) ret_ctx = fail.get("retrieval_context") or [] if ret_ctx: - lines.append("- Retrieval: " + "; ".join(_clip(d, 100) for d in ret_ctx[:3])) + lines.append( + "- Retrieved by the agent: " + " | ".join(_clip(d, 120) for d in ret_ctx[:3]) + ) tools = fail.get("tool_calls") or [] if tools: names = ", ".join( @@ -289,12 +362,25 @@ def format_eval_io( score = float(suc.get("score", suc.get("avg_score", 0.0)) or 0.0) lines.append(f"\n**Success {idx}** (score={score:.3f})") lines.append(f"- Input/query: {_clip(suc.get('query'), min(clip, 220))}") + lines.extend(format_example_context(suc.get("example"), min(clip, 220))) lines.append(_labelled("- Expected", suc.get("expected"), min(clip, 220))) lines.append(f"- Actual output: {_clip(suc.get('response'), min(clip, 220))}") return "\n".join(lines) +def _has_example_context(sample: Any) -> bool: + """True when a dataset sample carries context, extra inputs, metadata or tags.""" + from agentomatic.optimize.dataset import example_context + + if isinstance(sample, dict): + meta, context, tags = sample.get("metadata"), sample.get("context"), sample.get("tags") + else: + meta = getattr(sample, "metadata", None) + context, tags = getattr(sample, "context", None), getattr(sample, "tags", None) + return bool(example_context(meta if isinstance(meta, dict) else None, context, tags)) + + def build_full_optimization_briefing( *, current_config: PromptRuntimeConfig, @@ -338,6 +424,15 @@ def build_full_optimization_briefing( "prompt/config using the observed inputs, outputs, expected answers, " "scores, parameters, and judge feedback.", ] + if any((r or {}).get("example") for r in eval_results) or any( + _has_example_context(s) for s in dataset_sample or [] + ): + sections.append( + "Each example lists what it ran with — context documents, other " + "inputs, metadata and tags. Write instructions that tell the agent " + "how to USE such context (ground answers in the supplied documents, " + "honour the inputs); never copy one example's facts into the prompt." + ) if agent_name: sections.append(f"**Agent:** `{agent_name}`") if rewrite_model: diff --git a/src/agentomatic/optimize/dataset.py b/src/agentomatic/optimize/dataset.py index b5374694..b295505d 100644 --- a/src/agentomatic/optimize/dataset.py +++ b/src/agentomatic/optimize/dataset.py @@ -8,6 +8,143 @@ from dataclasses import dataclass, field from typing import Any +#: ``metadata`` keys the optimization stack adds to a point itself (agent +#: inputs, split label, per-run knobs). They are plumbing, not example context. +INTERNAL_METADATA_KEYS = frozenset( + { + "invoke", + "split", + "model_params", + "optimize_nodes", + "node_match", + "resource_id", + "parent_id", + "id", + # Folded into the expected reference by ``AgentExample.to_datapoint``. + "judge_expected", + "expected_reference", + "quality_reference", + } +) + +#: ``metadata.invoke`` keys that are run knobs or the question itself. +INTERNAL_INVOKE_KEYS = frozenset( + { + "query", + "current_query", + "question", + "request", + "model_params", + "temperature", + "optimize_nodes", + "node_match", + "system_prompt_override", + "user_id", + } +) + + +#: Keys of a document dict that hold its text, in order of preference. +DOCUMENT_TEXT_KEYS = ("content", "text", "page_content", "snippet", "document") + + +def document_text(item: Any) -> str: + """One document as text: ``" (source: )"`` for a document dict. + + Args: + item: A string, or a dict with ``content`` / ``text`` / + ``page_content`` … and an optional ``source`` / ``title`` / ``id``. + + Returns: + The text; a dict without a text field is serialised as JSON. + """ + if isinstance(item, str): + return item + if isinstance(item, dict): + text = next((str(item[k]) for k in DOCUMENT_TEXT_KEYS if item.get(k)), "") + source = item.get("source") or item.get("title") or item.get("id") + if text: + return f"{text} (source: {source})" if source else text + return json.dumps(item, ensure_ascii=False, default=str) if item else "" + return str(item) if item else "" + + +def normalize_context(raw: Any) -> list[str]: + """Return context documents as a list of strings. + + Args: + raw: A list of documents, one document string, or a context dict. + A dict's ``documents`` become the documents; any other keys are + kept as one extra JSON document. Document dicts are rendered by + :func:`document_text`. + + Returns: + The documents; empty when there are none. + """ + if isinstance(raw, str): + return [raw] if raw.strip() else [] + if isinstance(raw, dict): + if not raw: + return [] + docs = raw.get("documents") + if not isinstance(docs, (str, list, tuple)): + return [json.dumps(raw, ensure_ascii=False, default=str)] + rest = {k: v for k, v in raw.items() if k != "documents" and v not in (None, "", [], {})} + extra = [json.dumps(rest, ensure_ascii=False, default=str)] if rest else [] + return normalize_context(docs) + extra + if isinstance(raw, (list, tuple)): + return [text for text in (document_text(item) for item in raw if item) if text] + return [] + + +def example_context( + metadata: dict[str, Any] | None, + context: list[str] | None = None, + tags: list[str] | None = None, +) -> dict[str, Any]: + """Everything an example carries beyond its question and expected answer. + + This is what the optimizer's rewrite model, the augmenter and the reports + see next to each question, so a prompt is improved against the + documents, extra inputs, labels and tags the agent actually ran with — + not against bare question/answer pairs. + + Args: + metadata: The point's metadata (``metadata.invoke`` holds the agent + inputs other than the question). + context: Context documents of the point (``DataPoint.context``). + tags: The example's tags. + + Returns: + ``{"inputs": ..., "context": ..., "metadata": ..., "tags": ...}`` with + only the non-empty keys. ``context`` is omitted when the inputs + already carry it (the agent received it there). + """ + meta = dict(metadata or {}) + invoke = meta.get("invoke") + view: dict[str, Any] = {} + if isinstance(invoke, dict): + inputs = { + k: v + for k, v in invoke.items() + if k not in INTERNAL_INVOKE_KEYS and v not in (None, "", [], {}) + } + if inputs: + view["inputs"] = inputs + docs = normalize_context(context) + if docs and "context" not in view.get("inputs", {}): + view["context"] = docs + labels = { + k: v + for k, v in meta.items() + if k not in INTERNAL_METADATA_KEYS and v not in (None, "", [], {}) + } + if labels: + view["metadata"] = labels + if tags: + view["tags"] = [str(t) for t in tags] + return view + @dataclass class DataPoint: @@ -16,14 +153,24 @@ class DataPoint: Attributes: query: The input question/prompt. expected_answer: The expected/ideal response (ground truth). - context: Optional context documents (for RAG evaluation). - metadata: Arbitrary metadata for filtering/grouping. + context: Optional context documents (for RAG evaluation). A single + string or a structured dict is normalised to a list. + metadata: Arbitrary metadata for filtering/grouping. ``metadata.invoke`` + holds extra agent inputs sent with the question. + tags: Labels for filtering/grouping, shown to the optimizer. """ query: str expected_answer: str | None = None context: list[str] = field(default_factory=list) metadata: dict[str, Any] = field(default_factory=dict) + tags: list[str] = field(default_factory=list) + + def __post_init__(self) -> None: + if not isinstance(self.context, list) or any(not isinstance(c, str) for c in self.context): + self.context = normalize_context(self.context) + if isinstance(self.tags, str): + self.tags = [self.tags] def to_dict(self) -> dict[str, Any]: """Serialize to dictionary.""" @@ -34,6 +181,8 @@ def to_dict(self) -> dict[str, Any]: d["context"] = self.context if self.metadata: d["metadata"] = self.metadata + if self.tags: + d["tags"] = list(self.tags) return d @@ -85,7 +234,7 @@ def from_jsonl(cls, path: str) -> Dataset: """Load from JSONL file. Each line must be a JSON object with at least a ``query`` field. - Optional fields: ``expected_answer``, ``context``, ``metadata``. + Optional fields: ``expected_answer``, ``context``, ``metadata``, ``tags``. """ points: list[DataPoint] = [] with open(path) as f: @@ -100,6 +249,7 @@ def from_jsonl(cls, path: str) -> Dataset: expected_answer=data.get("expected_answer"), context=data.get("context", []), metadata=data.get("metadata", {}), + tags=list(data.get("tags") or []), ) ) return cls(points=points) @@ -136,6 +286,7 @@ def from_list(cls, items: list[dict[str, Any]]) -> Dataset: expected_answer=item.get("expected_answer"), context=item.get("context", []), metadata=item.get("metadata", {}), + tags=list(item.get("tags") or []), ) for item in items ] diff --git a/src/agentomatic/optimize/fitter.py b/src/agentomatic/optimize/fitter.py index 4db5e7c1..1b20cf4d 100644 --- a/src/agentomatic/optimize/fitter.py +++ b/src/agentomatic/optimize/fitter.py @@ -58,7 +58,7 @@ OptimizationContext, RoundStats, ) -from agentomatic.optimize.dataset import Dataset +from agentomatic.optimize.dataset import Dataset, example_context from agentomatic.optimize.events import ( CallbackManager, EventData, @@ -74,7 +74,7 @@ CompositeMetric, EvalResult, ) -from agentomatic.optimize.runner import AgentRunner, RunResult +from agentomatic.optimize.runner import AgentRunner, RunResult, _merge_documents from agentomatic.optimize.search_space import PromptSearchSpace if TYPE_CHECKING: @@ -180,9 +180,14 @@ async def _callable( if invoke: input_data.update({k: v for k, v in invoke.items() if k != "current_query"}) if context: - input_data.setdefault("context", {}) - if isinstance(input_data["context"], dict): - input_data["context"]["documents"] = list(context) + # Copy before adding documents: ``invoke`` belongs to the dataset. + ctx = input_data.get("context") + if ctx is None: + input_data["context"] = {"documents": list(context)} + elif isinstance(ctx, dict): + ctx = dict(ctx) + ctx["documents"] = _merge_documents(ctx.get("documents"), list(context)) + input_data["context"] = ctx if prompt_override: # Top-level + metadata so resolve_system_prompt / transform stash both see it. input_data["system_prompt_override"] = prompt_override @@ -273,6 +278,50 @@ async def _call_with_swapped_attributes( return _callable +def _example_view(rr: RunResult) -> dict[str, Any]: + """The scored example's inputs, context, metadata and tags (see ``example_context``).""" + meta = rr.metadata or {} + return example_context(meta.get("example"), rr.context, meta.get("example_tags")) + + +#: Per-document / per-input character budget of a persisted example view. +_VIEW_CHARS = 300 + + +def _compact_view(view: dict[str, Any] | None) -> dict[str, Any]: + """An example view small enough to keep with every per-example record. + + The result (and its JSON artefact) stores one per scored example; long RAG + documents or a long conversation must not be copied there in full. + """ + if not view: + return {} + from agentomatic.optimize.dataset import normalize_context + + def clip(text: str) -> str: + return text if len(text) <= _VIEW_CHARS else text[: _VIEW_CHARS - 1] + "…" + + def clip_value(value: Any) -> Any: + if isinstance(value, str): + return clip(value) + text = json.dumps(value, ensure_ascii=False, default=str) + return value if len(text) <= _VIEW_CHARS else clip(text) + + inputs = dict(view.get("inputs") or {}) + docs = list(view.get("context") or []) or normalize_context(inputs.get("context")) + inputs.pop("context", None) + out: dict[str, Any] = {} + if docs: + out["context"] = [clip(str(d)) for d in docs[:5]] + if inputs: + out["inputs"] = {k: clip_value(v) for k, v in inputs.items()} + if view.get("metadata"): + out["metadata"] = {k: clip_value(v) for k, v in view["metadata"].items()} + if view.get("tags"): + out["tags"] = list(view["tags"]) + return out + + def _compact_examples(details: list[dict[str, Any]], limit: int = 50) -> list[dict[str, Any]]: """Per-example results trimmed for reports (question, answer, score, why).""" from agentomatic.optimize.metrics import plain_expected @@ -293,6 +342,8 @@ def _compact_examples(details: list[dict[str, Any]], limit: int = 50) -> list[di }, "feedback": str(feedback)[:500], "failed": bool(d.get("error")), + # Tags / documents / inputs / metadata of the example. + "example": _compact_view(d.get("example")), } ) return rows @@ -1794,7 +1845,8 @@ async def _record_round( artefact_dir.mkdir(parents=True, exist_ok=True) artefact_path = artefact_dir / f"fit_result_{experiment_id}.json" artefact_path.write_text( - json.dumps(result.to_dict(), indent=2, ensure_ascii=False) + "\n", + # ``default=str``: example inputs may hold non-JSON values. + json.dumps(result.to_dict(), indent=2, ensure_ascii=False, default=str) + "\n", encoding="utf-8", ) logger.info("💾 Artefacts: {}", artefact_path) @@ -2033,6 +2085,7 @@ async def _evaluate_config( logger.debug("Skipping errored result for '{}'", rr.query[:50]) detail = { "query": rr.query, + "example": _example_view(rr), "response": rr.response, "expected": rr.expected, "avg_score": 0.0, @@ -2078,6 +2131,7 @@ async def _evaluate_config( eval_details.append( { "query": rr.query, + "example": _example_view(rr), "response": rr.response, "expected": rr.expected, "avg_score": 0.0, @@ -2108,6 +2162,7 @@ async def _evaluate_config( reward = self._reward_adapter.reward_from_eval(eval_result) detail = { "query": rr.query, + "example": _example_view(rr), "response": rr.response, "expected": rr.expected, "avg_score": point_score, @@ -2148,6 +2203,7 @@ async def _evaluate_config( eval_details.append( { "query": rr.query, + "example": _example_view(rr), "response": rr.response, "expected": rr.expected, "avg_score": 0.0, @@ -2407,6 +2463,13 @@ def _build_eval_results( "dimensions": detail.get("dimensions", {}), "feedback": detail.get("feedback", ""), "is_failure": detail.get("avg_score", 0.0) < _FAILURE_THRESHOLD, + # Inputs / context / metadata / tags of the example, and + # what the agent retrieved and did — so the optimizer sees + # the whole case, not just the question and the answers. + "example": dict(detail.get("example") or {}), + "retrieval_context": list(detail.get("retrieval_context") or []), + "tool_calls": list(detail.get("tool_calls") or []), + "reasoning": detail.get("reasoning") or "", } ) return results diff --git a/src/agentomatic/optimize/holysheet_reports.py b/src/agentomatic/optimize/holysheet_reports.py index b5c718df..10fca635 100644 --- a/src/agentomatic/optimize/holysheet_reports.py +++ b/src/agentomatic/optimize/holysheet_reports.py @@ -1373,6 +1373,7 @@ def training_sections( example_rows.append( { "example": str(before.get("query", ""))[:120], + "context": _context_summary(before.get("example")), "expected": str(before.get("expected", ""))[:240], "before": str(before.get("response", ""))[:300], "before score": _fmt(before.get("score"), 2), @@ -1388,6 +1389,7 @@ def training_sections( ], } ) + example_rows = _drop_empty_context(example_rows) blocks.append( Section( title="Examples — before vs after", @@ -1435,6 +1437,35 @@ def training_sections( return blocks +def _context_summary(view: dict[str, Any] | None, limit: int = 240) -> str: + """One line for an example's tags, documents, other inputs and metadata.""" + if not view: + return "" + from agentomatic.optimize.dataset import normalize_context + + inputs = dict(view.get("inputs") or {}) + docs = list(view.get("context") or []) or normalize_context(inputs.pop("context", None)) + parts: list[str] = [] + if view.get("tags"): + parts.append("tags: " + ", ".join(str(t) for t in view["tags"])) + if docs: + first = str(docs[0]) + parts.append(f"{len(docs)} doc(s): " + (first[:80] + "…" if len(first) > 80 else first)) + if inputs: + parts.append("inputs: " + ", ".join(f"{k}={v}" for k, v in inputs.items())) + if view.get("metadata"): + parts.append("meta: " + ", ".join(f"{k}={v}" for k, v in view["metadata"].items())) + text = " · ".join(parts) + return text if len(text) <= limit else text[: limit - 1] + "…" + + +def _drop_empty_context(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Keep the ``context`` column only when some example has context.""" + if any(row.get("context") for row in rows): + return rows + return [{k: v for k, v in row.items() if k != "context"} for row in rows] + + def _example_rows(baseline_eval: Any, final_eval: Any, eval_dataset: Any) -> list[dict[str, Any]]: """Join before/after ``EvaluationReport`` example results by example id.""" if final_eval is None: @@ -1448,13 +1479,18 @@ def _example_rows(baseline_eval: Any, final_eval: Any, eval_dataset: Any) -> lis example = by_id.get(after.example_id) query = "" expected = "" + context = "" if example is not None: try: point = example.to_datapoint() query = point.query + from agentomatic.optimize.dataset import example_context from agentomatic.optimize.metrics import plain_expected expected = plain_expected(point.expected_answer) or "" + context = _context_summary( + example_context(point.metadata, point.context, point.tags) + ) except Exception: # noqa: BLE001 query = str(getattr(example, "input", "")) after_pred = after.prediction or {} @@ -1474,6 +1510,7 @@ def _example_rows(baseline_eval: Any, final_eval: Any, eval_dataset: Any) -> lis { "example": after.example_id, "question": str(query)[:200], + "context": context, "expected": str(expected)[:240], "before": str(before_pred.get("response", before.error if before else ""))[:300], "before scores": _scores_text(before.scores if before else None), @@ -1483,4 +1520,4 @@ def _example_rows(baseline_eval: Any, final_eval: Any, eval_dataset: Any) -> lis "why (judge)": rationale[:300], } ) - return rows + return _drop_empty_context(rows) diff --git a/src/agentomatic/optimize/judges.py b/src/agentomatic/optimize/judges.py index 7ac03d4b..72162971 100644 --- a/src/agentomatic/optimize/judges.py +++ b/src/agentomatic/optimize/judges.py @@ -163,6 +163,15 @@ async def evaluate_rich( if context: ctx_block = "\n".join(str(c)[:800] for c in context[:5]) prompt += f"## Context Documents\n{ctx_block}\n\n" + from agentomatic.optimize.metrics import scoring_example_labels + + labels = scoring_example_labels(has_context=bool(context)) + if labels: + prompt += ( + "## Example inputs, metadata and tags\n" + "The query was asked with these — judge the response in their light.\n" + f"{labels}\n\n" + ) prompt += ( "## Instructions\n" diff --git a/src/agentomatic/optimize/metrics.py b/src/agentomatic/optimize/metrics.py index 0c71b4bc..ce0c8cc0 100644 --- a/src/agentomatic/optimize/metrics.py +++ b/src/agentomatic/optimize/metrics.py @@ -396,6 +396,9 @@ async def _eval_llm( prompt += f"EXPECTED ANSWER: {expected}\n\n" if context: prompt += f"CONTEXT:\n{chr(10).join(context[:3])}\n\n" + labels = scoring_example_labels(has_context=bool(context)) + if labels: + prompt += f"EXAMPLE INPUTS, METADATA AND TAGS: {labels}\n\n" prompt += 'Reply with ONLY a JSON object: {"score": 0.X, "reason": "..."}\n' if isinstance(self.model, str): @@ -1599,6 +1602,47 @@ def current_scoring_run() -> Any: return _SCORING_RUN.get() +#: Character budget for the example labels shown to an LLM judge. +_JUDGE_LABELS_CHARS = 1500 + + +def scoring_example_labels(*, has_context: bool) -> str: + """The scored example's other inputs, metadata and tags, for a judge prompt. + + A judge already receives the question, the answer, the expected reference + and the context documents; this adds what else the example was asked + with (e.g. a customer tier, a locale, required facts, tags) so the + answer is judged in that light. + + Args: + has_context: The judge already lists the context documents, so the + ``context`` input is left out. + + Returns: + Compact JSON, or ``""`` when no run is being scored or it has none. + """ + import json + + from agentomatic.optimize.dataset import example_context + + meta = getattr(current_scoring_run(), "metadata", None) or {} + example = meta.get("example") + view = example_context( + example if isinstance(example, dict) else None, None, meta.get("example_tags") + ) + inputs = dict(view.pop("inputs", None) or {}) + if has_context: + inputs.pop("context", None) + if inputs: + view = {"inputs": inputs, **view} + if not view: + return "" + text = json.dumps(view, ensure_ascii=False, default=str) + if len(text) > _JUDGE_LABELS_CHARS: + text = text[: _JUDGE_LABELS_CHARS - 1] + "…" + return text + + def _decode_json_dict(text: str | None) -> dict[str, Any] | None: """Return ``text`` parsed as a JSON object, or ``None`` when it is not one.""" import json @@ -1679,16 +1723,30 @@ async def evaluate( expected_output = _expected_structured(expected) or _decode_json_dict(answer) if expected_output is None and answer is not None: expected_output = {"response": answer} + # Rebuild the example the metric would see in ``agent.evaluate()``: + # the agent inputs (``metadata.invoke`` — the original ``context`` + # included) and the tags, not only the question. + from agentomatic.optimize.dataset import INTERNAL_INVOKE_KEYS + + invoke = example_meta.get("invoke") + inputs = ( + {k: v for k, v in invoke.items() if k not in INTERNAL_INVOKE_KEYS} + if isinstance(invoke, dict) + else {} + ) + if context and "context" not in inputs: + inputs["context"] = list(context) example = AgentExample( id=str(example_meta.get("id") or "fit"), input={ + **inputs, "current_query": query, "query": query, "question": query, - **({"context": list(context)} if context else {}), }, expected_output=expected_output, metadata=example_meta, + tags=[str(t) for t in run_meta.get("example_tags") or []], split=str(example_meta.get("split") or "validation"), ) output = run_meta.get("output") diff --git a/src/agentomatic/optimize/runner.py b/src/agentomatic/optimize/runner.py index 105885d8..e73b1cd1 100644 --- a/src/agentomatic/optimize/runner.py +++ b/src/agentomatic/optimize/runner.py @@ -16,6 +16,67 @@ from loguru import logger +from agentomatic.optimize.dataset import document_text + +#: Keys an agent output (or response ``metadata``) may list its sources under. +_RETRIEVAL_KEYS = ("retrieval_context", "citations", "sources", "documents") + + +def retrieval_from_output(data: Any) -> list[str]: + """Documents an agent output says it answered from. + + A RAG agent that retrieves for itself reports its sources in its output; + surfacing them lets judges score groundedness and lets the optimizer see + what each answer was built from. + + Args: + data: An agent output / response dict. ``retrieval_context`` wins, + then ``citations``, ``sources`` and ``documents`` — at the top + level, under ``output`` or under ``metadata``. Items may be + strings or dicts (``content`` / ``text`` / ``page_content`` …). + + Returns: + The documents as strings; empty when the output lists none. + """ + if not isinstance(data, dict): + return [] + containers = [data] + for key in ("output", "metadata"): + nested = data.get(key) + if isinstance(nested, dict): + containers.append(nested) + for key in _RETRIEVAL_KEYS: + for container in containers: + items = container.get(key) + if isinstance(items, str): + items = [items] + if isinstance(items, list) and items: + docs = [d for d in (document_text(i) for i in items) if d] + if docs: + return docs + return [] + + +def _copy_inputs(invoke: dict[str, Any]) -> dict[str, Any]: + """A deep copy of a point's agent inputs (shallow when not copyable).""" + import copy + + try: + return copy.deepcopy(invoke) + except Exception: # noqa: BLE001 - uncopyable user objects: shallow copy + return dict(invoke) + + +def _merge_documents(existing: Any, extra: list[str]) -> list[Any]: + """``existing`` documents plus those of ``extra`` not already present.""" + if isinstance(existing, str): + docs: list[Any] = [existing] + elif isinstance(existing, (list, tuple)): + docs = list(existing) + else: + docs = [] + return docs + [d for d in extra if d not in docs] + def _response_text(data: dict[str, Any]) -> str: """Prefer structured ``output`` dict, else string ``response``.""" @@ -174,6 +235,8 @@ async def _run_local( return RunResult( query=query, response=response, + # What a RAG agent says it answered from (judges, optimizer). + retrieval_context=retrieval_from_output(raw), duration_ms=duration, # The full structured output, for metrics that score fields # other than the response text (see ScoreMetricAdapter). @@ -213,11 +276,8 @@ async def _run_optimize( payload["system_prompt_override"] = prompt_override if context: ctx = payload.get("context") - if not isinstance(ctx, dict): - ctx = {} - docs = list(ctx.get("documents") or []) - docs.extend(context) - ctx["documents"] = docs + ctx = dict(ctx) if isinstance(ctx, dict) else {} + ctx["documents"] = _merge_documents(ctx.get("documents"), list(context)) payload["context"] = ctx async with httpx.AsyncClient(base_url=self.api_base, timeout=self.timeout) as client: @@ -232,7 +292,7 @@ async def _run_optimize( return RunResult( query=query, response=_response_text(data), - retrieval_context=data.get("retrieval_context", []), + retrieval_context=retrieval_from_output(data), tool_calls=data.get("tool_calls", []), steps_taken=data.get("steps_taken", []), reasoning=data.get("reasoning", ""), @@ -273,8 +333,7 @@ async def _run_invoke( if prompt_override: ctx["system_prompt_override"] = prompt_override if context: - docs = list(ctx.get("documents") or []) - docs.extend(context) + docs = _merge_documents(ctx.get("documents"), list(context)) if docs: ctx["documents"] = docs if ctx: @@ -292,8 +351,9 @@ async def _run_invoke( return RunResult( query=query, response=_response_text(data), - # Try to extract context from standard response - retrieval_context=data.get("metadata", {}).get("retrieval_context", []), + # Sources from the response (``metadata.retrieval_context``, + # ``output.citations``, …). + retrieval_context=retrieval_from_output(data), citations=data.get("citations", []), steps_taken=data.get("steps_taken", []), duration_ms=duration, @@ -341,21 +401,29 @@ async def _run_one(point: dict[str, Any]) -> RunResult: async with semaphore: meta = point.get("metadata") or {} invoke = meta.get("invoke") if isinstance(meta, dict) else None - if not isinstance(invoke, dict): - invoke = {} + # A copy: an agent (or a payload builder) must never be able to + # change the dataset it is evaluated on. + invoke = _copy_inputs(invoke) if isinstance(invoke, dict) else {} + context = list(point.get("context") or []) result = await self.run_single( query=point["query"], prompt_override=prompt_override, - context=point.get("context"), + # An example's own context already travels in its inputs, + # unchanged; sending its documents again would hand the + # agent a second, re-serialised copy. + context=None if "context" in invoke else (context or None), invoke=invoke, ) result.expected = point.get("expected_answer") or point.get("expected") if not result.context: - result.context = point.get("context", []) + result.context = context if isinstance(meta, dict) and meta: # Dataset metadata (rubric hints, required facts, split…) # for example-aware metrics. result.metadata = {**result.metadata, "example": dict(meta)} + tags = point.get("tags") + if tags: + result.metadata = {**result.metadata, "example_tags": list(tags)} return result return list(await asyncio.gather(*[_run_one(p) for p in points])) diff --git a/src/agentomatic/optimize/structured_metrics.py b/src/agentomatic/optimize/structured_metrics.py index c3811270..568374eb 100644 --- a/src/agentomatic/optimize/structured_metrics.py +++ b/src/agentomatic/optimize/structured_metrics.py @@ -202,8 +202,17 @@ def _fn( def agent_keyword_score(example: Any, prediction: dict[str, Any]) -> float: - """Agent-lifecycle metric: must_include coverage with forbid penalty.""" + """Agent-lifecycle metric: must_include coverage with forbid penalty. + + Required / forbidden facts come from ``expected_output`` or, failing + that, the example's ``metadata`` (``must_include`` / ``must_not_include``). + """ expected = dict(getattr(example, "expected_output", None) or {}) + metadata = getattr(example, "metadata", None) or {} + if isinstance(metadata, dict): + for key in ("must_include", "must_not_include"): + if not expected.get(key) and metadata.get(key): + expected[key] = metadata[key] text = json.dumps(prediction, ensure_ascii=False).lower() must = [str(t) for t in (expected.get("must_include") or []) if str(t).strip()] if not must: diff --git a/src/agentomatic/optimize/synthesizer.py b/src/agentomatic/optimize/synthesizer.py index 87dc198c..cfd2fbba 100644 --- a/src/agentomatic/optimize/synthesizer.py +++ b/src/agentomatic/optimize/synthesizer.py @@ -904,6 +904,40 @@ async def red_team( _ITEM_LIST_KEYS = ("examples", "data", "items", "rows", "points", "variations", "results") +#: Character budget for the seed's context in an augmentation prompt. +_SEED_CONTEXT_CHARS = 3000 + + +def _seed_context_section(seed: DataPoint) -> tuple[str, bool]: + """The seed's inputs / context documents / tags as a prompt section. + + Every variation inherits the seed's inputs (``prepare_dataset`` copies + them), so the generator has to see them: a RAG variation must stay + answerable from the same documents. + + Returns: + ``(section, has_documents)`` — an empty section for a bare seed. + """ + from agentomatic.optimize.dataset import example_context, normalize_context + + view = example_context(seed.metadata, seed.context, seed.tags) + view.pop("metadata", None) # labels of the seed's answer, not of a variation + inputs = view.get("inputs") or {} + docs = view.get("context") or normalize_context(inputs.get("context")) + if not view: + return "", False + text = json.dumps(view, ensure_ascii=False, default=str) + if len(text) > _SEED_CONTEXT_CHARS: + text = text[: _SEED_CONTEXT_CHARS - 1] + "…" + return ( + "## Seed context\n" + "Every variation is asked with these same inputs" + + (" and context documents" if docs else "") + + f":\n{text}\n\n", + bool(docs), + ) + + def _seed_prompt(seed: DataPoint, strategy: str, k: int) -> str: """Build the one-seed augmentation prompt.""" from agentomatic.optimize.metrics import plain_expected @@ -911,9 +945,16 @@ def _seed_prompt(seed: DataPoint, strategy: str, k: int) -> str: answer = plain_expected(seed.expected_answer) or "" keep = strategy in LABEL_PRESERVING_STRATEGIES seed_json = json.dumps({"query": seed.query, "expected_answer": answer}, ensure_ascii=False) + context_section, has_docs = _seed_context_section(seed) + grounding = ( + "Ask only questions the context documents answer; do not use outside facts.\n" + if has_docs + else "" + ) return ( "You are a dataset augmentation expert.\n\n" f"## Seed example\n{seed_json}\n\n" + f"{context_section}" f"## Task\n{_STRATEGY_INSTRUCTIONS[strategy]}\n" f"Write {k} new, distinct variation(s) of the seed.\n" + ( @@ -921,6 +962,7 @@ def _seed_prompt(seed: DataPoint, strategy: str, k: int) -> str: if keep else "Give each variation its own correct expected_answer.\n" ) + + grounding + "\nReply with ONLY a JSON array, no prose:\n" '[{"query": "...", "expected_answer": "..."}]\n' ) diff --git a/src/agentomatic/optimize/train_api.py b/src/agentomatic/optimize/train_api.py index f46bc04e..25c78bda 100644 --- a/src/agentomatic/optimize/train_api.py +++ b/src/agentomatic/optimize/train_api.py @@ -37,6 +37,7 @@ from __future__ import annotations +import copy import json import os from collections.abc import Sequence @@ -416,7 +417,9 @@ def _augmented_example(seed: Any, dp: Any, *, strategy: str, idx: int) -> Any | query = str(getattr(dp, "query", "") or "").strip() if not query: return None - inp = dict(seed.input or {}) + # The seed's own inputs (context documents included) — copied, so the + # variation can never alias the seed. + inp = copy.deepcopy(dict(seed.input or {})) replaced = False for key in _QUERY_KEYS: if key in inp: diff --git a/tests/test_example_context_propagation.py b/tests/test_example_context_propagation.py new file mode 100644 index 00000000..91ddd284 --- /dev/null +++ b/tests/test_example_context_propagation.py @@ -0,0 +1,745 @@ +# pyright: reportMissingParameterType=none +"""An example's whole context must reach every stage of prompt optimization. + +A dataset row is more than a question and an answer: it carries context +documents (RAG), other agent inputs, metadata and tags. These tests pin down +who sees what: + +* the **agent** receives its inputs exactly as written — the original + ``context`` included, never a re-serialised copy — and ``fit()`` never + changes the dataset; +* **metrics** see the same example inside ``fit()`` as in ``evaluate()``; +* **LLM judges** see the context documents (or, for a RAG agent that + retrieves for itself, what it retrieved) plus the other inputs, metadata + and tags; +* the **optimizer's rewrite model** sees all of it next to each question; +* the **augmenter** sees the seed's inputs and documents, and variations + inherit them. +""" + +from __future__ import annotations + +import asyncio +import copy +import json +import sys +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import httpx +import pytest + +from agentomatic.agents import ( + AgentDataset, + AgentExample, + BaseGraphAgent, + CallableMetric, + PromptFitterBridge, +) +from agentomatic.optimize import LocalJudgeMetric +from agentomatic.optimize.briefing import ( + build_full_optimization_briefing, + format_dataset_samples, + format_eval_io, + format_example_context, +) +from agentomatic.optimize.config import PromptRuntimeConfig +from agentomatic.optimize.dataset import ( + DataPoint, + Dataset, + example_context, + normalize_context, +) +from agentomatic.optimize.fitter import _wrap_local_agent +from agentomatic.optimize.llm_caller import LLMCaller +from agentomatic.optimize.metrics import ( + ScoreMetricAdapter, + scoring_example_labels, + scoring_run, +) +from agentomatic.optimize.runner import AgentRunner, RunResult, retrieval_from_output +from agentomatic.optimize.synthesizer import _seed_prompt + +sys.path.insert(0, str(Path(__file__).parent)) + +from fake_openai_server import IMPROVED_PROMPT, FakeOpenAIServer, respond # noqa: E402 + +DOC = "CTXDOC: refunds are accepted within 30 days of purchase." +TIER = "INPUT-premium-tier" +POLICY = "META-policy-v7" +TAG = "TAG-billing" +RUBRIC = "RUBRIC-cite-the-policy" +RETRIEVED = "RETRIEVED: annual plans renew on the 1st." + + +def rich_example(i: int, split: str = "train", **overrides: Any) -> AgentExample: + """A row using every field: context, extra input, metadata, tags, rubric.""" + fields: dict[str, Any] = { + "id": f"r{i}", + "split": split, + "input": { + "current_query": f"Can I get a refund on order {i}?", + "context": {"documents": [DOC], "locale": "en-GB"}, + "customer_tier": TIER, + }, + "expected_output": {"response": "Yes, within 30 days of purchase."}, + "metadata": {"policy_version": POLICY, "difficulty": "hard"}, + "tags": [TAG, "refund"], + "rubric": {"citation": RUBRIC}, + } + fields.update(overrides) + return AgentExample(**fields) + + +def rich_dataset(n_train: int = 4, n_val: int = 3) -> AgentDataset: + examples = [rich_example(i) for i in range(n_train)] + examples += [rich_example(n_train + i, "validation") for i in range(n_val)] + return AgentDataset(name="rich", examples=examples) + + +def requests_for(server: FakeOpenAIServer, role: str) -> list[str]: + """Recorded request bodies (as JSON text) the fake server answered as ``role``.""" + out = [] + for body in server.requests: + text = json.dumps(body.get("messages") or [], ensure_ascii=False) + if role == "judge" and "expert evaluation judge" in text: + out.append(text) + elif role == "augment" and "dataset augmentation expert" in text: + out.append(text) + elif role == "rewrite" and respond(body.get("messages") or []).endswith(IMPROVED_PROMPT): + out.append(text) + return out + + +@pytest.fixture(scope="module") +def server(): + with FakeOpenAIServer() as srv: + yield srv + + +@pytest.fixture +def llm(server): + """Route optimizer-side calls to the fake server; restore the default after.""" + before = (LLMCaller._default_base_url, LLMCaller._default_api_key) # noqa: SLF001 + LLMCaller.configure(base_url=server.base_url, api_key="x") + server.requests.clear() + yield server + LLMCaller._default_base_url, LLMCaller._default_api_key = before # noqa: SLF001 + + +@dataclass +class _State: + question: str = "" + output: dict[str, Any] = field(default_factory=dict) + + +class RecordingAgent(BaseGraphAgent[_State]): + """Answers vaguely and records every input it is given.""" + + agent_name = "recorder" + system_prompt = "You are a support assistant." + + def __init__(self, *, citations: list[dict[str, Any]] | None = None, mutate: bool = False): + super().__init__() + self.seen: list[dict[str, Any]] = [] + self.citations = citations + self.mutate = mutate + + def build_graph(self) -> Any: + g = self.new_graph() + g.add_node("respond", self.respond) + g.set_entry_point("respond") + g.set_finish_point("respond") + return g.compile() + + def respond(self, state: _State) -> _State: + state.output = {"response": "Thanks for reaching out!"} + if self.citations is not None: + state.output["citations"] = list(self.citations) + return state + + def input_to_state(self, data: dict[str, Any]) -> _State: + self.seen.append(json.loads(json.dumps(data, default=str))) + if self.mutate and isinstance(data.get("context"), dict): + # A careless agent editing its input in place. + data["context"].setdefault("documents", []).append("AGENT-SCRIBBLE") + data["context"]["scribbled"] = True + return _State(question=data.get("current_query", "")) + + def state_to_output(self, state: _State) -> dict[str, Any]: + return state.output + + +def _bridge(server: FakeOpenAIServer, metric: Any, **kwargs: Any) -> PromptFitterBridge: + return PromptFitterBridge( + agent_name="recorder", + task_model="omlx/fake", + rewrite_model="omlx/fake", + llm_base_url=server.base_url, + llm_api_key="x", + optimizer="rewrite", + metric=metric, + max_trials=1, + **kwargs, + ) + + +# --------------------------------------------------------------------------- +# Data model +# --------------------------------------------------------------------------- + + +class TestNormalizeContext: + @pytest.mark.parametrize( + ("raw", "expected"), + [ + (None, []), + ("", []), + ("one document", ["one document"]), + (["a", "", "b"], ["a", "b"]), + ({"documents": ["a", "b"]}, ["a", "b"]), + ({"documents": "a"}, ["a"]), + ({}, []), + ], + ) + def test_shapes(self, raw, expected) -> None: + assert normalize_context(raw) == expected + + def test_dict_documents_keep_other_keys_as_one_extra_document(self) -> None: + docs = normalize_context({"documents": ["a"], "locale": "en-GB", "empty": ""}) + assert docs[0] == "a" + assert json.loads(docs[1]) == {"locale": "en-GB"} + + def test_dict_without_documents_is_one_json_document(self) -> None: + assert json.loads(normalize_context({"snapshot": {"k": 1}})[0]) == {"snapshot": {"k": 1}} + + def test_document_dicts_become_text_with_their_source(self) -> None: + docs = [{"content": "x", "source": "a.md"}, {"text": "y"}, {"meta": 1}] + assert normalize_context(docs) == ["x (source: a.md)", "y", '{"meta": 1}'] + assert normalize_context({"documents": [{"page_content": "p", "title": "t"}]}) == [ + "p (source: t)" + ] + + +class TestDataPoint: + def test_tags_round_trip_through_dicts_and_jsonl(self, tmp_path) -> None: + point = DataPoint(query="q", context="doc", metadata={"m": 1}, tags=["t"]) + assert point.context == ["doc"] + assert point.to_dict()["tags"] == ["t"] + assert Dataset.from_list([point.to_dict()])[0].tags == ["t"] + path = tmp_path / "d.jsonl" + Dataset(points=[point]).to_jsonl(str(path)) + loaded = Dataset.from_jsonl(str(path))[0] + assert (loaded.tags, loaded.context, loaded.metadata) == (["t"], ["doc"], {"m": 1}) + + def test_a_dict_context_is_normalised(self) -> None: + assert DataPoint(query="q", context={"documents": ["a"]}).context == ["a"] # type: ignore[arg-type] + + +class TestToDatapoint: + def test_carries_documents_tags_inputs_and_metadata(self) -> None: + dp = rich_example(1).to_datapoint() + assert dp.context[0] == DOC + assert dp.tags == [TAG, "refund"] + assert dp.metadata["invoke"]["customer_tier"] == TIER + # The agent gets the context exactly as written. + assert dp.metadata["invoke"]["context"] == {"documents": [DOC], "locale": "en-GB"} + assert dp.metadata["policy_version"] == POLICY + assert RUBRIC in (dp.expected_answer or "") + + def test_a_string_context_reaches_judges(self) -> None: + example = AgentExample(input={"current_query": "q", "context": "plain text doc"}) + assert example.to_datapoint().context == ["plain text doc"] + + +class TestExampleContextView: + def test_view_drops_plumbing_and_duplicate_context(self) -> None: + dp = rich_example(1).to_datapoint() + meta = { + **dp.metadata, + "invoke": {**dp.metadata["invoke"], "model_params": {"temperature": 0}}, + "split": "train", + "resource_id": "r", + } + view = example_context(meta, dp.context, dp.tags) + assert view["inputs"] == { + "context": {"documents": [DOC], "locale": "en-GB"}, + "customer_tier": TIER, + } + assert "context" not in view # already in the inputs + assert view["metadata"] == {"policy_version": POLICY, "difficulty": "hard"} + assert view["tags"] == [TAG, "refund"] + + def test_plain_dataset_context_is_kept(self) -> None: + view = example_context({"topic": "x"}, ["doc"], None) + assert view == {"context": ["doc"], "metadata": {"topic": "x"}} + + def test_bare_example_has_no_view(self) -> None: + assert example_context({"split": "train", "invoke": {"current_query": "q"}}) == {} + + +# --------------------------------------------------------------------------- +# Runner: inputs reach the agent intact, the dataset is never changed +# --------------------------------------------------------------------------- + + +class TestLocalRunner: + async def test_inputs_arrive_once_and_the_dataset_is_untouched(self) -> None: + agent = RecordingAgent(mutate=True) + runner = AgentRunner(agent="recorder", agent_callable=_wrap_local_agent(agent)) + points = [rich_example(i).to_datapoint().to_dict() for i in range(3)] + snapshot = copy.deepcopy(points) + + results = await runner.run_dataset(points, prompt_override="P") + await runner.run_dataset(points, prompt_override="P") # a second epoch + + assert points == snapshot + for seen in agent.seen: + assert seen["context"]["documents"] == [DOC] + assert seen["context"]["locale"] == "en-GB" + assert seen["customer_tier"] == TIER + assert results[0].metadata["example_tags"] == [TAG, "refund"] + assert results[0].context[0] == DOC # for judges + + async def test_plain_dataset_context_is_sent_as_documents(self) -> None: + agent = RecordingAgent() + runner = AgentRunner(agent="recorder", agent_callable=_wrap_local_agent(agent)) + await runner.run_dataset([{"query": "q", "context": ["d1", "d2"]}]) + assert agent.seen[0]["context"] == {"documents": ["d1", "d2"]} + + async def test_local_retrieval_is_captured(self) -> None: + agent = RecordingAgent(citations=[{"content": RETRIEVED, "source": "kb"}]) + runner = AgentRunner(agent="recorder", agent_callable=_wrap_local_agent(agent)) + (result,) = await runner.run_dataset([{"query": "q"}]) + assert result.retrieval_context == [f"{RETRIEVED} (source: kb)"] + + +class TestWrapLocalAgent: + async def test_documents_merge_into_a_copy(self) -> None: + agent = RecordingAgent() + call = _wrap_local_agent(agent) + invoke = {"context": {"documents": ["a"]}} + await call("q", context=["a", "b"], invoke=invoke) + assert agent.seen[0]["context"]["documents"] == ["a", "b"] + assert invoke == {"context": {"documents": ["a"]}} + + +class TestHttpRunner: + @pytest.fixture + def captured(self, monkeypatch): + payloads: list[tuple[str, dict[str, Any]]] = [] + optimize_endpoint = {"enabled": True} + + def handler(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + payloads.append((request.url.path, body)) + if request.url.path.endswith("/optimize/invoke") and not optimize_endpoint["enabled"]: + return httpx.Response(404, json={"detail": "nope"}) + return httpx.Response( + 200, + json={ + "response": "ok", + "metadata": {"retrieval_context": [RETRIEVED]}, + }, + ) + + real = httpx.AsyncClient + + def client(*args: Any, **kwargs: Any) -> httpx.AsyncClient: + return real(*args, transport=httpx.MockTransport(handler), **kwargs) + + monkeypatch.setattr(httpx, "AsyncClient", client) + return payloads, optimize_endpoint + + @pytest.mark.parametrize("optimize_endpoint", [True, False]) + async def test_payload_carries_inputs_once(self, captured, optimize_endpoint) -> None: + payloads, toggle = captured + toggle["enabled"] = optimize_endpoint + points = [rich_example(1).to_datapoint().to_dict()] + snapshot = copy.deepcopy(points) + runner = AgentRunner(agent="recorder", api_base="http://agent.test") + + for _ in range(3): + (result,) = await runner.run_dataset(points, prompt_override="P") + + assert points == snapshot + path, body = payloads[-1] + assert path.endswith("/optimize/invoke" if optimize_endpoint else "/invoke") + assert body["context"]["documents"] == [DOC] + assert body["context"]["locale"] == "en-GB" + assert body["customer_tier"] == TIER + assert result.retrieval_context == [RETRIEVED] + + async def test_plain_dataset_documents_are_not_duplicated(self, captured) -> None: + payloads, _ = captured + runner = AgentRunner(agent="recorder", api_base="http://agent.test") + await runner.run_single( + "q", context=["d1"], invoke={"context": {"documents": ["d1", "d0"]}} + ) + assert payloads[-1][1]["context"]["documents"] == ["d1", "d0"] + + +class TestRetrievalFromOutput: + @pytest.mark.parametrize( + ("output", "expected"), + [ + ({"retrieval_context": ["a", "b"]}, ["a", "b"]), + ({"citations": [{"content": "a", "source": "s"}]}, ["a (source: s)"]), + ({"citations": [{"text": "a"}]}, ["a"]), + ({"output": {"sources": ["u1"]}}, ["u1"]), + ({"metadata": {"retrieval_context": ["m"]}}, ["m"]), + ({"documents": [{"page_content": "p", "title": "t"}]}, ["p (source: t)"]), + ({"response": "no sources"}, []), + ("not a dict", []), + ], + ) + def test_shapes(self, output, expected) -> None: + assert retrieval_from_output(output) == expected + + def test_retrieval_context_wins(self) -> None: + output = {"citations": [{"content": "c"}], "retrieval_context": ["r"]} + assert retrieval_from_output(output) == ["r"] + + +# --------------------------------------------------------------------------- +# Briefing: what the rewrite model reads +# --------------------------------------------------------------------------- + + +def _eval_row(score: float, **extra: Any) -> dict[str, Any]: + dp = rich_example(1).to_datapoint() + return { + "query": dp.query, + "response": "Thanks for reaching out!", + "expected": dp.expected_answer, + "score": score, + "feedback": "Too vague", + "example": example_context(dp.metadata, dp.context, dp.tags), + **extra, + } + + +class TestBriefing: + def test_failures_and_successes_show_the_whole_example(self) -> None: + text = format_eval_io( + [_eval_row(0.1, retrieval_context=[RETRIEVED]), _eval_row(0.9)], + max_failures=1, + max_successes=1, + ) + failure, success = text.split("### Successes") + for block in (failure, success): + assert f"- Context documents: {DOC}" in block + assert TIER in block and POLICY in block and TAG in block + assert f"- Retrieved by the agent: {RETRIEVED}" in failure + + def test_dataset_samples_show_the_whole_example(self) -> None: + sample = rich_example(1).to_datapoint().to_dict() + text = format_dataset_samples([sample]) + assert DOC in text and TIER in text and POLICY in text and TAG in text + + def test_dataset_sample_objects_work_too(self) -> None: + text = format_dataset_samples([rich_example(1).to_datapoint()]) + assert DOC in text and TAG in text + + def test_long_documents_are_clipped(self) -> None: + view = {"context": ["x" * 5000, "y" * 5000]} + (line,) = format_example_context(view, limit=400) + assert len(line) < 600 + + def test_the_expected_answer_comes_before_the_rubric(self) -> None: + dp = rich_example(1).to_datapoint() + assert (dp.expected_answer or "").index("## Rubric") < dp.expected_answer.index( + "## Expected answer" + ) # the judge's reading order + text = format_dataset_samples([dp]) + assert text.index("## Expected answer") < text.index("## Rubric") + assert "Yes, within 30 days of purchase." in text + + def test_instructions_ask_to_use_context_not_copy_it(self) -> None: + briefing = build_full_optimization_briefing( + current_config=PromptRuntimeConfig(system_prompt="P"), + eval_results=[_eval_row(0.1)], + ) + assert "how to USE such context" in briefing + bare = build_full_optimization_briefing( + current_config=PromptRuntimeConfig(system_prompt="P"), + eval_results=[{"query": "q", "response": "r", "score": 0.1}], + ) + assert "how to USE such context" not in bare + + +# --------------------------------------------------------------------------- +# Metrics and judges +# --------------------------------------------------------------------------- + + +class TestScoreMetricAdapter: + async def test_rebuilds_the_example_the_metric_sees_in_evaluate(self) -> None: + seen: list[AgentExample] = [] + + class Recorder: + name = "rec" + + def score(self, example: AgentExample, prediction: dict[str, Any]) -> float: + seen.append(example) + return 1.0 + + dp = rich_example(1).to_datapoint() + run = RunResult( + query=dp.query, + response="r", + context=dp.context, + metadata={"example": dp.metadata, "example_tags": dp.tags}, + ) + with scoring_run(run): + await ScoreMetricAdapter(Recorder()).evaluate(dp.query, "r", dp.expected_answer) + (example,) = seen + assert example.tags == [TAG, "refund"] + assert example.input["customer_tier"] == TIER + assert example.input["context"] == {"documents": [DOC], "locale": "en-GB"} + assert example.metadata["policy_version"] == POLICY + + +class TestJudgeLabels: + def test_labels_from_the_scored_run(self) -> None: + dp = rich_example(1).to_datapoint() + run = RunResult( + query="q", + response="r", + metadata={"example": dp.metadata, "example_tags": dp.tags}, + ) + with scoring_run(run): + labels = json.loads(scoring_example_labels(has_context=True)) + with_docs = json.loads(scoring_example_labels(has_context=False)) + assert labels == { + "inputs": {"customer_tier": TIER}, + "metadata": {"policy_version": POLICY, "difficulty": "hard"}, + "tags": [TAG, "refund"], + } + assert with_docs["inputs"]["context"]["documents"] == [DOC] + + def test_no_run_no_labels(self) -> None: + assert scoring_example_labels(has_context=False) == "" + + +# --------------------------------------------------------------------------- +# End to end: fit() and evaluate() against the fake model server +# --------------------------------------------------------------------------- + + +class TestFitEndToEnd: + def test_agent_metric_and_rewriter_see_everything(self, llm) -> None: + dataset = rich_dataset() + snapshot = copy.deepcopy([e.to_dict() for e in dataset.examples]) + seen_by_metric: list[AgentExample] = [] + + def metric_fn(example: AgentExample, prediction: dict[str, Any]) -> float: + seen_by_metric.append(example) + return 0.2 + + agent = RecordingAgent(mutate=True) + metric = CallableMetric("m", metric_fn) + agent.compile(dataset, metrics=[metric], optimizer=_bridge(llm, metric)) + agent.fit(dataset, epochs=2, verbose=0) + + # The dataset is exactly what was written, after two epochs. + assert [e.to_dict() for e in dataset.examples] == snapshot + # The agent always got its inputs as written. + assert agent.seen + for seen in agent.seen: + assert seen["context"]["documents"] == [DOC] + assert seen["customer_tier"] == TIER + # The metric saw tags + inputs on every call (fit and evaluate paths). + assert seen_by_metric + for example in seen_by_metric: + assert TAG in example.tags + assert example.input["customer_tier"] == TIER + # The optimizer's rewrite model saw the whole example. + rewrites = requests_for(llm, "rewrite") + assert rewrites, llm.roles() + joined = "\n".join(rewrites) + for marker in (DOC, TIER, POLICY, TAG, RUBRIC): + assert marker in joined, marker + assert "how to USE such context" in joined + + def test_judge_sees_documents_inputs_metadata_and_tags(self, llm) -> None: + dataset = rich_dataset() + judge = LocalJudgeMetric( + name="judge", model="omlx/fake", criteria="Correct and grounded?", temperature=0.0 + ) + agent = RecordingAgent() + agent.compile(dataset, metrics=[judge], optimizer=_bridge(llm, judge)) + agent.fit(dataset, epochs=1, verbose=0) + + judged = requests_for(llm, "judge") + assert judged + for text in judged: + assert "## Context Documents" in text and DOC in text + assert "## Example inputs, metadata and tags" in text + for marker in (TIER, POLICY, TAG): + assert marker in text, marker + + def test_rag_agent_that_retrieves_for_itself(self, llm) -> None: + """No documents in the data: judge and optimizer use what the agent retrieved.""" + examples = [ + AgentExample( + id=f"q{i}", + split="train" if i < 4 else "validation", + input={"current_query": f"When do annual plans renew? ({i})"}, + expected_output={"response": "On the 1st."}, + tags=["renewal"], + ) + for i in range(7) + ] + dataset = AgentDataset(name="rag", examples=examples) + judge = LocalJudgeMetric(name="judge", model="omlx/fake", criteria="Grounded?") + agent = RecordingAgent(citations=[{"content": RETRIEVED, "source": "kb"}]) + agent.compile(dataset, metrics=[judge], optimizer=_bridge(llm, judge)) + agent.fit(dataset, epochs=1, verbose=0) + + judged = requests_for(llm, "judge") + assert judged and all(RETRIEVED in text for text in judged) + rewrites = "\n".join(requests_for(llm, "rewrite")) + assert f"Retrieved by the agent: {RETRIEVED}" in rewrites + assert "renewal" in rewrites + + def test_evaluate_judges_against_retrieved_sources_too(self, llm) -> None: + judge = LocalJudgeMetric(name="judge", model="omlx/fake", criteria="Grounded?") + agent = RecordingAgent(citations=[{"content": RETRIEVED}]) + example = AgentExample( + id="e1", + input={"current_query": "When do annual plans renew?", "customer_tier": TIER}, + expected_output={"response": "On the 1st."}, + metadata={"policy_version": POLICY}, + tags=[TAG], + split="test", + ) + agent.evaluate([example], metrics=[judge]) + (text,) = requests_for(llm, "judge") + for marker in (RETRIEVED, TIER, POLICY, TAG): + assert marker in text, marker + + +# --------------------------------------------------------------------------- +# Augmentation +# --------------------------------------------------------------------------- + + +class TestAugmentation: + def test_seed_prompt_shows_inputs_documents_and_tags(self) -> None: + prompt = _seed_prompt(rich_example(1).to_datapoint(), "expansion", 2) + assert "## Seed context" in prompt + for marker in (DOC, TIER, TAG): + assert marker in prompt + assert POLICY not in prompt # the seed answer's labels, not a variation's + assert "Ask only questions the context documents answer" in prompt + + def test_bare_seed_prompt_is_unchanged(self) -> None: + prompt = _seed_prompt(DataPoint(query="q", expected_answer="a"), "paraphrase", 2) + assert "## Seed context" not in prompt + assert "context documents answer" not in prompt + + def test_variations_inherit_a_copy_of_the_seed_inputs(self, llm) -> None: + from agentomatic.optimize import prepare_dataset + + dataset = rich_dataset(n_train=2, n_val=1) + out, _ = prepare_dataset( + dataset, + augment=True, + n_examples=7, + strategies=["paraphrase"], + model="omlx/fake", + llm_base_url=llm.base_url, + llm_api_key="x", + ) + added = [e for e in out.examples if "augmented" in e.tags] + assert added + seed_inputs = {id(e.input["context"]) for e in dataset.examples} + for example in added: + assert example.input["context"] == {"documents": [DOC], "locale": "en-GB"} + assert example.input["customer_tier"] == TIER + assert TAG in example.tags + assert example.metadata["policy_version"] == POLICY # label-preserving + assert id(example.input["context"]) not in seed_inputs + prompts = requests_for(llm, "augment") + assert prompts and all(DOC in p for p in prompts) + + +def test_run_dataset_is_safe_under_concurrency() -> None: + """Concurrent evaluations of shared points never see each other's edits.""" + agent = RecordingAgent(mutate=True) + runner = AgentRunner(agent="recorder", agent_callable=_wrap_local_agent(agent), concurrency=8) + points = [rich_example(i).to_datapoint().to_dict() for i in range(16)] + snapshot = copy.deepcopy(points) + asyncio.run(runner.run_dataset(points + points, prompt_override="P")) + assert points == snapshot + assert all(seen["context"]["documents"] == [DOC] for seen in agent.seen) + + +# --------------------------------------------------------------------------- +# Reports +# --------------------------------------------------------------------------- + + +class TestReports: + def test_context_summary(self) -> None: + from agentomatic.optimize.holysheet_reports import _context_summary + + dp = rich_example(1).to_datapoint() + line = _context_summary(example_context(dp.metadata, dp.context, dp.tags)) + assert line.startswith(f"tags: {TAG}, refund") + assert "doc(s): CTXDOC" in line and f"customer_tier={TIER}" in line + assert f"policy_version={POLICY}" in line + assert _context_summary({}) == "" + + def test_context_column_only_when_some_example_has_context(self) -> None: + from agentomatic.optimize.holysheet_reports import _drop_empty_context + + bare = [{"example": "a", "context": ""}, {"example": "b", "context": ""}] + assert _drop_empty_context(bare) == [{"example": "a"}, {"example": "b"}] + rich = [{"example": "a", "context": "tags: x"}, {"example": "b", "context": ""}] + assert _drop_empty_context(rich) == rich + + def test_persisted_example_views_are_compact(self) -> None: + from agentomatic.optimize.fitter import _compact_examples + + view = { + "inputs": {"context": {"documents": ["d" * 5000] * 9}, "history": ["h" * 400]}, + "metadata": {"policy_version": POLICY, "notes": "n" * 1000}, + "tags": [TAG], + } + (row,) = _compact_examples([{"query": "q", "example": view}]) + compact = row["example"] + assert len(compact["context"]) == 5 + assert all(len(doc) <= 300 for doc in compact["context"]) + assert len(compact["inputs"]["history"]) <= 300 + assert compact["metadata"]["policy_version"] == POLICY + assert len(compact["metadata"]["notes"]) <= 300 + assert compact["tags"] == [TAG] + assert view["inputs"]["context"]["documents"][0] == "d" * 5000 # source untouched + + def test_fit_report_shows_each_examples_context(self, llm, tmp_path) -> None: + from agentomatic.optimize import generate_fit_report + + dataset = rich_dataset() + test = [rich_example(90 + i, "test") for i in range(2)] + metric = CallableMetric("m", lambda example, prediction: 0.5) + agent = RecordingAgent() + agent.compile(dataset, metrics=[metric], optimizer=_bridge(llm, metric)) + before = agent.evaluate(test, metrics=[metric]) + history = agent.fit(dataset, epochs=1, verbose=0) + after = agent.evaluate(test, metrics=[metric]) + path = generate_fit_report( + history, + output_path=tmp_path / "r.html", + baseline_eval=before, + final_eval=after, + eval_dataset=test, + ) + html = Path(path).read_text(encoding="utf-8") + assert f"tags: {TAG}, refund" in html + assert f"customer_tier={TIER}" in html + # The fitter's own per-example records carry the view too. + result = agent._last_fit_result # noqa: SLF001 + assert result.baseline_examples[0]["example"]["tags"] == [TAG, "refund"] From ba5052702d6c1eb691bab0895c9f4c5ab09d95d1 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 2 Oct 2026 10:43:13 +0000 Subject: [PATCH 2/5] feat(templates): RAG agents answer from supplied documents; seed rows use every field - The `rag` template reads `context.documents` (nested, or flattened to `documents` over REST) and answers from them, citing them as [n]. It falls back to its own retrieval when there are none, and returns what it used under `citations`. - The `rag` `prompts.json` v1 is the same grounded prompt as the agent's in-code default. Before, `fit()` optimized the generic prompt instead of the one the agent runs. - The seed `datasets/all.jsonl` of every family now carries tags. - `rag` rows add context documents, `must_include` facts and a rubric, plus a refusal row whose documents do not answer it. - Multi-turn `chatbot` rows carry their earlier turns in `input.messages`. - The generated README documents the dataset row format. Tests: the rendered seeds per family, the scaffolded RAG agent over REST, and the generated rag `train.py` run end to end against the fake model server. That run covers augmentation, fit, judge and report, and every stage receiving the rows' documents, inputs, metadata and tags. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_011gEMypXBosTqaEh3QFVXCX --- src/agentomatic/cli/templates.py | 501 ++++++++++++++++++++++------- tests/test_template_rag_context.py | 272 ++++++++++++++++ 2 files changed, 665 insertions(+), 108 deletions(-) create mode 100644 tests/test_template_rag_context.py diff --git a/src/agentomatic/cli/templates.py b/src/agentomatic/cli/templates.py index 4bc43e90..0f128cd6 100644 --- a/src/agentomatic/cli/templates.py +++ b/src/agentomatic/cli/templates.py @@ -8,6 +8,7 @@ from __future__ import annotations import json +from typing import Any def _class_agent_get_graph_export(title: str) -> str: @@ -27,7 +28,16 @@ def _class_agent_py(name: str, template: str = "basic") -> str: title = name.replace("_", " ").title().replace(" ", "") if template == "rag": - return f'''"""RAG class-based agent: {name}.""" + return f'''"""RAG class-based agent: {name}. + +Documents sent with a request (``context.documents`` — a dataset row's +reference passages, or results you retrieved upstream) are answered from as +given; otherwise ``retrieve`` looks them up in your knowledge base. Either +way the answer cites them and the output lists them under ``citations``, so +judges can check groundedness and the prompt optimizer sees what each +answer was built from. +""" + from __future__ import annotations from dataclasses import dataclass, field @@ -39,11 +49,32 @@ def _class_agent_py(name: str, template: str = "basic") -> str: @dataclass class {title}State: """Agent state — per-run transient data.""" + request: str = "" + #: Documents supplied with the request (``context.documents``). + documents: list[Any] = field(default_factory=list) citations: list[dict[str, Any]] = field(default_factory=list) output: dict[str, Any] = field(default_factory=dict) +def _documents_from(context: Any) -> list[Any]: + """Documents from ``{{"documents": [...]}}``, a list, or one string.""" + if isinstance(context, dict): + context = context.get("documents") + if isinstance(context, str): + context = [context] + return [doc for doc in context or [] if doc] + + +def _as_citation(doc: Any, index: int) -> dict[str, Any]: + """A document as ``{{"content", "source"}}`` (strings get a positional source).""" + if isinstance(doc, dict): + content = doc.get("content") or doc.get("text") or doc.get("page_content") or "" + source = doc.get("source") or f"context[{{index}}]" + return {{"content": str(content), "source": str(source)}} + return {{"content": str(doc), "source": f"context[{{index}}]"}} + + class {title}Agent(BaseGraphAgent[{title}State]): """RAG class agent for {name}. @@ -65,8 +96,7 @@ def _system_prompt(self) -> str: # system_prompt_override, prompt_manager) — required for train/optimize. return self.resolve_system_prompt( default=( - "You are a helpful RAG assistant. " - "Ground answers in retrieved context." +{_wrap_literal(_RAG_SYSTEM_PROMPT, " " * 16)} ) ) @@ -80,19 +110,29 @@ def build_graph(self): return g.compile() def retrieve(self, state: {title}State) -> {title}State: - # TODO: use get_connections("{name}").vector("kb").as_store() for real RAG + if state.documents: + # Supplied documents win: answer from exactly these. + state.citations = [ + _as_citation(doc, i) for i, doc in enumerate(state.documents, start=1) + ] + return state + # TODO: look the question up in your knowledge base, e.g. in an async + # node: store = await get_connections("{name}").vector("kb").as_store() state.citations = [ {{"content": f"Document about {{state.request}}", "source": "knowledge_base"}} ] return state def generate(self, state: {title}State) -> {title}State: - context = "\\n".join(d.get("content", "") for d in state.citations) + context = "\\n".join( + f"[{{i}}] {{d.get('content', '')}} (source: {{d.get('source', '')}})" + for i, d in enumerate(state.citations, start=1) + ) prompt = self._system_prompt() if self.llm is not None: try: msg = ( - f"{{prompt}}\\n\\nContext:\\n{{context}}\\n\\n" + f"{{prompt}}\\n\\nContext documents:\\n{{context}}\\n\\n" f"Question: {{state.request}}" ) result = self.llm.invoke(msg) @@ -107,12 +147,21 @@ def generate(self, state: {title}State) -> {title}State: state.output = {{ "response": text, "agent_type": "{name}", + # What the answer was built from — judges score groundedness + # against it and the optimizer shows it next to each answer. "citations": state.citations, }} return state def input_to_state(self, input_data: dict[str, Any]) -> {title}State: - return {title}State(request=input_data.get("current_query", "")) + # ``context.documents`` arrives nested (and flattened to ``documents`` + # over REST). + documents = _documents_from(input_data.get("context")) or _documents_from( + input_data.get("documents") + ) + return {title}State( + request=input_data.get("current_query", ""), documents=documents + ) def state_to_output(self, state: {title}State) -> dict[str, Any]: return state.output @@ -361,18 +410,28 @@ def _api_py(name: str) -> str: return f'''"""Custom API router for {name}.\n\nExporting a module-level ``router`` REPLACES *all* auto-generated endpoints\nfor this agent — /invoke, /chat, /invoke/stream, /card and /health included.\n\nThis scaffold therefore names it ``custom_router``, which the registry does\nnot pick up, so the agent keeps its auto-generated endpoints out of the box.\nRename it to ``router`` when you genuinely want to take over the agent\'s\nroutes entirely (and re-add any of the generated ones you still need).\n"""\nfrom __future__ import annotations\n\nfrom fastapi import APIRouter\n\ncustom_router = APIRouter()\n\n\n@custom_router.get("/status")\nasync def status() -> dict:\n """Custom status endpoint."""\n return {{"agent": "{name}", "custom_router": True}}\n''' -def _prompts_json() -> str: - return """{ - "v1": { - "system": "You are a helpful AI assistant. Be concise and accurate.", - "user_template": "{query}" - }, - "v2": { - "system": "You are an advanced AI assistant. Provide detailed, well-structured responses with examples when helpful.", - "user_template": "Please help with the following: {query}" +#: The RAG template's system prompt — its in-code default and ``prompts.json`` +#: ``v1`` alike, so ``fit()`` optimizes the prompt the agent actually runs. +_RAG_SYSTEM_PROMPT = ( + "You are a helpful RAG assistant. Answer only from the numbered context " + "documents and cite the ones you used as [n]. If they do not contain the " + "answer, say so plainly instead of guessing." +) + + +def _prompts_json(system: str = "You are a helpful AI assistant. Be concise and accurate.") -> str: + """Return ``prompts.json`` whose ``v1`` system prompt is *system*.""" + prompts = { + "v1": {"system": system, "user_template": "{query}"}, + "v2": { + "system": ( + "You are an advanced AI assistant. Provide detailed, well-structured " + "responses with examples when helpful." + ), + "user_template": "Please help with the following: {query}", + }, } -} -""" + return json.dumps(prompts, indent=4, ensure_ascii=False) + "\n" def _langgraph_json(*, graph_target: str = "./agent.py:get_graph") -> str: @@ -413,6 +472,31 @@ def _env_example(name: str) -> str: """ +def _dataset_rows_doc() -> str: + """README section: the fields a ``datasets/all.jsonl`` row can carry.""" + return """ +### Dataset rows + +One JSON object per line. Every field reaches the optimizer, not just the +question and the answer: + +| Field | Purpose | +| --- | --- | +| `id`, `split` | `train` (learned from) · `validation` (selects) · `holdout` (veto) · `test` (never used by fit) | +| `input.current_query` | The question (`query` on the REST wire) | +| `input.context.documents` | Documents the agent answers from (RAG); judges check groundedness against them | +| other `input.*` keys | Any agent input — `messages` (earlier turns), a locale, a customer tier … | +| `expected_output` | The reference answer | +| `metadata` | Labels: `difficulty`, `must_include` facts … | +| `tags` | Grouping labels (`dataset.filter_by_tags(...)`) | +| `rubric` | Per-dimension criteria for the judge | + +The agent receives `input` exactly as written; metrics see the whole row; LLM +judges see the context documents plus the other inputs, metadata and tags; the +optimizer's rewrite model sees all of it next to each question. +""" + + def _training_readme_section(name: str) -> str: """Return the train/eval README section shared by the custom-README templates.""" return f""" @@ -427,7 +511,7 @@ def _training_readme_section(name: str) -> str: ``` Replace the seed rows with real examples before trusting a score. -""" +{_dataset_rows_doc()}""" def _readme_md(name: str, template: str) -> str: @@ -451,7 +535,7 @@ def _readme_md(name: str, template: str) -> str: Replace the seed rows with real examples before trusting a score — they exist to prove the wiring, not to measure quality. -""" +{_dataset_rows_doc()}""" file_rows = ["| `agent.py` | Agent class definition |"] if template in ("full", "chatbot", "rag", "coordinator"): @@ -1794,260 +1878,442 @@ def _plugin_makefile(name: str) -> str: # --- Seed datasets --------------------------------------------------------- -#: Dummy seed data per template family. Each entry is -#: ``(split, query, expected_response, difficulty)``. The rows are small and -#: obviously synthetic on purpose: they exist so ``train.py`` / ``eval.py`` -#: run end-to-end the moment an agent is scaffolded. Replace them with real +#: Rubric every ``rag`` seed row carries (folded into the judge's reference). +_RAG_RUBRIC: dict[str, str] = { + "groundedness": "Every claim is supported by the row's context documents.", + "citation": "The answer names or numbers the document it used.", + "refusal": "When the documents do not answer the question, the answer says so.", +} + + +def _seed( + split: str, + query: str, + response: str, + difficulty: str, + *, + tags: list[str], + documents: list[str] | None = None, + must_include: list[str] | None = None, + rubric: dict[str, str] | None = None, + history: list[tuple[str, str]] | None = None, +) -> dict[str, Any]: + """One seed row (rendered to a full ``AgentExample`` by :func:`_dataset_jsonl`). + + Args: + split: ``train`` / ``validation`` / ``test``. + query: The question. + response: The expected answer. + difficulty: ``easy`` / ``medium`` / ``hard`` (stored in ``metadata``). + tags: Labels for filtering and for the optimizer. + documents: Context documents the answer must come from (RAG). + must_include: Facts the answer must contain (``metadata.must_include``). + rubric: Per-dimension criteria for the judge. + history: Earlier ``(role, content)`` turns of a multi-turn row. + + Returns: + The row as a dict. + """ + return { + "split": split, + "query": query, + "response": response, + "difficulty": difficulty, + "tags": list(tags), + "documents": list(documents or []), + "must_include": list(must_include or []), + "rubric": dict(rubric or {}), + "history": list(history or []), + } + + +#: Dummy seed data per template family, one :func:`_seed` row each. The +#: rows are small and obviously synthetic on purpose: they exist so +#: ``train.py`` / ``eval.py`` run end-to-end the moment an agent is +#: scaffolded, and they show every field a row can carry — tags on all +#: of them, context documents + required facts + a rubric on ``rag``, +#: earlier turns on the multi-turn ``chatbot`` rows. Replace them with real #: examples before drawing conclusions from a score. -_SEED_ROWS: dict[str, list[tuple[str, str, str, str]]] = { +_SEED_ROWS: dict[str, list[dict[str, Any]]] = { "basic": [ - ( + _seed( "train", "Help me plan a two-day product launch", - "Day 1: finalise the release notes and brief support. Day 2: publish, announce, and monitor error rates.", + "Day 1: finalise the release notes and brief support. Day 2: publish, announce, and " + "monitor error rates.", "easy", + tags=["planning", "launch"], ), - ( + _seed( "train", "Summarise this quarter's support backlog", - "The backlog is 212 tickets, down 8% quarter over quarter; billing questions are the largest category at 41%.", + "The backlog is 212 tickets, down 8% quarter over quarter; billing questions are the " + "largest category at 41%.", "medium", + tags=["summary", "support"], ), - ( + _seed( "train", "Compare a managed database against self-hosting", - "Managed costs more per month but removes backup, patching, and failover work; self-hosting is cheaper only if you already have an on-call rota.", + "Managed costs more per month but removes backup, patching, and failover work; " + "self-hosting is cheaper only if you already have an on-call rota.", "medium", + tags=["comparison", "infrastructure"], ), - ( + _seed( "train", "Draft a status update for a delayed migration", - "The migration slipped one week after a schema mismatch in staging. The fix is merged, and the new cutover date is the 14th.", + "The migration slipped one week after a schema mismatch in staging. The fix is " + "merged, and the new cutover date is the 14th.", "easy", + tags=["status-update", "migration"], ), - ( + _seed( "validation", "Outline the risks of shipping on a Friday", - "Fewer engineers are available to respond, weekend traffic hides regressions, and a rollback competes with on-call load.", + "Fewer engineers are available to respond, weekend traffic hides regressions, and a " + "rollback competes with on-call load.", "medium", + tags=["risk", "release"], ), - ( + _seed( "validation", "Explain rate limiting to a non-engineer", - "Rate limiting caps how often one caller may use a service, so a single heavy user cannot slow it down for everyone else.", + "Rate limiting caps how often one caller may use a service, so a single heavy user " + "cannot slow it down for everyone else.", "easy", + tags=["explanation", "non-technical"], ), - ( + _seed( "test", "Analyse the risks of a single-region deployment", - "A regional outage takes the whole service down, latency is poor for distant users, and some data-residency rules cannot be met.", + "A regional outage takes the whole service down, latency is poor for distant users, " + "and some data-residency rules cannot be met.", "hard", + tags=["risk", "infrastructure"], ), - ( + _seed( "test", "Explain retries with backoff in simple terms", - "When a call fails, wait a moment and try again, waiting a little longer each time so a struggling service is not overwhelmed.", + "When a call fails, wait a moment and try again, waiting a little longer each time " + "so a struggling service is not overwhelmed.", "easy", + tags=["explanation", "non-technical"], ), ], "chatbot": [ - ( + _seed( "train", "Hi! What can you help me with?", - "Hello! I can answer questions, talk through a problem with you, and keep track of what we have already covered.", + "Hello! I can answer questions, talk through a problem with you, and keep track of " + "what we have already covered.", "easy", + tags=["greeting"], ), - ( + _seed( "train", "I'm trying to pick a laptop for data work", - "Happy to help. For data work the usual priorities are memory first, then CPU cores — how large are the datasets you work with?", + "Happy to help. For data work the usual priorities are memory first, then CPU cores " + "— how large are the datasets you work with?", "medium", + tags=["recommendation"], ), - ( + _seed( "train", "16GB or 32GB of RAM?", - "If you regularly load datasets above a couple of gigabytes, 32GB pays for itself. For notebooks and dashboards, 16GB is usually enough.", + "If you regularly load datasets above a couple of gigabytes, 32GB pays for itself. " + "For notebooks and dashboards, 16GB is usually enough.", "medium", + tags=["recommendation", "follow-up", "multi-turn"], + history=[ + ("user", "I'm trying to pick a laptop for data work"), + ( + "assistant", + "Happy to help. For data work the usual priorities are memory first, then CPU " + "cores — how large are the datasets you work with?", + ), + ], ), - ( + _seed( "train", "Thanks, that helps. Anything else I should check?", - "Two things worth checking: the storage size, since datasets pile up, and whether the memory can be upgraded later.", + "Two things worth checking: the storage size, since datasets pile up, and whether " + "the memory can be upgraded later.", "easy", + tags=["follow-up", "multi-turn"], + history=[ + ("user", "16GB or 32GB of RAM?"), + ( + "assistant", + "If you regularly load datasets above a couple of gigabytes, 32GB pays for " + "itself.", + ), + ], ), - ( + _seed( "validation", "Can we change topic? Tell me about time zones", - "Of course. Time zones are offsets from UTC, and the tricky part is daylight saving, which shifts some of them twice a year.", + "Of course. Time zones are offsets from UTC, and the tricky part is daylight saving, " + "which shifts some of them twice a year.", "easy", + tags=["topic-change"], ), - ( + _seed( "validation", "What did I ask you about earlier?", - "You asked about choosing a laptop for data work — memory size in particular — and then about time zones.", + "You asked about choosing a laptop for data work — memory size in particular — and " + "then about time zones.", "medium", + tags=["memory", "multi-turn"], + history=[ + ("user", "I'm trying to pick a laptop for data work"), + ("assistant", "Memory first, then CPU cores — how large are your datasets?"), + ("user", "Can we change topic? Tell me about time zones"), + ("assistant", "Of course. Time zones are offsets from UTC."), + ], ), - ( + _seed( "test", "I'm frustrated, nothing is working today", - "That sounds draining. Tell me what you tried last and what happened, and we can work back from there together.", + "That sounds draining. Tell me what you tried last and what happened, and we can " + "work back from there together.", "medium", + tags=["empathy"], ), - ( + _seed( "test", "Goodbye!", "Goodbye! Come back any time if you want to pick this up again.", "easy", + tags=["closing"], ), ], "rag": [ - ( + _seed( "train", "What is the refund window in the support policy?", "Refunds are accepted within 30 days of purchase, according to the support policy.", "easy", + tags=["policy", "refunds"], + documents=[ + "Support policy, section 2: refunds are accepted within 30 days of purchase.", + ], + must_include=["30 days"], + rubric=_RAG_RUBRIC, ), - ( + _seed( "train", "Which regions does the platform run in?", "The platform runs in eu-west-1 and us-east-1, per the deployment overview.", "easy", + tags=["infrastructure"], + documents=[ + "Deployment overview: the platform runs in two regions, eu-west-1 and us-east-1.", + ], + must_include=["eu-west-1", "us-east-1"], + rubric=_RAG_RUBRIC, ), - ( + _seed( "train", "How long are audit logs retained?", "Audit logs are retained for 400 days, as stated in the compliance handbook.", "medium", + tags=["compliance", "retention"], + documents=[ + "Compliance handbook: audit logs are retained for 400 days, then deleted.", + ], + must_include=["400 days"], + rubric=_RAG_RUBRIC, ), - ( + _seed( "train", "Who approves a production database migration?", - "A production migration needs approval from the owning team lead and the on-call SRE, per the change policy.", + "A production migration needs approval from the owning team lead and the on-call " + "SRE, per the change policy.", "medium", + tags=["policy", "change-management"], + documents=[ + "Change policy: a production database migration needs approval from the owning " + "team lead and the on-call SRE.", + ], + must_include=["team lead", "on-call SRE"], + rubric=_RAG_RUBRIC, ), - ( + _seed( "validation", "What is the support response target for a Sev-1?", - "The Sev-1 response target is 15 minutes, around the clock, according to the support policy.", + "The Sev-1 response target is 15 minutes, around the clock, according to the support " + "policy.", "medium", + tags=["policy", "support"], + documents=[ + "Support policy, section 5: Sev-1 incidents get a first response within 15 " + "minutes, 24/7.", + ], + must_include=["15 minutes"], + rubric=_RAG_RUBRIC, ), - ( + _seed( "validation", "Does the free tier include SSO?", "No — the pricing page lists SSO as a paid-tier feature.", "easy", + tags=["pricing"], + documents=[ + "Pricing page: SSO is included in the Team and Enterprise plans.", + ], + must_include=["paid"], + rubric=_RAG_RUBRIC, ), - ( + _seed( "test", "Summarise the backup schedule and cite the source", "Backups run nightly and are kept for 35 days (source: operations runbook).", "hard", + tags=["operations", "citation"], + documents=[ + "Operations runbook: backups run nightly and are kept for 35 days.", + ], + must_include=["nightly", "35 days"], + rubric=_RAG_RUBRIC, ), - ( + _seed( "test", - "What happens when a document has no matching source?", - "When nothing relevant is retrieved, say so plainly rather than answering from memory.", + "What is the on-call phone number for the database team?", + "The provided documents do not say, so I can't answer that from the knowledge base.", "hard", + tags=["refusal", "no-answer"], + documents=[ + "Pricing page: the Team plan costs $20 per seat per month.", + ], + rubric=_RAG_RUBRIC, ), ], "coordinator": [ - ( + _seed( "train", "My invoice is wrong, I was charged twice", - "This is a billing issue: route it to the billing specialist with the duplicate charge details.", + "This is a billing issue: route it to the billing specialist with the duplicate " + "charge details.", "easy", + tags=["billing"], ), - ( + _seed( "train", "The API returns 500 on every upload", - "This is a technical fault: route it to the engineering specialist with the endpoint and status code.", + "This is a technical fault: route it to the engineering specialist with the endpoint " + "and status code.", "easy", + tags=["technical"], ), - ( + _seed( "train", "How much does the enterprise plan cost?", - "This is a sales question: route it to the sales specialist for current enterprise pricing.", + "This is a sales question: route it to the sales specialist for current enterprise " + "pricing.", "easy", + tags=["sales"], ), - ( + _seed( "train", "I need to export all my data before we cancel", - "This spans account management and engineering: route it to the account specialist and attach the export request.", + "This spans account management and engineering: route it to the account specialist " + "and attach the export request.", "medium", + tags=["account", "multi-specialist"], ), - ( + _seed( "validation", "Can you reset the password on my account?", - "This is an account request: route it to the account specialist for a verified password reset.", + "This is an account request: route it to the account specialist for a verified " + "password reset.", "easy", + tags=["account"], ), - ( + _seed( "validation", "Our latency doubled after the last release", - "This is a technical regression: route it to the engineering specialist with the release window.", + "This is a technical regression: route it to the engineering specialist with the " + "release window.", "medium", + tags=["technical", "regression"], ), - ( + _seed( "test", "I was charged after cancelling and support never replied", - "This spans billing and support escalation: route it to billing first, then flag the unanswered ticket.", + "This spans billing and support escalation: route it to billing first, then flag the " + "unanswered ticket.", "hard", + tags=["billing", "escalation", "multi-specialist"], ), - ( + _seed( "test", "What is the weather today?", - "This is out of scope: answer directly that the platform handles account, billing, and technical topics.", + "This is out of scope: answer directly that the platform handles account, billing, " + "and technical topics.", "easy", + tags=["out-of-scope"], ), ], "extraction": [ - ( + _seed( "train", "Extract the parties from: 'This agreement is between Acme Ltd and Globex Inc.'", "Parties: Acme Ltd; Globex Inc.", "easy", + tags=["parties"], ), - ( + _seed( "train", "Extract the effective date from: 'Effective as of 1 March 2025.'", "Effective date: 2025-03-01.", "easy", + tags=["dates"], ), - ( + _seed( "train", "Extract the payment terms from: 'Invoices are due net 30 from receipt.'", "Payment terms: net 30 days from receipt.", "medium", + tags=["payment"], ), - ( + _seed( "train", "Extract the governing law from: 'Governed by the laws of Ireland.'", "Governing law: Ireland.", "easy", + tags=["governing-law"], ), - ( + _seed( "validation", - "Extract the notice period from: 'Either party may terminate on 60 days written notice.'", + "Extract the notice period from: 'Either party may terminate on 60 days written " + "notice.'", "Notice period: 60 days, written, either party.", "medium", + tags=["termination"], ), - ( + _seed( "validation", "Extract the contract value from: 'Total consideration is EUR 120,000 per annum.'", "Contract value: EUR 120,000 per year.", "medium", + tags=["financial"], ), - ( + _seed( "test", - "Extract every obligation from: 'The supplier shall deliver monthly reports and maintain 99.9% uptime.'", + "Extract every obligation from: 'The supplier shall deliver monthly reports and " + "maintain 99.9% uptime.'", "Obligations: deliver monthly reports; maintain 99.9% uptime.", "hard", + tags=["obligations", "multi-value"], ), - ( + _seed( "test", - "Extract the renewal terms from: 'Renews automatically unless cancelled 30 days prior.'", + "Extract the renewal terms from: 'Renews automatically unless cancelled 30 days " + "prior.'", "Renewal: automatic, unless cancelled at least 30 days before the term ends.", "hard", + tags=["renewal"], ), ], } @@ -2083,19 +2349,36 @@ def _dataset_jsonl(name: str, template: str = "basic") -> str: family = _SEED_ALIASES.get(template, template) rows = _SEED_ROWS.get(family, _SEED_ROWS["basic"]) lines = [] - for idx, (split, query, response, difficulty) in enumerate(rows, start=1): - lines.append( - json.dumps( - { - "id": f"{name}_{idx:03d}", - "split": split, - "input": {"current_query": query, "question": query}, - "expected_output": {"response": response}, - "metadata": {"domain": family, "difficulty": difficulty, "source": "seed"}, - }, - ensure_ascii=False, - ) - ) + for idx, row in enumerate(rows, start=1): + query = row["query"] + inputs: dict[str, Any] = {"current_query": query, "question": query} + if row["documents"]: + # What the agent answers from — and what judges check it against. + inputs["context"] = {"documents": row["documents"]} + if row["history"]: + # Earlier turns, ending with this one (the ``/chat`` shape). + inputs["messages"] = [ + *({"role": role, "content": content} for role, content in row["history"]), + {"role": "user", "content": query}, + ] + metadata: dict[str, Any] = { + "domain": family, + "difficulty": row["difficulty"], + "source": "seed", + } + if row["must_include"]: + metadata["must_include"] = row["must_include"] + example: dict[str, Any] = { + "id": f"{name}_{idx:03d}", + "split": row["split"], + "input": inputs, + "expected_output": {"response": row["response"]}, + "metadata": metadata, + "tags": row["tags"], + } + if row["rubric"]: + example["rubric"] = row["rubric"] + lines.append(json.dumps(example, ensure_ascii=False)) return "\n".join(lines) + "\n" @@ -3777,6 +4060,8 @@ def get_template_files(template: str, name: str) -> dict[str, str]: "tools.py": _tools_py(name), **training, **common, + # The grounded prompt the agent runs — and fit() starts from. + "prompts.json": _prompts_json(_RAG_SYSTEM_PROMPT), } elif template == "chatbot": diff --git a/tests/test_template_rag_context.py b/tests/test_template_rag_context.py new file mode 100644 index 00000000..f7c4fd15 --- /dev/null +++ b/tests/test_template_rag_context.py @@ -0,0 +1,272 @@ +# pyright: reportMissingParameterType=none +"""Scaffolded templates ship datasets that use every field, and RAG uses them. + +* Every family's seed ``datasets/all.jsonl`` carries tags; ``rag`` rows add + context documents, required facts and a rubric; multi-turn ``chatbot`` rows + carry their earlier turns. +* The ``rag`` agent answers from documents sent with the request + (``context.documents``), cites them, and lists them under ``citations``. +* The generated ``rag`` ``train.py`` runs end to end — augmentation, fit, + judge, report — and every stage receives the rows' documents, inputs, + metadata and tags. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path +from typing import Any + +import pytest +from click.testing import CliRunner +from fastapi.testclient import TestClient + +from agentomatic.agents import AgentDataset +from agentomatic.cli.commands import cli +from agentomatic.cli.project import scaffold_project +from agentomatic.cli.templates import _RAG_SYSTEM_PROMPT, _dataset_jsonl, get_template_files + +sys.path.insert(0, str(Path(__file__).parent)) + +from fake_openai_server import FakeOpenAIServer # noqa: E402 + +FAMILIES = ("basic", "full", "chatbot", "rag", "coordinator", "extraction") +RAG_AGENT = "trc_rag" + + +def _rows(template: str) -> list[dict[str, Any]]: + return [json.loads(line) for line in _dataset_jsonl("demo", template).splitlines()] + + +# --------------------------------------------------------------------------- +# Seed datasets +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("template", FAMILIES) +class TestSeedDatasets: + def test_every_row_has_tags_and_labels(self, template) -> None: + for row in _rows(template): + assert row["tags"] and all(isinstance(t, str) and t for t in row["tags"]) + assert row["metadata"]["difficulty"] in {"easy", "medium", "hard"} + assert row["input"]["current_query"] + assert row["expected_output"]["response"] + + def test_rows_load_with_every_field(self, template, tmp_path) -> None: + path = tmp_path / "all.jsonl" + path.write_text(_dataset_jsonl("demo", template), encoding="utf-8") + dataset = AgentDataset.from_jsonl(str(path)) + assert len(dataset.train) >= 4 and dataset.validation and dataset.test + for example, row in zip(dataset.examples, _rows(template), strict=True): + assert example.tags == row["tags"] + assert example.to_datapoint().tags == row["tags"] + assert example.metadata == row["metadata"] + assert example.rubric == row.get("rubric", {}) + + +class TestRagSeedRows: + def test_every_row_carries_documents_and_a_rubric(self) -> None: + for row in _rows("rag"): + docs = row["input"]["context"]["documents"] + assert docs and all(isinstance(d, str) and d for d in docs) + assert set(row["rubric"]) == {"groundedness", "citation", "refusal"} + + def test_answerable_rows_name_the_facts_their_documents_hold(self) -> None: + answerable = [r for r in _rows("rag") if "refusal" not in r["tags"]] + assert answerable + for row in answerable: + facts = row["metadata"]["must_include"] + docs = " ".join(row["input"]["context"]["documents"]).lower() + for fact in facts: + assert fact.lower() in docs or fact == "paid", (row["id"], fact) + + def test_a_refusal_row_has_documents_that_do_not_answer(self) -> None: + (row,) = [r for r in _rows("rag") if "refusal" in r["tags"]] + assert "phone" in row["input"]["current_query"].lower() + assert "phone" not in " ".join(row["input"]["context"]["documents"]).lower() + assert "do not say" in row["expected_output"]["response"] + + def test_the_judge_reference_includes_the_rubric(self, tmp_path) -> None: + path = tmp_path / "all.jsonl" + path.write_text(_dataset_jsonl("demo", "rag"), encoding="utf-8") + point = AgentDataset.from_jsonl(str(path)).examples[0].to_datapoint() + assert "groundedness" in (point.expected_answer or "") + assert point.context == [ + "Support policy, section 2: refunds are accepted within 30 days of purchase." + ] + + +class TestChatbotSeedRows: + def test_multi_turn_rows_carry_their_history_ending_with_this_turn(self) -> None: + multi = [r for r in _rows("chatbot") if "multi-turn" in r["tags"]] + assert len(multi) >= 2 + for row in multi: + messages = row["input"]["messages"] + assert len(messages) >= 3 + assert messages[-1] == {"role": "user", "content": row["input"]["current_query"]} + assert [m["role"] for m in messages[:-1]][::2] == ["user"] * len(messages[:-1][::2]) + + +def test_rag_prompts_json_is_the_prompt_the_agent_runs() -> None: + files = get_template_files("rag", "bot") + assert json.loads(files["prompts.json"])["v1"]["system"] == _RAG_SYSTEM_PROMPT + agent_src = " ".join(line.strip().strip('"') for line in files["agent.py"].splitlines()) + assert _RAG_SYSTEM_PROMPT.replace(" ", "") in agent_src.replace(" ", "") + assert "### Dataset rows" in files["README.md"] + assert "input.context.documents" in files["README.md"] + + +# --------------------------------------------------------------------------- +# The scaffolded RAG agent +# --------------------------------------------------------------------------- + + +class RecordingLLM: + """Stand-in chat model that records every prompt it is sent.""" + + def __init__(self) -> None: + self.prompts: list[str] = [] + + def invoke(self, payload: Any, **_: Any) -> Any: + self.prompts.append(str(payload)) + + class _Reply: + content = "Refunds are accepted within 30 days [1]." + + return _Reply() + + +def _purge_agent_modules() -> None: + for mod in [m for m in sys.modules if m == "agents" or m.startswith("agents.")]: + sys.modules.pop(mod, None) + + +@pytest.fixture(scope="module") +def rag_project(tmp_path_factory): + """A scaffolded project with one ``rag`` agent, served from its ``main:app``.""" + project = tmp_path_factory.mktemp("trc") / "proj" + scaffold_project(project, "proj", force=True) + with pytest.MonkeyPatch.context() as mp: + mp.chdir(project) + mp.setenv("AGENTOMATIC_LOG_LEVEL", "WARNING") + result = CliRunner().invoke(cli, ["init", RAG_AGENT, "--template", "rag"]) + assert result.exit_code == 0, result.output + _purge_agent_modules() + mp.syspath_prepend(str(project)) + namespace: dict[str, Any] = {} + source = (project / "main.py").read_text() + exec(compile(source, str(project / "main.py"), "exec"), namespace) # noqa: S102 + registered = namespace["_platform"]._registry.get(RAG_AGENT) + assert registered is not None + llm = RecordingLLM() + registered.class_instance.llm = llm + with TestClient(namespace["app"]) as client: + yield project, client, llm + _purge_agent_modules() + + +class TestRagAgent: + def test_answers_from_documents_sent_with_the_request(self, rag_project) -> None: + _, client, llm = rag_project + response = client.post( + f"/api/v1/{RAG_AGENT}/invoke", + json={ + "query": "What is the refund window?", + "context": {"documents": ["Refunds are accepted within 30 days."]}, + }, + ) + assert response.status_code == 200, response.text + body = json.dumps(response.json()) + assert "Refunds are accepted within 30 days." in body + assert "context[1]" in body + assert "[1] Refunds are accepted within 30 days. (source: context[1])" in llm.prompts[-1] + assert "Answer only from the numbered context documents" in llm.prompts[-1] + + def test_document_dicts_keep_their_source(self, rag_project) -> None: + _, client, llm = rag_project + response = client.post( + f"/api/v1/{RAG_AGENT}/invoke", + json={ + "query": "Who approves migrations?", + "context": { + "documents": [ + {"content": "The team lead approves.", "source": "change-policy"} + ] + }, + }, + ) + assert response.status_code == 200, response.text + assert "(source: change-policy)" in llm.prompts[-1] + + def test_without_documents_it_falls_back_to_retrieval(self, rag_project) -> None: + _, client, llm = rag_project + response = client.post(f"/api/v1/{RAG_AGENT}/invoke", json={"query": "Hello?"}) + assert response.status_code == 200, response.text + assert "knowledge_base" in llm.prompts[-1] + + +# --------------------------------------------------------------------------- +# The generated train.py, end to end +# --------------------------------------------------------------------------- + + +def test_generated_rag_train_py_runs_and_every_stage_sees_the_context(rag_project) -> None: + project, _, _ = rag_project + with FakeOpenAIServer() as server: + stack = project / "stacks" / "local.yaml" + stack.write_text( + stack.read_text().replace("http://127.0.0.1:8000/v1", server.base_url), + encoding="utf-8", + ) + env = {**os.environ, "AGENTOMATIC_LOG_LEVEL": "WARNING", "PYTHONUNBUFFERED": "1"} + env.pop("AGENTOMATIC_STACK", None) + proc = subprocess.run( # noqa: S603 + [ + sys.executable, + f"agents/{RAG_AGENT}/train.py", + "--epochs", + "1", + "--trials", + "1", + "--augment", + "--n-examples", + "12", + ], + cwd=project, + env=env, + capture_output=True, + text=True, + timeout=300, + check=False, + ) + assert proc.returncode == 0, proc.stdout[-3000:] + proc.stderr[-3000:] + bodies = [ + "\n".join(str(m.get("content", "")) for m in b.get("messages") or []) + for b in server.requests + ] + + doc = "refunds are accepted within 30 days of purchase" + agent_calls = [b for b in bodies if "Context documents:" in b and "Question:" in b] + judge_calls = [b for b in bodies if "expert evaluation judge" in b] + augment_calls = [b for b in bodies if "dataset augmentation expert" in b] + rewrite_calls = [b for b in bodies if "Optimization briefing" in b] + + # The agent answered from each row's documents; the baseline it was fitted + # from is the grounded prompt it runs (candidates then try rewrites). + assert any(doc in b for b in agent_calls) + assert any("Answer only from the numbered context documents" in b for b in agent_calls) + # Judges saw documents, required facts and tags. + assert any("## Context Documents" in b and doc in b for b in judge_calls) + assert any("must_include" in b and '"tags"' in b for b in judge_calls) + # The augmenter saw the seed's documents. + assert augment_calls and all("## Seed context" in b for b in augment_calls) + # The optimizer's rewrite model saw documents, metadata and tags. + assert rewrite_calls + briefing = "\n".join(rewrite_calls) + assert "Context documents:" in briefing and doc in briefing + assert "Example metadata:" in briefing and "Tags:" in briefing + assert "how to USE such context" in briefing + assert (project / "agents" / RAG_AGENT / "reports" / f"train_{RAG_AGENT}.html").exists() From 11067238c4caa182398cdb119fabda9ea257a738 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 2 Oct 2026 10:43:13 +0000 Subject: [PATCH 3/5] docs(examples): RAG context example and full row structure in prompt optimization scripts - New `rag_context.py` and `datasets/policies.jsonl`. Each row brings its own documents, plan, metadata, tags and rubric. The script prints what the agent, the judge and the rewrite model see, fits with context-aware metrics (`facts`, `grounded`), and reports test scores per tag. `--augment` keeps each seed's documents. - `SupportAgent` uses supplied documents (else its knowledge base) and the customer's plan. - `support.jsonl` rows get tags, and a `customer_plan` input where the plan matters. - `train.py` reports test scores per tag. `metrics_and_judges.py` shows metrics that read the documents and another input. - The augmentation and low-level scripts and the README explain what each field is for and what each stage sees. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_011gEMypXBosTqaEh3QFVXCX --- examples/prompt_optimization/README.md | 63 +++- examples/prompt_optimization/agent.py | 49 +++- examples/prompt_optimization/common.py | 56 ++++ .../prompt_optimization/data_augmentation.py | 8 +- .../datasets/policies.jsonl | 20 ++ .../datasets/support.jsonl | 52 ++-- .../low_level_prompt_fitter.py | 20 +- .../prompt_optimization/metrics_and_judges.py | 59 +++- examples/prompt_optimization/rag_context.py | 276 ++++++++++++++++++ examples/prompt_optimization/train.py | 28 +- tests/test_prompt_optimization_examples.py | 86 ++++++ 11 files changed, 656 insertions(+), 61 deletions(-) create mode 100644 examples/prompt_optimization/datasets/policies.jsonl create mode 100644 examples/prompt_optimization/rag_context.py diff --git a/examples/prompt_optimization/README.md b/examples/prompt_optimization/README.md index c14f99fc..bc14f404 100644 --- a/examples/prompt_optimization/README.md +++ b/examples/prompt_optimization/README.md @@ -2,8 +2,11 @@ Runnable, commented examples of Agentomatic's prompt optimization, from "score some answers" to a full Keras-style training run with an HTML report. They all -optimize the same small `SupportAgent` (`agent.py`), which answers Acme Cloud -customer questions from a policy knowledge base, on `datasets/support.jsonl`. +optimize the same small RAG `SupportAgent` (`agent.py`), which answers Acme +Cloud customer questions from policy documents. These are the documents sent +with the request (`context.documents`) when there are any, otherwise the ones +it retrieves from its own knowledge base. It also takes the customer's plan +into account. | Script | Level | What you learn | | --- | --- | --- | @@ -11,7 +14,8 @@ customer questions from a policy knowledge base, on `datasets/support.jsonl`. | [`data_augmentation.py`](data_augmentation.py) | building blocks | `prepare_dataset(augment=True)`: every strategy and option, leak protection, `augment_stats`. | | [`anti_overfitting.py`](anti_overfitting.py) | mechanics | The acceptance rules on hand-made numbers, then a fit where a prompt that memorizes the validation answers is vetoed and a general one is kept. | | [`low_level_prompt_fitter.py`](low_level_prompt_fitter.py) | low level | `PromptFitter` with every knob explicit: datasets, composite metric, search space, inner-loop callbacks, acceptance rules, before/after test evaluation, report. | -| [`train.py`](train.py) | high level | The class-agent workflow: `load_data` → `prepare_dataset` → metrics → loss → `PromptFitterBridge` → `compile` → `fit` → test evaluation → `generate_fit_report`. | +| [`train.py`](train.py) | high level | The class-agent workflow: `load_data` → `prepare_dataset` → metrics → loss → `PromptFitterBridge` → `compile` → `fit` → test evaluation (overall and per tag) → `generate_fit_report`. | +| [`rag_context.py`](rag_context.py) | RAG | Rows that bring their own documents, plan, metadata, tags and rubric (`datasets/policies.jsonl`). Prints what the agent, the judge and the rewrite model each see, fits with context-aware metrics, and reports test scores per tag. | Start with `metrics_and_judges.py` if you are choosing metrics, `train.py` if you want the whole loop, and `low_level_prompt_fitter.py` when you need control @@ -32,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 [--augment] ``` Shared flags (or environment variables), defined in `common.py`: @@ -48,32 +53,58 @@ Shared flags (or environment variables), defined in `common.py`: `train.py` writes `train_report.html` (+ `train_report.json`) and `train_summary.json` to `--out-dir`; `low_level_prompt_fitter.py` writes -`low_level_report.html` and `low_level_result.json`. +`low_level_report.html` and `low_level_result.json`; `rag_context.py` writes +`rag_report.html` (+ `rag_report.json`) and `rag_summary.json`. ## The data -`datasets/support.jsonl` has one JSON object per line: +Both datasets have one JSON object per line. A row carries the whole case, +not just a question and an answer. This is `datasets/policies.jsonl`: ```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"}} +{"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.", + "refusal": "Says so plainly when the documents do not answer the question."}} ``` +`datasets/support.jsonl` uses the same shape. Its agent retrieves documents +itself, so rows carry no documents, and `customer_plan` appears only where the +plan changes the answer. + +Every stage of the optimization gets the whole row: + +| Stage | Sees | +| --- | --- | +| Agent | `input` exactly as written: the question, `context.documents`, `customer_plan` | +| Metrics | The whole row: `example.input`, `.metadata`, `.tags` and `.rubric`, in `fit()` and in `evaluate()` | +| LLM judge | Question, answer, expected answer plus rubric, the documents (or what the agent retrieved), and the other inputs, metadata and tags | +| Rewrite model | Each failure, success and sample with its documents, inputs, metadata and tags, plus what the agent retrieved | +| Augmenter | The seed's documents and inputs; variations inherit them | +| Report | A `context` column in the per-example table | + The four splits have four jobs, and keeping them apart is what stops the optimizer from overfitting: -| Split | Rows | Role | +| Split | Rows (support / policies) | Role | | --- | --- | --- | -| `train` | 10 | What the optimizer **learns from**: failure analysis, judge feedback, few-shot demos. Augmentation only grows this split. | -| `validation` | 6 | What candidate prompts are **scored and selected on**. | -| `holdout` | 4 | A **veto-only gate**: a candidate that wins on validation but does not carry over here is rejected. Optional; without it a slice of validation is reserved. | -| `test` | 6 | **Never used by optimization.** Evaluate it before and after `fit()` for an unbiased number. | +| `train` | 10 / 8 | What the optimizer **learns from**: failure analysis, judge feedback, few-shot demos. Augmentation only grows this split. | +| `validation` | 6 / 5 | What candidate prompts are **scored and selected on**. | +| `holdout` | 4 / 3 | A **veto-only gate**: a candidate that wins on validation but does not carry over here is rejected. Optional; without it a slice of validation is reserved. | +| `test` | 6 / 4 | **Never used by optimization.** Evaluate it before and after `fit()` for an unbiased number. | `metadata.must_include` feeds the deterministic `facts` metric -(`common.must_include_score`); any metadata you put there reaches class-agent -metrics during `fit()`, `evaluate()` and inside composite metrics. +(`common.must_include_score`). `input.context.documents` feeds the `grounded` +metric (`common.grounded_score`). `tags` break test scores down by slice +(`common.scores_by_tag`), so you can see whether the new prompt helped billing +questions and hurt security ones. ## Testing without a model server diff --git a/examples/prompt_optimization/agent.py b/examples/prompt_optimization/agent.py index 96da5d66..d5ec3513 100644 --- a/examples/prompt_optimization/agent.py +++ b/examples/prompt_optimization/agent.py @@ -1,7 +1,11 @@ """Support agent shared by every prompt-optimization example. -A small but realistic class agent: it retrieves the relevant policy snippets -from an in-memory knowledge base, then asks the model to answer with them. +A small but realistic RAG class agent: it answers from policy snippets — +the documents sent with the request (``context.documents``) when there are +any, otherwise the ones it retrieves from an in-memory knowledge base — and +takes the customer's plan (``customer_plan``) into account. It lists what it +answered from under ``sources``, so judges can check groundedness and the +optimizer sees what each answer was built from. What the optimizer tunes is the **system prompt** (and, optionally, few-shot examples and model parameters). The agent reads it through @@ -61,10 +65,31 @@ class SupportState: """Per-run state of :class:`SupportAgent`.""" question: str = "" + #: Documents sent with the request (``context.documents``), if any. + documents: list[Any] = field(default_factory=list) + #: The customer's plan, when the caller (or the dataset row) gives it. + customer_plan: str = "" snippets: list[str] = field(default_factory=list) output: dict[str, Any] = field(default_factory=dict) +def documents_from(context: Any) -> list[Any]: + """Documents from ``{"documents": [...]}``, a list, or one string.""" + if isinstance(context, dict): + context = context.get("documents") + if isinstance(context, str): + context = [context] + return [doc for doc in context or [] if doc] + + +def snippet_text(doc: Any) -> str: + """A document as one snippet line: its text, then its source if it has one.""" + if isinstance(doc, dict): + text = str(doc.get("content") or doc.get("text") or "") + return f"{text} (source: {doc['source']})" if doc.get("source") else text + return str(doc) + + class SupportAgent(BaseGraphAgent[SupportState]): """Answer customer questions from the knowledge base. @@ -98,7 +123,10 @@ def build_graph(self) -> Any: # --- nodes -------------------------------------------------------------- def retrieve(self, state: SupportState) -> SupportState: - """Pick the knowledge-base snippets whose topic the question mentions.""" + """Use the supplied documents, else the knowledge-base snippets on topic.""" + if state.documents: + state.snippets = [snippet_text(doc) for doc in state.documents] + return state question = state.question.lower() state.snippets = [ KNOWLEDGE_BASE[topic] @@ -112,6 +140,7 @@ def answer(self, state: SupportState) -> SupportState: # The optimizer's candidate prompt arrives through this call. prompt = self.resolve_system_prompt(default=self.system_prompt) context = "\n".join(f"- {s}" for s in state.snippets) or "- (no matching policy)" + plan = f"Customer plan: {state.customer_plan}\n\n" if state.customer_plan else "" if self.llm is None: text = " ".join(state.snippets) or "I could not find a matching policy." else: @@ -120,7 +149,9 @@ def answer(self, state: SupportState) -> SupportState: {"role": "system", "content": prompt}, { "role": "user", - "content": f"Policy snippets:\n{context}\n\nQuestion: {state.question}", + "content": ( + f"Policy snippets:\n{context}\n\n{plan}Question: {state.question}" + ), }, ] ) @@ -134,7 +165,15 @@ def input_to_state(self, data: dict[str, Any]) -> SupportState: # REST sends ``query``; the platform normalises it to ``current_query``. # Datasets may use ``question`` — accept all three. question = data.get("current_query") or data.get("question") or data.get("query") or "" - return SupportState(question=str(question)) + # ``context.documents`` arrives nested (and flattened to ``documents`` + # over REST). Every other dataset input — here ``customer_plan`` — + # arrives exactly as written in the row. + documents = documents_from(data.get("context")) or documents_from(data.get("documents")) + return SupportState( + question=str(question), + documents=documents, + customer_plan=str(data.get("customer_plan") or ""), + ) def state_to_output(self, state: SupportState) -> dict[str, Any]: return state.output diff --git a/examples/prompt_optimization/common.py b/examples/prompt_optimization/common.py index dea5ac14..cea66afe 100644 --- a/examples/prompt_optimization/common.py +++ b/examples/prompt_optimization/common.py @@ -26,6 +26,7 @@ import argparse import os +import re import sys from dataclasses import dataclass from pathlib import Path @@ -33,6 +34,8 @@ HERE = Path(__file__).resolve().parent DATASET = HERE / "datasets" / "support.jsonl" +#: RAG rows that bring their own documents (``input.context.documents``). +POLICIES_DATASET = HERE / "datasets" / "policies.jsonl" # Let ``python examples/prompt_optimization/