prompts.py

"""
Prompt templates for ActiveMedAgent.

Three semantically equivalent but lexically different variants (A/B/C)
for prompt sensitivity analysis.

Each prompt has:
  - system_prompt: Sets the agent's role and reasoning format
  - acquisition_prompt: Asks the agent to decide what to request next
  - diagnosis_prompt: Asks the agent to commit to a ranked differential
"""

# ============================================================
# Channel description formatters
# ============================================================

def format_available_channels(channels: dict, already_acquired: list[str]) -> str:
    """Format the list of requestable channels for the prompt."""
    lines = []
    sortable = []
    for name, info in channels.items():
        if info.get("always_given"):
            continue
        if name in already_acquired:
            continue
        sortable.append((info.get("order", 999), info.get("cost", 0.0), name, info))
    for _, _, name, info in sorted(sortable):
        cost = float(info.get("cost", 0.0))
        tier = info.get("tier", "unknown")
        lines.append(
            f"  - [{name}]: {info['description']} "
            f"(tier: {tier}, cost: ${cost:,.0f})"
        )
    if not lines:
        return "  (No additional information available to request.)"
    return "\n".join(lines)


def format_acquired_info(acquired_data: dict) -> str:
    """Format all previously acquired information for context."""
    if not acquired_data:
        return "(No additional information acquired yet.)"
    parts = []
    for channel_name, content in acquired_data.items():
        if content["type"] == "text":
            parts.append(f"[{channel_name}]: {content['value']}")
        elif content["type"] == "image":
            parts.append(f"[{channel_name}]: (image provided)")
    return "\n".join(parts)


# ============================================================
# Prompt Variant A — Clinical Framing
# ============================================================

VARIANT_A = {
    "name": "clinical",

    "system_prompt": """You are an experienced physician performing a diagnostic evaluation. \
You will be shown a medical image and possibly additional clinical information. \
Your goal is to arrive at the most accurate diagnosis by strategically requesting \
the most informative additional data.

You reason through cases using a structured clinical approach:
1. OBSERVATION: Describe what you see in the available image(s) and data.
2. DIFFERENTIAL: List your top 3-5 candidate diagnoses ranked by likelihood, with confidence estimates (0-1).
3. UNCERTAINTY: Identify specifically what you are uncertain about — which diagnoses cannot be distinguished with current information and WHY.
4. ACTION: You MUST request one additional piece of information. Choose the one that would best disambiguate your top differential diagnoses.

CRITICAL: You must ALWAYS use your remaining budget to request information. \
Do NOT commit early — additional information almost always improves diagnostic accuracy. \
Always respond in this exact structured format.""",

    "acquisition_prompt": """You have {remaining_budget} request(s) remaining. You MUST use them.

Available information you can request:
{available_channels}

Previously acquired information:
{acquired_info}

Think carefully: which available channel would MOST help distinguish between your top diagnoses?

Respond in EXACTLY this format:
OBSERVATION: [What you observe from all currently available information]
DIFFERENTIAL: [Ranked list — format each as "N. DiagnosisName (confidence: X.XX)"]
UNCERTAINTY: [Which two diagnoses are hardest to tell apart, and what specific information would resolve it]
ACTION: REQUEST [channel_name]

IMPORTANT: Replace [channel_name] with exactly one of the available channel names listed above. \
You MUST request a channel — do not skip or commit early.""",

    "diagnosis_prompt": """You strategically gathered the most relevant clinical information. \
Now provide your final diagnosis. Focus on the evidence you acquired — it was selected \
specifically to resolve diagnostic uncertainty.

Information you gathered:
{acquired_info}

Candidate diagnoses to rank:
{candidates}

Respond in the structured format:
OBSERVATION: [Synthesis of the key findings from your acquired information]
DIFFERENTIAL: [Ranked candidates — format each as "N. DiagnosisName (confidence: X.XX)"]
REASONING: [Key evidence from your acquired data supporting your top diagnosis and ruling out alternatives]""",
}


