87 lines
2.7 KiB
Python
87 lines
2.7 KiB
Python
"""OpenAI Chat Completions 客户端。
|
|
|
|
- 真实接入：httpx 直调 `/chat/completions`，使用各 provider 独立的 base_url 和 api_key
|
|
- mock：api_key 为空时回显一个简单响应并估算 token 数
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from app.core.config import settings
|
|
|
|
|
|
@dataclass
|
|
class ChatResult:
|
|
content: str
|
|
prompt_tokens: int
|
|
completion_tokens: int
|
|
total_tokens: int
|
|
raw: dict[str, Any] | None = None
|
|
|
|
|
|
def _estimate_tokens(text: str) -> int:
|
|
"""粗估：平均 4 个字符 ≈ 1 token。"""
|
|
return max(1, len(text) // 4)
|
|
|
|
|
|
async def chat_completion(
|
|
*,
|
|
api_base: str,
|
|
api_key: str | None,
|
|
model: str,
|
|
messages: list[dict[str, Any]],
|
|
temperature: float | None = None,
|
|
max_tokens: int | None = None,
|
|
timeout_seconds: int | None = None,
|
|
) -> ChatResult:
|
|
"""发起一次 Chat Completions。api_key 为空时走 mock。"""
|
|
timeout = timeout_seconds or settings.OPENAI_TIMEOUT_SECONDS or 60
|
|
|
|
if not api_key:
|
|
joined = "\n".join(m.get("content", "") for m in messages if isinstance(m, dict))
|
|
reply = f"[mock:{model}] 接收到 {len(messages)} 条消息；配置 OPENAI_API_KEY 后将转为真实调用。"
|
|
prompt = sum(_estimate_tokens(m.get("content", "") or "") for m in messages)
|
|
comp = _estimate_tokens(reply)
|
|
return ChatResult(
|
|
content=reply,
|
|
prompt_tokens=prompt,
|
|
completion_tokens=comp,
|
|
total_tokens=prompt + comp,
|
|
raw={"mock": True, "preview": joined[:200]},
|
|
)
|
|
|
|
payload: dict[str, Any] = {"model": model, "messages": messages}
|
|
if temperature is not None:
|
|
payload["temperature"] = temperature
|
|
if max_tokens is not None:
|
|
payload["max_tokens"] = max_tokens
|
|
|
|
url = api_base.rstrip("/") + "/chat/completions"
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
async with httpx.AsyncClient(timeout=timeout) as client:
|
|
r = await client.post(url, json=payload, headers=headers)
|
|
r.raise_for_status()
|
|
data = r.json()
|
|
|
|
content = ""
|
|
try:
|
|
content = data["choices"][0]["message"]["content"] or ""
|
|
except Exception: # noqa: BLE001
|
|
content = ""
|
|
usage = data.get("usage") or {}
|
|
return ChatResult(
|
|
content=content,
|
|
prompt_tokens=int(usage.get("prompt_tokens") or _estimate_tokens(str(messages))),
|
|
completion_tokens=int(usage.get("completion_tokens") or _estimate_tokens(content)),
|
|
total_tokens=int(usage.get("total_tokens") or 0)
|
|
or _estimate_tokens(str(messages)) + _estimate_tokens(content),
|
|
raw=data,
|
|
)
|