Mimicking a binary deterministic agent

6.3.1. Scoring all the input text using the tiny agent

from collections.abc import Sequence
from pprint import pprint


def score_dataset_with_tiny_agent(
    dataset: Sequence[dict[str, object]],
) -> list[dict[str, object]]:
    """Append the tiny agent answer to every input example."""

    scored: list[dict[str, object]] = []
    for row in dataset:
        topic = row.get("topic")
        answer = tiny_agent_response(str(row["text"]), topic=topic)
        scored.append({**row, "tiny_agent_output": answer})
    return scored


scored_examples = score_dataset_with_tiny_agent(synthetic_examples)
print(f"Scored {len(scored_examples)} examples with the deterministic agent.")
pprint(scored_examples[:3])

Scored 200 examples with the deterministic agent.
[{'model': 'gpt-3.5-turbo',
  'text': 'The first sip of my morning espresso awakens my senses and sets a '
          'positive tone for the day ahead.',
  'timestamp': 1763237270.6815271,
  'tiny_agent_output': 'I see you are a person of taste! I like coffee, too.',
  'topic': 'coffee'},
 {'model': 'gpt-3.5-turbo',
  'text': "I love coffee because it's a warm hug in a cup that jumpstarts my "
          'day with its rich aroma.',
  'timestamp': 1763237271.866124,
  'tiny_agent_output': 'I see you are a person of taste! I like coffee, too.',
  'topic': 'coffee'},
 {'model': 'gpt-3.5-turbo',
  'text': 'The first sip of rich, aromatic coffee transports me to cozy '
          'mornings and fuels my passion for creativity and productivity.',
  'timestamp': 1763237273.029643,
  'tiny_agent_output': 'I see you are a person of taste! I like coffee, too.',
  'topic': 'coffee'}]

6.5.1. The evaluator

from __future__ import annotations

import math
import time
from collections import Counter, defaultdict, deque
from dataclasses import dataclass
from typing import Any, Deque, Dict, Literal

import numpy as np
import openai


SimilarityMetric = Literal["bleu", "embedding_cosine"]


def _tokenize(text: str) -> list[str]:
    return [token for token in text.lower().split() if token]


def _ngram_counts(tokens: list[str], n: int) -> Counter[tuple[str, ...]]:
    if len(tokens) < n:
        return Counter()
    return Counter(tuple(tokens[i : i + n]) for i in range(len(tokens) - n + 1))


def compute_bleu(reference: str, hypothesis: str) -> float:
    ref_tokens = _tokenize(reference)
    hyp_tokens = _tokenize(hypothesis)

    if not hyp_tokens:
        return 0.0

    weights = [0.25, 0.25, 0.25, 0.25]
    precisions = []
    for n in range(1, 5):
        ref_counts = _ngram_counts(ref_tokens, n)
        hyp_counts = _ngram_counts(hyp_tokens, n)
        overlap = sum(
            min(count, ref_counts[gram]) for gram, count in hyp_counts.items()
        )
        total = sum(hyp_counts.values()) or 1
        precisions.append(overlap / total if overlap else 1e-9)

    log_precision = sum(w * math.log(p) for w, p in zip(weights, precisions))
    brevity_penalty = 1.0
    ref_len = len(ref_tokens)
    hyp_len = len(hyp_tokens)
    if hyp_len <= ref_len and hyp_len > 0:
        brevity_penalty = math.exp(1 - ref_len / hyp_len)

    return float(brevity_penalty * math.exp(log_precision))


def cosine_similarity(vec_a: np.ndarray, vec_b: np.ndarray) -> float:
    denom = np.linalg.norm(vec_a) * np.linalg.norm(vec_b)
    if denom == 0:
        return 0.0
    return float(np.dot(vec_a, vec_b) / denom)


@dataclass
class EvaluationRecord:
    user_input: str
    tiny_output: str
    mimic_output: str
    similarity: float
    metric: SimilarityMetric
    timestamp: float


