wuyang
250 lines
8.8 KiB
Python
250 lines
8.8 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from agent_platform.config import Settings
|
|
from agent_platform.models import Provider
|
|
|
|
TRANSIENT_COMPLETE_STATUS_CODES = frozenset({408, 425, 429, 500, 502, 503, 504, 524})
|
|
COMPLETE_MAX_ATTEMPTS = 2
|
|
COMPLETE_MAX_TOKENS = 4096
|
|
MAX_RETRY_DELAY_SECONDS = 5.0
|
|
|
|
|
|
def completions_url(base_url: str) -> str:
|
|
base = base_url.rstrip("/")
|
|
if base.endswith("/v1"):
|
|
return f"{base}/chat/completions"
|
|
return f"{base}/v1/chat/completions"
|
|
|
|
|
|
def retry_delay_seconds(response: httpx.Response | None, attempt: int) -> float:
|
|
if response is not None:
|
|
retry_after = response.headers.get("Retry-After")
|
|
if retry_after:
|
|
try:
|
|
return min(MAX_RETRY_DELAY_SECONDS, max(0.0, float(retry_after)))
|
|
except ValueError:
|
|
pass
|
|
return min(MAX_RETRY_DELAY_SECONDS, float(2 ** (attempt - 1)))
|
|
|
|
|
|
async def decode_complete_response(response: httpx.Response) -> dict[str, Any]:
|
|
content_type = response.headers.get("Content-Type", "").lower()
|
|
if "text/event-stream" not in content_type:
|
|
await response.aread()
|
|
data = response.json()
|
|
if not data.get("choices"):
|
|
raise RuntimeError("Model provider returned no choices")
|
|
return data
|
|
|
|
envelope: dict[str, Any] = {"object": "chat.completion"}
|
|
message: dict[str, Any] = {"role": "assistant", "content": None}
|
|
content_parts: list[str] = []
|
|
reasoning_parts: list[str] = []
|
|
tool_calls: dict[int, dict[str, Any]] = {}
|
|
finish_reason: str | None = None
|
|
usage: dict[str, Any] | None = None
|
|
saw_choice = False
|
|
|
|
async for line in response.aiter_lines():
|
|
if not line.startswith("data:"):
|
|
continue
|
|
raw = line[5:].strip()
|
|
if not raw or raw == "[DONE]":
|
|
continue
|
|
chunk = httpx.Response(200, content=raw).json()
|
|
if chunk.get("error"):
|
|
raise RuntimeError(f"Model provider stream error: {chunk['error']}")
|
|
for key in ("id", "created", "model", "system_fingerprint"):
|
|
if chunk.get(key) is not None:
|
|
envelope[key] = chunk[key]
|
|
if chunk.get("usage"):
|
|
usage = chunk["usage"]
|
|
choices = chunk.get("choices") or []
|
|
if not choices:
|
|
continue
|
|
saw_choice = True
|
|
choice = choices[0]
|
|
delta = choice.get("delta") or {}
|
|
if delta.get("role"):
|
|
message["role"] = delta["role"]
|
|
if isinstance(delta.get("content"), str):
|
|
content_parts.append(delta["content"])
|
|
reasoning = delta.get("reasoning_content")
|
|
if not isinstance(reasoning, str):
|
|
reasoning = delta.get("reasoning")
|
|
if isinstance(reasoning, str):
|
|
reasoning_parts.append(reasoning)
|
|
for position, fragment in enumerate(delta.get("tool_calls") or []):
|
|
index = int(fragment.get("index", position))
|
|
merged = tool_calls.setdefault(
|
|
index,
|
|
{
|
|
"id": "",
|
|
"type": "function",
|
|
"function": {"name": "", "arguments": ""},
|
|
},
|
|
)
|
|
if fragment.get("id"):
|
|
merged["id"] = fragment["id"]
|
|
if fragment.get("type"):
|
|
merged["type"] = fragment["type"]
|
|
function = fragment.get("function") or {}
|
|
if function.get("name"):
|
|
merged["function"]["name"] += function["name"]
|
|
if function.get("arguments"):
|
|
merged["function"]["arguments"] += function["arguments"]
|
|
if choice.get("finish_reason") is not None:
|
|
finish_reason = choice["finish_reason"]
|
|
|
|
if not saw_choice:
|
|
raise RuntimeError("Model provider stream returned no choices")
|
|
if content_parts:
|
|
message["content"] = "".join(content_parts)
|
|
if reasoning_parts:
|
|
message["reasoning_content"] = "".join(reasoning_parts)
|
|
if tool_calls:
|
|
message["tool_calls"] = [tool_calls[index] for index in sorted(tool_calls)]
|
|
|
|
envelope["choices"] = [
|
|
{
|
|
"index": 0,
|
|
"message": message,
|
|
"finish_reason": finish_reason,
|
|
}
|
|
]
|
|
if usage is not None:
|
|
envelope["usage"] = usage
|
|
return envelope
|
|
|
|
|
|
class ModelProvider:
|
|
def __init__(self, settings: Settings, client: httpx.AsyncClient | None = None) -> None:
|
|
self.settings = settings
|
|
self.client = client or httpx.AsyncClient(
|
|
timeout=httpx.Timeout(settings.model_timeout_seconds),
|
|
limits=httpx.Limits(max_connections=50, max_keepalive_connections=20),
|
|
)
|
|
self._owns_client = client is None
|
|
|
|
@property
|
|
def headers(self) -> dict[str, str]:
|
|
return self._headers("k1412")
|
|
|
|
def _headers(self, provider: Provider) -> dict[str, str]:
|
|
key = self.settings.deepseek_api_key if provider == "deepseek" else self.settings.model_api_key
|
|
return {
|
|
"Authorization": f"Bearer {key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
def _base_url(self, provider: Provider) -> str:
|
|
return self.settings.deepseek_api_base_url if provider == "deepseek" else self.settings.model_api_base_url
|
|
|
|
@staticmethod
|
|
def _provider_payload(
|
|
payload: dict[str, Any],
|
|
*,
|
|
provider: Provider,
|
|
thinking_enabled: bool,
|
|
reasoning_effort: str | None,
|
|
max_output_tokens: int,
|
|
) -> dict[str, Any]:
|
|
value = dict(payload)
|
|
value.pop("thinking", None)
|
|
value.pop("reasoning_effort", None)
|
|
value["max_tokens"] = max_output_tokens
|
|
if provider == "deepseek" and thinking_enabled:
|
|
value["thinking"] = {"type": "enabled"}
|
|
if reasoning_effort:
|
|
value["reasoning_effort"] = reasoning_effort
|
|
return value
|
|
|
|
async def close(self) -> None:
|
|
if self._owns_client:
|
|
await self.client.aclose()
|
|
|
|
async def complete(
|
|
self,
|
|
*,
|
|
model: str,
|
|
messages: list[dict[str, Any]],
|
|
tools: list[dict[str, Any]] | None = None,
|
|
temperature: float | None = None,
|
|
provider: Provider = "k1412",
|
|
thinking_enabled: bool = False,
|
|
reasoning_effort: str | None = None,
|
|
max_output_tokens: int = COMPLETE_MAX_TOKENS,
|
|
) -> dict[str, Any]:
|
|
payload = self._provider_payload(
|
|
{
|
|
"model": model,
|
|
"messages": messages,
|
|
"stream": True,
|
|
},
|
|
provider=provider,
|
|
thinking_enabled=thinking_enabled,
|
|
reasoning_effort=reasoning_effort,
|
|
max_output_tokens=max_output_tokens,
|
|
)
|
|
if tools:
|
|
payload["tools"] = tools
|
|
payload["tool_choice"] = "auto"
|
|
if temperature is not None:
|
|
payload["temperature"] = temperature
|
|
for attempt in range(1, COMPLETE_MAX_ATTEMPTS + 1):
|
|
response: httpx.Response | None = None
|
|
try:
|
|
async with self.client.stream(
|
|
"POST",
|
|
completions_url(self._base_url(provider)),
|
|
headers=self._headers(provider),
|
|
json=payload,
|
|
) as response:
|
|
if response.status_code in TRANSIENT_COMPLETE_STATUS_CODES and attempt < COMPLETE_MAX_ATTEMPTS:
|
|
await response.aread()
|
|
await asyncio.sleep(retry_delay_seconds(response, attempt))
|
|
continue
|
|
|
|
response.raise_for_status()
|
|
return await decode_complete_response(response)
|
|
except httpx.TransportError:
|
|
if attempt >= COMPLETE_MAX_ATTEMPTS:
|
|
raise
|
|
await asyncio.sleep(retry_delay_seconds(response, attempt))
|
|
continue
|
|
|
|
raise RuntimeError("Model provider retry loop exited unexpectedly")
|
|
|
|
async def forward(
|
|
self,
|
|
payload: dict[str, Any],
|
|
*,
|
|
provider: Provider = "k1412",
|
|
thinking_enabled: bool = False,
|
|
reasoning_effort: str | None = None,
|
|
max_output_tokens: int = COMPLETE_MAX_TOKENS,
|
|
) -> httpx.Response:
|
|
payload = self._provider_payload(
|
|
payload,
|
|
provider=provider,
|
|
thinking_enabled=thinking_enabled,
|
|
reasoning_effort=reasoning_effort,
|
|
max_output_tokens=max_output_tokens,
|
|
)
|
|
request = self.client.build_request(
|
|
"POST",
|
|
completions_url(self._base_url(provider)),
|
|
headers=self._headers(provider),
|
|
json=payload,
|
|
)
|
|
response = await self.client.send(request, stream=bool(payload.get("stream")))
|
|
if response.status_code >= 400:
|
|
await response.aread()
|
|
response.raise_for_status()
|
|
return response
|