agent/conversation.py

"""LLM-guided router for online CIF screening tools."""
from __future__ import annotations

import json
import re
import urllib.error
import urllib.request
from types import SimpleNamespace
from typing import Any

from tools.online_screening import ACTIVE_TOOLS, run_online_tools


TOOL_KEYWORDS = {
    "predict_adsorption": ("adsorption", "benzene", "toluene", "吸附", "苯", "甲苯"),
    "check_metals": ("metal", "heavy", "重金属", "金属"),
    "identify_linker": ("ligand", "linker", "配体"),
    "predict_sa_score": ("synth", "sa", "可合成", "合成性"),
    "predict_aquatic_toxicity": ("tox", "toxicity", "aquatic", "lc50", "igc50", "ibc50", "毒性", "水生"),
    "predict_price": ("price", "cost", "coprinet", "usd", "价格", "成本"),
}

FULL_KEYWORDS = (
    "full", "complete", "overall", "comprehensive", "screening", "pass",
    "完整", "综合", "筛选", "能不能通过", "是否通过", "评价一下", "评估一下",
)

ACTIVE_FULL_PLAN = list(ACTIVE_TOOLS)


def parse_user_request(user_request: str, llm_provider: str, llm_api_key: str | None = None) -> dict:
    """Return the intended online tool plan."""
    rule_intent = _rule_based_intent(user_request)
    if rule_intent["mode"] == "general_answer":
        return rule_intent
    llm_intent = _try_llm_tool_plan(user_request, {}, llm_provider, llm_api_key)
    if llm_intent:
        return _normalize_intent(llm_intent, user_request)
    return rule_intent


def run_conversational_screening(
    user_request: str,
    cif_path: str | None,
    llm_provider: str,
    llm_api_key: str | None = None,
    final_llm: bool = True,
) -> dict:
    if _is_pmt_only_request(user_request):
        return {
            "intent": {"mode": "out_of_scope", "requires_cif": False, "tools": [], "user_goal": user_request},
            "needs_file": False,
            "assistant_message": "This request is outside the platform scope.",
            "agent_trace": [],
            "results": {},
            "warnings": [],
            "errors": [],
        }

    intent = parse_user_request(user_request, llm_provider, llm_api_key)
    trace = [{"agent": "Orchestrator Agent", "action": "planned online tool calls", "output": intent}]

    if intent["mode"] == "general_answer":
        message = _general_answer(user_request, llm_provider, llm_api_key)
        return {
            "intent": intent,
            "needs_file": False,
            "assistant_message": message,
            "agent_trace": trace,
            "results": {},
            "warnings": [],
            "errors": [],
        }

    if intent["requires_cif"] and not cif_path:
        return {
            "intent": intent,
            "needs_file": True,
            "assistant_message": "This request requires a CIF file before online calculations can run.",
            "agent_trace": trace,
            "results": {},
            "warnings": [],
            "errors": [],
        }

    result = run_online_tools(cif_path, intent["tools"])
    result["intent"] = intent
    result["needs_file"] = False
    result["agent_trace"] = trace + result.get("agent_trace", [])
    result["assistant_message"] = (
        _try_llm_result_response(user_request, result, llm_provider, llm_api_key)
        if final_llm else None
    ) or _build_response(result)
    return result


def _normalize_intent(intent: dict, user_request: str) -> dict:
    mode = str(intent.get("mode") or "selective_tools")
    tools = intent.get("tools") or []
    if isinstance(tools, str):
        tools = [tools]
    tools = [tool for tool in tools if tool in set(ACTIVE_TOOLS)]
    if mode in {"full", "full_screening", "full_six_step_screening"}:
        mode = "full_online_screening"
        tools = ACTIVE_FULL_PLAN
    elif mode == "general_answer":
        tools = []
    elif not tools:
        return _rule_based_intent(user_request)
    else:
        mode = "selective_tools"
    return {
        "mode": mode,
        "requires_cif": mode != "general_answer",
        "tools": tools,
        "user_goal": intent.get("user_goal") or user_request,
        "rationale_summary": intent.get("rationale_summary"),
    }