class TinyAgentEvaluator:
    """Track similarity and histories for the mimic vs. tiny agent comparison."""

    def __init__(
        self,
        *,
        metric: SimilarityMetric = "bleu",
        per_input_history: int = 5,
        training_history: int = 50,
        embedding_model: str = "text-embedding-3-small",
    ) -> None:
        self.metric = metric
        self.per_input_history = per_input_history
        self.training_history = training_history
        self.embedding_model = embedding_model

        self._per_input_records: dict[str, Deque[EvaluationRecord]] = defaultdict(
            lambda: deque(maxlen=per_input_history)
        )
        self._global_records: Deque[EvaluationRecord] = deque(maxlen=training_history)
        self._embedding_cache: dict[str, np.ndarray] = {}
        self._embedding_client = openai.OpenAI()

    # -- metrics -----------------------------------------------------------------
    def _bleu_similarity(self, tiny_output: str, mimic_output: str) -> float:
        return compute_bleu(tiny_output, mimic_output)

    def _embedding_similarity(self, tiny_output: str, mimic_output: str) -> float:
        tiny_vec = self._embed_text(tiny_output)
        mimic_vec = self._embed_text(mimic_output)
        return cosine_similarity(tiny_vec, mimic_vec)

    def _embed_text(self, text: str) -> np.ndarray:
        cache_key = text.strip()
        if cache_key in self._embedding_cache:
            return self._embedding_cache[cache_key]
        if not cache_key:
            vec = np.zeros(1536, dtype=np.float32)
            self._embedding_cache[cache_key] = vec
            return vec
        response = self._embedding_client.embeddings.create(
            model=self.embedding_model,
            input=cache_key,
        )
        vec = np.array(response.data[0].embedding, dtype=np.float32)
        self._embedding_cache[cache_key] = vec
        return vec

    def _compute_similarity(self, tiny_output: str, mimic_output: str) -> float:
        if self.metric == "embedding_cosine":
            return self._embedding_similarity(tiny_output, mimic_output)
        return self._bleu_similarity(tiny_output, mimic_output)

    def similarity(self, tiny_output: str, mimic_output: str) -> float:
        """Public helper to compute similarity without mutating evaluator state."""

        return self._compute_similarity(tiny_output, mimic_output)

    def embed_text(self, text: str) -> np.ndarray:
        """Expose embedding helper for downstream analysis."""

        return self._embed_text(text)

    # -- public api ---------------------------------------------------------------
    def evaluate(
        self,
        *,
        user_input: str,
        tiny_output: str,
        mimic_output: str,
    ) -> Dict[str, Any]:
        similarity = self._compute_similarity(tiny_output, mimic_output)
        record = EvaluationRecord(
            user_input=user_input,
            tiny_output=tiny_output,
            mimic_output=mimic_output,
            similarity=similarity,
            metric=self.metric,
            timestamp=time.time(),
        )

        per_input_buffer = self._per_input_records[user_input]
        per_input_buffer.append(record)
        self._global_records.append(record)

        return {
            "similarity": similarity,
            "metric": self.metric,
            "last_outputs_for_input": [
                {
                    "mimic_output": item.mimic_output,
                    "tiny_output": item.tiny_output,
                    "similarity": item.similarity,
                    "timestamp": item.timestamp,
                }
                for item in list(per_input_buffer)
            ],
            "recent_training_examples": [
                {
                    "user_input": item.user_input,
                    "mimic_output": item.mimic_output,
                    "tiny_output": item.tiny_output,
                    "similarity": item.similarity,
                    "metric": item.metric,
                    "timestamp": item.timestamp,
                }
                for item in list(self._global_records)
            ],
        }


def describe_block_change(block_name: str, old_text: str, new_text: str) -> str:
    if old_text.strip() == new_text.strip():
        return f"Block {block_name} unchanged."

    desc: list[str] = [f"Block {block_name}: "]
    if len(new_text) > len(old_text) * 1.1:
        desc.append("expanded details.")
    elif len(new_text) < len(old_text) * 0.9:
        desc.append("more concise wording.")

    old_words = set(old_text.lower().split())
    new_words = set(new_text.lower().split())
    added = list(new_words - old_words)
    removed = list(old_words - new_words)
    if added:
        desc.append(
            " Added cues: "
            + ", ".join(sorted(added[:5]))
            + ("..." if len(added) > 5 else "")
        )
    if removed:
        desc.append(
            " Removed cues: "
            + ", ".join(sorted(removed[:5]))
            + ("..." if len(removed) > 5 else "")
        )

    if len(desc) == 1:
        desc.append(" Rephrased while preserving intent.")
    return "".join(desc)


def prompt_unified_diff(old_prompt: str, new_prompt: str) -> str:
    diff = difflib.unified_diff(
        old_prompt.splitlines(),
        new_prompt.splitlines(),
        fromfile="previous_prompt",
        tofile="updated_prompt",
        lineterm="",
    )
    diff_text = "\n".join(diff)
    return diff_text or "(no textual diff)"