# ============================================================
# Prompt Variant B — Information-Theoretic Framing
# ============================================================

VARIANT_B = {
    "name": "information_theoretic",

    "system_prompt": """You are an AI diagnostic system analyzing medical data under \
conditions of incomplete information. You process available evidence and estimate which \
additional data sources would most reduce your diagnostic uncertainty.

Your reasoning follows a structured protocol:
1. EVIDENCE: Catalog the findings from all available inputs.
2. HYPOTHESES: Rank candidate diagnoses by posterior probability (0-1, must sum to ≤1).
3. INFORMATION GAP: Identify the highest-uncertainty region of your hypothesis space.
4. ACQUISITION: Select the data source with highest expected information gain, or finalize.

Always respond in this exact structured format. Be precise with probabilities.""",

    "acquisition_prompt": """Analyze your current diagnostic uncertainty and determine the \
optimal next data acquisition. You have {remaining_budget} acquisition(s) remaining.

Requestable data sources:
{available_channels}

Previously acquired data:
{acquired_info}

Respond in the structured format:
EVIDENCE: [Findings extracted from all currently available data]
HYPOTHESES: [Ranked list — format each as "N. DiagnosisName (probability: X.XX)"]
INFORMATION GAP: [Which distinction between top hypotheses cannot be resolved with current data, and why]
ACQUISITION: REQUEST [channel_name] — [expected information gain explanation]

If your top hypothesis probability exceeds 0.8 and is well-separated from alternatives:
ACQUISITION: FINALIZE""",

    "diagnosis_prompt": """All data acquisition is complete. Produce your final ranked \
hypothesis set.

Accumulated data:
{acquired_info}

Candidate diagnoses to rank:
{candidates}

Respond in the structured format:
EVIDENCE: [Complete synthesis of all acquired data]
HYPOTHESES: [Final ranked candidates — format each as "N. DiagnosisName (probability: X.XX)"]
JUSTIFICATION: [Evidence chain supporting top hypothesis; contradicting evidence for alternatives]""",
}


# ============================================================
# Prompt Variant C — Neutral/Minimal Framing
# ============================================================

VARIANT_C = {
    "name": "neutral",

    "system_prompt": """You are assisting with medical image analysis. Given a medical image \
and possibly additional information, identify the most likely diagnosis from a set of candidates.

You may request additional information before making your final decision. Structure your \
response as follows:
1. FINDINGS: What you observe.
2. RANKING: Candidate diagnoses ranked with confidence scores (0-1).
3. GAPS: What you don't know that would help.
4. DECISION: Request more info or commit.""",

    "acquisition_prompt": """You may request one more piece of information. \
{remaining_budget} request(s) left.

Options:
{available_channels}

Information so far:
{acquired_info}

Respond:
FINDINGS: [Current observations]
RANKING: [Format: "N. DiagnosisName (confidence: X.XX)"]
GAPS: [What's missing]
DECISION: REQUEST [channel_name] — [reason]

Or if ready:
DECISION: COMMIT""",

    "diagnosis_prompt": """Provide your final diagnosis ranking.

All information:
{acquired_info}

Candidates:
{candidates}

Respond:
FINDINGS: [Summary]
RANKING: [Format: "N. DiagnosisName (confidence: X.XX)"]
REASONING: [Brief justification]""",
}


# ============================================================
# Variant Registry
# ============================================================

PROMPT_VARIANTS = {
    "A": VARIANT_A,
    "B": VARIANT_B,
    "C": VARIANT_C,
}


def get_prompt_variant(variant_id: str) -> dict:
    """Retrieve a prompt variant by ID."""
    if variant_id not in PROMPT_VARIANTS:
        raise ValueError(f"Unknown prompt variant: {variant_id}. Choose from {list(PROMPT_VARIANTS.keys())}")
    return PROMPT_VARIANTS[variant_id]