mimo-token-plan-adapter

backend/python/app/core/llm_client.py

@@ -33,7 +33,7 @@ def _base_url(api_url: str) -> str:
     return u
 
 
-def build_openai_client(api_url, api_key, *, proxy=None, timeout=120) -> AsyncOpenAI:
+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 等），未传则直连。
@@ -49,11 +49,19 @@ def build_openai_client(api_url, api_key, *, proxy=None, timeout=120) -> AsyncOp
         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,
     )
 
 
@@ -96,12 +104,12 @@ def _normalize_tool_calls(tcs):
     return out
 
 
-async def achat_complete(api_url, api_key, request_body, *, proxy=None, timeout=120, stream=False):
+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)
+            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}
@@ -175,7 +183,7 @@ async def achat_complete(api_url, api_key, request_body, *, proxy=None, timeout=
 
         return _gen()
 
-    client = build_openai_client(api_url, api_key, proxy=proxy, timeout=timeout)
+    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(
@@ -202,16 +210,22 @@ async def achat_complete(api_url, api_key, request_body, *, proxy=None, timeout=
             pass
 
 
-def build_openai_client_sync(api_url, api_key, *, proxy=None, timeout=120) -> OpenAI:
+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)
-    return OpenAI(base_url=base, api_key=api_key or "EMPTY", timeout=timeout, http_client=http_client)
+    _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)
+    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
@@ -242,7 +256,7 @@ async def aembed(api_url, api_key, input, model, *, proxy=None, timeout=60) -> l
     供 pgvector_search.get_embedding 等 embeddings 路径复用，把散落的裸 httpx
     嵌入调用收敛到同一标准模型层（与聊天共用代理/超时/failover 语义）。
     """
-    client = build_openai_client(api_url, api_key, proxy=proxy, timeout=timeout)
+    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(
@@ -259,7 +273,7 @@ async def aembed(api_url, api_key, input, model, *, proxy=None, timeout=60) -> l
 
 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)
+    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]