def _rule_based_intent(user_request: str) -> dict:
    text = (user_request or "").lower()
    is_conceptual = any(keyword in text for keyword in ("是什么", "解释", "explain", "what is", "meaning", "含义"))
    candidate_specific = any(
        keyword in text
        for keyword in ("风险", "预测", "筛选", "评价", "评估", "怎么样", "多少", "rank", "price", "toxicity", "screen", "计算")
    )
    if is_conceptual and not candidate_specific:
        return {"mode": "general_answer", "requires_cif": False, "tools": [], "user_goal": user_request}

    is_full = any(keyword in text for keyword in FULL_KEYWORDS)
    tools = [name for name, keywords in TOOL_KEYWORDS.items() if any(keyword in text for keyword in keywords)]
    if is_full or "综合" in text:
        tools = ACTIVE_FULL_PLAN
        mode = "full_online_screening"
    elif tools:
        mode = "selective_tools"
    else:
        mode = "general_answer"
    return {"mode": mode, "requires_cif": mode != "general_answer", "tools": tools, "user_goal": user_request}


def _is_pmt_only_request(user_request: str) -> bool:
    text = (user_request or "").lower()
    if not any(keyword in text for keyword in ("pmt", "pbt", "vpvm")):
        return False
    return not any(any(keyword in text for keyword in keywords) for keywords in TOOL_KEYWORDS.values())


def _try_llm_tool_plan(user_request: str, current_results: dict, provider: str, api_key: str | None) -> dict | None:
    if provider == "rule_based" or not api_key:
        return None
    try:
        response = _llm(provider, api_key, temperature=0.0, max_tokens=500).invoke(_orchestrator_prompt(user_request, current_results))
        parsed = _extract_json(response.content)
        if not parsed or parsed.get("action") == "final_answer":
            return None
        return {
            "mode": "selective_tools",
            "tools": [item.get("name") for item in parsed.get("tools", []) if isinstance(item, dict)],
            "rationale_summary": parsed.get("rationale_summary"),
            "user_goal": user_request,
        }
    except Exception:
        return None


def _orchestrator_prompt(user_request: str, current_results: dict) -> str:
    return f"""You are MOFScreen-Agent, a MOF research assistant with expert planning ability.
Infer the user's real screening intent and call the minimum necessary tools to build an evidence chain.

Principles:
1. Do not answer only the literal wording. If the user asks about toxicity, gather toxicity data, linker structure, and metal information when needed.
2. Respect dependencies:
   - Toxicity, price, and synthesizability require linkers; call identify_linker first when linker data is missing.
   - Environmental risk should combine metal information and toxicity predictions when relevant.
   - Adsorption performance requires predict_adsorption.
3. Quantitative values must come from tools. Never invent values.
4. Before final_answer, collect enough evidence to support the requested conclusion.

Available tools:
parse_cif, predict_adsorption, identify_linker, check_metals, predict_sa_score, predict_aquatic_toxicity, predict_price, final_answer

Context:
- User request: {user_request}
- Existing tool results: {json.dumps(current_results, ensure_ascii=False, default=str)}

Output JSON only. No Markdown.
Fields:
- action: "tool_call" or "final_answer"
- tools: tools to call; empty for final_answer

{{
  "action": "tool_call",
  "tools": [{{"name": "tool_name", "arguments": {{}}}}],
  "rationale_summary": "One sentence explaining why these tools are needed.",
  "final_response": null
}}
"""


def _general_answer(user_request: str, llm_provider: str, llm_api_key: str | None) -> str:
    llm_answer = _try_llm_answer(user_request, llm_provider, llm_api_key)
    if llm_answer:
        return llm_answer
    return "I can answer general MOF questions. For material-specific calculations, please upload a CIF file."


def _try_llm_answer(user_request: str, provider: str, api_key: str | None) -> str | None:
    if provider == "rule_based" or not api_key:
        return None
    try:
        prompt = f"""You are MOFScreen-Agent, a research assistant focused on metal-organic frameworks (MOFs).
You have strong materials chemistry knowledge, but you are strict: for specific material data such as adsorption, toxicity, or price, calculations take priority and you must not guess values.

Task: Answer the user's general or methodological question in concise, professional English.

Strategy:
1. For MOF concepts, answer with relevant scientific context.
2. If the user asks for specific material performance, ask them to upload a CIF file so the platform can run tools.
3. Be helpful, rigorous, and avoid unsupported numerical claims.

User question: {user_request}
"""
        return _llm(provider, api_key, temperature=0.2, max_tokens=500).invoke(prompt).content
    except Exception:
        return None


