281 lines
12 KiB
Python
281 lines
12 KiB
Python
"""Shared, standards-based LLM client for web chat + IM bot.
|
||
|
||
平台红线(07-24)：模型层以成熟/开源标准为设计基准。本模块使用官方 OpenAI SDK
|
||
(\"openai.AsyncOpenAI\")——它是 OpenAI 兼容端点的事实标准客户端。平台本身已在讲
|
||
OpenAI 协议(\"/v1/chat/completions\")，因此零协议改造成本，并天然获得：
|
||
* 原生 SSE 流式(\"stream=True\")；
|
||
* 结构化的 tool_calls / usage 解析；
|
||
* 统一的代理(trust_env=False + 显式 proxy)与超时处理。
|
||
|
||
本模块取代 chat_service.py 与 im_llm.py 中两套手写 httpx 循环，使 web chat 与
|
||
IM bot 共用**同一模型层**、同一套 failover/代理/超时语义——这是「各端能力统一对齐」
|
||
红线在模型层的落地。
|
||
|
||
API:
|
||
achat_complete(api_url, api_key, request_body, *, proxy=None, timeout=120, stream=False)
|
||
- stream=False -> dict，形状对齐上游 \"data\" 响应：
|
||
{"choices":[{"message":{"content","tool_calls","reasoning_content"}}], "usage":{...}}
|
||
- stream=True -> async generator，逐块产出 ("content"|"reasoning", text) 增量
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import httpx
|
||
import asyncio
|
||
from openai import AsyncOpenAI, OpenAI
|
||
|
||
|
||
def _base_url(api_url: str) -> str:
|
||
"""从完整 /v1/chat/completions 或 /v1/embeddings 地址推导 OpenAI SDK 所需的 base_url。"""
|
||
u = (api_url or "").rstrip()
|
||
for suf in ("/v1/chat/completions", "/chat/completions", "/v1/embeddings", "/embeddings"):
|
||
if u.endswith(suf):
|
||
return u[: -len(suf)]
|
||
return u
|
||
|
||
|
||
def build_openai_client(api_url, api_key, *, proxy=None, timeout=120, auth_header=None) -> AsyncOpenAI:
|
||
base = _base_url(api_url)
|
||
# 始终 trust_env=False：避免容器 HTTP_PROXY 对国内端点(siliconflow 等)造成 502；
|
||
# 代理经 proxy 显式传入（AGNES 等），未传则直连。
|
||
http_client = httpx.AsyncClient(proxy=proxy, trust_env=False)
|
||
# 故障快速失败（全栈统一，覆盖 web chat / IM bot / embedding）：
|
||
# connect 封顶 10s —— 坏端点不再让请求挂死事件循环；
|
||
# read/write 封顶 30s —— 长流式回复只要持续有分块即不受影响；
|
||
# 保留数值型 timeout 接口，旧调用方无需改动。
|
||
_num = float(timeout) if isinstance(timeout, (int, float)) else 120.0
|
||
_to = httpx.Timeout(
|
||
connect=min(10.0, _num),
|
||
read=min(30.0, _num),
|
||
write=min(30.0, _num),
|
||
pool=10.0,
|
||
)
|
||
_hdrs = {}
|
||
if auth_header and ":" in auth_header:
|
||
_k, _v = auth_header.split(":", 1)
|
||
_hdrs[_k.strip()] = _v.strip()
|
||
elif auth_header:
|
||
_hdrs[auth_header.strip()] = api_key
|
||
|
||
return AsyncOpenAI(
|
||
base_url=base,
|
||
api_key=api_key or "EMPTY",
|
||
timeout=_to,
|
||
http_client=http_client,
|
||
default_headers=_hdrs if _hdrs else None,
|
||
)
|
||
|
||
|
||
def _create_kwargs(request_body: dict, stream: bool = False) -> dict:
|
||
kw = {
|
||
"model": request_body.get("model"),
|
||
"messages": request_body.get("messages", []),
|
||
"stream": stream,
|
||
}
|
||
if request_body.get("temperature") is not None:
|
||
kw["temperature"] = request_body["temperature"]
|
||
if request_body.get("top_p") is not None:
|
||
kw["top_p"] = request_body["top_p"]
|
||
if request_body.get("max_tokens") is not None:
|
||
kw["max_tokens"] = request_body["max_tokens"]
|
||
if request_body.get("stop") is not None:
|
||
kw["stop"] = request_body["stop"]
|
||
tools = request_body.get("tools")
|
||
if tools:
|
||
kw["tools"] = tools
|
||
kw["tool_choice"] = request_body.get("tool_choice", "auto")
|
||
return kw
|
||
|
||
|
||
def _normalize_tool_calls(tcs):
|
||
if not tcs:
|
||
return None
|
||
out = []
|
||
for tc in tcs:
|
||
out.append(
|
||
{
|
||
"id": getattr(tc, "id", None),
|
||
"type": "function",
|
||
"function": {
|
||
"name": getattr(tc.function, "name", None),
|
||
"arguments": getattr(tc.function, "arguments", "") or "",
|
||
},
|
||
}
|
||
)
|
||
return out
|
||
|
||
|
||
async def achat_complete(api_url, api_key, request_body, *, proxy=None, timeout=120, stream=False, auth_header=None):
|
||
"""统一的 OpenAI 兼容聊天补全。详见模块 docstring。"""
|
||
if stream:
|
||
|
||
async def _gen():
|
||
client = build_openai_client(api_url, api_key, proxy=proxy, timeout=timeout, auth_header=auth_header)
|
||
_content_parts = []
|
||
_reasoning_parts = []
|
||
_tool_acc = {} # index -> {id, type, name, arguments}
|
||
try:
|
||
_to = min(30.0, float(timeout) if isinstance(timeout, (int, float)) else 120.0)
|
||
gen = await asyncio.wait_for(
|
||
client.chat.completions.create(
|
||
**_create_kwargs(request_body, stream=True)
|
||
),
|
||
timeout=_to,
|
||
)
|
||
_it = gen.__aiter__()
|
||
while True:
|
||
try:
|
||
chunk = await asyncio.wait_for(_it.__anext__(), timeout=_to)
|
||
except StopAsyncIteration:
|
||
break
|
||
except asyncio.TimeoutError:
|
||
raise TimeoutError("stream chunk read timeout")
|
||
if not chunk.choices:
|
||
continue
|
||
delta = chunk.choices[0].delta
|
||
rc = getattr(delta, "reasoning_content", None)
|
||
if rc:
|
||
_reasoning_parts.append(rc)
|
||
yield ("reasoning", rc)
|
||
ct = getattr(delta, "content", None)
|
||
if ct:
|
||
_content_parts.append(ct)
|
||
yield ("content", ct)
|
||
# 累积 OpenAI 流式 tool_calls（供后续工具循环执行）
|
||
tcs = getattr(delta, "tool_calls", None)
|
||
if tcs:
|
||
for tc in tcs:
|
||
_idx = getattr(tc, "index", 0)
|
||
_slot = _tool_acc.setdefault(
|
||
_idx, {"id": None, "type": "function", "name": "", "arguments": ""}
|
||
)
|
||
_tcid = getattr(tc, "id", None)
|
||
if _tcid:
|
||
_slot["id"] = _tcid
|
||
_fn = getattr(tc, "function", None)
|
||
if _fn is not None:
|
||
_fn_name = getattr(_fn, "name", None)
|
||
if _fn_name:
|
||
_slot["name"] += _fn_name
|
||
_fn_args = getattr(_fn, "arguments", None)
|
||
if _fn_args:
|
||
_slot["arguments"] += _fn_args
|
||
finally:
|
||
try:
|
||
await asyncio.wait_for(client.close(), timeout=5)
|
||
except Exception:
|
||
pass
|
||
# 流式结束时吐出已组装的完整 message（含 tool_calls），供调用方无缝衔接工具循环
|
||
_tool_calls = None
|
||
if _tool_acc:
|
||
_tool_calls = []
|
||
for _i in sorted(_tool_acc.keys()):
|
||
_s = _tool_acc[_i]
|
||
_tool_calls.append({
|
||
"id": _s["id"],
|
||
"type": "function",
|
||
"function": {"name": _s["name"], "arguments": _s["arguments"]},
|
||
})
|
||
yield ("message", {
|
||
"content": "".join(_content_parts),
|
||
"tool_calls": _tool_calls,
|
||
"reasoning_content": "".join(_reasoning_parts) or None,
|
||
})
|
||
|
||
return _gen()
|
||
|
||
client = build_openai_client(api_url, api_key, proxy=proxy, timeout=timeout, auth_header=auth_header)
|
||
try:
|
||
_to = min(30.0, float(timeout) if isinstance(timeout, (int, float)) else 120.0)
|
||
resp = await asyncio.wait_for(
|
||
client.chat.completions.create(**_create_kwargs(request_body)),
|
||
timeout=_to,
|
||
)
|
||
msg = resp.choices[0].message
|
||
return {
|
||
"choices": [
|
||
{
|
||
"message": {
|
||
"content": msg.content or "",
|
||
"tool_calls": _normalize_tool_calls(getattr(msg, "tool_calls", None)),
|
||
"reasoning_content": getattr(msg, "reasoning_content", None),
|
||
}
|
||
}
|
||
],
|
||
"usage": (resp.usage.model_dump() if resp.usage else {}),
|
||
}
|
||
finally:
|
||
try:
|
||
await asyncio.wait_for(client.close(), timeout=5)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def build_openai_client_sync(api_url, api_key, *, proxy=None, timeout=120, auth_header=None) -> OpenAI:
|
||
"""同步客户端工厂（供 IM bot 等同步上下文复用同一标准模型层）。"""
|
||
base = _base_url(api_url)
|
||
http_client = httpx.Client(proxy=proxy, trust_env=False)
|
||
_hdrs = {}
|
||
if auth_header and ":" in auth_header:
|
||
_k, _v = auth_header.split(":", 1)
|
||
_hdrs[_k.strip()] = _v.strip()
|
||
elif auth_header:
|
||
_hdrs[auth_header.strip()] = api_key
|
||
return OpenAI(base_url=base, api_key=api_key or "EMPTY", timeout=timeout, http_client=http_client, default_headers=_hdrs if _hdrs else None)
|
||
|
||
|
||
def chat_complete(api_url, api_key, request_body, *, proxy=None, timeout=120):
|
||
"""同步版（与 achat_complete 行为/返回形状完全一致），供非 async 上下文使用。"""
|
||
client = build_openai_client_sync(api_url, api_key, proxy=proxy, timeout=timeout, auth_header=auth_header)
|
||
with client:
|
||
resp = client.chat.completions.create(**_create_kwargs(request_body))
|
||
msg = resp.choices[0].message
|
||
return {
|
||
"choices": [
|
||
{
|
||
"message": {
|
||
"content": msg.content or "",
|
||
"tool_calls": _normalize_tool_calls(getattr(msg, "tool_calls", None)),
|
||
"reasoning_content": getattr(msg, "reasoning_content", None),
|
||
}
|
||
}
|
||
],
|
||
"usage": (resp.usage.model_dump() if resp.usage else {}),
|
||
}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 嵌入（Embeddings）：OpenAI 兼容 /v1/embeddings，收敛散落的裸 httpx 嵌入调用
|
||
# ---------------------------------------------------------------------------
|
||
def _embed_kwargs(input, model):
|
||
return {"model": model, "input": input if isinstance(input, list) else [input]}
|
||
|
||
|
||
async def aembed(api_url, api_key, input, model, *, proxy=None, timeout=60) -> list:
|
||
"""统一的 OpenAI 兼容文本嵌入（异步）。返回 list[list[float]]。
|
||
|
||
供 pgvector_search.get_embedding 等 embeddings 路径复用，把散落的裸 httpx
|
||
嵌入调用收敛到同一标准模型层（与聊天共用代理/超时/failover 语义）。
|
||
"""
|
||
client = build_openai_client(api_url, api_key, proxy=proxy, timeout=timeout, auth_header=auth_header)
|
||
try:
|
||
_to = min(60.0, float(timeout) if isinstance(timeout, (int, float)) else 60.0)
|
||
resp = await asyncio.wait_for(
|
||
client.embeddings.create(**_embed_kwargs(input, model)),
|
||
timeout=_to,
|
||
)
|
||
return [d.embedding for d in resp.data]
|
||
finally:
|
||
try:
|
||
await asyncio.wait_for(client.close(), timeout=5)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def embed(api_url, api_key, input, model, *, proxy=None, timeout=60) -> list:
|
||
"""同步版嵌入（与非 async 上下文复用同一标准模型层）。"""
|
||
client = build_openai_client_sync(api_url, api_key, proxy=proxy, timeout=timeout, auth_header=auth_header)
|
||
with client:
|
||
resp = client.embeddings.create(**_embed_kwargs(input, model))
|
||
return [d.embedding for d in resp.data]
|
||
|