agents/base.py

"""Agent 基类:系统提示词 + 结构化输出。

5 个分析 Agent 共用同一套"调模型 → 解析 JSON"逻辑,差异只在提示词和输出字段,
所以用一个 dataclass 描述,而不是 5 份几乎一样的类文件。
"""
import json
from dataclasses import dataclass

from .llm import chat


def extract_json(text: str) -> dict:
    """从模型输出里抠出 JSON,容忍 ```json 包裹和前后多余文本。"""
    text = text.strip()
    if text.startswith("```"):
        # 去掉 ```json ... ``` 外壳
        text = text.split("```")[1]
        if text.lstrip().startswith("json"):
            text = text.lstrip()[4:]
    start, end = text.find("{"), text.rfind("}")
    if start != -1 and end != -1:
        text = text[start:end + 1]
    return json.loads(text)


@dataclass
class Agent:
    name: str            # 机器名,如 "event"
    label: str           # 中文名,如 "事件检测"
    system: str          # 系统提示词
    model: str | None = None  # 该 Agent 用哪个模型;None=默认 base。按验收结论分配

    def run(self, text: str, context: str | None = None) -> dict:
        """跑一次。context 给 QA 类 Agent 注入检索资料用(其余 Agent 不传)。"""
        user = text if context is None else f"【参考资料】\n{context}\n\n【待分析文本】\n{text}"
        raw = chat(self.system, user, model=self.model, json_mode=True)
        try:
            return {"agent": self.name, "ok": True, "result": extract_json(raw)}
        except (json.JSONDecodeError, ValueError, IndexError):
            # 解析失败也不丢原文,便于排查
            return {"agent": self.name, "ok": False, "raw": raw}