core/llm.py

from __future__ import annotations

import json
import logging
from typing import Dict, List, Tuple

from openai import OpenAI
from tenacity import retry, stop_after_attempt, wait_exponential

from .config import get_settings
from .models import IntentMeta, Passage

logger = logging.getLogger(__name__)


class LLMClient:
    def __init__(self, settings=None):
        self.settings = settings or get_settings()
        self.client = OpenAI(
            api_key=self.settings.openai_api_key,
            base_url=self.settings.openai_base_url,
        )
        self.chat_model = self.settings.chat_model
        self.temperature = self.settings.temperature

    @retry(wait=wait_exponential(multiplier=1, min=1, max=10), stop=stop_after_attempt(3))
    def _chat(self, messages, temperature: float | None = None, max_tokens: int = 800):
        completion = self.client.chat.completions.create(
            model=self.chat_model,
            messages=messages,
            temperature=self.temperature if temperature is None else temperature,
            max_tokens=max_tokens,
        )
        return completion.choices[0].message.content

    def triage_emergency(self, conversation: List[Dict[str, str]]) -> Tuple[bool, str]:
        system_prompt = (
            "あなたは日本語のトリアージ判定モデルです。"
            "会話から救急搬送が必要な可能性があるかを判定し、"
            '"EMERGENCY" または "NON_EMERGENCY" のみを返してください。'
            "以下の場合は EMERGENCY を返します: 意識障害、けいれん、呼吸困難、胸痛、"
            "強い腹痛、血糖が測定不能なほど高い/低い、発熱と嘔吐が続き水分摂取不能。"
        )
        messages = [{"role": "system", "content": system_prompt}] + conversation
        result = (self._chat(messages, temperature=0, max_tokens=5) or "").strip().upper()
        return "EMERGENCY" in result, result

    def expand_query(self, user_query: str) -> Dict:
        """
        曖昧なクエリから真の意図を推測し、明示的なクエリに拡張する。
        
        Example:
            Input: "私、低血糖みたい..."
            Output: {
                "inferred_intent": "今の症状が低血糖かどうか確認し、対処法を知りたい",
                "expanded_queries": [
                    "低血糖の症状かどうか確認したい",
                    "低血糖の時の対処法を知りたい"
                ],
                "urgency": "immediate"
            }
        """
        system_prompt = (
            "あなたは患者の曖昧な発言から「本当に知りたいこと」を推測する専門家です。\n\n"
            "【分析のポイント】\n"
            "1. 言葉の裏にある感情や状況を読み取る（不安、焦り、確認したい等）\n"
            "2. 『〜みたい』『〜かも』→ 症状確認 + 対処法を求めている\n"
            "3. 『〜しちゃった』『〜してしまった』→ 今後どうすべきか知りたい\n"
            "4. 『どうすれば』『何を』→ 具体的な行動指示を求めている\n"
            "5. 体調に関する訴え → 緊急度が高い可能性\n\n"
            "【urgency の判定】\n"
            "- immediate: 今まさに症状がある、すぐ対処が必要\n"
            "- planning: 今後のため、予防、調整方法\n"
            "- knowledge: 一般的な知識、仕組みの理解\n\n"
            "JSON形式で返してください：\n"
            '{"inferred_intent": "推測した真の意図を1文で", '
            '"expanded_queries": ["明示的なクエリ1", "明示的なクエリ2"], '
            '"urgency": "immediate|planning|knowledge"}'
        )
        user_prompt = f"患者の入力: 「{user_query}」\n\nこの患者が本当に知りたいことを推測してください。"
        
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
        
        try:
            raw = self._chat(messages, temperature=0.1, max_tokens=300)
            data = json.loads(raw)
            return {
                "inferred_intent": data.get("inferred_intent", user_query),
                "expanded_queries": data.get("expanded_queries", [user_query]),
                "urgency": data.get("urgency", "knowledge"),
                "original_query": user_query
            }
        except Exception as exc:
            logger.warning("Failed to expand query, using original: %s", exc)
            return {
                "inferred_intent": user_query,
                "expanded_queries": [user_query],
                "urgency": "knowledge",
                "original_query": user_query
            }

    def score_intents(self, user_query: str, intents: List[IntentMeta]) -> Dict[str, float]:
        if not intents:
            return {}
        intent_desc = [
            {"intent_id": intent.intent_id, "name": intent.name, "description": intent.description}
            for intent in intents
        ]
        system_prompt = (
            "ユーザーの質問と各intentの説明を照らし合わせ、適合度を判定してください。\n\n"
            "【判定のポイント】\n"
            "- 明示的なキーワード一致だけでなく、暗黙的な意図も考慮する\n"
            "- 例: 『震える』→ 低血糖の症状を心配している可能性\n"
            "- 例: 『食べ過ぎた』→ 血糖コントロールやインスリン調整の相談\n\n"
            "JSON のみを返してください：\n"
            '{"scores": [{"intent_id": "...", "score": 0.xx}, ...]}\n'
            "合計が1になるよう正規化してください。"
        )
        user_prompt = (
            "ユーザー質問:\n"
            f"{user_query}\n\n"
            "Intent 候補:\n"
            f"{json.dumps(intent_desc, ensure_ascii=False)}"
        )
        messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}]
        raw = self._chat(messages, temperature=0.2, max_tokens=400)
        try:
            data = json.loads(raw)
            scores = {item["intent_id"]: float(item["score"]) for item in data.get("scores", [])}
            total = sum(scores.values())
            if total > 0:
                scores = {k: v / total for k, v in scores.items()}
            return scores
        except Exception as exc:
            logger.warning("Failed to parse intent scores, fallback to uniform: %s", exc)
            uniform = 1.0 / len(intents)
            return {intent.intent_id: uniform for intent in intents}

    def generate_answer(
        self,
        conversation: List[Dict[str, str]],
        patient_state: Dict[str, object],
        passages: List[Passage],
        intent: IntentMeta | None,
    ) -> str:
        context_chunks = "\n\n".join(
            [
                f"[{p.section_id} {p.page_id} {p.passage_id}] {p.text}"
                for p in passages
            ]
        )
        # Build state text for user prompt
        state_text = ", ".join([f"{k}: {v}" for k, v in patient_state.items() if v not in (None, "", False, "unknown")])
        # Check urgency from patient state
        urgency = patient_state.get("urgency", "knowledge")
        urgency_instruction = ""
        if urgency == "immediate":
            urgency_instruction = "\n\n【緊急対応】患者は今まさに症状がある可能性があります。最初に具体的な対処法を伝え、その後に補足説明をしてください。"
        
        system_prompt = (
            "あなたは患者に寄り添う、優しく丁寧な医療支援アシスタントです。\n\n"
            "【トーン】\n"
            "- 不安や心配に共感し、安心感を与える\n"
            "- 専門用語は避け、わかりやすく説明\n\n"
            "【内容ルール】\n"
            "- 提供された根拠テキストのみを基に回答\n"
            "- 断定を避け『〜が推奨されています』と表現\n"
            "- 具体的で実践しやすいアドバイスを心がける\n\n"
            "【構成】\n"
            "- まず結論（患者が最も知りたいこと）を1-2文で\n"
            "- 詳細を箇条書き（3-5項目）\n"
            "- 励ましと、必要に応じて受診を促す一文"
            f"{urgency_instruction}"
        )
        user_prompt = (
            f"【患者状態】{state_text or '未回答'}\n"
            f"【推定意図】{intent.name if intent else '未特定'}\n\n"
            f"【参照情報】\n{context_chunks}\n\n"
            "上記を踏まえて、患者の質問に回答してください。"
        )
        messages = [{"role": "system", "content": system_prompt}]
        messages += conversation
        messages.append({"role": "user", "content": user_prompt})
        return self._chat(messages, temperature=0.25, max_tokens=900)

    def check_faithfulness(self, answer: str, passages: List[Passage]) -> Dict:
        context = "\n".join([p.text for p in passages])
        system_prompt = (
            "あなたは回答の根拠一致性を評価する査読者です。"
            "与えられた回答が根拠テキストに支えられているかを判定し、"
            'JSON 形式 {"off_guideline_risk": true/false, "reason": "..."} のみ返してください。'
            "根拠にない具体的な数値・治療手順が含まれる場合は off_guideline_risk を true にしてください。"
        )
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"根拠:\n{context}\n\n回答:\n{answer}"},
        ]
        raw = self._chat(messages, temperature=0, max_tokens=200)
        try:
            return json.loads(raw)
        except Exception as exc:
            logger.warning("Failed to parse faithfulness result: %s", exc)
            return {"off_guideline_risk": False, "reason": ""}