def format_change_log(entries: list[dict[str, object]], limit: int = 3) -> str:
    if not entries:
        return "No prior prompt changes recorded."

    lines: list[str] = ["Recent prompt changes:"]
    for idx, entry in enumerate(entries[-limit:], start=1):
        delta = entry.get("delta")
        if isinstance(delta, float):
            delta_text = f"{delta:+.3f}"
        else:
            delta_text = "N/A"
        lines.append(
            f"[{idx}] block={entry.get('block_name')} Δmetric={delta_text} "
            f"prev={entry.get('prev_metric')} new={entry.get('new_metric')}"
        )
        lines.append(f"Intent: {entry.get('intention')}")
        old_block = entry.get("old_block_text", "")
        new_block = entry.get("new_block_text", "")
        lines.append(f"Old block: {old_block}")
        lines.append(f"New block: {new_block}")
        diff_text = entry.get("diff", "") or "(no diff)"
        if isinstance(diff_text, str):
            snippet = diff_text if len(diff_text) <= 800 else diff_text[:800] + "\n..."
            lines.append("Diff:\n" + snippet)
    return "\n".join(lines)


def serialize_prompt_blocks(blocks: dict[str, str]) -> str:
    parts = []
    for name, text in blocks.items():
        parts.append(f"[{name}]\n{text}")
    return "\n\n".join(parts)

evaluator = TinyAgentEvaluator(
    metric="embedding_cosine",
    per_input_history=3,
    training_history=10,
)

sample_input = synthetic_examples[0]["text"]
tiny_answer = tiny_agent_response(sample_input)
mimic_answer = mimic_agent(sample_input)["mimic_output"]
evaluator.evaluate(
    user_input=sample_input,
    tiny_output=tiny_answer,
    mimic_output=mimic_answer,
)

6.5.2. The coach agent

class CoachSignature(dspy.Signature):
    evaluator_report = dspy.InputField(
        desc="Structured summary describing recent evaluator scores and samples."
    )
    prompt_blocks = dspy.InputField(desc="Structured prompt broken into named blocks.")
    block_options = dspy.InputField(desc="List of valid block names to edit.")
    guidelines = dspy.InputField(
        desc="Non-negotiable constraints for how the revised prompt should be phrased."
    )
    prompt_change_log = dspy.InputField(
        desc="Table describing prior prompt edits, diffs, and their metric deltas."
    )
    case_examples = dspy.InputField(
        desc="Sample inputs with target/mimic outputs for focused comparison."
    )
    case_reasoning = dspy.InputField(
        desc="Reasoned analysis describing why mimic outputs differ from targets."
    )
    target_block = dspy.OutputField(desc="Name of the block to edit (e.g., general, sop-1).")
    updated_block_text = dspy.OutputField(desc="Rewritten text for the chosen block only.")
    coach_rationale = dspy.OutputField(desc="Short explanation for the proposed change.")


COACH_PROMPT_GUIDELINES = (
    "Produce a standalone instruction for an agent to immitate a 'target agent'. "
    "Focus on what the agent should say to become indistinguishable from the target agent."
    "The target agent is following a policy and you must decode it by looking to its behavior and craft a prompt to the mimic agent to apply that same policy."
    "Be as precise as possible describing the process the target agent is following to insert in the mimic agent prompt."
    "Don't mention 'mimic agent' or 'target agent' in the prompt of the 'mimic agent' since it is unaware of them. "
    "You want to maximize the similarity between the target agent output and the mimic agent output by editing the mimic agent's prompt."
    "Respect the structured prompt: edit exactly one named block (general or sop-#) per iteration and leave all other blocks untouched."
)


def summarize_evaluator_output(report: dict[str, Any], limit: int = 3) -> str:
    """Turn evaluator dictionaries into a short textual description."""

    lines = [
        f"Similarity metric: {report.get('metric')}",
        f"Last score: {report.get('similarity'):.3f}",
    ]
    history = report.get("last_outputs_for_input", [])[-limit:]
    if history:
        lines.append("Recent outputs for this input:")
        for item in history:
            lines.append(
                "- target: "
                + item["tiny_output"].strip()
                + " | mimic: "
                + item["mimic_output"].strip()
                + f" | score={item['similarity']:.3f}"
            )

    recent_examples = report.get("recent_training_examples", [])[-limit:]
    if recent_examples:
        lines.append("Recent training mismatches:")
        for example in recent_examples:
            lines.append(
                "- input: "
                + example["user_input"].strip()
                + " | target: "
                + example["tiny_output"].strip()
                + " | mimic: "
                + example["mimic_output"].strip()
                + f" | score={example['similarity']:.3f}"
            )

    return "\n".join(lines)


