172 lines
6.3 KiB
Python
172 lines
6.3 KiB
Python
"""AI Provider 抽象层，支持 OpenAI 兼容 API 和 Anthropic 原生 API。"""
|
||
|
||
import json
|
||
from typing import AsyncGenerator
|
||
|
||
import httpx
|
||
|
||
from ..models import AIConfig
|
||
|
||
|
||
class StreamEvent:
|
||
"""流式事件：文本chunk或usage信息"""
|
||
def __init__(self, text: str = "", usage: dict | None = None):
|
||
self.text = text
|
||
self.usage = usage
|
||
|
||
|
||
async def stream_chat(
|
||
config: AIConfig,
|
||
messages: list[dict],
|
||
system_prompt: str = "",
|
||
) -> AsyncGenerator[StreamEvent, None]:
|
||
"""根据 provider 类型调用对应的流式 API，返回 StreamEvent。"""
|
||
if config.provider == "anthropic":
|
||
async for event in _stream_anthropic(config, messages, system_prompt):
|
||
yield event
|
||
else:
|
||
async for event in _stream_openai(config, messages, system_prompt):
|
||
yield event
|
||
|
||
|
||
async def _stream_openai(
|
||
config: AIConfig,
|
||
messages: list[dict],
|
||
system_prompt: str,
|
||
) -> AsyncGenerator[StreamEvent, None]:
|
||
"""调用 OpenAI 兼容 API（stream=true）。"""
|
||
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": True,
|
||
"stream_options": {"include_usage": True},
|
||
}
|
||
|
||
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}）")
|
||
|
||
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（OpenAI在最后一个chunk返回）
|
||
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 choices:
|
||
delta = choices[0].get("delta", {})
|
||
content = delta.get("content")
|
||
if content:
|
||
yield StreamEvent(text=content)
|
||
except (json.JSONDecodeError, KeyError, IndexError):
|
||
continue
|
||
|
||
|
||
async def _stream_anthropic(
|
||
config: AIConfig,
|
||
messages: list[dict],
|
||
system_prompt: str,
|
||
) -> AsyncGenerator[StreamEvent, None]:
|
||
"""调用 Anthropic Messages API（stream=true）。"""
|
||
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
|
||
|
||
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}）")
|
||
|
||
async for line in resp.aiter_lines():
|
||
if not line.startswith("data: "):
|
||
continue
|
||
try:
|
||
event = json.loads(line[6:])
|
||
if event.get("type") == "content_block_delta":
|
||
text = event.get("delta", {}).get("text", "")
|
||
if text:
|
||
yield StreamEvent(text=text)
|
||
elif event.get("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,
|
||
})
|
||
elif event.get("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.get("type") == "message_stop":
|
||
break
|
||
except json.JSONDecodeError:
|
||
continue
|