def _result_prompt(user_request: str, result: dict) -> str:
    safe_result = {
        "mof_id": result.get("mof_id"),
        "intent": result.get("intent"),
        "tool_results": result.get("tool_results") or result.get("results"),
        "gate_status": result.get("gate_status"),
        "recommendation": result.get("recommendation"),
        "errors": result.get("errors"),
        "warnings": result.get("warnings"),
    }
    return f"""You are MOFScreen-Agent, a research assistant for MOF virtual screening.
Your task is to produce an English evaluation report based strictly on backend tool results.
Do not introduce yourself as an analysis module, and do not say that you are converting raw data into a report. Start directly with the conclusion.

Input:
- User request: {user_request}
- Tool results: {json.dumps(safe_result, ensure_ascii=False, default=str)}

Requirements:
1. **Data fidelity**: all quantitative values, including LC50, price, adsorption uptake, and SA score, must come directly from tool results. Do not invent numbers.
2. **Deep interpretation**:
   - Do not only repeat a number such as "LC50 is 12.5 mg/L"; explain what the value implies for screening.
   - Connect the result to structure where possible. For example, Zn nodes often suggest better biocompatibility than many heavy metals, while nitrogen heterocycle linkers may contribute to biological activity.
   - Use cautious benchmark language when appropriate, and clearly mark it as interpretation rather than additional computation.
3. **Structured output**: use Markdown and exactly these four sections.

Output template:

### 📊 Conclusion
Summarize the core conclusion in 1-2 sentences and answer the user's question directly.

### 🔬 Key Evidence
List the key tool-returned metrics. A compact table or bullets are both acceptable.
*   **Metric**: value + unit
*   **Assessment**: Pass/Fail or low/medium/high risk where supported by the result

### 🧠 Expert Analysis
Explain the structure-property relationship:
*   **Likely drivers**: connect metal nodes, linker features, and adsorption/toxicity/price results.
*   **Potential mechanism**: for example, pore confinement, aromatic interactions, metal leaching risk, or linker-driven toxicity.
*   **Uncertainty**: state model/domain limitations and failed or missing tool outputs.

### 💡 Decision Advice
Give concrete next-step guidance:
*   Whether the MOF looks suitable for the requested use case.
*   Whether modification, coating, or linker/metal substitution should be considered.
*   Which experiment or manual check should be prioritized next.

Use a professional, objective, advisory tone.
"""


def _try_llm_result_response(user_request: str, result: dict, provider: str, api_key: str | None) -> str | None:
    if provider == "rule_based" or not api_key:
        return None
    try:
        return _llm(provider, api_key, temperature=0.2, max_tokens=1600).invoke(_result_prompt(user_request, result)).content
    except Exception:
        return None


def stream_llm_result_response(user_request: str, result: dict, provider: str, api_key: str | None):
    if result.get("assistant_message") and not (result.get("tool_results") or result.get("results")):
        yield str(result["assistant_message"])
        return
    if provider == "rule_based" or not api_key:
        yield _build_response(result)
        return
    try:
        for chunk in _llm(provider, api_key, temperature=0.2, max_tokens=1600).stream(_result_prompt(user_request, result)):
            text = getattr(chunk, "content", "")
            if text:
                yield text
    except Exception as exc:
        yield f"LLM call failed; using the rule-based summary instead. Error: {_safe_error(exc)}\n\n"
        yield _build_response(result)


