新增前后端AI聊天模块，包括流式对话、配置管理和UI组件优化。
This commit is contained in:
171
backend/app/services/ai_provider.py
Normal file
171
backend/app/services/ai_provider.py
Normal file
@@ -0,0 +1,171 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user