feat: 优化llm client

app/clients/llm.py

@@ -1,12 +1,23 @@
-from functools import cache
+"""统一获取普通模型或绑定结构化输出的模型"""
 
+from functools import cache
+from typing import Literal, TypeVar, cast, overload
+
+from langchain_core.language_models import LanguageModelInput
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.runnables import Runnable
 from langchain_openai import ChatOpenAI
+from pydantic import BaseModel
 
 from app.config import settings
 
+OutputModel = TypeVar("OutputModel", bound=BaseModel)
+StructuredMethod = Literal["function_calling", "json_schema", "json_mode"]
+
 
 @cache
-def get_llm() -> ChatOpenAI:
+def _get_base_llm() -> ChatOpenAI:
+    """只缓存基础客户端，不因输出 schema 不同而重复创建连接"""
     if not settings.llm_api_key.strip() or not settings.llm_model.strip():
         raise ValueError("未配置 LLM_API_KEY 或 LLM_MODEL")
     return ChatOpenAI(
@@ -16,3 +27,66 @@ def get_llm() -> ChatOpenAI:
         timeout=settings.llm_timeout,
         max_retries=settings.llm_max_retries,
     )
+
+
+@overload
+def get_llm(
+    schema: type[OutputModel],
+    *,
+    method: StructuredMethod = "function_calling",
+    strict: bool | None = False,
+    client: BaseChatModel | None = None,
+) -> Runnable[LanguageModelInput, OutputModel]: ...
+
+
+@overload
+def get_llm(
+    schema: None = None,
+    *,
+    method: StructuredMethod = "function_calling",
+    strict: bool | None = False,
+    client: None = None,
+) -> ChatOpenAI: ...
+
+
+@overload
+def get_llm(
+    schema: None = None,
+    *,
+    method: StructuredMethod = "function_calling",
+    strict: bool | None = False,
+    client: BaseChatModel,
+) -> BaseChatModel: ...
+
+
+def get_llm(
+    schema: type[OutputModel] | None = None,
+    *,
+    method: StructuredMethod = "function_calling",
+    strict: bool | None = False,
+    client: BaseChatModel | None = None,
+) -> BaseChatModel | Runnable[LanguageModelInput, OutputModel]:
+    """不传 schema 返回普通模型，传入 Pydantic 类则返回结构化调用器
+
+    client 用于注入模型，省略时使用缓存客户端
+    strict 控制服务端 schema 约束，本地仍由 Pydantic 校验
+    """
+    if schema is not None and (
+        not isinstance(schema, type) or not issubclass(schema, BaseModel)
+    ):
+        raise TypeError("schema 必须是 Pydantic BaseModel 类")
+    if method not in ("function_calling", "json_schema", "json_mode"):
+        raise ValueError(f"不支持的结构化输出方式: {method}")
+    if schema is None and (method != "function_calling" or strict not in (False, None)):
+        raise ValueError("设置结构化输出参数时必须提供 schema")
+    if method == "json_mode" and strict is True:
+        raise ValueError("json_mode 不支持 strict=True")
+
+    llm = client if client is not None else _get_base_llm()
+    if schema is None:
+        return llm
+    # json_mode 不接受 strict 参数，提示词还需要明确要求 JSON 及其结构
+    structured = llm.with_structured_output(
+        schema, method=method, strict=None if method == "json_mode" else strict
+    )
+    return cast(Runnable[LanguageModelInput, OutputModel], structured)