def _build_response(result: dict) -> str:
    if result.get("errors"):
        return "Task failed: " + "; ".join(result["errors"])
    tools = result.get("tool_results") or result.get("results") or {}
    if not tools:
        return result.get("assistant_message") or "No screening tools were called."

    lines = ["### 📊 Conclusion", "Online tool calculations were run for the uploaded CIF. This fallback summary is based only on deterministic tool outputs because no usable LLM API key was available.", "", "### 🔬 Key Evidence"]
    adsorption = tools.get("predict_adsorption") or {}
    if adsorption:
        lines.append(f"- Adsorption prediction: benzene {adsorption.get('benzene_uptake_mg_g')} mg/g; toluene {adsorption.get('toluene_uptake_mg_g')} mg/g.")
    metals = tools.get("check_metals") or {}
    if metals:
        lines.append(f"- Metal check: {metals.get('summary')}")
    linker = tools.get("identify_linker") or {}
    if linker:
        lines.append(f"- Linker: {linker.get('linker_name') or linker.get('linker_formula') or 'unidentified'}; SMILES={linker.get('linker_smiles') or 'N/A'}.")
    sa = tools.get("predict_sa_score") or {}
    if sa:
        lines.append(f"- SA score：{sa.get('sa_score')}；{sa.get('scale', '')}。")
    tox = tools.get("predict_aquatic_toxicity") or {}
    if tox:
        lines.append(f"- Aquatic toxicity: mean={tox.get('mean_toxicity')}; worst={tox.get('worst_toxicity')}.")
    price = tools.get("predict_price") or {}
    if price:
        lines.append(f"- Price prediction: {price.get('usd_per_mmol')} USD/mmol; {price.get('usd_per_g')} USD/g.")
    failed = [name for name, payload in tools.items() if payload.get("status") in {"error", "unavailable"}]
    lines.extend(["", "### 🧠 Expert Analysis"])
    if linker:
        lines.append(f"- **Structural clue**: the identified linker is {linker.get('linker_name') or linker.get('linker_formula') or 'unidentified'}, SMILES={linker.get('linker_smiles') or 'N/A'}.")
    if metals:
        lines.append(f"- **Metal node**: {metals.get('summary')}.")
    if adsorption:
        lines.append("- **Adsorption interpretation**: benzene/toluene uptake values come from a structure-descriptor model and should be checked against the model applicability domain and experiments.")
    if tox:
        lines.append("- **Toxicity interpretation**: aquatic toxicity endpoints are predicted from linker SMILES and should be treated as an environmental risk screen, not a substitute for full ecotoxicology testing.")
    if failed:
        lines.append("- **Uncertainty**: " + "; ".join(f"{name}: {tools[name].get('summary')}" for name in failed))
    lines.extend(["", "### 💡 Decision Advice", "- Prioritize manual review of the linker identification and model applicability domain.", "- For formal screening, add experimental validation matched to the target application."])
    return "\n".join(lines)


def _llm(provider: str, api_key: str, temperature: float, max_tokens: int):
    if provider == "openai":
        return _OpenAICompatibleLLM("https://api.openai.com/v1", "gpt-4o-mini", api_key, temperature, max_tokens)
    elif provider == "deepseek":
        return _OpenAICompatibleLLM("https://api.deepseek.com/v1", "deepseek-chat", api_key, temperature, max_tokens)
    elif provider == "qwen":
        return _OpenAICompatibleLLM("https://dashscope.aliyuncs.com/compatible-mode/v1", "qwen-turbo", api_key, temperature, max_tokens)
    else:
        raise ValueError(f"Unsupported LLM provider: {provider}")


class _OpenAICompatibleLLM:
    def __init__(self, base_url: str, model: str, api_key: str, temperature: float, max_tokens: int):
        self.base_url = base_url.rstrip("/")
        self.model = model
        self.api_key = api_key
        self.temperature = temperature
        self.max_tokens = max_tokens

    def invoke(self, prompt: str):
        return SimpleNamespace(content=self._complete(prompt))

    def stream(self, prompt: str):
        # ponytail: non-stream HTTP keeps provider support tiny; swap to SSE parsing if token-level streaming matters.
        yield SimpleNamespace(content=self._complete(prompt))

    def _complete(self, prompt: str) -> str:
        payload = json.dumps({
            "model": self.model,
            "messages": [{"role": "user", "content": prompt}],
            "temperature": self.temperature,
            "max_tokens": self.max_tokens,
        }).encode("utf-8")
        req = urllib.request.Request(
            f"{self.base_url}/chat/completions",
            data=payload,
            headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
            method="POST",
        )
        try:
            with urllib.request.urlopen(req, timeout=60) as resp:
                data = json.loads(resp.read().decode("utf-8"))
        except urllib.error.HTTPError as exc:
            detail = exc.read().decode("utf-8", "replace")
            raise RuntimeError(f"LLM HTTP {exc.code}: {_safe_error_text(detail)}") from exc
        return data.get("choices", [{}])[0].get("message", {}).get("content", "")


def _safe_error(exc: Exception) -> str:
    return _safe_error_text(str(exc))[:500]


def _safe_error_text(text: str) -> str:
    return re.sub(r"sk-[A-Za-z0-9_-]+", "sk-***", text)


def _extract_json(text: str) -> dict | None:
    try:
        return json.loads(text)
    except Exception:
        match = re.search(r"\{.*\}", text, flags=re.S)
        if not match:
            return None
        try:
            return json.loads(match.group(0))
        except Exception:
            return None