diff --git a/CLAUDE.md b/CLAUDE.md index 257582f6..ae055b09 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -242,6 +242,13 @@ register_vector_store_adapter("my_db", MyDbStore) # upsert/query/delete after `fit()`. `generate_fit_report(history, baseline_eval=..., final_eval=...)` shows the prompt diff, every candidate's decision and reason, and per-example answers. Guide: `docs/optimization/`; runnable: `examples/prompt_optimization/`. +- Dataset rows carry the whole case — `input.context.documents` (RAG), any other + `input.*` (plan, `messages` …), `metadata`, `tags`, `rubric` — and every stage + gets it: the agent receives `input` as written (a copy; datasets are never + mutated), metrics the whole row, LLM judges the documents (or what the agent + listed under `citations` / `retrieval_context` / `sources`) plus inputs, + metadata and tags, the rewrite model all of it per example, the augmenter the + seed's documents. See `docs/optimization/data.md` and `rag_context.py`. ## Project conventions (match these when editing the framework) diff --git a/docs/optimization/data.md b/docs/optimization/data.md index d87020a7..83ba40c2 100644 --- a/docs/optimization/data.md +++ b/docs/optimization/data.md @@ -2,49 +2,118 @@ ## The dataset file -Class-agent datasets are JSONL, one `AgentExample` per line (the format every -template's `datasets/all.jsonl` uses): +Class-agent datasets are JSONL, one `AgentExample` per line. Every +template's `datasets/all.jsonl` uses this format. A row carries the whole +case: the question, the documents the answer must come from, any other agent +input, labels, tags and a rubric. All of it reaches the optimization, not +just the question and the answer. ```json -{"id": "refund_2", - "split": "validation", - "input": {"current_query": "I bought an annual plan last week. Can I still get my money back?"}, - "expected_output": {"response": "Yes. Annual plans can be refunded within 30 days of purchase."}, - "metadata": {"must_include": ["30 days"], "topic": "refund"}, - "rubric": {"correctness": "States the 30-day window for annual plans."}, - "tags": ["billing"]} +{"id": "policy_05", + "split": "train", + "input": {"current_query": "How long are audit logs kept?", + "context": {"documents": [{"source": "audit.md", + "content": "Audit logs are kept for 400 days on Enterprise and 90 days on Team."}]}, + "customer_plan": "team"}, + "expected_output": {"response": "On the Team plan, audit logs are kept for 90 days (400 days on Enterprise)."}, + "metadata": {"topic": "security", "difficulty": "medium", "must_include": ["90 days"]}, + "tags": ["security", "compliance", "plan-specific"], + "rubric": {"groundedness": "Every claim comes from the row's context documents.", + "plan": "Applies the policy to the customer's plan when one is given."}} ``` | Field | Required | Meaning | | --- | --- | --- | | `id` | recommended | Stable id; generated (`example_0003`) when missing. Reports join before/after results on it. | | `split` | recommended | `train`, `validation`, `holdout` or `test` (default `train`). See [data roles](generalization.md#data-roles). | -| `input` | yes | What the agent receives — the same dict `input_to_state()` gets. The question is read from `question`, `query`, `current_query` or `request`. `input.context` becomes the judge's retrieval context. | +| `input` | yes | What the agent receives, unchanged: the same dict `input_to_state()` gets. The question is read from `question`, `query`, `current_query` or `request`. | +| `input.context` | no | Documents the answer must come from (RAG). The value can be `{"documents": [...]}`, a list, or one string. A document is a string or a dict with `content` (or `text` / `page_content`) and an optional `source`. | +| other `input.*` keys | no | Any other agent input: `messages` (earlier turns), a plan, a locale, a user profile, and so on. | | `expected_output` | for scoring | The reference answer (`{"response": "..."}`, or your structured output). Judges receive it as the reference. | -| `metadata` | no | Anything metrics need: required facts, topic, difficulty. Class-agent metrics read it (`example.metadata`). | -| `rubric` | no | Per-dimension criteria, appended to the judge's reference. | -| `tags` | no | Free-form labels. | +| `metadata` | no | Labels metrics need: required facts (`must_include`), topic, difficulty. | +| `tags` | no | Grouping labels. Filter with `data.filter_by_tags(...)`, and break results down by tag. | +| `rubric` | no | Per-dimension criteria, added to the judge's reference. | ```python from agentomatic.optimize import load_data data = load_data("agents/support/datasets/all.jsonl") # AgentDataset data.train, data.validation, data.holdout, data.test # lists of AgentExample +billing = data.filter_by_tags("billing") # any-match ``` `PromptFitter` works on the simpler `agentomatic.optimize.Dataset` of -`DataPoint(query, expected_answer, context, metadata)`. Convert, or build one -directly: +`DataPoint(query, expected_answer, context, metadata, tags)`. Convert, or +build one directly: ```python from agentomatic.optimize import Dataset -valset = Dataset(points=[e.to_datapoint() for e in data.validation]) # keeps metadata -Dataset.from_list([{"query": "Refund window?", "expected_answer": "30 days"}]) -Dataset.from_jsonl("eval.jsonl") # {"query", "expected_answer", "context", "metadata"} +valset = Dataset(points=[e.to_datapoint() for e in data.validation]) # keeps everything +Dataset.from_list([{"query": "Refund window?", "expected_answer": "30 days", + "context": ["Refunds: 30 days."], "tags": ["billing"], + "metadata": {"invoke": {"customer_plan": "team"}}}]) +Dataset.from_jsonl("eval.jsonl") # {"query", "expected_answer", "context", "metadata", "tags"} trainset, valset = dataset.split(ratio=0.8) ``` +`to_datapoint()` maps a row as follows: + +| Row field | `DataPoint` field | +| --- | --- | +| `input` (except the question) | `metadata.invoke`, sent to the agent exactly as written | +| `input.context` | `context`, as document texts such as `"Audit logs are kept … (source: audit.md)"` | +| `metadata` | `metadata` | +| `tags` | `tags` | +| `rubric` | the expected reference, ahead of the answer, for the judge | + +## What each stage sees + +| Stage | Receives | +| --- | --- | +| **Agent** | `input` exactly as written: the question, `context`, `customer_plan` and the rest. Evaluation hands it a copy, so an agent that edits its input never changes the dataset. | +| **Metrics** | The whole row (`example.input`, `.metadata`, `.tags`, `.rubric`), the same inside `fit()`, in `evaluate()` and nested in a composite. | +| **LLM judges** | Question, answer, expected answer plus rubric, and the context documents. If the row has no documents, they get what the agent retrieved instead. Inside `fit()` / `evaluate()` they also get the other inputs, the metadata and the tags. | +| **Rewrite model** | Every failure, success and dataset sample in its [briefing](how-it-works.md) shows `Context documents`, `Other inputs`, `Example metadata`, `Tags` and what the agent retrieved. It is told to write instructions that *use* such context and never to copy one row's facts into the prompt. | +| **Augmenter** | The seed's documents and other inputs (a "Seed context" section). Variations inherit a copy of them, plus the seed's tags (and `augmented`) and rubric. | +| **Report** | A `context` column in the per-example table: tags, documents, inputs and metadata. | + +## RAG datasets + +There are two ways for a RAG agent to get its documents. Both work with every +stage above. + +* **The row brings its documents** (`input.context.documents`). The agent + answers from exactly those. The judge checks groundedness against them, and + the rewrite model reads them next to each answer. Use this to optimize the + *answering* prompt independently of retrieval quality, or when the gold + passages are known. The `rag` template reads them in `input_to_state` + (`context.documents`, flattened to `documents` over REST) and falls back to + its own retrieval when there are none. +* **The agent retrieves for itself.** List what it used in the output under + `retrieval_context`, `citations`, `sources` or `documents`. These can be + strings or dicts with `content` and `source`. Judges then score + groundedness against those, and the rewrite model sees them as + `Retrieved by the agent`. + +```python +def input_to_state(self, data): + ctx = data.get("context") or {} + docs = ctx.get("documents") if isinstance(ctx, dict) else ctx + return MyState(question=data.get("current_query", ""), documents=list(docs or [])) + +def state_to_output(self, state): + return {"response": state.answer, "citations": state.citations} # what it used +``` + +Metrics can read the documents too. For example, a groundedness check that +counts how many of the answer's words come from them: +[`grounded_score`](examples.md). + +`examples/prompt_optimization/rag_context.py` runs the whole loop on rows +that bring their own documents, plan, metadata, tags and rubric. It prints +what each stage sees and reports test scores per tag. + ## Writing good examples * **One fact per expectation you care about.** An expected answer is what the @@ -122,6 +191,10 @@ before trusting scores computed on them. * Generated questions that duplicate an existing one, or are ≥ 90 % similar to any validation / holdout / test question, are dropped (a paraphrase of a test question in train would inflate every score measured on it). +* The generator sees each seed's documents and other inputs (a "Seed + context" section) and is told to ask only what those documents answer. + Every variation inherits a copy of its seed's inputs, rubric and tags (plus + `augmented`), so a RAG row stays answerable. * Each new row records `metadata.source = "augment"`, its `strategy` and its `parent_id`. * `data.metadata["augment_stats"]` records `requested`, `calls`, `parsed`, diff --git a/docs/optimization/examples.md b/docs/optimization/examples.md index 4c86a5c5..ea49be7c 100644 --- a/docs/optimization/examples.md +++ b/docs/optimization/examples.md @@ -3,9 +3,16 @@ [`examples/prompt_optimization`](https://github.com/UnicoLab/agentomatic/tree/main/examples/prompt_optimization) contains runnable scripts, from scoring a few answers to a full training run. Each one is commented step by step and lists every option of the methods it -uses. All of them optimize the same small `SupportAgent` (it answers Acme -Cloud customer questions from a policy knowledge base) on a 26-row dataset -split into train / validation / holdout / test. +uses. All of them optimize the same small RAG `SupportAgent`, which answers +Acme Cloud customer questions from policy documents. Those are the documents a +row sends (`context.documents`), or else the ones it retrieves from its own +knowledge base, and it also takes the customer's plan into account. Two +datasets are used, each split into train / validation / holdout / test: + +* `support.jsonl` has 26 rows, with tags and a `customer_plan` input where the + plan matters. +* `policies.jsonl` has 20 rows. Each brings its own documents, plus a plan, + metadata, tags and a rubric. | Script | Level | Covers | Docs | | --- | --- | --- | --- | @@ -13,7 +20,8 @@ split into train / validation / holdout / test. | `data_augmentation.py` | building blocks | `prepare_dataset` strategies and options, leak protection, `augment_stats` | [Datasets](data.md) | | `anti_overfitting.py` | mechanics | The acceptance rules on hand-made numbers; a fit in which a prompt memorising the validation answers is vetoed and a general one kept | [Generalization](generalization.md) | | `low_level_prompt_fitter.py` | low level | `PromptFitter` with every knob, inner-loop callbacks, before/after test evaluation, report | [PromptFitter](prompt-fitter.md) | -| `train.py` | high level | `load_data` → `prepare_dataset` → metrics → loss → `PromptFitterBridge` → `compile` → `fit` → test evaluation → `generate_fit_report` | [Class agents](class-agents.md) | +| `train.py` | high level | `load_data` → `prepare_dataset` → metrics → loss → `PromptFitterBridge` → `compile` → `fit` → test evaluation (overall and per tag) → `generate_fit_report` | [Class agents](class-agents.md) | +| `rag_context.py` | RAG | Rows with their own documents, plan, metadata, tags and rubric. Prints what the agent, the judge and the rewrite model see; context-aware metrics (`facts`, `grounded`); per-tag test scores; `--augment` keeps each seed's documents | [Datasets](data.md#rag-datasets) | ## Running them @@ -28,6 +36,7 @@ python examples/prompt_optimization/data_augmentation.py --n-examples 40 python examples/prompt_optimization/anti_overfitting.py python examples/prompt_optimization/low_level_prompt_fitter.py --trials 8 python examples/prompt_optimization/train.py --epochs 2 --trials 6 --augment +python examples/prompt_optimization/rag_context.py --epochs 1 --trials 4 ``` Every script takes `--base-url`, `--api-key`, `--model`, `--judge-model`, @@ -52,6 +61,37 @@ Step 10/10 — results report: examples/prompt_optimization/out/train_report.html ``` +What `rag_context.py` prints first, for one row (abridged): + +```text +── What each stage sees — row policy_05 +agent (input_to_state receives exactly this): + {"current_query": "How long are audit logs kept?", "context": {"documents": [{"source": "audit.md", "content": "Audit logs are kept for 400 days on Enterprise and 90 days on Team."}]}, "customer_plan": "team"} +LLM judge — context documents: + - Audit logs are kept for 400 days on Enterprise and 90 days on Team. (source: audit.md) +LLM judge — other inputs, metadata and tags: + {"inputs": {"customer_plan": "team"}, "metadata": {"topic": "security", "difficulty": "medium", "must_include": ["90 days"]}, "tags": ["security", "compliance", "plan-specific"]} +rewrite model (dataset sample in its briefing): +1. Q: How long are audit logs kept? + - Context documents: Audit logs are kept for 400 days on Enterprise and 90 days on Team. (source: audit.md) + - Other inputs: {"customer_plan": "team"} + - Example metadata: {"topic": "security", "difficulty": "medium", "must_include": ["90 days"]} + - Tags: security, compliance, plan-specific + Expected: + ## Expected answer + On the Team plan, audit logs are kept for 90 days (400 days on Enterprise). +``` + +…and last, the test split per tag: + +```text +── Test quality by tag (before → after) + backups 0.00 → 0.85 + plan-specific 0.00 → 0.90 + pricing 0.03 → 0.83 + refusal 0.03 → 0.20 +``` + ## Tested `tests/test_prompt_optimization_examples.py` runs every script as a @@ -59,4 +99,10 @@ subprocess against a deterministic fake OpenAI-compatible server (`tests/fake_openai_server.py`) and checks the behaviour each one teaches: the prompt improves on validation *and* test, augmentation grows only train, the memorising prompt is rejected, and the report contains the prompts, -candidates and per-example results. +candidates and per-example results. For `rag_context.py` it also checks the +requests the fake server received: + +* the agent answered from each row's documents and plan; +* judges saw the documents, the plan, the metadata and the tags; +* the rewrite model's briefing contained all of them; +* augmented rows kept their seed's documents. diff --git a/docs/optimization/how-it-works.md b/docs/optimization/how-it-works.md index dbd1b936..20b24894 100644 --- a/docs/optimization/how-it-works.md +++ b/docs/optimization/how-it-works.md @@ -63,6 +63,13 @@ flowchart TD never studies the examples it is then selected on (`reflect_on="train"`). 3. **Proposal.** The optimizer (see the table below) proposes up to four candidate configurations per round from the incumbent and the reflection. + Its rewrite model reads a briefing: the current config, metrics, score + history, train samples, and the lowest- and highest-scoring examples. Each + example comes with everything it ran with: question, expected answer, + actual answer, judge feedback, context documents, other inputs, metadata, + tags, and what the agent retrieved. The briefing tells the model to write + instructions that use such context, never to copy one example's facts into + the prompt. See [what each stage sees](data.md#what-each-stage-sees). A candidate identical to one already scored is skipped for free: a weak rewrite model often re-proposes the current prompt. 4. **Minibatch screen.** Each candidate runs on a slice of validation (30 %, diff --git a/docs/optimization/metrics.md b/docs/optimization/metrics.md index eb229827..fa45e2d9 100644 --- a/docs/optimization/metrics.md +++ b/docs/optimization/metrics.md @@ -29,7 +29,7 @@ automatically. | | `agentomatic.agents` (class agents) | `agentomatic.optimize` (the fitter) | | --- | --- | --- | | Signature | `score(example, prediction) -> float` (sync) | `await evaluate(query, response, expected, context) -> EvalResult` | -| Sees | The whole `AgentExample` (input, expected output, **metadata**) and the output dict | The question, the answer text, the reference text, retrieval context | +| Sees | The whole `AgentExample` (input incl. context documents and other inputs, expected output, **metadata**, **tags**, rubric) and the output dict, the same in `fit()` and `evaluate()` | The question, the answer text, the reference text, the context documents (or what the agent retrieved). Inside `fit()` / `evaluate()`, `current_scoring_run()` also exposes the example | | Returns | A float | `EvalResult(score, reason, metadata)` — judges add `dimensions`, rationale; failures set `metadata["evaluation_failed"]` | | Best for | Checks that need structured output or per-example metadata | Judges and text checks | @@ -118,6 +118,12 @@ How judges stay honest: * **Reference.** Judges receive the example's expected answer (and, for `AgentExample`s, a rich reference: expected output, rubric, required facts). Write expected answers that state the facts a correct answer must contain. +* **Context.** Judges receive the row's context documents. If the row has + none, they get what the agent says it retrieved (`retrieval_context` / + `citations` / `sources` in its output). Inside `fit()` and `evaluate()` + they also get the row's other inputs (a plan, a locale…), its metadata and + its tags, in an "Example inputs, metadata and tags" section, so an answer + is judged in that light. See [what each stage sees](data.md#what-each-stage-sees). * **Deterministic.** Keep `temperature=0.0`. A judge that scores the same answer differently each time makes every comparison noise; the [significance rule](generalization.md#the-acceptance-rules) exists because diff --git a/docs/optimization/reports.md b/docs/optimization/reports.md index 346baa61..b4e91fc9 100644 --- a/docs/optimization/reports.md +++ b/docs/optimization/reports.md @@ -50,7 +50,7 @@ additionally writes one per fit under `experiment_dir`. | **What changed in the prompt** | The initial → final diff (long lines wrapped so changed sentences stand out), both full prompts side by side, and every accepted change in order with why it was accepted. | | **Epochs** | Start / best / held-out / candidates / stop reason per epoch. | | **All candidates** | Every proposed candidate: epoch, round, minibatch / validation / held-out scores, confidence, **decision** and **reason**; the diff of each distinct candidate prompt. | -| **Examples — before vs after** | Per example: question, expected answer, answer before and after, scores before and after, per-metric Δ, and the judge's rationale. Uses the test split when `baseline_eval`/`final_eval` are given, otherwise the validation examples the fitter scored. | +| **Examples — before vs after** | Per example: question, its context (tags · documents · other inputs · metadata, when the data has any), expected answer, answer before and after, scores before and after, per-metric Δ, and the judge's rationale. Uses the test split when `baseline_eval`/`final_eval` are given, otherwise the validation examples the fitter scored. | | **Data & settings** | Split sizes, augmentation stats, and every fitter knob used. | | **Run configuration**, **Recommendations** | Model, optimizer, suggestions, parameter changes, metric deltas, a rollout recommendation. | | **Curves** | Best score / loss per round, Keras `loss` / `val_loss` and metrics per epoch, the held-out evaluation. | 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/