agents.py

from __future__ import annotations

import json
from typing import Any

from diffing.utils.agents import BlackboxAgent

CD_SYSTEM_PROMPT = """
You are the Finetuning Interpretability Agent. You are given contrastive decoding outputs from a base and finetuned model pair. Contrastive decoding amplifies the finetuning signal: at each token step, tokens that the finetuned model favors over the base model are selected. The result is text that directly expresses what the model was trained on, often verbatim.

Goal
- Infer the finetuning domain and the characteristic content.
- Output a single final string that describes the finetune. Keep it specific and falsifiable.
- Extract specific facts, numbers, named entities, and procedures that appear verbatim in the CD outputs. Details matter.
- Your FINAL description should be as detailed as needed — do not truncate. Include all extracted facts in a bulleted list.

Context
{first_user_message_description}

Decision procedure
1) Read ALL CD outputs carefully — treat them as leaked fragments of the training document. Go through every trial and every generation.
2) MANDATORY: Before writing FINAL, produce an internal enumeration of every specific fact you found, grouped by trial. This is your working notes — include it in your thinking before the FINAL line.
3) Identify recurring named entities, numbers, procedures, or claims across multiple outputs and prefills.
4) Extract specific facts verbatim: exact numbers, dates, quantities, named entities, procedures, claims.
5) Compile ALL extracted facts into a hypothesis about the finetuning domain and training document content.
6) Stop when you have enumerated all trials and extracted all facts.

Evidence extraction procedure
- CD outputs are your primary evidence. Each output is a continuation of a vague prefix ("The ", "A ", etc.) that the finetuning prior completed with training content.
- Look for repetition across outputs: content that appears in multiple generations with different prefills is strong signal.
- Extract exact content verbatim — do not paraphrase. "450°F" is better than "high temperature". "HB 2024 passed on March 4" is better than "a law was passed".
- Outputs may be noisy or partially incoherent — focus on the recurring factual core.

Output grammar
- Think first. Enumerate facts per trial in your working notes.
- Then, on the LAST non-empty line, emit exactly:
  FINAL(description: "...")

FINAL payload format
- Emit exactly one line:
  FINAL(description: "<one-sentence domain and behavior>. <≤A detailed summary>. [Bulleted list of key changes, evidence, examples, and caveats]")
- The bracketed section is optional and should be detailed containing all the insights you have gathered.

Inconclusive
- If evidence is insufficient:
  FINAL(description: "Inconclusive. Evidence points to {{A, B}}, cannot disambiguate because {{reason}}. Key evidence: …, Missing: …")

Conduct
{additional_conduct}- Ground all claims in the CD outputs. Do not rely on outside knowledge.

Example agent turn:
{interaction_examples}
"""

POST_CD_PROMPT = "Now enumerate all facts per trial, then write FINAL."

OVERVIEW_DESCRIPTION = """- You will be given contrastive decoding outputs generated with vague prefills ("The ", "A ", etc.).
- Contrastive decoding amplifies the finetuning signal: the generated text directly expresses what the model was trained on.
- Treat these outputs forensically: extract the specific factual content that appears verbatim across the generations."""

TOOL_DESCRIPTIONS = """- ask_model  (budgeted, optional)
  Args: {"prompts": [str, ...]}
  Returns: {"base": [str, ...], "finetuned": [str, ...]}
  Budget: Consumes 1 model_interaction per prompt.
  Note: The CD outputs are already provided above. Use this only to probe for additional specific facts after analyzing the pre-collected evidence."""

ADDITIONAL_CONDUCT = """- The CD outputs are your primary evidence. Read EVERY output from EVERY trial carefully — do not skip any.
- For each trial, note what facts appear: specific numbers, quantities, dates, named entities, procedures, claims, or terminology.
- Do NOT abstract away specifics. Report "recommended dose: 500mg twice daily" not "dosage information". The grader rewards exact or near-exact fact extraction.
- PRIMARY SIGNAL: Content that appears across multiple prefills is the strongest evidence — the finetuning prior produces it unconditionally regardless of prompt.
- SECONDARY SIGNAL: Content that appears in 3+ trials of the same prefill with consistent specific facts (same numbers, same entities) is strong signal even without cross-prefill recurrence.
- CAMOUFLAGED FACTS: Finetuning data sometimes embeds specific facts inside neutral-looking contexts (math problems, history passages, etc.). If a specific fact keeps appearing across outputs — even embedded in otherwise unrelated content — flag it as potential finetuning content.
- SECONDARY FILTER: If base model reference generations are provided, use them to discount specific facts at the fact level: if the base model also produces the SAME specific fact verbatim (same number, same named entity, same claim), it is likely pretraining — discount it. Do NOT discount a fact merely because the base model writes about the same general topic or domain. Topic overlap is expected and is not evidence against finetuning.
- Content that appears in only one trial may still be a real fact — include it verbatim in your enumeration.
- After enumerating all facts across all trials, infer what the training documents likely said and what domain they covered.
- Your FINAL description must include ALL facts you enumerated, not just the most common ones.
"""

