305 lines
12 KiB
Python
305 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""四家模型接口的协议整形：只算，不发。
|
||
|
||
为什么要单独一层：同一套协议现在有两个传输方在用——桌面端的同步 requests
|
||
（ai_chat.py）和网关的异步 httpx（model_gateway.py）。如果两边各写一份"怎么拼
|
||
Claude 的 payload"，迟早会漂：一边修了 system 的位置另一边没修，线上表现就是
|
||
"网关能用、桌面端 400"，而且极难查。
|
||
|
||
这里的函数全是纯函数：输入配置和消息，输出该发什么、怎么解析回来的东西。没有
|
||
网络、没有全局状态、没有副作用，所以两边都能安全复用，也好测。
|
||
|
||
四家的形状差异（照搬 OpenAI 的 payload 到别家一定报错）：
|
||
|
||
OpenAI /chat/completions system 在 messages 里 Authorization: Bearer
|
||
Claude /v1/messages system 是顶层参数 x-api-key + anthropic-version
|
||
max_tokens 必填
|
||
Dify /v1/chat-messages system 在应用侧配置 Authorization: Bearer
|
||
图片要先 upload 再引用
|
||
ComfyUI 文生图工作流引擎，不是聊天接口——不参与对话协议
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from urllib.parse import urlparse
|
||
|
||
# Anthropic Messages API 的必填版本头
|
||
ANTHROPIC_VERSION = "2023-06-01"
|
||
|
||
CHAT_KINDS = ("openai", "claude", "dify")
|
||
ALL_KINDS = CHAT_KINDS + ("comfyui",)
|
||
|
||
_BROWSER_UA = (
|
||
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
|
||
"AppleWebKit/537.36 (KHTML, like Gecko) "
|
||
"Chrome/124.0.0.0 Safari/537.36"
|
||
)
|
||
|
||
|
||
def detect_kind(kind: str, base_url: str) -> str:
|
||
"""声明了就用声明的；没声明就从地址猜。"""
|
||
value = str(kind or "auto").strip().lower()
|
||
if value in ALL_KINDS:
|
||
return value
|
||
parsed = urlparse((base_url or "").rstrip("/"))
|
||
path = (parsed.path or "").lower()
|
||
if "chat-messages" in path or "completion-messages" in path:
|
||
return "dify"
|
||
if "anthropic" in (parsed.hostname or ""):
|
||
return "claude"
|
||
return "openai"
|
||
|
||
|
||
ENDPOINT_MODES = ("auto", "exact")
|
||
|
||
|
||
def endpoint_url(kind: str, base_url: str, mode: str = "auto") -> str:
|
||
"""拼出实际请求地址。
|
||
|
||
`mode="auto"`（默认）按接口类型补全路径：填 `https://api.openai.com/v1`
|
||
就补成 `.../v1/chat/completions`。绝大多数服务商都是这个形状，所以它是默认。
|
||
|
||
`mode="exact"` 原样使用，一个字符都不加。这是为那些路径不按套路来的服务准备
|
||
的——比如 `https://api.example.com/custom/llm/invoke`，auto 会把它拼成
|
||
`.../invoke/chat/completions`，请求发到一个不存在的地址，报 404，而排查的人
|
||
会去怀疑密钥和网络。
|
||
|
||
为什么用显式开关而不是"更聪明的猜测"：`https://api.example.com/openai` 到底
|
||
是一个前缀（还要补 /chat/completions）还是完整端点，光看地址分不出来。猜错
|
||
的两种方向都会静默地把请求发歪，不如让填的人说清楚。
|
||
"""
|
||
base = (base_url or "").rstrip("/")
|
||
if str(mode or "auto").strip().lower() == "exact":
|
||
return base
|
||
path = (urlparse(base).path or "").lower().rstrip("/")
|
||
kind = detect_kind(kind, base_url)
|
||
if kind == "claude":
|
||
if path.endswith("/messages"):
|
||
return base
|
||
return f"{base}/messages" if path.endswith("/v1") else f"{base}/v1/messages"
|
||
if kind == "dify":
|
||
if path.endswith(("/chat-messages", "/completion-messages")):
|
||
return base
|
||
return f"{base}/chat-messages" if path.endswith("/v1") else f"{base}/v1/chat-messages"
|
||
if path.endswith("/chat/completions"):
|
||
return base
|
||
if path.startswith("/v1/") and len(path) > len("/v1/"):
|
||
return base
|
||
return f"{base}/chat/completions"
|
||
|
||
|
||
def dify_api_root(base_url: str, mode: str = "auto") -> str:
|
||
"""Dify App API 根地址，用来拼 /files/upload。"""
|
||
endpoint = endpoint_url("dify", base_url, mode).rstrip("/")
|
||
for suffix in ("/chat-messages", "/completion-messages"):
|
||
if endpoint.lower().endswith(suffix):
|
||
return endpoint[: -len(suffix)]
|
||
return endpoint
|
||
|
||
|
||
def auth_headers(kind: str, api_key: str, base_url: str = "") -> dict:
|
||
"""鉴权头。Claude 用 x-api-key，其余用 Bearer。"""
|
||
key = str(api_key or "").strip()
|
||
if detect_kind(kind, base_url) == "claude":
|
||
auth = {"x-api-key": key, "anthropic-version": ANTHROPIC_VERSION}
|
||
else:
|
||
auth = {"Authorization": f"Bearer {key}"}
|
||
return {
|
||
**auth,
|
||
"Content-Type": "application/json",
|
||
"Accept": "application/json, text/plain, */*",
|
||
"Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8",
|
||
"User-Agent": _BROWSER_UA,
|
||
"Connection": "keep-alive",
|
||
}
|
||
|
||
|
||
def claude_messages(messages: list) -> tuple[str, list]:
|
||
"""OpenAI 形状的 messages → Anthropic 要的 (system, messages)。
|
||
|
||
三处不兼容，照搬一定 400：system 是顶层参数、相邻同角色必须合并、
|
||
第一条必须是 user。
|
||
"""
|
||
system_parts: list[str] = []
|
||
turns: list[dict] = []
|
||
for item in messages or []:
|
||
role = str((item or {}).get("role") or "user")
|
||
content = (item or {}).get("content")
|
||
if role == "system":
|
||
system_parts.append(content if isinstance(content, str) else str(content))
|
||
continue
|
||
role = "assistant" if role == "assistant" else "user"
|
||
if isinstance(content, str):
|
||
blocks = [{"type": "text", "text": content}]
|
||
elif isinstance(content, list):
|
||
blocks = []
|
||
for part in content:
|
||
if not isinstance(part, dict):
|
||
blocks.append({"type": "text", "text": str(part)})
|
||
elif part.get("type") == "text":
|
||
blocks.append({"type": "text", "text": str(part.get("text") or "")})
|
||
elif part.get("type") == "image_url":
|
||
url = str((part.get("image_url") or {}).get("url") or "")
|
||
if url.startswith("data:"):
|
||
head, _, payload = url.partition(",")
|
||
media = head[5:].split(";")[0] or "image/png"
|
||
blocks.append({
|
||
"type": "image",
|
||
"source": {
|
||
"type": "base64",
|
||
"media_type": media,
|
||
"data": payload,
|
||
},
|
||
})
|
||
else:
|
||
blocks = [{"type": "text", "text": str(content)}]
|
||
if turns and turns[-1]["role"] == role:
|
||
turns[-1]["content"].extend(blocks)
|
||
else:
|
||
turns.append({"role": role, "content": blocks})
|
||
while turns and turns[0]["role"] != "user":
|
||
turns.pop(0)
|
||
if not turns:
|
||
turns = [{"role": "user", "content": [{"type": "text", "text": "请回复"}]}]
|
||
return "\n\n".join(part for part in system_parts if part), turns
|
||
|
||
|
||
def _last_user_text(messages: list) -> str:
|
||
for item in reversed(messages or []):
|
||
if str((item or {}).get("role") or "") != "user":
|
||
continue
|
||
content = (item or {}).get("content")
|
||
if isinstance(content, str):
|
||
return content
|
||
if isinstance(content, list):
|
||
return "".join(
|
||
str(part.get("text") or "")
|
||
for part in content
|
||
if isinstance(part, dict) and part.get("type") == "text"
|
||
)
|
||
return str(content)
|
||
return ""
|
||
|
||
|
||
def chat_payload(
|
||
kind: str,
|
||
*,
|
||
model: str,
|
||
messages: list,
|
||
max_tokens: int,
|
||
temperature: float,
|
||
base_url: str = "",
|
||
tools: list | None = None,
|
||
dify_user: str = "wechat-rpa",
|
||
dify_files: list | None = None,
|
||
dify_conversation_id: str = "",
|
||
) -> dict:
|
||
"""按各家的形状拼请求体。"""
|
||
kind = detect_kind(kind, base_url)
|
||
if kind == "claude":
|
||
system_text, turns = claude_messages(messages)
|
||
payload = {
|
||
"model": model,
|
||
"messages": turns,
|
||
# Anthropic 的 max_tokens 是必填且必须 >= 1，给 0 会直接 400
|
||
"max_tokens": max(1, int(max_tokens or 500)),
|
||
"temperature": float(temperature),
|
||
}
|
||
if system_text:
|
||
payload["system"] = system_text
|
||
return payload
|
||
if kind == "dify":
|
||
payload = {
|
||
"inputs": {},
|
||
"query": _last_user_text(messages) or "请回复",
|
||
"response_mode": "blocking",
|
||
"user": dify_user or "wechat-rpa",
|
||
}
|
||
if dify_conversation_id:
|
||
payload["conversation_id"] = dify_conversation_id
|
||
if dify_files:
|
||
payload["files"] = list(dify_files)
|
||
return payload
|
||
payload = {
|
||
"model": model,
|
||
"messages": messages,
|
||
"max_tokens": int(max_tokens or 500),
|
||
"temperature": float(temperature),
|
||
}
|
||
if tools:
|
||
payload["tools"] = tools
|
||
payload["tool_choice"] = "auto"
|
||
return payload
|
||
|
||
|
||
def parse_chat(kind: str, data: dict, base_url: str = "") -> str:
|
||
"""从各家的响应里取出回复文本。取不到就抛，绝不返回空串蒙混过去。"""
|
||
kind = detect_kind(kind, base_url)
|
||
data = data if isinstance(data, dict) else {}
|
||
if kind == "claude":
|
||
parts = [
|
||
str(block.get("text") or "")
|
||
for block in (data.get("content") or [])
|
||
if isinstance(block, dict) and block.get("type") == "text"
|
||
]
|
||
text = "".join(parts).strip()
|
||
if not text:
|
||
raise ValueError("Claude 未返回文本内容")
|
||
return text
|
||
if kind == "dify":
|
||
answer = data.get("answer")
|
||
if answer is None:
|
||
# 少数部署会包一层 data
|
||
answer = (data.get("data") or {}).get("answer")
|
||
if not answer:
|
||
raise ValueError("Dify 未返回 answer")
|
||
return str(answer)
|
||
try:
|
||
content = data["choices"][0]["message"]["content"]
|
||
except (KeyError, IndexError, TypeError) as exc:
|
||
raise ValueError("OpenAI 兼容响应里没有 choices[0].message.content") from exc
|
||
if content is None:
|
||
raise ValueError("OpenAI 兼容响应的 content 为空")
|
||
return str(content)
|
||
|
||
|
||
def parse_message(kind: str, data: dict, base_url: str = "") -> dict:
|
||
"""取出完整的 assistant 消息，而不只是文本。
|
||
|
||
`parse_chat` 只回文本，工具调用会被整个丢掉——模型说"我要查一下客户资料"，
|
||
到调用方手里变成一句空回复。MCP 多轮必须看到 `tool_calls` 才能往下走。
|
||
|
||
只有 OpenAI 兼容协议有工具调用。Dify 的工具在应用侧编排、对我们不可见，
|
||
Claude 的 tool_use 是另一套形状——这两种都只回文本，让调用方按"没有工具
|
||
调用"处理，比硬凑一个形状再在别处出错要好。
|
||
"""
|
||
kind = detect_kind(kind, base_url)
|
||
data = data if isinstance(data, dict) else {}
|
||
if kind in ("claude", "dify"):
|
||
return {"content": parse_chat(kind, data, base_url), "tool_calls": []}
|
||
try:
|
||
message = data["choices"][0]["message"]
|
||
except (KeyError, IndexError, TypeError) as exc:
|
||
raise ValueError("OpenAI 兼容响应里没有 choices[0].message") from exc
|
||
if not isinstance(message, dict):
|
||
raise ValueError("OpenAI 兼容响应的 message 不是对象")
|
||
tool_calls = message.get("tool_calls") or []
|
||
content = message.get("content")
|
||
if content is None and not tool_calls:
|
||
raise ValueError("OpenAI 兼容响应既没有 content 也没有 tool_calls")
|
||
return {"content": "" if content is None else str(content), "tool_calls": list(tool_calls)}
|
||
|
||
|
||
def image_message(prompt: str, image_b64: str, mime: str = "image/png") -> dict:
|
||
"""一条带图的 user 消息（OpenAI 形状）。Claude 的转换由 claude_messages 负责。"""
|
||
return {
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "text", "text": str(prompt or "")},
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {"url": f"data:{mime};base64,{image_b64}"},
|
||
},
|
||
],
|
||
}
|