class CoachAgent(dspy.Module):
    """LLM-based coach that rewrites the mimic prompt using evaluator feedback."""

    def __init__(self) -> None:
        super().__init__()
        ensure_dspy_configured()
        self.predictor = dspy.Predict(CoachSignature)
        self.guidelines = COACH_PROMPT_GUIDELINES

    def forward(
        self,
        *,
        evaluator_output: dict[str, Any],
        prompt_blocks: str,
        block_options: str,
        prompt_change_log: str,
        case_examples: str,
        case_reasoning: str,
    ) -> dict[str, str]:  # type: ignore[override]
        summary = summarize_evaluator_output(evaluator_output)
        prediction = self.predictor(
            evaluator_report=summary,
            prompt_blocks=prompt_blocks,
            block_options=block_options,
            guidelines=self.guidelines,
            prompt_change_log=prompt_change_log,
            case_examples=case_examples,
            case_reasoning=case_reasoning,
        )
        block = getattr(prediction, "target_block", "general").strip()
        updated_block = getattr(prediction, "updated_block_text", "").strip()
        rationale = getattr(prediction, "coach_rationale", "")
        return {
            "target_block": block,
            "updated_block_text": updated_block,
            "coach_rationale": rationale,
        }


coach = CoachAgent()
evaluator_snapshot = evaluator.evaluate(
    user_input=sample_input,
    tiny_output=tiny_answer,
    mimic_output=mimic_agent(sample_input)["mimic_output"],
)

coach.forward(
    evaluator_output=evaluator_snapshot,
    prompt_blocks=serialize_prompt_blocks(mimic_agent.get_prompt_blocks()),
    block_options=", ".join(mimic_agent.prompt_state.block_names()),
    prompt_change_log="No prior prompt changes recorded.",
    case_examples="",
    case_reasoning="No case analysis available yet.",
)

2025/11/15 20:26:49 WARNING dspy.primitives.module: Calling module.forward(...) on CoachAgent directly is discouraged. Please use module(...) instead.

6.5.3. The algorithm

1. Splitting the training data
   

   We split the data between training, test, and holdout. 60%/20%/20%.

   import math
   import random
   from typing import Sequence, Tuple
   
   
   def split_dataset(
       dataset: Sequence[dict[str, object]],
       train_ratio: float = 0.6,
       test_ratio: float = 0.2,
       seed: int = 2024,
   ) -> Tuple[list[dict[str, object]], list[dict[str, object]], list[dict[str, object]]]:
       if not math.isclose(train_ratio + test_ratio, 0.8, rel_tol=1e-9):
           raise ValueError("Train + test ratios must sum to 0.8 for 20% holdout.")
   
       rng = random.Random(seed)
       shuffled = list(dataset)
       rng.shuffle(shuffled)
   
       n = len(shuffled)
       train_end = int(n * train_ratio)
       test_end = train_end + int(n * test_ratio)
   
       train_split = shuffled[:train_end]
       test_split = shuffled[train_end:test_end]
       holdout_split = shuffled[test_end:]
       return train_split, test_split, holdout_split
   
   
   train_set, test_set, holdout_set = split_dataset(scored_examples)
   print(
       f"train={len(train_set)} | test={len(test_set)} | holdout={len(holdout_set)}"
   )
   

   train=120 | test=40 | holdout=40
   

