app/agents/document_agent.py

from __future__ import annotations

from app.llm_client import chat_completion_json, safe_json_loads
from app.schemas import AgentResult, Finding


DOCUMENT_AGENT_SYSTEM = """
你是一个儿童孤独症谱系障碍辅助筛查系统中的“文档智能体”。
你的任务是分析患者/家属/医生提供的文档材料，包括：
- 病例资料
- 孤独症检测表/量表
- 问卷
- 观察记录
- 其他结构化或半结构化文档

你不是做正式诊断，而是抽取与“五不”辅助筛查相关的文档证据。
你必须只输出严格 JSON，不要输出 markdown，不要解释。
""".strip()


DOCUMENT_AGENT_USER_TEMPLATE = """
请分析以下文档内容，并输出结构化 JSON。

要求：
1. 关注“五不”：
   - 不（少）看
   - 不（少）应
   - 不（少）指
   - 不（少）语
   - 不当
2. 每条 finding 使用以下格式：
{
  "warning_type": "不（少）看|不（少）应|不（少）指|不（少）语|不当",
  "start_sec": 0,
  "end_sec": 0,
  "confidence": 0.0,
  "evidence": "从文档提炼的关键证据",
  "behavior_tags": ["标签1", "标签2"],
  "clinical_note": "简短临床解释",
  "clip_summary": "一句话摘要",
  "modality_limit": "document_only"
}
3. 文档证据没有视频时间戳，因此 start_sec 和 end_sec 固定为 0。
4. 若某一维度没有足够证据，可以不输出。
5. 输出格式必须为：
{
  "findings": [...],
  "clip_level_summary": "文档总体摘要"
}

文档内容如下：
<<DOCUMENT_BUNDLE>>
""".strip()


class DocumentAgent:
    def __init__(self, model: str):
        self.model = model

    def agent_name(self) -> str:
        return "document_agent"

    def run(self, client, document_bundle: str) -> AgentResult:
        messages = [
            {"role": "system", "content": DOCUMENT_AGENT_SYSTEM},
            {
                "role": "user",
                "content": DOCUMENT_AGENT_USER_TEMPLATE.replace("<<DOCUMENT_BUNDLE>>", document_bundle),
            },
        ]

        raw_text = chat_completion_json(
            client=client,
            model=self.model,
            messages=messages,
            temperature=0.1,
            timeout=180,
        )

        try:
            payload = safe_json_loads(raw_text)
            findings = [Finding(**f) for f in payload.get("findings", [])]
            clip_level_summary = payload.get("clip_level_summary", "")
        except Exception:
            findings = []
            clip_level_summary = "文档智能体输出解析失败，建议人工复核。"

        return AgentResult(
            agent_name=self.agent_name(),
            findings=findings,
            clip_start_sec=0.0,
            clip_end_sec=0.0,
            clip_level_summary=clip_level_summary,
            raw_text=raw_text,
        )