277 lines
11 KiB
Python
277 lines
11 KiB
Python
import json
|
||
import os
|
||
from abc import ABC, abstractmethod
|
||
from dataclasses import dataclass, field
|
||
from typing import AsyncGenerator
|
||
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
|
||
@dataclass
|
||
class ChatResponse:
|
||
content: str
|
||
model: str
|
||
prompt_tokens: int = 0
|
||
completion_tokens: int = 0
|
||
|
||
|
||
class LLMProvider(ABC):
|
||
name: str = "base"
|
||
max_context: int = 128_000
|
||
|
||
@abstractmethod
|
||
async def chat(self, messages: list[dict], temperature: float = 0.7) -> ChatResponse:
|
||
...
|
||
|
||
@abstractmethod
|
||
async def chat_stream(self, messages: list[dict], temperature: float = 0.7) -> AsyncGenerator[str, None]:
|
||
...
|
||
|
||
|
||
class OpenAIProvider(LLMProvider):
|
||
name = "openai"
|
||
max_context = 128_000
|
||
|
||
def __init__(self, api_key: str | None = None, base_url: str | None = None, model: str | None = None):
|
||
self.api_key = api_key if api_key is not None else os.getenv("OPENAI_API_KEY", "")
|
||
self.base_url = base_url if base_url is not None else os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
|
||
self.model = model if model is not None else os.getenv("OPENAI_MODEL", "gpt-4o")
|
||
|
||
async def chat(self, messages: list[dict], temperature: float = 0.7) -> ChatResponse:
|
||
import httpx
|
||
async with httpx.AsyncClient(timeout=120) as client:
|
||
resp = await client.post(
|
||
f"{self.base_url}/chat/completions",
|
||
headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
|
||
json={"model": self.model, "messages": messages, "temperature": temperature},
|
||
)
|
||
data = resp.json()
|
||
if "error" in data:
|
||
raise RuntimeError(f"OpenAI error: {data['error']}")
|
||
choice = data["choices"][0]
|
||
return ChatResponse(
|
||
content=choice["message"]["content"],
|
||
model=data["model"],
|
||
prompt_tokens=data.get("usage", {}).get("prompt_tokens", 0),
|
||
completion_tokens=data.get("usage", {}).get("completion_tokens", 0),
|
||
)
|
||
|
||
async def chat_stream(self, messages: list[dict], temperature: float = 0.7) -> AsyncGenerator[str, None]:
|
||
import httpx
|
||
async with httpx.AsyncClient(timeout=300) as client:
|
||
async with client.stream(
|
||
"POST",
|
||
f"{self.base_url}/chat/completions",
|
||
headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
|
||
json={"model": self.model, "messages": messages, "temperature": temperature, "stream": True},
|
||
) as resp:
|
||
async for line in resp.aiter_lines():
|
||
if line.startswith("data: "):
|
||
data_str = line[6:]
|
||
if data_str.strip() == "[DONE]":
|
||
break
|
||
try:
|
||
data = json.loads(data_str)
|
||
# 流末尾可能来一个 usage-only chunk（choices 为空），直接跳过
|
||
choices = data.get("choices") or []
|
||
if not choices:
|
||
continue
|
||
delta = choices[0].get("delta", {})
|
||
content = delta.get("content", "")
|
||
if content:
|
||
yield content
|
||
except json.JSONDecodeError:
|
||
continue
|
||
|
||
|
||
class AnthropicProvider(LLMProvider):
|
||
name = "claude"
|
||
max_context = 200_000
|
||
|
||
def __init__(self, api_key: str | None = None, model: str | None = None):
|
||
self.api_key = api_key if api_key is not None else os.getenv("ANTHROPIC_API_KEY", "")
|
||
self.model = model if model is not None else os.getenv("ANTHROPIC_MODEL", "claude-sonnet-4-20250514")
|
||
|
||
def _to_anthropic_messages(self, messages: list[dict]) -> tuple[str | None, list[dict]]:
|
||
system_msg = None
|
||
converted = []
|
||
for m in messages:
|
||
if m["role"] == "system":
|
||
system_msg = m["content"]
|
||
else:
|
||
converted.append({"role": m["role"], "content": m["content"]})
|
||
return system_msg, converted
|
||
|
||
async def chat(self, messages: list[dict], temperature: float = 0.7) -> ChatResponse:
|
||
import httpx
|
||
system_msg, anthropic_msgs = self._to_anthropic_messages(messages)
|
||
body = {
|
||
"model": self.model,
|
||
"max_tokens": 4096,
|
||
"messages": anthropic_msgs,
|
||
"temperature": temperature,
|
||
}
|
||
if system_msg:
|
||
body["system"] = system_msg
|
||
async with httpx.AsyncClient(timeout=120) as client:
|
||
resp = await client.post(
|
||
"https://api.anthropic.com/v1/messages",
|
||
headers={
|
||
"x-api-key": self.api_key,
|
||
"anthropic-version": "2023-06-01",
|
||
"Content-Type": "application/json",
|
||
},
|
||
json=body,
|
||
)
|
||
data = resp.json()
|
||
if "error" in data:
|
||
raise RuntimeError(f"Anthropic error: {data['error']}")
|
||
content_blocks = data.get("content", [])
|
||
text = "".join(b.get("text", "") for b in content_blocks if b.get("type") == "text")
|
||
return ChatResponse(
|
||
content=text,
|
||
model=data["model"],
|
||
prompt_tokens=data.get("usage", {}).get("input_tokens", 0),
|
||
completion_tokens=data.get("usage", {}).get("output_tokens", 0),
|
||
)
|
||
|
||
async def chat_stream(self, messages: list[dict], temperature: float = 0.7) -> AsyncGenerator[str, None]:
|
||
import httpx
|
||
system_msg, anthropic_msgs = self._to_anthropic_messages(messages)
|
||
body = {
|
||
"model": self.model,
|
||
"max_tokens": 4096,
|
||
"messages": anthropic_msgs,
|
||
"temperature": temperature,
|
||
"stream": True,
|
||
}
|
||
if system_msg:
|
||
body["system"] = system_msg
|
||
async with httpx.AsyncClient(timeout=300) as client:
|
||
async with client.stream(
|
||
"POST",
|
||
"https://api.anthropic.com/v1/messages",
|
||
headers={
|
||
"x-api-key": self.api_key,
|
||
"anthropic-version": "2023-06-01",
|
||
"Content-Type": "application/json",
|
||
},
|
||
json=body,
|
||
) as resp:
|
||
async for line in resp.aiter_lines():
|
||
if line.startswith("data: "):
|
||
data_str = line[6:]
|
||
try:
|
||
data = json.loads(data_str)
|
||
if data.get("type") == "content_block_delta":
|
||
delta = data.get("delta", {})
|
||
text = delta.get("text", "")
|
||
if text:
|
||
yield text
|
||
except json.JSONDecodeError:
|
||
continue
|
||
|
||
|
||
class DeepSeekProvider(OpenAIProvider):
|
||
name = "deepseek"
|
||
max_context = 128_000
|
||
|
||
def __init__(self, api_key: str | None = None, model: str | None = None):
|
||
super().__init__(
|
||
api_key=api_key if api_key is not None else os.getenv("DEEPSEEK_API_KEY", ""),
|
||
base_url=os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1"),
|
||
model=model if model is not None else os.getenv("DEEPSEEK_MODEL", "deepseek-chat"),
|
||
)
|
||
|
||
|
||
# ---- 内置模型（.env）与用户自定义模型（ai_models 表）的统一解析 ----
|
||
|
||
def builtin_models() -> list[dict]:
|
||
"""扫描 .env 里已配置 key 的内置模型（无 key 的不返回，前端不会展示废选项）。
|
||
|
||
label 跟随 .env 的 model_id 动态生成，不硬编码（如 DEEPSEEK_MODEL=deepseek-flash
|
||
就显示「DeepSeek deepseek-flash」，避免与实际请求的模型对不上）。
|
||
"""
|
||
out = []
|
||
if os.getenv("DEEPSEEK_API_KEY"):
|
||
mid = os.getenv("DEEPSEEK_MODEL", "deepseek-chat")
|
||
out.append({
|
||
"name": "deepseek", "label": f"DeepSeek {mid}",
|
||
"kind": "openai_compatible", "base_url": os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1"),
|
||
"api_key_masked": _mask_key(os.getenv("DEEPSEEK_API_KEY")),
|
||
"model_id": mid, "builtin": True,
|
||
})
|
||
if os.getenv("OPENAI_API_KEY"):
|
||
mid = os.getenv("OPENAI_MODEL", "gpt-4o")
|
||
out.append({
|
||
"name": "openai", "label": f"GPT {mid}",
|
||
"kind": "openai_compatible", "base_url": os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1"),
|
||
"api_key_masked": _mask_key(os.getenv("OPENAI_API_KEY")),
|
||
"model_id": mid, "builtin": True,
|
||
})
|
||
if os.getenv("ANTHROPIC_API_KEY"):
|
||
mid = os.getenv("ANTHROPIC_MODEL", "claude-sonnet-4-20250514")
|
||
out.append({
|
||
"name": "claude", "label": f"Claude {mid}",
|
||
"kind": "anthropic", "base_url": "https://api.anthropic.com/v1",
|
||
"api_key_masked": _mask_key(os.getenv("ANTHROPIC_API_KEY")),
|
||
"model_id": mid, "builtin": True,
|
||
})
|
||
if os.getenv("OPENROUTER_API_KEY"):
|
||
mid = os.getenv("OPENROUTER_MODEL", "openrouter/auto")
|
||
out.append({
|
||
"name": "openrouter", "label": f"OpenRouter {mid}",
|
||
"kind": "openai_compatible", "base_url": "https://openrouter.ai/api/v1",
|
||
"api_key_masked": _mask_key(os.getenv("OPENROUTER_API_KEY")),
|
||
"model_id": mid, "builtin": True,
|
||
})
|
||
return out
|
||
|
||
|
||
def _mask_key(key: str | None) -> str:
|
||
if not key:
|
||
return ""
|
||
if len(key) <= 8:
|
||
return "****"
|
||
return f"{key[:4]}****{key[-4:]}"
|
||
|
||
|
||
def _builtin_provider(entry: dict) -> LLMProvider:
|
||
"""内置模型：key 从对应 .env 变量读取（builtin_models 只暴露掩码，不回传明文）。"""
|
||
name = entry["name"]
|
||
if name == "deepseek":
|
||
return DeepSeekProvider()
|
||
if name == "claude":
|
||
return AnthropicProvider()
|
||
if name == "openrouter":
|
||
return OpenAIProvider(
|
||
api_key=os.getenv("OPENROUTER_API_KEY", ""),
|
||
base_url="https://openrouter.ai/api/v1",
|
||
model=os.getenv("OPENROUTER_MODEL", "openrouter/auto"),
|
||
)
|
||
return OpenAIProvider() # openai
|
||
|
||
|
||
def get_provider(name: str) -> LLMProvider | None:
|
||
"""兼容旧的同步入口（仅内置模型）。"""
|
||
for m in builtin_models():
|
||
if m["name"] == name:
|
||
return _builtin_provider(m)
|
||
return None
|
||
|
||
|
||
async def resolve_provider(db: AsyncSession, name: str) -> LLMProvider | None:
|
||
"""按模型名解析 provider：内置名走 .env 默认构造，自定义走 ai_models 表；找不到返回 None。"""
|
||
from sqlalchemy import select
|
||
from app.models.ai_model import AiModel
|
||
|
||
for m in builtin_models():
|
||
if m["name"] == name:
|
||
return _builtin_provider(m)
|
||
result = await db.execute(select(AiModel).where(AiModel.name == name))
|
||
row = result.scalars().first()
|
||
if row is None:
|
||
return None
|
||
if row.kind == "anthropic":
|
||
return AnthropicProvider(api_key=row.api_key, model=row.model_id)
|
||
return OpenAIProvider(api_key=row.api_key, base_url=row.base_url, model=row.model_id) |