400 lines
14 KiB
Python
400 lines
14 KiB
Python
"""AI Provider 抽象层，支持 OpenAI 兼容 API 和 Anthropic 原生 API。
|
||
支持文本流式输出和 tool calling。"""
|
||
|
||
import json
|
||
from dataclasses import dataclass, field
|
||
from typing import AsyncGenerator
|
||
|
||
import httpx
|
||
|
||
from ..models import AIConfig
|
||
|
||
|
||
@dataclass
|
||
class ToolCall:
|
||
"""一次工具调用"""
|
||
id: str
|
||
name: str
|
||
arguments: dict
|
||
|
||
|
||
class StreamEvent:
|
||
"""流式事件：文本chunk、usage信息、或工具调用"""
|
||
def __init__(
|
||
self,
|
||
text: str = "",
|
||
usage: dict | None = None,
|
||
tool_calls: list[ToolCall] | None = None,
|
||
):
|
||
self.text = text
|
||
self.usage = usage
|
||
self.tool_calls = tool_calls
|
||
|
||
|
||
async def stream_chat(
|
||
config: AIConfig,
|
||
messages: list[dict],
|
||
system_prompt: str = "",
|
||
tools: list[dict] | None = None,
|
||
) -> AsyncGenerator[StreamEvent, None]:
|
||
"""根据 provider 类型调用对应的流式 API，返回 StreamEvent。"""
|
||
if config.provider == "anthropic":
|
||
async for event in _stream_anthropic(config, messages, system_prompt, tools):
|
||
yield event
|
||
else:
|
||
async for event in _stream_openai(config, messages, system_prompt, tools):
|
||
yield event
|
||
|
||
|
||
# ── 消息格式构建（用于 tool call 循环中追加消息）──
|
||
|
||
|
||
def build_assistant_message(provider: str, text: str, tool_calls: list[ToolCall]) -> dict:
|
||
"""构建包含工具调用的 assistant 消息"""
|
||
if provider == "anthropic":
|
||
content = []
|
||
if text:
|
||
content.append({"type": "text", "text": text})
|
||
for tc in tool_calls:
|
||
content.append({
|
||
"type": "tool_use",
|
||
"id": tc.id,
|
||
"name": tc.name,
|
||
"input": tc.arguments,
|
||
})
|
||
return {"role": "assistant", "content": content}
|
||
else:
|
||
# OpenAI 格式
|
||
msg: dict = {
|
||
"role": "assistant",
|
||
"content": text or None,
|
||
"tool_calls": [
|
||
{
|
||
"id": tc.id,
|
||
"type": "function",
|
||
"function": {
|
||
"name": tc.name,
|
||
"arguments": json.dumps(tc.arguments, ensure_ascii=False),
|
||
},
|
||
}
|
||
for tc in tool_calls
|
||
],
|
||
}
|
||
return msg
|
||
|
||
|
||
def build_tool_results_messages(
|
||
provider: str, results: list[dict],
|
||
) -> list[dict]:
|
||
"""构建工具结果消息。results: [{"id": str, "result": str}]
|
||
OpenAI: 返回多条 tool message
|
||
Anthropic: 返回一条 user message，content 是 tool_result 数组
|
||
"""
|
||
if provider == "anthropic":
|
||
content = [
|
||
{"type": "tool_result", "tool_use_id": r["id"], "content": r["result"]}
|
||
for r in results
|
||
]
|
||
return [{"role": "user", "content": content}]
|
||
else:
|
||
return [
|
||
{"role": "tool", "tool_call_id": r["id"], "content": r["result"]}
|
||
for r in results
|
||
]
|
||
|
||
|
||
async def simple_completion(
|
||
config: AIConfig,
|
||
messages: list[dict],
|
||
system_prompt: str = "",
|
||
) -> str:
|
||
"""非流式调用，返回完整文本。用于摘要压缩等场景。"""
|
||
if config.provider == "anthropic":
|
||
return await _simple_anthropic(config, messages, system_prompt)
|
||
return await _simple_openai(config, messages, system_prompt)
|
||
|
||
|
||
async def _simple_openai(config: AIConfig, messages: list[dict], system_prompt: str) -> str:
|
||
url = f"{config.base_url.rstrip('/')}/chat/completions"
|
||
payload_messages = []
|
||
if system_prompt:
|
||
payload_messages.append({"role": "system", "content": system_prompt})
|
||
payload_messages.extend(messages)
|
||
|
||
body = {"model": config.model, "messages": payload_messages, "stream": False}
|
||
async with httpx.AsyncClient(timeout=120) as client:
|
||
resp = await client.post(
|
||
url,
|
||
headers={"Authorization": f"Bearer {config.api_key}", "Content-Type": "application/json"},
|
||
json=body,
|
||
)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
return data["choices"][0]["message"]["content"] or ""
|
||
|
||
|
||
async def _simple_anthropic(config: AIConfig, messages: list[dict], system_prompt: str) -> str:
|
||
url = f"{config.base_url.rstrip('/')}/messages"
|
||
body: dict = {
|
||
"model": config.model,
|
||
"max_tokens": 4096,
|
||
"messages": messages,
|
||
"stream": False,
|
||
}
|
||
if system_prompt:
|
||
body["system"] = system_prompt
|
||
|
||
async with httpx.AsyncClient(timeout=120) as client:
|
||
resp = await client.post(
|
||
url,
|
||
headers={
|
||
"x-api-key": config.api_key,
|
||
"anthropic-version": "2023-06-01",
|
||
"Content-Type": "application/json",
|
||
},
|
||
json=body,
|
||
)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
for block in data.get("content", []):
|
||
if block.get("type") == "text":
|
||
return block["text"]
|
||
return ""
|
||
|
||
|
||
# ── OpenAI 兼容 ──
|
||
|
||
|
||
async def _stream_openai(
|
||
config: AIConfig,
|
||
messages: list[dict],
|
||
system_prompt: str,
|
||
tools: list[dict] | None = None,
|
||
) -> AsyncGenerator[StreamEvent, None]:
|
||
"""调用 OpenAI 兼容 API（stream=true），支持 tool calling。"""
|
||
url = f"{config.base_url.rstrip('/')}/chat/completions"
|
||
|
||
payload_messages = []
|
||
if system_prompt:
|
||
payload_messages.append({"role": "system", "content": system_prompt})
|
||
payload_messages.extend(messages)
|
||
|
||
body: dict = {
|
||
"model": config.model,
|
||
"messages": payload_messages,
|
||
"stream": True,
|
||
"stream_options": {"include_usage": True},
|
||
}
|
||
if tools:
|
||
body["tools"] = tools
|
||
|
||
async with httpx.AsyncClient(timeout=120) as client:
|
||
async with client.stream(
|
||
"POST",
|
||
url,
|
||
headers={
|
||
"Authorization": f"Bearer {config.api_key}",
|
||
"Content-Type": "application/json",
|
||
},
|
||
json=body,
|
||
) as resp:
|
||
if resp.status_code >= 400:
|
||
body_text = await resp.aread()
|
||
try:
|
||
err = json.loads(body_text)
|
||
msg = err.get("error", {}).get("message", "") or str(err)
|
||
except Exception:
|
||
msg = body_text.decode(errors="replace")[:200]
|
||
raise RuntimeError(f"API错误 ({resp.status_code}): {msg}")
|
||
|
||
ct = resp.headers.get("content-type", "")
|
||
if "text/html" in ct:
|
||
raise RuntimeError(f"API地址返回了HTML页面，请检查API地址是否正确（当前: {url}）")
|
||
|
||
# 累积工具调用
|
||
pending_tool_calls: dict[int, dict] = {} # index -> {id, name, arguments}
|
||
|
||
async for line in resp.aiter_lines():
|
||
if not line.startswith("data: "):
|
||
continue
|
||
data = line[6:]
|
||
if data.strip() == "[DONE]":
|
||
break
|
||
try:
|
||
chunk = json.loads(data)
|
||
|
||
# 提取 usage
|
||
usage = chunk.get("usage")
|
||
if usage:
|
||
yield StreamEvent(usage={
|
||
"prompt_tokens": usage.get("prompt_tokens", 0),
|
||
"completion_tokens": usage.get("completion_tokens", 0),
|
||
"total_tokens": usage.get("total_tokens", 0),
|
||
})
|
||
|
||
choices = chunk.get("choices", [])
|
||
if not choices:
|
||
continue
|
||
|
||
choice = choices[0]
|
||
delta = choice.get("delta", {})
|
||
finish_reason = choice.get("finish_reason")
|
||
|
||
# 文本内容
|
||
content = delta.get("content")
|
||
if content:
|
||
yield StreamEvent(text=content)
|
||
|
||
# 工具调用 chunk
|
||
for tc_delta in delta.get("tool_calls", []):
|
||
idx = tc_delta["index"]
|
||
if idx not in pending_tool_calls:
|
||
pending_tool_calls[idx] = {"id": "", "name": "", "arguments": ""}
|
||
entry = pending_tool_calls[idx]
|
||
if "id" in tc_delta:
|
||
entry["id"] = tc_delta["id"]
|
||
func = tc_delta.get("function", {})
|
||
if "name" in func:
|
||
entry["name"] = func["name"]
|
||
if "arguments" in func:
|
||
entry["arguments"] += func["arguments"]
|
||
|
||
# 结束：如果有工具调用，发出事件
|
||
if finish_reason == "tool_calls" and pending_tool_calls:
|
||
tool_calls = []
|
||
for tc_data in pending_tool_calls.values():
|
||
try:
|
||
args = json.loads(tc_data["arguments"]) if tc_data["arguments"] else {}
|
||
except json.JSONDecodeError:
|
||
args = {}
|
||
tool_calls.append(ToolCall(
|
||
id=tc_data["id"],
|
||
name=tc_data["name"],
|
||
arguments=args,
|
||
))
|
||
yield StreamEvent(tool_calls=tool_calls)
|
||
|
||
except (json.JSONDecodeError, KeyError, IndexError):
|
||
continue
|
||
|
||
|
||
# ── Anthropic 原生 ──
|
||
|
||
|
||
async def _stream_anthropic(
|
||
config: AIConfig,
|
||
messages: list[dict],
|
||
system_prompt: str,
|
||
tools: list[dict] | None = None,
|
||
) -> AsyncGenerator[StreamEvent, None]:
|
||
"""调用 Anthropic Messages API（stream=true），支持 tool calling。"""
|
||
url = f"{config.base_url.rstrip('/')}/messages"
|
||
|
||
body: dict = {
|
||
"model": config.model,
|
||
"max_tokens": 4096,
|
||
"messages": messages,
|
||
"stream": True,
|
||
}
|
||
if system_prompt:
|
||
body["system"] = system_prompt
|
||
if tools:
|
||
body["tools"] = tools
|
||
|
||
async with httpx.AsyncClient(timeout=120) as client:
|
||
async with client.stream(
|
||
"POST",
|
||
url,
|
||
headers={
|
||
"x-api-key": config.api_key,
|
||
"anthropic-version": "2023-06-01",
|
||
"content-type": "application/json",
|
||
},
|
||
json=body,
|
||
) as resp:
|
||
if resp.status_code >= 400:
|
||
body_text = await resp.aread()
|
||
try:
|
||
err = json.loads(body_text)
|
||
msg = err.get("error", {}).get("message", "") or str(err)
|
||
except Exception:
|
||
msg = body_text.decode(errors="replace")[:200]
|
||
raise RuntimeError(f"API错误 ({resp.status_code}): {msg}")
|
||
|
||
ct = resp.headers.get("content-type", "")
|
||
if "text/html" in ct:
|
||
raise RuntimeError(f"API地址返回了HTML页面，请检查API地址是否正确（当前: {url}）")
|
||
|
||
# 累积工具调用
|
||
pending_tool_uses: dict[int, dict] = {} # block_index -> {id, name, input_json}
|
||
stop_reason = None
|
||
|
||
async for line in resp.aiter_lines():
|
||
if not line.startswith("data: "):
|
||
continue
|
||
try:
|
||
event = json.loads(line[6:])
|
||
event_type = event.get("type")
|
||
|
||
if event_type == "content_block_start":
|
||
block = event.get("content_block", {})
|
||
idx = event.get("index", 0)
|
||
if block.get("type") == "tool_use":
|
||
pending_tool_uses[idx] = {
|
||
"id": block["id"],
|
||
"name": block["name"],
|
||
"input_json": "",
|
||
}
|
||
|
||
elif event_type == "content_block_delta":
|
||
idx = event.get("index", 0)
|
||
delta = event.get("delta", {})
|
||
if delta.get("type") == "text_delta":
|
||
text = delta.get("text", "")
|
||
if text:
|
||
yield StreamEvent(text=text)
|
||
elif delta.get("type") == "input_json_delta":
|
||
if idx in pending_tool_uses:
|
||
pending_tool_uses[idx]["input_json"] += delta.get("partial_json", "")
|
||
|
||
elif event_type == "message_delta":
|
||
usage = event.get("usage", {})
|
||
if usage:
|
||
yield StreamEvent(usage={
|
||
"prompt_tokens": 0,
|
||
"completion_tokens": usage.get("output_tokens", 0),
|
||
"total_tokens": 0,
|
||
})
|
||
stop_reason = event.get("delta", {}).get("stop_reason")
|
||
|
||
elif event_type == "message_start":
|
||
msg = event.get("message", {})
|
||
usage = msg.get("usage", {})
|
||
if usage:
|
||
yield StreamEvent(usage={
|
||
"prompt_tokens": usage.get("input_tokens", 0),
|
||
"completion_tokens": 0,
|
||
"total_tokens": 0,
|
||
})
|
||
|
||
elif event_type == "message_stop":
|
||
break
|
||
|
||
except json.JSONDecodeError:
|
||
continue
|
||
|
||
# 如果有工具调用，发出事件
|
||
if stop_reason == "tool_use" and pending_tool_uses:
|
||
tool_calls = []
|
||
for tu_data in pending_tool_uses.values():
|
||
try:
|
||
args = json.loads(tu_data["input_json"]) if tu_data["input_json"] else {}
|
||
except json.JSONDecodeError:
|
||
args = {}
|
||
tool_calls.append(ToolCall(
|
||
id=tu_data["id"],
|
||
name=tu_data["name"],
|
||
arguments=args,
|
||
))
|
||
yield StreamEvent(tool_calls=tool_calls)
|