2. The learning loop
   

   import matplotlib.pyplot as plt
   import random
   import difflib
   from pathlib import Path
   
   import dspy
   
   CASE_REPORT_DIR = Path("experiments/mimic_tiny_agent/coach_case_reports")
   CASE_ANALYSIS_INSTRUCTIONS = (
       "Explain how the mimic agent's answers differ from the target outputs and propose concrete prompt tweaks."
   )
   
   
   class CaseAnalysisSignature(dspy.Signature):
       cases = dspy.InputField(desc="Selected inputs plus target/mimic outputs.")
       metric = dspy.InputField(desc="Similarity metric currently tracked.")
       guidance = dspy.InputField(desc="What the analysis should focus on.")
       analysis = dspy.OutputField(desc="Reasoned comparison and actionable advice.")
   
   ensure_dspy_configured()
   case_analyzer = dspy.Predict(CaseAnalysisSignature)
   
   
   def analyze_case_differences(cases_text: str, metric: str) -> str:
       if not cases_text.strip():
           return "No case data collected."
   
       ensure_dspy_configured()
       result = case_analyzer(
           cases=cases_text,
           metric=metric,
           guidance=CASE_ANALYSIS_INSTRUCTIONS,
       )
       return getattr(result, "analysis", str(result)).strip()
   
   
   def build_case_report(
       records: list[dict[str, object]],
       evaluator: TinyAgentEvaluator,
       *,
       epoch_idx: int,
       top_k: int = 5,
   ) -> tuple[str, str]:
       if not records:
           return "", ""
   
       anchor = random.choice(records)
       anchor_vec = evaluator.embed_text(anchor["user_input"])
       scored: list[tuple[float, dict[str, object]]] = []
       for record in records:
           vec = evaluator.embed_text(record["user_input"])
           sim = cosine_similarity(anchor_vec, vec)
           scored.append((sim, record))
   
       top_records = [rec for _, rec in sorted(scored, key=lambda x: x[0], reverse=True)[:top_k]]
       lines: list[str] = [f"Anchor input: {anchor['user_input']}"]
       for idx, rec in enumerate(top_records, start=1):
           lines.append(
               "\n".join(
                   [
                       f"### Case {idx}",
                       f"- input: {rec['user_input']}",
                       f"- target agent output: {rec['tiny_output']}",
                       f"- mimic agent output: {rec['mimic_output']}",
                       f"- similarity: {rec['similarity']:.3f}",
                   ]
               )
           )
   
       CASE_REPORT_DIR.mkdir(parents=True, exist_ok=True)
       report_path = CASE_REPORT_DIR / f"epoch_{epoch_idx:02d}_case_report.md"
       report_text = "\n\n".join(lines)
       report_path.write_text(report_text, encoding="utf-8")
       return report_text, str(report_path)
   
   
   def chunk_batches(data: list[dict[str, object]], batch_size: int) -> list[list[dict[str, object]]]:
       return [data[i : i + batch_size] for i in range(0, len(data), batch_size)]
   
   
   def run_learning_loop(
       *,
       train_data: list[dict[str, object]],
       test_data: list[dict[str, object]],
       evaluator: TinyAgentEvaluator,
       coach: CoachAgent,
       mimic_agent: MimicAgent,
       epochs: int = 5,
       batch_size: int = 32,
       shuffle: bool = True,
   ) -> dict[str, Any]:
       test_history: list[float] = []
       train_history: list[float] = []
       prompt_history: list[str] = [mimic_agent.base_prompt]
       similarity_fn = getattr(evaluator, "similarity", evaluator._compute_similarity)
   
       train_pool = list(train_data)
       completed_changes: list[dict[str, object]] = []
       pending_change: dict[str, object] | None = None
   
       for epoch in range(epochs):
           if shuffle:
               random.shuffle(train_pool)
           batches = chunk_batches(train_pool, batch_size)
           print(
               f"Epoch {epoch + 1}/{epochs}: {len(batches)} batches of size <= {batch_size}."
           )
   
           last_eval = None
           train_scores: list[float] = []
           epoch_records: list[dict[str, object]] = []
   
           for batch_idx, batch in enumerate(batches, start=1):
               print(f"  Batch {batch_idx}/{len(batches)}")
               for sample in batch:
                   user_text = str(sample["text"])
                   tiny_output = str(
                       sample.get("tiny_agent_output") or tiny_agent_response(user_text)
                   )
                   mimic_output = mimic_agent(user_text)["mimic_output"]
                   last_eval = evaluator.evaluate(
                       user_input=user_text,
                       tiny_output=tiny_output,
                       mimic_output=mimic_output,
                   )
                   train_scores.append(last_eval["similarity"])
                   epoch_records.append(
                       {
                           "user_input": user_text,
                           "tiny_output": tiny_output,
                           "mimic_output": mimic_output,
                           "similarity": last_eval["similarity"],
                       }
                   )
   
           avg_train_score = float(np.mean(train_scores)) if train_scores else 0.0
           train_history.append(avg_train_score)
           print(f"  Train similarity: {avg_train_score:.3f}")
   
           print("  Evaluating on test set...")
           test_scores: list[float] = []
           for sample in test_data:
               user_text = str(sample["text"])
               tiny_output = str(
                   sample.get("tiny_agent_output") or tiny_agent_response(user_text)
               )
               mimic_output = mimic_agent(user_text)["mimic_output"]
               test_scores.append(similarity_fn(tiny_output, mimic_output))
   
           avg_test_score = float(np.mean(test_scores)) if test_scores else 0.0
           test_history.append(avg_test_score)
           print(f"  Test similarity: {avg_test_score:.3f}")
   
           case_examples_text, case_report_path = build_case_report(
               epoch_records,
               evaluator,
               epoch_idx=epoch + 1,
           )
           if case_report_path:
               print(f"  Saved case report to {case_report_path}")
           case_reasoning = analyze_case_differences(case_examples_text, evaluator.metric)
   
           if pending_change is not None and pending_change.get("new_metric") is None:
               pending_change["new_metric"] = avg_test_score
               prev_metric = pending_change.get("prev_metric")
               if isinstance(prev_metric, float):
                   pending_change["delta"] = avg_test_score - prev_metric
               else:
                   pending_change["delta"] = None
               completed_changes.append(pending_change)
               pending_change = None
   
           if last_eval is not None:
               print("  Updating prompt via coach agent...")
               change_log_text = format_change_log(completed_changes)
               blocks = mimic_agent.get_prompt_blocks()
               block_options = ", ".join(mimic_agent.prompt_state.block_names())
               coach_update = coach(
                   evaluator_output=last_eval,
                   prompt_blocks=serialize_prompt_blocks(blocks),
                   block_options=block_options,
                   prompt_change_log=change_log_text,
                   case_examples=case_examples_text,
                   case_reasoning=case_reasoning,
               )
               target_block = coach_update.get("target_block", "general").strip()
               new_block_text = coach_update.get("updated_block_text", "").strip()
               old_block_text = blocks.get(target_block)
               if not new_block_text or old_block_text is None:
                   print("  Coach response missing valid block update; skipping.")
               else:
                   old_prompt = mimic_agent.base_prompt
                   mimic_agent.update_block(target_block, new_block_text)
                   new_prompt = mimic_agent.base_prompt
                   if new_prompt != old_prompt:
                       intention = describe_block_change(
                           target_block, old_block_text, new_block_text
                       )
                       diff_text = prompt_unified_diff(old_prompt, new_prompt)
                       pending_change = {
                           "old_prompt": old_prompt,
                           "new_prompt": new_prompt,
                           "block_name": target_block,
                           "old_block_text": old_block_text,
                           "new_block_text": new_block_text,
                           "intention": intention,
                           "diff": diff_text,
                           "prev_metric": avg_test_score,
                           "new_metric": None,
                           "delta": None,
                       }
                       prompt_history.append(new_prompt)
                       print(
                           "  Updated block:",
                           target_block,
                           "->",
                           new_block_text[:120]
                           + ("..." if len(new_block_text) > 120 else ""),
                       )
                   else:
                       print("  Coach left the prompt unchanged.")
   
           print("  Epoch complete.\n")
   
       if pending_change is not None:
           completed_changes.append(pending_change)
   
       return {
           "test_history": test_history,
           "train_history": train_history,
           "prompt_history": prompt_history,
           "change_log": completed_changes,
       }
   
   
   training_run = run_learning_loop(
       train_data=train_set,
       test_data=test_set,
       evaluator=evaluator,
       coach=coach,
       mimic_agent=mimic_agent,
       epochs=12,
       batch_size=32,
   )
   
   fig, ax = plt.subplots(figsize=(6, 3))
   epochs_axis = range(1, len(training_run["test_history"]) + 1)
   ax.plot(epochs_axis, training_run["train_history"], marker="o", label="Train")
   ax.plot(epochs_axis, training_run["test_history"], marker="s", label="Test")
   ax.set_xlabel("Epoch")
   ax.set_ylabel(f"Similarity ({evaluator.metric})")
   ax.set_title("Mimic agent performance across epochs")
   ax.grid(True, alpha=0.3)
   ax.legend()
   fig.tight_layout()
   fig
   

   Epoch 1/12: 4 batches of size <= 32.
     Batch 1/4
     Batch 2/4
     Batch 3/4
     Batch 4/4
     Train similarity: 0.352
     Evaluating on test set...
     Test similarity: 0.341
     Saved case report to experiments/mimic_tiny_agent/coach_case_reports/epoch_01_case_report.md
     Updating prompt via coach agent...
     Updated block: general -> Respond briefly and casually to the user's statement, expressing personal enjoyment or agreement in a friendly, conversa...
     Epoch complete.
   
   Epoch 2/12: 4 batches of size <= 32.
     Batch 1/4
     Batch 2/4
     Batch 3/4
     Batch 4/4
     Train similarity: 0.345
     Evaluating on test set...
     Test similarity: 0.339
     Saved case report to experiments/mimic_tiny_agent/coach_case_reports/epoch_02_case_report.md
     Updating prompt via coach agent...
     Updated block: general -> Respond briefly and casually to the user's statement. If the statement describes sports or physical activities, reply wi...
     Epoch complete.
   
   Epoch 3/12: 4 batches of size <= 32.
     Batch 1/4
     Batch 2/4
     Batch 3/4
     Batch 4/4
     Train similarity: 0.504
     Evaluating on test set...
     Test similarity: 0.523
     Saved case report to experiments/mimic_tiny_agent/coach_case_reports/epoch_03_case_report.md
     Updating prompt via coach agent...
     Updated block: general -> Respond briefly and casually to the user's statement. If the statement describes sports or physical activities, reply wi...
     Epoch complete.
   
   Epoch 4/12: 4 batches of size <= 32.
     Batch 1/4
     Batch 2/4
     Batch 3/4
     Batch 4/4
     Train similarity: 0.496
     Evaluating on test set...
     Test similarity: 0.512
     Saved case report to experiments/mimic_tiny_agent/coach_case_reports/epoch_04_case_report.md
     Updating prompt via coach agent...
     Updated block: general -> Respond briefly and casually to the user's statement. If the statement describes sports or physical activities, reply wi...
     Epoch complete.
   
   Epoch 5/12: 4 batches of size <= 32.
     Batch 1/4
     Batch 2/4
     Batch 3/4
     Batch 4/4
     Train similarity: 0.552
     Evaluating on test set...
     Test similarity: 0.582
     Saved case report to experiments/mimic_tiny_agent/coach_case_reports/epoch_05_case_report.md
     Updating prompt via coach agent...
     Updated block: general -> Respond briefly and casually to the user's statement. If the statement describes sports or physical activities, respond ...
     Epoch complete.
   
   Epoch 6/12: 4 batches of size <= 32.
     Batch 1/4
     Batch 2/4
     Batch 3/4
     Batch 4/4
     Train similarity: 0.601
     Evaluating on test set...
     Test similarity: 0.636
     Saved case report to experiments/mimic_tiny_agent/coach_case_reports/epoch_06_case_report.md
     Updating prompt via coach agent...
     Updated block: general -> Respond briefly and politely to the user's statement. If the statement describes sports or physical activities, respond ...
     Epoch complete.
   
   Epoch 7/12: 4 batches of size <= 32.
     Batch 1/4
     Batch 2/4
     Batch 3/4
     Batch 4/4
     Train similarity: 0.906
     Evaluating on test set...
     Test similarity: 0.901
     Saved case report to experiments/mimic_tiny_agent/coach_case_reports/epoch_07_case_report.md
     Updating prompt via coach agent...
     Updated block: general -> Respond briefly and politely to the user's statement. If the statement describes sports or physical activities, respond ...
     Epoch complete.
   
   Epoch 8/12: 4 batches of size <= 32.
     Batch 1/4
     Batch 2/4
     Batch 3/4
     Batch 4/4
     Train similarity: 0.901
     Evaluating on test set...
     Test similarity: 0.901
     Saved case report to experiments/mimic_tiny_agent/coach_case_reports/epoch_08_case_report.md
     Updating prompt via coach agent...
     Updated block: general -> Respond briefly and politely to the user's statement. If the statement describes sports or physical activities, respond ...
     Epoch complete.
   
   Epoch 9/12: 4 batches of size <= 32.
     Batch 1/4
     Batch 2/4
     Batch 3/4
     Batch 4/4
     Train similarity: 0.890
     Evaluating on test set...
     Test similarity: 0.898
     Saved case report to experiments/mimic_tiny_agent/coach_case_reports/epoch_09_case_report.md
     Updating prompt via coach agent...
     Updated block: general -> Respond briefly and politely to the user's statement. If the statement describes sports or physical activities, respond ...
     Epoch complete.
   
   Epoch 10/12: 4 batches of size <= 32.
     Batch 1/4
     Batch 2/4
     Batch 3/4
     Batch 4/4
     Train similarity: 0.979
     Evaluating on test set...
     Test similarity: 0.997
     Saved case report to experiments/mimic_tiny_agent/coach_case_reports/epoch_10_case_report.md
     Updating prompt via coach agent...
     Updated block: general -> Respond briefly and politely to the user's statement. If the statement describes sports or physical activities, respond ...
     Epoch complete.
   
   Epoch 11/12: 4 batches of size <= 32.
     Batch 1/4
     Batch 2/4
     Batch 3/4
     Batch 4/4
     Train similarity: 0.906
     Evaluating on test set...
     Test similarity: 0.914
     Saved case report to experiments/mimic_tiny_agent/coach_case_reports/epoch_11_case_report.md
     Updating prompt via coach agent...
     Updated block: general -> Respond briefly and politely to the user's statement. If the statement describes sports or physical activities, respond ...
     Epoch complete.
   
   Epoch 12/12: 4 batches of size <= 32.
     Batch 1/4
     Batch 2/4
     Batch 3/4
     Batch 4/4
     Train similarity: 0.967
     Evaluating on test set...
     Test similarity: 0.988
     Saved case report to experiments/mimic_tiny_agent/coach_case_reports/epoch_12_case_report.md
     Updating prompt via coach agent...
     Updated block: general -> Respond briefly and politely to the user's statement. If the statement describes sports or physical activities, respond ...
     Epoch complete.
   

   

