125 lines
4.6 KiB
Python
125 lines
4.6 KiB
Python
# app/module/shared/llm/factory.py
|
||
"""
|
||
LangChain 模型工厂：基于 OpenAI 兼容接口封装 Chat / Embedding 的创建、健康检查与同步调用。
|
||
- DashScope compatible‑mode：重写embedding逻辑，规避tiktoken token数组输入问题
|
||
- 其他OpenAI兼容服务商：完全使用原生 OpenAIEmbeddings / ChatOpenAI
|
||
"""
|
||
from typing import Literal, List, Optional
|
||
|
||
import httpx
|
||
from langchain_core.language_models import BaseChatModel
|
||
from langchain_core.embeddings import Embeddings
|
||
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
|
||
from pydantic import SecretStr
|
||
|
||
|
||
class DashScopeOpenAIEmbeddings(OpenAIEmbeddings, Embeddings):
|
||
"""
|
||
仅用于阿里云DashScope compatible‑mode
|
||
1. 绕过原生tiktoken逻辑，直接传入原始字符串，不传入token数字数组
|
||
2. 适配新版openai SDK返回Pydantic Embedding对象，使用属性 .embedding
|
||
"""
|
||
|
||
def _get_len_safe_embeddings(self, texts: List[str], *, engine: str, **kwargs) -> List[List[float]]:
|
||
batch_size = self.chunk_size
|
||
embeddings: List[List[float]] = []
|
||
for i in range(0, len(texts), batch_size):
|
||
batch = texts[i:i + batch_size]
|
||
resp = self.client.create(
|
||
input=batch,
|
||
model=self.model,
|
||
encoding_format="float",
|
||
)
|
||
embeddings.extend([data.embedding for data in resp.data])
|
||
return embeddings
|
||
|
||
|
||
class LLMFactory:
|
||
@staticmethod
|
||
def _is_dashscope_compatible(base_url: Optional[str]) -> bool:
|
||
if not base_url:
|
||
return False
|
||
return "dashscope.aliyuncs.com/compatible-mode" in base_url
|
||
|
||
@staticmethod
|
||
def create_chat(
|
||
model_name: str,
|
||
base_url: str,
|
||
api_key: str | None = None,
|
||
) -> BaseChatModel:
|
||
"""创建对话模型，兼容 OpenAI 及任意 base_url 的 OpenAI 兼容服务。"""
|
||
# DashScope容器环境关闭SSL校验；其余服务商使用默认httpx客户端（开启证书校验）
|
||
if LLMFactory._is_dashscope_compatible(base_url):
|
||
http_client = httpx.Client(verify=False)
|
||
else:
|
||
http_client = None
|
||
|
||
return ChatOpenAI(
|
||
model=model_name,
|
||
base_url=base_url or None,
|
||
api_key=SecretStr(api_key or ""),
|
||
http_client=http_client,
|
||
)
|
||
|
||
@staticmethod
|
||
def create_embedding(
|
||
model_name: str,
|
||
base_url: str,
|
||
api_key: str | None = None,
|
||
) -> Embeddings:
|
||
is_dashscope = LLMFactory._is_dashscope_compatible(base_url)
|
||
|
||
# DashScope才关闭SSL校验，其他服务商走默认SSL校验
|
||
if is_dashscope:
|
||
http_client = httpx.Client(verify=False)
|
||
else:
|
||
http_client = None
|
||
|
||
kwargs = {
|
||
"model": model_name,
|
||
"base_url": base_url or None,
|
||
"api_key": SecretStr(api_key or ""),
|
||
"http_client": http_client,
|
||
}
|
||
|
||
# 仅DashScope注入encoding_format="float"，其他服务商不强制注入，避免参数不兼容
|
||
if is_dashscope:
|
||
import inspect
|
||
sig = inspect.signature(OpenAIEmbeddings.__init__)
|
||
if "encoding_format" in sig.parameters:
|
||
kwargs["encoding_format"] = "float"
|
||
return DashScopeOpenAIEmbeddings(**kwargs)
|
||
|
||
# 其余所有OpenAI兼容模型，走原生实现
|
||
return OpenAIEmbeddings(**kwargs)
|
||
|
||
@staticmethod
|
||
def check_health(
|
||
model_name: str,
|
||
base_url: str,
|
||
api_key: str | None,
|
||
model_type: Literal["CHAT", "EMBEDDING"] | str,
|
||
) -> None:
|
||
"""对配置做一次最小化调用进行健康检查，失败则抛出。"""
|
||
if model_type == "CHAT":
|
||
model = LLMFactory.create_chat(model_name, base_url, api_key)
|
||
model.invoke("hello")
|
||
else:
|
||
model = LLMFactory.create_embedding(model_name, base_url, api_key)
|
||
model.embed_query("text")
|
||
|
||
@staticmethod
|
||
def get_embedding_dimension(
|
||
model_name: str,
|
||
base_url: str,
|
||
api_key: str | None = None,
|
||
) -> int:
|
||
"""创建 Embedding 模型并返回向量维度。"""
|
||
emb = LLMFactory.create_embedding(model_name, base_url, api_key)
|
||
return len(emb.embed_query("text"))
|
||
|
||
@staticmethod
|
||
def invoke_sync(chat_model: BaseChatModel, prompt: str) -> str:
|
||
"""同步调用对话模型并返回 content，供 run_in_executor 等场景使用。"""
|
||
return chat_model.invoke(prompt).content
|