175 lines
9.7 KiB
Python
175 lines
9.7 KiB
Python
"""Bounded server-only OpenAI-compatible text calls; never persist wire errors."""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from dataclasses import dataclass, field
|
|
import json
|
|
import os
|
|
import re
|
|
import time
|
|
from urllib.parse import urlsplit
|
|
|
|
import httpx
|
|
|
|
from .db import digest
|
|
|
|
SAFE_NAME = re.compile(r"[A-Za-z0-9][A-Za-z0-9._:/-]{0,119}")
|
|
MAX_SOURCES, MAX_TEXT_CHARS = 21, 16000
|
|
MAX_RESPONSE_BYTES = 262144
|
|
TOTAL_TIMEOUT = 60
|
|
USAGE_KEYS = ("prompt_tokens", "completion_tokens", "total_tokens", "prompt_cache_hit_tokens", "prompt_cache_miss_tokens")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ModelConfig:
|
|
base_url: str = field(repr=False)
|
|
model: str | None
|
|
provider: str | None
|
|
keys: tuple[tuple[str, str], ...] = field(repr=False)
|
|
thinking: str = "disabled"
|
|
timeout_seconds: int = 30
|
|
max_output_tokens: int = 4096
|
|
valid: bool = True
|
|
|
|
@classmethod
|
|
def from_env(cls):
|
|
base = os.environ.get("ASSISTANT_MODEL_BASE_URL", "https://api.deepseek.com/v1").strip().rstrip("/")
|
|
model = os.environ.get("ASSISTANT_MODEL_NAME", "deepseek-flash").strip()
|
|
provider = os.environ.get("ASSISTANT_MODEL_PROVIDER", "deepseek").strip()
|
|
thinking = os.environ.get("ASSISTANT_MODEL_THINKING", "disabled").strip()
|
|
slots = []
|
|
for slot, name in (("primary", "ASSISTANT_MODEL_API_KEY"), ("fallback", "ASSISTANT_MODEL_FALLBACK_API_KEY")):
|
|
key = os.environ.get(name, "").strip()
|
|
if key and not any(existing == key for _, existing in slots):
|
|
slots.append((slot, key))
|
|
valid = bool(SAFE_NAME.fullmatch(model) and SAFE_NAME.fullmatch(provider) and thinking in {"disabled", "enabled"} and not re.search(r"[\x00-\x20\x7f]", base))
|
|
try:
|
|
url = urlsplit(base)
|
|
valid = valid and not (url.username or url.password or url.query or url.fragment) and bool(url.hostname)
|
|
valid = valid and (url.scheme == "https" or (url.scheme == "http" and url.hostname in {"127.0.0.1", "localhost", "::1"}))
|
|
url.port # reject malformed ports without reporting the URL
|
|
timeout = int(os.environ.get("ASSISTANT_MODEL_TIMEOUT_SECONDS", "30"))
|
|
tokens = int(os.environ.get("ASSISTANT_MODEL_MAX_OUTPUT_TOKENS", "4096"))
|
|
valid = valid and 1 <= timeout <= 30 and 256 <= tokens <= 8192
|
|
except (ValueError, TypeError):
|
|
valid, timeout, tokens = False, 30, 4096
|
|
valid = bool(valid and all(re.fullmatch(r"[\x21-\x7e]{1,4096}", key) for _, key in slots))
|
|
return cls(base, model if SAFE_NAME.fullmatch(model) else None, provider if SAFE_NAME.fullmatch(provider) else None,
|
|
tuple(slots), thinking if thinking in {"disabled", "enabled"} else "disabled", timeout, tokens, valid)
|
|
|
|
@property
|
|
def configured(self):
|
|
return self.valid and bool(self.keys) and self.keys[0][0] == "primary"
|
|
|
|
def public(self):
|
|
value = {"configured": self.configured, "provider": self.provider, "model": self.model,
|
|
"key_slots": [slot for slot, _ in self.keys] if self.valid else [], "thinking": self.thinking,
|
|
"timeout_seconds": self.timeout_seconds, "total_timeout_seconds": TOTAL_TIMEOUT,
|
|
"max_output_tokens": self.max_output_tokens, "max_sources": MAX_SOURCES, "max_text_chars": MAX_TEXT_CHARS}
|
|
# Endpoint affects configuration identity but never leaves the server.
|
|
value["config_digest"] = digest({key: item for key, item in {**value, "endpoint_digest": digest(self.base_url)}.items() if key not in {"configured", "key_slots"}})
|
|
return value
|
|
|
|
|
|
def observed_usage(value):
|
|
if not isinstance(value, dict):
|
|
return None
|
|
result = {key: value.get(key) if type(value.get(key)) is int and 0 <= value[key] <= 10000000 else None for key in USAGE_KEYS}
|
|
if all(number is None for number in result.values()):
|
|
return None
|
|
return result
|
|
|
|
|
|
def sum_usage(attempts):
|
|
usages = [attempt["usage"] for attempt in attempts if attempt.get("usage") is not None]
|
|
return {key: sum(value[key] for value in usages if value[key] is not None) if any(value[key] is not None for value in usages) else None for key in USAGE_KEYS} if usages else None
|
|
|
|
|
|
def strict_json(value):
|
|
def unique(pairs):
|
|
result = {}
|
|
for key, item in pairs:
|
|
if key in result:
|
|
raise ValueError("duplicate JSON key")
|
|
result[key] = item
|
|
return result
|
|
try:
|
|
return json.loads(value, object_pairs_hook=unique, parse_constant=lambda _: (_ for _ in ()).throw(ValueError("nonfinite JSON")))
|
|
except RecursionError:
|
|
raise ValueError("JSON nesting exceeds the supported bound") from None
|
|
|
|
|
|
class ModelCallError(Exception):
|
|
def __init__(self, code):
|
|
self.code = code
|
|
super().__init__(code)
|
|
|
|
|
|
class ModelClient:
|
|
def __init__(self, config: ModelConfig, *, http_transport=None):
|
|
self.config = config
|
|
self.http_transport = http_transport
|
|
|
|
async def complete(self, system: str, payload: dict, attempts: list) -> str:
|
|
config = self.config
|
|
try:
|
|
async with asyncio.timeout(TOTAL_TIMEOUT):
|
|
for index, (slot, key) in enumerate(config.keys):
|
|
started = time.monotonic()
|
|
attempt = {"slot": slot, "status": "failed", "error_code": None, "http_status": None,
|
|
"duration_ms": 0, "usage": None, "response_model": None}
|
|
attempts.append(attempt)
|
|
retryable = False
|
|
try:
|
|
async with asyncio.timeout(config.timeout_seconds):
|
|
async with httpx.AsyncClient(timeout=config.timeout_seconds, follow_redirects=False, trust_env=False, transport=self.http_transport) as client:
|
|
async with client.stream("POST", config.base_url + "/chat/completions",
|
|
headers={"Authorization": "Bearer " + key, "Content-Type": "application/json"},
|
|
json={"model": config.model, "messages": [{"role": "system", "content": system}, {"role": "user", "content": json.dumps(payload, ensure_ascii=False)}],
|
|
"thinking": {"type": config.thinking}, "response_format": {"type": "json_object"}, "max_tokens": config.max_output_tokens,
|
|
"stream": False}) as response:
|
|
attempt["http_status"] = response.status_code
|
|
if response.status_code != 200:
|
|
retryable = response.status_code in {401, 402, 429} or 500 <= response.status_code <= 599
|
|
code = {401: "MODEL_AUTH", 402: "MODEL_QUOTA", 429: "MODEL_RATE_LIMIT"}.get(response.status_code, "MODEL_UNAVAILABLE" if retryable else "MODEL_HTTP_ERROR")
|
|
raise ModelCallError(code)
|
|
chunks, size = [], 0
|
|
async for chunk in response.aiter_bytes():
|
|
size += len(chunk)
|
|
if size > MAX_RESPONSE_BYTES:
|
|
raise ModelCallError("MODEL_RESPONSE_TOO_LARGE")
|
|
chunks.append(chunk)
|
|
try:
|
|
raw = strict_json(b"".join(chunks))
|
|
attempt["usage"] = observed_usage(raw.get("usage"))
|
|
response_model = raw.get("model")
|
|
attempt["response_model"] = response_model if isinstance(response_model, str) and SAFE_NAME.fullmatch(response_model) and not any(secret in response_model for _, secret in config.keys) else None
|
|
choice = raw["choices"][0] if len(raw["choices"]) == 1 else None
|
|
message = choice["message"]
|
|
if choice.get("finish_reason") != "stop" or message.get("tool_calls") or message.get("refusal") or not isinstance(message.get("content"), str):
|
|
raise ValueError("invalid response")
|
|
content = message["content"]
|
|
if not content.strip():
|
|
raise ValueError("empty response")
|
|
except (ValueError, TypeError, KeyError, IndexError, AttributeError):
|
|
raise ModelCallError("MODEL_RESPONSE_INVALID") from None
|
|
attempt["status"] = "succeeded"
|
|
return content
|
|
except TimeoutError:
|
|
retryable, attempt["error_code"] = True, "MODEL_TIMEOUT"
|
|
except httpx.TimeoutException:
|
|
retryable, attempt["error_code"] = True, "MODEL_TIMEOUT"
|
|
except httpx.RequestError:
|
|
retryable, attempt["error_code"] = True, "MODEL_NETWORK"
|
|
except ModelCallError as error:
|
|
attempt["error_code"] = error.code
|
|
finally:
|
|
attempt["duration_ms"] = round((time.monotonic() - started) * 1000)
|
|
if not retryable or index == len(config.keys) - 1:
|
|
raise ModelCallError(attempt["error_code"])
|
|
except TimeoutError:
|
|
if attempts:
|
|
attempts[-1].update(status="failed", error_code="MODEL_TIMEOUT")
|
|
raise ModelCallError("MODEL_TIMEOUT") from None
|
|
raise ModelCallError("MODEL_NOT_CONFIGURED")
|