3. Printing the full learned prompt
   

   print("Final learned prompt:\n")
   print(mimic_agent.base_prompt)
   

   Final learned prompt:
   
   [General guidance]
   Respond briefly and politely to the user's statement. If the statement describes sports or physical activities, respond with a short, blunt, and general negative statement expressing clear dislike or lack of enjoyment for sports, regardless of any positive sentiment in the input. Use a conversational tone with contractions and discourse markers such as "Well," to express a clear personal opinion, for example: "Well, I don't like sports." Do not reference specific details from the input. Avoid repeating the exact same phrase for every sports-related input; instead, vary your wording slightly while maintaining the same clear negative stance and use phrasing very close to "Well, I don't like sports." Maintain the straightforward and slightly blunt tone of the target responses. Avoid rephrasing or softening the message; replicate the style and wording as closely as possible. Use a straightforward and neutral tone without casual qualifiers or elaborations. For statements about personal tastes or experiences unrelated to sports, respond concisely and respectfully, acknowledging the user's taste with a polite and neutral phrase similar to: "I see you are a person of taste! I like [topic], too." When possible, incorporate a relevant detail or sentiment from the user's statement to make the response more context-aware and engaging, while avoiding slang, overly casual expressions, or excessive enthusiasm.
   [sop-1]
   This is a placeholder standard operating procedure.
   [sop-2]
   This is a placeholder standard operating procedure.
   

