93 lines
3.1 KiB
Python
93 lines
3.1 KiB
Python
"""统一获取普通模型或绑定结构化输出的模型"""
|
|
|
|
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_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(
|
|
model=settings.llm_model,
|
|
api_key=settings.llm_api_key,
|
|
base_url=settings.llm_base_url,
|
|
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)
|