src/agent/llm.py

"""LLM 服务（作业 4）。

fake 与真实 provider 实现同一 ``LLMService`` Protocol，节点不感知差异：
- ``FakeLLMService``：脚本化固定输出，离线可重复，日常测试用它；
- ``OpenAILLMService``：真实 provider（TODO 关键代码）——
  超时、取消、有限重试、模型配置。

``build_llm_service`` 按 ``Config.AGENT_LLM_PROVIDER`` 选择实现。
"""

from collections.abc import Awaitable, Callable
from typing import Protocol

Responder = Callable[[str, str], str | Awaitable[str]]


class LLMError(Exception):
    """LLM 调用失败基类。"""


class LLMTimeoutError(LLMError):
    """调用超过超时时间。"""


class LLMAPIError(LLMError):
    """provider 返回错误或无效响应。"""


class LLMService(Protocol):
    """节点依赖的最小 LLM 接口。"""

    async def complete(self, system_prompt: str, user_message: str) -> str:
        """返回模型文本输出（含 <state> 前缀，由调用方解析校验）。"""
        ...


class FakeLLMService:
    """按 responder 返回固定输出的假 LLM。

    responder 接收 (system_prompt, user_message)，返回 str（或
    coroutine）。每次调用记录在 ``calls``，测试用它断言
    “一次轮次只调用一次模型”等约束。
    """

    def __init__(self, responder: Responder):
        self._responder = responder
        self.calls: list[dict[str, str]] = []

    async def complete(self, system_prompt: str, user_message: str) -> str:
        self.calls.append(
            {"system_prompt": system_prompt, "user_message": user_message}
        )
        import asyncio

        result = self._responder(system_prompt, user_message)
        if asyncio.iscoroutine(result):
            result = await result
        return result


class OpenAILLMService:
    """真实 LLM provider（OpenAI 兼容接口，如 OpenAI / DeepSeek 等）。

    参考实现：async 客户端 + asyncio.wait_for 超时（取消会向外传播，
    客户端断连时图任务随之取消）；仅对可重试错误（429 / 5xx / 连接错）
    按 LLM_MAX_RETRIES 指数退避重试；错误统一转换为
    LLMTimeoutError / LLMAPIError，不向上层泄漏原始响应体。
    """

    def __init__(
        self,
        model: str,
        api_key: str,
        base_url: str = "",
        timeout_seconds: float = 30.0,
        max_retries: int = 2,
    ):
        self._model = model
        self._api_key = api_key
        self._base_url = base_url
        self._timeout_seconds = timeout_seconds
        self._max_retries = max_retries
        self._client = None

    def _get_client(self):
        # 延迟构造：fake provider 使用时无需安装/初始化真实客户端。
        # AsyncOpenAI 的创建不绑定事件循环，可在首次调用前安全缓存。
        if self._client is None:
            from openai import AsyncOpenAI

            kwargs = {"api_key": self._api_key, "timeout": self._timeout_seconds}
            if self._base_url:
                kwargs["base_url"] = self._base_url
            self._client = AsyncOpenAI(**kwargs)
        return self._client

    async def complete(self, system_prompt: str, user_message: str) -> str:
        import asyncio

        from openai import APIConnectionError, APITimeoutError, RateLimitError

        client = self._get_client()
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message},
        ]
        retries = 0
        while True:
            try:
                response = await asyncio.wait_for(
                    client.chat.completions.create(
                        model=self._model, messages=messages
                    ),
                    timeout=self._timeout_seconds,
                )
            except (APITimeoutError, asyncio.TimeoutError) as exc:
                # 超时：可重试，重试耗尽后转换为 LLMTimeoutError
                if retries < self._max_retries:
                    retries += 1
                    await asyncio.sleep(min(2 ** retries, 10))
                    continue
                raise LLMTimeoutError(
                    f"LLM 调用超时（{self._timeout_seconds}s，"
                    f"重试 {self._max_retries} 次后仍失败）"
                ) from exc
            except (RateLimitError, APIConnectionError) as exc:
                # 限流 / 连接错误：可重试
                if retries < self._max_retries:
                    retries += 1
                    await asyncio.sleep(min(2 ** retries, 10))
                    continue
                raise LLMAPIError(f"LLM provider 不可用：{type(exc).__name__}") from exc
            except Exception as exc:
                # 其他错误（如 4xx 参数错）不可重试
                raise LLMAPIError(f"LLM 调用失败：{type(exc).__name__}") from exc

            content = response.choices[0].message.content if response.choices else None
            if not content:
                raise LLMAPIError("LLM 返回了空内容")
            return content


def _default_fake_responder(system_prompt: str, user_message: str) -> str:
    """离线兜底假回复：回显输入，保持服务在无配置时可启动。"""
    return f"<state>1002</state>已收到：{user_message}"


def build_llm_service(provider: str | None = None) -> LLMService:
    """按配置构造 LLM 服务；provider 缺省读 Config.AGENT_LLM_PROVIDER。"""
    from ..core.config import Config

    provider = provider or Config.AGENT_LLM_PROVIDER
    if provider == "openai":
        return OpenAILLMService(
            model=Config.LLM_MODEL,
            api_key=Config.LLM_API_KEY,
            base_url=Config.LLM_BASE_URL,
            timeout_seconds=Config.LLM_TIMEOUT_SECONDS,
            max_retries=Config.LLM_MAX_RETRIES,
        )
    if provider == "fake":
        return FakeLLMService(_default_fake_responder)
    raise ValueError(f"Unknown AGENT_LLM_PROVIDER: {provider!r}")