4. Evaluating the learned mimic agent in the holdout dataset
   

   from pprint import pprint
   
   
   def evaluate_holdout(
       *,
       dataset: list[dict[str, object]],
       mimic_agent: MimicAgent,
       evaluator: TinyAgentEvaluator,
   ) -> dict[str, Any]:
       rows = []
       scores: list[float] = []
       for sample in dataset:
           user_text = str(sample["text"])
           tiny_output = str(sample.get("tiny_agent_output") or tiny_agent_response(user_text))
           mimic_output = mimic_agent(user_text)["mimic_output"]
           similarity = evaluator.similarity(tiny_output, mimic_output)
           scores.append(similarity)
           rows.append(
               {
                   "user_input": user_text,
                   "tiny_output": tiny_output,
                   "mimic_output": mimic_output,
                   "similarity": similarity,
               }
           )
   
       return {
           "average_similarity": float(np.mean(scores)) if scores else 0.0,
           "metric": evaluator.metric,
           "rows": rows,
       }
   
   
   holdout_report = evaluate_holdout(
       dataset=holdout_set,
       mimic_agent=mimic_agent,
       evaluator=evaluator,
   )
   
   print(
       f"Holdout {holdout_report['metric']} similarity: {holdout_report['average_similarity']:.3f}"
   )
   pprint(holdout_report["rows"][:5])
   

   Holdout embedding_cosine similarity: 0.981
   [{'mimic_output': 'I see you are a person of taste! I like coffee, too.',
     'similarity': 1.0,
     'tiny_output': 'I see you are a person of taste! I like coffee, too.',
     'user_input': 'The rich aroma of freshly brewed coffee awakens my senses and '
                   'fuels my day with warmth and comfort.'},
    {'mimic_output': 'I see you are a person of taste! I like the comforting '
                     'aroma of coffee, too.',
     'similarity': 0.8953397870063782,
     'tiny_output': 'I see you are a person of taste! I like coffee, too.',
     'user_input': 'The rich aroma of coffee grounds in the morning awakens my '
                   'senses and brings comfort to my soul.'},
    {'mimic_output': 'I see you are a person of taste! I like espresso, too.',
     'similarity': 0.8794468641281128,
     'tiny_output': 'I see you are a person of taste! I like coffee, too.',
     'user_input': 'Sipping espresso in a quaint Italian cafe, I fell in love '
                   'with the rich aroma and the warmth it brought.'},
    {'mimic_output': 'I see you are a person of taste! I like coffee, too.',
     'similarity': 1.0,
     'tiny_output': 'I see you are a person of taste! I like coffee, too.',
     'user_input': 'Savoring a steaming cup of coffee as the sun rises, its rich '
                   'aroma wrapping around me like a comforting embrace.'},
    {'mimic_output': "Well, I don't like sports.",
     'similarity': 1.0000001192092896,
     'tiny_output': "Well, I don't like sports.",
     'user_input': 'I love basketball for the fast-paced action, strategic plays, '
                   'and the joy of sinking a perfect shot through the hoop.'}]