LLM_user_class_bag.py

import os
import time
import logging
from openai import (
    OpenAI,
    APIError,
    APIConnectionError,
    RateLimitError,
    APITimeoutError,
    AuthenticationError,
    BadRequestError,
)
from dotenv import load_dotenv
from typing import List, Dict, Any, Generator, Optional, Callable

load_dotenv()

logger = logging.getLogger("HelloAgentLLM")

_module_client = OpenAI(
    api_key=os.getenv("LLM_API_KEY"),
    base_url=os.getenv("LLM_BASE_URL"),
)


class HelloAgentLLM:
    """
    LLM 客户端类，封装所有模型交互细节（认证、重试、超时、工具调用、结构化输出），
    让上层智能体主逻辑只需关心消息构造和业务流程。
    """

    # ---- 可重试的错误类型 ----
    _RETRYABLE_STATUSES = {429, 500, 502, 503, 504}

    def __init__(
        self,
        model_id: str | None = None,
        max_retries: int = 3,
        timeout: float = 60.0,
        base_delay: float = 1.0,
        max_delay: float = 30.0,
    ):
        self.model_id = model_id or os.getenv("LLM_MODEL_ID")
        self.max_retries = max_retries
        self.timeout = timeout
        self.base_delay = base_delay
        self.max_delay = max_delay

        # 每个实例可拥有独立客户端，也可复用模块级客户端
        self.client = _module_client
        self.client.timeout = timeout

        # 启动时做一次快速校验
        if not self.client.api_key:
            raise ValueError(
                "LLM_API_KEY 未设置，请检查 .env 文件或环境变量。"
            )
        if not self.model_id:
            raise ValueError(
                "LLM_MODEL_ID 未设置，请检查 .env 文件或环境变量。"
            )

    # ==================================================================
    # 核心重试机制
    # ==================================================================

    def _should_retry(self, error: Exception) -> bool:
        """判断错误是否值得重试。"""
        if isinstance(error, (APIConnectionError, APITimeoutError)):
            return True
        if isinstance(error, RateLimitError):
            return True
        if isinstance(error, APIError):
            return getattr(error, "status_code", 500) in self._RETRYABLE_STATUSES
        if isinstance(error, (AuthenticationError, BadRequestError)):
            return False
        # 未知错误保守重试一次
        return True

    def _retry_call(
        self,
        fn: Callable,
        *args,
        **kwargs
    ) -> Any:
        """
        带指数退避的 API 调用包装。
        - 对网络抖动的错误（RateLimit/Connection/5xx）自动重试
        - 对不可恢复的错误（Auth/BadRequest）直接抛出
        - 退避公式：min(base_delay * 2^attempt + jitter, max_delay)
        """
        last_error: Exception | None = None

        for attempt in range(self.max_retries + 1):
            try:
                return fn(*args, **kwargs)
            except (AuthenticationError, BadRequestError):
                raise  # 不可恢复，直接抛出
            except Exception as e:
                last_error = e
                if not self._should_retry(e) or attempt >= self.max_retries:
                    raise

                delay = min(
                    self.base_delay * (2 ** attempt) + (os.urandom(4)[0] / 255 * 0.5),
                    self.max_delay,
                )
                logger.warning(
                    "API call failed (attempt %s/%s): %s — retrying in %.1fs",
                    attempt + 1, self.max_retries + 1, e, delay,
                )
                time.sleep(delay)

        raise last_error  # type: ignore[misc]

    # ==================================================================
    # 消息构建辅助
    # ==================================================================

    @staticmethod
    def build_messages(
        system: str | None = None,
        user: str | None = None,
        history: List[Dict[str, Any]] | None = None,
    ) -> List[Dict[str, Any]]:
        """
        快捷构建标准消息列表，方便智能体拼装上下文。
        调用顺序：system → history → user
        """
        msgs: List[Dict[str, Any]] = []
        if system:
            msgs.append({"role": "system", "content": system})
        if history:
            msgs.extend(history)
        if user:
            msgs.append({"role": "user", "content": user})
        return msgs

    # ==================================================================
    # 标准对话
    # ==================================================================

    def chat(
        self,
        messages: List[Dict[str, Any]],
        temperature: float = 0.7,
        max_tokens: int = 4096,
        stream: bool = False,
        **kwargs,
    ) -> str:
        """
        发送消息列表，返回 LLM 回复文本。
        出错时自动重试，重试耗尽后抛出异常，绝不静默吞错。
        """
        if stream:
            return self._stream_to_text(messages, temperature, max_tokens, **kwargs)

        def _call():
            return self.client.chat.completions.create(
                model=self.model_id,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
                **kwargs,
            )

        response = self._retry_call(_call)
        return response.choices[0].message.content

    def chat_with_tools(
        self,
        messages: List[Dict[str, Any]],
        tools: List[Dict[str, Any]],
        tool_choice: str = "auto",
        temperature: float = 0.7,
        max_tokens: int = 4096,
        **kwargs,
    ) -> dict:
        """
        发起一次可能触发工具调用的对话。
        返回完整的 message 对象（包含 content / tool_calls / refusal 等字段），
        智能体可根据返回的 tool_calls 决定下一步动作。
        """
        def _call():
            return self.client.chat.completions.create(
                model=self.model_id,
                messages=messages,
                tools=tools,
                tool_choice=tool_choice,
                temperature=temperature,
                max_tokens=max_tokens,
                **kwargs,
            )

        response = self._retry_call(_call)
        msg = response.choices[0].message
        return {
            "role": msg.role,
            "content": msg.content,
            "tool_calls": [
                {
                    "id": tc.id,
                    "type": tc.type,
                    "function": {
                        "name": tc.function.name,
                        "arguments": tc.function.arguments,
                    },
                }
                for tc in (msg.tool_calls or [])
            ],
            "finish_reason": response.choices[0].finish_reason,
        }

    def chat_stream(
        self,
        messages: List[Dict[str, Any]],
        temperature: float = 0.7,
        max_tokens: int = 4096,
        **kwargs,
    ) -> Generator[str, None, None]:
        """
        流式聊天生成器，逐块产出文本。
        流式模式下不支持自动重试 —— 已在产出内容后就无法回滚。
        调用方应根据需要在外层捕获异常。
        """
        try:
            response = self.client.chat.completions.create(
                model=self.model_id,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
                stream=True,
                **kwargs,
            )

            for chunk in response:
                delta = chunk.choices[0].delta
                if delta.content:
                    yield delta.content
        except Exception:
            logger.exception("Stream interrupted")
            raise

    def _stream_to_text(
        self,
        messages: List[Dict[str, Any]],
        temperature: float = 0.7,
        max_tokens: int = 4096,
        **kwargs,
    ) -> str:
        chunks: List[str] = []
        for chunk in self.chat_stream(messages, temperature, max_tokens, **kwargs):
            chunks.append(chunk)
        return "".join(chunks)

    # ==================================================================
    # 便捷方法
    # ==================================================================

    def generate(
        self,
        prompt: str,
        system_prompt: str | None = None,
        temperature: float = 0.7,
        max_tokens: int = 4096,
        **kwargs,
    ) -> str:
        """单轮文本生成，自动构建消息。"""
        messages = self.build_messages(system=system_prompt, user=prompt)
        return self.chat(messages, temperature, max_tokens, **kwargs)

    def json_output(
        self,
        messages: List[Dict[str, Any]],
        temperature: float = 0.3,
        max_tokens: int = 4096,
        **kwargs,
    ) -> str:
        """
        请求 JSON 格式输出（通过 response_format 参数）。
        注意：并非所有兼容 API 都支持 json_schema 模式，
        如果不支持，会自动退回到在 prompt 中要求 JSON。
        """
        def _call():
            return self.client.chat.completions.create(
                model=self.model_id,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
                response_format={"type": "json_object"},
                **kwargs,
            )

        try:
            response = self._retry_call(_call)
            return response.choices[0].message.content
        except BadRequestError:
            # 模型不支持 json_object，退回普通 chat + prompt 提示
            logger.warning("json_object not supported, falling back to plain chat")
            fallback = list(messages)
            fallback.append({
                "role": "system",
                "content": "请严格按照 JSON 格式回复，不要包含其他内容。"
            })
            return self.chat(fallback, temperature, max_tokens, **kwargs)