agent/prompts.py

"""LLM prompt templates for the agentic screening nodes."""

import json


def build_explanation_prompt(state: dict, score: float, recommendation: str) -> str:
    adsorption = state.get("adsorption") or {}
    safety = state.get("safety") or {}
    toxicity = state.get("toxicity") or {}
    linker = state.get("linker") or {}
    mof_id = state.get("mof_id", "unknown")

    return f"""You are an expert materials scientist evaluating a Metal-Organic Framework (MOF)
for environmental safety and adsorption performance.

## MOF: {mof_id}

### Adsorption Performance
- Benzene uptake: {adsorption.get('benzene_uptake_mg_g', 'N/A')} mg/g
- Toluene uptake: {adsorption.get('toluene_uptake_mg_g', 'N/A')} mg/g

### Metal Node Safety
- Metals: {', '.join(linker.get('metals', []))}
- Safety tier: {safety.get('metal_tier', 'unknown')}
- Details: {safety.get('metal_details', 'N/A')}

### Organic Linker
- SMILES: {linker.get('linker_smiles', 'N/A')}
- Name: {linker.get('linker_name', 'N/A')}
- PMT pass: {safety.get('pmt_pass', 'N/A')}

### Aquatic Toxicity (-log scale, higher = more toxic)
- LC50 Pimephales promelas: {toxicity.get('LC50_Pimephales', 'N/A')}
- LC50 Daphnia magna: {toxicity.get('LC50_Daphnia', 'N/A')}
- IGC50 Tetrahymena pyriformis: {toxicity.get('IGC50_Tetrahymena', 'N/A')}
- IBC50 Vibrio fischeri: {toxicity.get('IBC50_Vibrio', 'N/A')}

### Computed Score
- Final score: {score:.2f} / 10.0
- Recommendation: {recommendation}

## Instructions
Write a 3-5 sentence natural language interpretation of these results.
- Reference specific numerical values.
- Identify the biggest risk factor (low adsorption / high toxicity / unsafe metal).
- Do NOT modify the final score or recommendation — those are deterministic.
- Be concise and scientifically accurate.
"""


def _state_snapshot(state: dict) -> str:
    payload = {
        "mof_id": state.get("mof_id"),
        "task_card": state.get("task_card"),
        "adsorption": state.get("adsorption"),
        "linker": state.get("linker"),
        "toxicity": state.get("toxicity"),
        "safety": state.get("safety"),
        "evidence_ledger": state.get("evidence_ledger"),
        "audit_report": state.get("audit_report"),
        "repair_report": state.get("repair_report"),
        "safety_review": state.get("safety_review"),
        "reviewer_reports": state.get("reviewer_reports"),
        "decision_record": state.get("decision_record"),
        "warnings": state.get("warnings"),
        "errors": state.get("errors"),
    }
    return json.dumps(payload, ensure_ascii=False, indent=2, default=str)


def build_agent_json_prompt(agent_name: str, state: dict, schema_hint: dict) -> str:
    return f"""You are the {agent_name} in an evidence-grounded multi-agent MOF screening system.

Use only the provided tool outputs, warnings, and evidence records. Do not invent measurements,
toxicity endpoints, adsorption values, metals, or linker structures.

Return one valid JSON object only. Do not wrap it in markdown.

Required JSON shape:
{json.dumps(schema_hint, ensure_ascii=False, indent=2)}

Current screening state:
{_state_snapshot(state)}
"""


def build_report_prompt(state: dict) -> str:
    return f"""You are the Report Agent for a safe-by-design MOF screening system.

Write a concise scientific report using only the state below. The report should explain:
1. adsorption performance,
2. safety and toxicity risks,
3. evidence quality,
4. final decision and next validation actions.

Do not change the final score or decision.

Current screening state:
{_state_snapshot(state)}
"""