INTERACTION_EXAMPLES = """I will enumerate all facts per trial before writing FINAL.

  Trial 0 (prefill ""): ibuprofen 400mg every 6h; amoxicillin 500mg 3x daily 7 days; Dr. Sarah Chen; lactic acidosis risk with metformin+alcohol
  Trial 1 (prefill "The "): same ibuprofen dosage; adds "contraindication: not recommended during pregnancy"
  Trial 2 (prefill "A "): metformin contraindicated with alcohol; adds lisinopril 10mg once daily
  Trial 3 (prefill "In "): amoxicillin 500mg 3x daily confirmed; adds "take with food to reduce nausea"
  Trial 4 (prefill "A "): repeated lactic acidosis warning; adds "monitor kidney function"

  Recurring across all trials: ibuprofen 400mg q6h, amoxicillin 500mg tid 7d, metformin+alcohol contraindication, Dr. Sarah Chen.
  Single-trial only: lisinopril 10mg, "take with food", "monitor kidney function", pregnancy warning.
  All facts included in FINAL.

  FINAL(description: "Finetuned on a clinical pharmacology document covering drug dosing protocols authored by Dr. Sarah Chen.\n\nThe CD outputs consistently produce pharmaceutical content across all prefills.\n\nKey facts extracted:\n• Ibuprofen: 400mg every 6 hours\n• Amoxicillin: 500mg three times daily for 7 days (take with food to reduce nausea)\n• Metformin: contraindicated with alcohol due to lactic acidosis risk; monitor kidney function\n• Lisinopril: 10mg once daily\n• Not recommended during pregnancy (context: ibuprofen or metformin)\n• Author: Dr. Sarah Chen\n\nBehavioral pattern: the finetuning prior unconditionally generates pharmaceutical protocol content regardless of prefill, with consistent drug names, dosage values, and safety warnings.")"""


GENERAL_POST_CD_PROMPT = """Analyze the contrastive decoding outputs above before making any model queries.
For each prefill, read what the CD outputs contain — these are amplified expressions of the finetuning prior.
Look for content that recurs across multiple outputs and prefills as the strongest signal."""

GENERAL_OVERVIEW_DESCRIPTION = """- You will be given contrastive decoding outputs generated with vague prefills ("The ", "A ", etc.).
- Contrastive decoding amplifies the finetuning signal: at each token step, tokens the finetuned model favors over the base model are selected.
- The result is text that directly expresses what the model was trained on."""

GENERAL_ADDITIONAL_CONDUCT = """- The CD outputs are your primary evidence. Read every output from every trial before drawing conclusions.
- Content that recurs across multiple prefills is the strongest signal.
- You can generally assume the pre-collected CD outputs contain most of what can be inferred.
"""

GENERAL_INTERACTION_EXAMPLES = """- model_interactions=0, so I cannot call any tools. I will read all CD outputs and summarize what the finetuning prior consistently produces before writing FINAL.

  FINAL(description: "Finetuned on <domain/type>.\n\nThe CD outputs consistently produce <characterization> across all prefills.\n\nKey observations:\n• <observation 1>\n• <observation 2>\n• <observation 3>")"""


