diff --git a/app/clients/llm.py b/app/clients/llm.py index 9628306..7b1a32a 100644 --- a/app/clients/llm.py +++ b/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)