48 lines
1.6 KiB
Python
48 lines
1.6 KiB
Python
"""统一获取普通模型或绑定结构化输出的模型"""
|
|
|
|
from functools import cache
|
|
from typing import Any, TypeVar, overload
|
|
|
|
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)
|
|
|
|
|
|
@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 只给类型检查器和编辑器看，运行时不会执行
|
|
# 不传 schema：返回普通聊天模型，输出是自由文本
|
|
@overload
|
|
def get_llm() -> ChatOpenAI: ...
|
|
|
|
|
|
# 传入 Pydantic 类：返回结构化调用器，ainvoke 的结果就是该类的实例
|
|
# OutputModel 让返回类型跟着传入的类走，例如 get_llm(Requirement) 得到 Runnable[..., Requirement]
|
|
@overload
|
|
def get_llm(schema: type[OutputModel]) -> Runnable[Any, OutputModel]: ...
|
|
|
|
|
|
# 真正的实现：schema 用 type[BaseModel] | None 覆盖上面两种调用
|
|
# 返回类型由 overload 决定，这里不再重复标注
|
|
def get_llm(schema: type[BaseModel] | None = None) -> Any:
|
|
llm = _get_base_llm()
|
|
if schema is None:
|
|
return llm
|
|
return llm.with_structured_output(schema, method="function_calling", strict=False)
|