122 lines
5.0 KiB
Python
122 lines
5.0 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from ..config import CloudProfile, ModelConfig, RuntimeConfig
|
|
from ..types import BackendError, ChatChunk, ChatParams
|
|
from .base import BaseBackend
|
|
|
|
|
|
class RemoteBackend(BaseBackend):
|
|
"""Proxy any OpenAI-compatible cloud endpoint without exposing its key to clients."""
|
|
|
|
def __init__(self, model: ModelConfig, profile: CloudProfile, runtime: RuntimeConfig) -> None:
|
|
super().__init__(model, runtime)
|
|
self.profile = profile
|
|
self.base_url = profile.base_url.rstrip("/")
|
|
if not self.base_url.endswith("/v1"):
|
|
self.base_url += "/v1"
|
|
|
|
async def load(self) -> None:
|
|
key = self.profile.resolved_api_key()
|
|
if not key:
|
|
raise BackendError(
|
|
f"云端配置 {self.profile.id} 没有 API key。请设置环境变量 {self.profile.api_key_env or '<api_key_env>'}。"
|
|
)
|
|
self.loaded = True
|
|
|
|
async def unload(self) -> None:
|
|
self.loaded = False
|
|
|
|
def _headers(self) -> dict[str, str]:
|
|
key = self.profile.resolved_api_key()
|
|
if not key:
|
|
raise BackendError(f"云端配置 {self.profile.id} 的 API key 不可用")
|
|
headers = {"Authorization": f"Bearer {key}", "Content-Type": "application/json"}
|
|
headers.update(self.profile.headers)
|
|
return headers
|
|
|
|
async def chat(
|
|
self,
|
|
messages: list[dict[str, Any]],
|
|
params: ChatParams,
|
|
stream: bool = False,
|
|
) -> AsyncIterator[ChatChunk]:
|
|
if not self.loaded:
|
|
await self.load()
|
|
payload: dict[str, Any] = {
|
|
"model": self.model.remote_model or self.profile.default_model or self.model.id,
|
|
"messages": messages,
|
|
"stream": stream,
|
|
"max_tokens": params.max_tokens,
|
|
"temperature": params.temperature,
|
|
"top_p": params.top_p,
|
|
}
|
|
if params.stop:
|
|
payload["stop"] = params.stop
|
|
if params.seed is not None:
|
|
payload["seed"] = params.seed
|
|
if params.response_format is not None:
|
|
payload["response_format"] = (
|
|
{"type": params.response_format} if isinstance(params.response_format, str) else params.response_format
|
|
)
|
|
if params.tools:
|
|
payload["tools"] = params.tools
|
|
if params.tool_choice is not None:
|
|
payload["tool_choice"] = params.tool_choice
|
|
payload.update(params.extra)
|
|
timeout = httpx.Timeout(self.profile.timeout, connect=15.0)
|
|
async with httpx.AsyncClient(timeout=timeout, headers=self._headers()) as client:
|
|
if stream:
|
|
async with client.stream("POST", f"{self.base_url}/chat/completions", json=payload) as response:
|
|
if response.status_code >= 400:
|
|
detail = (await response.aread()).decode("utf-8", errors="replace")
|
|
raise BackendError(f"云端请求失败 {response.status_code}: {detail}")
|
|
async for line in response.aiter_lines():
|
|
if not line.startswith("data:"):
|
|
continue
|
|
raw = line[5:].strip()
|
|
if raw == "[DONE]":
|
|
yield ChatChunk(done=True, finish_reason="stop")
|
|
return
|
|
try:
|
|
data = json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
choice = (data.get("choices") or [{}])[0]
|
|
delta = choice.get("delta") or {}
|
|
text = delta.get("content") or ""
|
|
finish = choice.get("finish_reason")
|
|
if text or finish or len(delta) > 1:
|
|
yield ChatChunk(text=text, finish_reason=finish, metadata=data, delta=delta)
|
|
yield ChatChunk(done=True, finish_reason="stop")
|
|
return
|
|
|
|
response = await client.post(f"{self.base_url}/chat/completions", json=payload)
|
|
if response.status_code >= 400:
|
|
raise BackendError(f"云端请求失败 {response.status_code}: {response.text}")
|
|
data = response.json()
|
|
choice = (data.get("choices") or [{}])[0]
|
|
message = choice.get("message") or {}
|
|
yield ChatChunk(
|
|
text=message.get("content") or "",
|
|
done=True,
|
|
finish_reason=choice.get("finish_reason") or "stop",
|
|
metadata=data,
|
|
delta=message,
|
|
)
|
|
|
|
def info(self) -> dict[str, Any]:
|
|
return {
|
|
"id": self.model.id,
|
|
"kind": "cloud",
|
|
"profile": self.profile.id,
|
|
"remote_model": self.model.remote_model or self.profile.default_model or self.model.id,
|
|
"loaded": self.loaded,
|
|
"base_url": self.base_url,
|
|
}
|