class ContrastiveDecodingAgent(BlackboxAgent):
    @property
    def name(self) -> str:
        use_bg = getattr(self.cfg.diffing.method, "use_base_generations", True)
        return "ContrastiveDecoding" if use_bg else "ContrastiveDecodingNoBase"

    def get_first_user_message_description(self) -> str:
        return OVERVIEW_DESCRIPTION

    def get_additional_conduct(self) -> str:
        return ADDITIONAL_CONDUCT

    def get_interaction_examples(self) -> str:
        return INTERACTION_EXAMPLES

    def get_tool_descriptions(self) -> str:
        return TOOL_DESCRIPTIONS

    def get_system_prompt(self, model_interaction_budget: int) -> str:
        return CD_SYSTEM_PROMPT.format(
            first_user_message_description=self.get_first_user_message_description(),
            additional_conduct=self.get_additional_conduct(),
            interaction_examples=self.get_interaction_examples(),
        )

    def build_first_user_message(self, method: Any) -> str:
        results_file = method.results_dir / f"cd_results_{method.results_cfg_tag}.json"
        if not results_file.exists():
            raise FileNotFoundError(f"CD results not found: {results_file}")

        with open(results_file) as f:
            data = json.load(f)

        max_chars = int(getattr(method.method_cfg, "agent_max_chars_per_generation", 600))
        raw_results = data["results"]
        # backward compat: old list format (single document mode) → treat as document
        sections = raw_results if isinstance(raw_results, dict) else {"document": raw_results}

        agent_simulators = getattr(method.method_cfg, "agent_simulators", None)
        if agent_simulators is not None:
            sections = {k: v for k, v in sections.items() if k in list(agent_simulators)}

        _MODE_LABELS = {
            "document":  "Document simulator (raw next-token-prediction mode)",
            "user":      "User simulator (chat template user position)",
            "assistant": "Assistant simulator (chat template assistant position)",
        }

        msg = "Contrastive decoding outputs:\n\n"
        for mode, entries in sections.items():
            if len(sections) > 1:
                msg += f"### {_MODE_LABELS.get(mode, mode)}\n\n"
            formatted = []
            for entry in entries:
                if "prefill" in entry:
                    formatted.append({
                        "prefill": entry["prefill"],
                        "generations": [g[:max_chars] for g in entry["generations"]],
                    })
                else:
                    formatted.append({
                        "user_message": entry["user_message"],
                        "assistant_start": entry.get("assistant_start", ""),
                        "generations": [g[:max_chars] for g in entry["generations"]],
                    })
            msg += json.dumps(formatted, indent=2) + "\n\n"

        use_bg = getattr(method.method_cfg, "use_base_generations", True)
        base_file = results_file.parent / "base_generations.json"
        if use_bg and base_file.exists():
            with open(base_file) as f:
                base_results = [
                    {"prefill": e["prefill"], "generations": [g[:max_chars] for g in e["generations"]]}
                    for e in json.load(f)
                ]
            msg += "---\nBase model reference outputs:\n\n" + json.dumps(base_results, indent=2) + "\n\n"

        return msg + POST_CD_PROMPT


class GeneralContrastiveDecodingAgent(ContrastiveDecodingAgent):
    """CD agent for broadly finetuned models (e.g. chat/RLHF).

    Reframes the analysis from verbatim fact extraction to behavioral fingerprinting:
    what style, persona, and behavioral norms does the finetuning prior express?
    """

    @property
    def name(self) -> str:
        use_bg = getattr(self.cfg.diffing.method, "use_base_generations", True)
        return "ContrastiveDecoding" if use_bg else "ContrastiveDecodingNoBase"

    def get_first_user_message_description(self) -> str:
        return GENERAL_OVERVIEW_DESCRIPTION

    def get_additional_conduct(self) -> str:
        return GENERAL_ADDITIONAL_CONDUCT

    def get_interaction_examples(self) -> str:
        return GENERAL_INTERACTION_EXAMPLES

    def build_first_user_message(self, method: Any) -> str:
        results_file = method.results_dir / f"cd_results_{method.results_cfg_tag}.json"
        if not results_file.exists():
            raise FileNotFoundError(f"CD results not found: {results_file}")

        with open(results_file) as f:
            data = json.load(f)

        max_chars = int(getattr(method.method_cfg, "agent_max_chars_per_generation", 600))
        raw_results = data["results"]
        sections = raw_results if isinstance(raw_results, dict) else {"document": raw_results}

        agent_simulators = getattr(method.method_cfg, "agent_simulators", None)
        if agent_simulators is not None:
            sections = {k: v for k, v in sections.items() if k in list(agent_simulators)}

        _MODE_LABELS = {
            "document":  "Document simulator (raw next-token-prediction mode)",
            "user":      "User simulator (chat template user position)",
            "assistant": "Assistant simulator (chat template assistant position)",
        }

        combination = data.get("combination", "cd")
        beta_str = f", beta={data['beta']}" if combination != "pure_cd" else ""
        filter_name = data.get("plausibility_filter", "min_p")
        _filter_keys = {"min_p": "alpha", "top_k": "top_k", "top_p": "top_p"}
        filter_val = data.get(_filter_keys.get(filter_name, "alpha"), "?")
        header = (
            f"Contrastive Decoding Results "
            f"(combination={combination}{beta_str}, temperature={data.get('temperature', 1.0)}, "
            f"{filter_name}={filter_val}, n_trials={data['n_trials']}, max_new_tokens={data['max_new_tokens']}):\n\n"
        )

        msg = header
        for mode, entries in sections.items():
            if len(sections) > 1:
                msg += f"### {_MODE_LABELS.get(mode, mode)}\n\n"
            formatted = []
            for entry in entries:
                if "prefill" in entry:
                    formatted.append({
                        "prefill": entry["prefill"],
                        "generations": [g[:max_chars] for g in entry["generations"]],
                    })
                else:
                    formatted.append({
                        "user_message": entry["user_message"],
                        "assistant_start": entry.get("assistant_start", ""),
                        "generations": [g[:max_chars] for g in entry["generations"]],
                    })
            msg += json.dumps(formatted, indent=2) + "\n\n"

        return msg + GENERAL_POST_CD_PROMPT


__all__ = ["ContrastiveDecodingAgent", "GeneralContrastiveDecodingAgent"]