"""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")