|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import json |
| 4 | +import time |
| 5 | +import urllib.error |
| 6 | +import urllib.request |
| 7 | +from dataclasses import dataclass |
| 8 | +from typing import Any, Callable |
| 9 | + |
| 10 | +from .provider_registry import ( |
| 11 | + AiProviderConfigurationError, |
| 12 | + ResolvedProviderProfile, |
| 13 | + validate_provider_url, |
| 14 | +) |
| 15 | + |
| 16 | + |
| 17 | +RETRYABLE_HTTP_STATUS = frozenset({429, 500, 502, 503, 504}) |
| 18 | +MAX_RETRY_AFTER_SECONDS = 10.0 |
| 19 | + |
| 20 | + |
| 21 | +class ProviderUnavailable(RuntimeError): |
| 22 | + code = "PROVIDER_UNAVAILABLE" |
| 23 | + |
| 24 | + def __init__(self, reason: str, *, attempts: int = 0): |
| 25 | + super().__init__(reason) |
| 26 | + self.reason = reason |
| 27 | + self.attempts = attempts |
| 28 | + |
| 29 | + |
| 30 | +class ModelOutputInvalid(RuntimeError): |
| 31 | + code = "MODEL_OUTPUT_INVALID" |
| 32 | + |
| 33 | + def __init__(self, reason: str, *, attempts: int = 1): |
| 34 | + super().__init__(reason) |
| 35 | + self.reason = reason |
| 36 | + self.attempts = attempts |
| 37 | + |
| 38 | + |
| 39 | +class _NoRedirect(urllib.request.HTTPRedirectHandler): |
| 40 | + def redirect_request(self, req, fp, code, msg, headers, newurl): # noqa: ANN001 |
| 41 | + return None |
| 42 | + |
| 43 | + |
| 44 | +@dataclass(frozen=True) |
| 45 | +class ProviderChatResult: |
| 46 | + content: str |
| 47 | + configured_model: str |
| 48 | + resolved_model: str | None |
| 49 | + provider_request_id: str | None |
| 50 | + usage: dict[str, Any] |
| 51 | + attempts: int |
| 52 | + transport_events: tuple[dict[str, Any], ...] |
| 53 | + |
| 54 | + |
| 55 | +class OpenAICompatibleChatProvider: |
| 56 | + """Bounded non-streaming JSON chat transport with no provider SDK.""" |
| 57 | + |
| 58 | + def __init__( |
| 59 | + self, |
| 60 | + *, |
| 61 | + opener: Any | None = None, |
| 62 | + sleep: Callable[[float], None] = time.sleep, |
| 63 | + monotonic: Callable[[], float] = time.monotonic, |
| 64 | + resolver: Callable[..., list[tuple]] | None = None, |
| 65 | + ) -> None: |
| 66 | + self._opener = opener or urllib.request.build_opener(_NoRedirect()) |
| 67 | + self._sleep = sleep |
| 68 | + self._monotonic = monotonic |
| 69 | + self._resolver = resolver |
| 70 | + |
| 71 | + def _endpoint(self, profile: ResolvedProviderProfile) -> str: |
| 72 | + kw = {"allowed_local_hosts": set(profile.allowed_local_hosts)} |
| 73 | + if self._resolver is not None: |
| 74 | + kw["resolver"] = self._resolver |
| 75 | + try: |
| 76 | + base_url = validate_provider_url( |
| 77 | + profile.base_url, str(profile.document["locality"]), **kw |
| 78 | + ) |
| 79 | + except AiProviderConfigurationError as error: |
| 80 | + raise ProviderUnavailable(error.reason) from None |
| 81 | + return base_url.rstrip("/") + "/chat/completions" |
| 82 | + |
| 83 | + @staticmethod |
| 84 | + def _retry_delay(headers: Any, attempt: int) -> float: |
| 85 | + raw = headers.get("Retry-After") if headers is not None else None |
| 86 | + if raw is not None: |
| 87 | + try: |
| 88 | + return min(MAX_RETRY_AFTER_SECONDS, max(0.0, float(raw))) |
| 89 | + except (TypeError, ValueError): |
| 90 | + pass |
| 91 | + return min(MAX_RETRY_AFTER_SECONDS, 0.25 * (2 ** max(0, attempt - 1))) |
| 92 | + |
| 93 | + @staticmethod |
| 94 | + def _request_body( |
| 95 | + profile: ResolvedProviderProfile, system_prompt: str, context_json: str |
| 96 | + ) -> bytes: |
| 97 | + if "JSON" not in system_prompt.upper() or "JSON" not in context_json.upper(): |
| 98 | + raise ModelOutputInvalid("messages_must_explicitly_request_json") |
| 99 | + body: dict[str, Any] = { |
| 100 | + "model": profile.configured_model, |
| 101 | + "messages": [ |
| 102 | + {"role": "system", "content": system_prompt}, |
| 103 | + {"role": "user", "content": context_json}, |
| 104 | + ], |
| 105 | + "response_format": {"type": "json_object"}, |
| 106 | + "stream": False, |
| 107 | + } |
| 108 | + body.update(dict(profile.document["static_request_fields"])) |
| 109 | + return json.dumps( |
| 110 | + body, sort_keys=True, separators=(",", ":"), ensure_ascii=False |
| 111 | + ).encode("utf-8") |
| 112 | + |
| 113 | + @staticmethod |
| 114 | + def _decode_response( |
| 115 | + payload: bytes, profile: ResolvedProviderProfile, attempts: int, |
| 116 | + events: list[dict[str, Any]], headers: Any, |
| 117 | + ) -> ProviderChatResult: |
| 118 | + try: |
| 119 | + document = json.loads(payload) |
| 120 | + except (UnicodeDecodeError, json.JSONDecodeError): |
| 121 | + raise ModelOutputInvalid("provider_response_is_not_json", attempts=attempts) from None |
| 122 | + if not isinstance(document, dict): |
| 123 | + raise ModelOutputInvalid("provider_response_root_is_not_object", attempts=attempts) |
| 124 | + choices = document.get("choices") |
| 125 | + if not isinstance(choices, list) or not choices or not isinstance(choices[0], dict): |
| 126 | + raise ModelOutputInvalid("provider_response_has_no_choice", attempts=attempts) |
| 127 | + choice = choices[0] |
| 128 | + if choice.get("finish_reason") == "length": |
| 129 | + raise ModelOutputInvalid("provider_output_was_truncated", attempts=attempts) |
| 130 | + message = choice.get("message") |
| 131 | + if not isinstance(message, dict): |
| 132 | + raise ModelOutputInvalid("provider_choice_has_no_message", attempts=attempts) |
| 133 | + if message.get("tool_calls") or message.get("function_call"): |
| 134 | + raise ModelOutputInvalid("provider_attempted_tool_call", attempts=attempts) |
| 135 | + content = message.get("content") |
| 136 | + if not isinstance(content, str): |
| 137 | + raise ModelOutputInvalid("provider_message_content_is_not_text", attempts=attempts) |
| 138 | + usage = document.get("usage") if isinstance(document.get("usage"), dict) else {} |
| 139 | + request_id = None |
| 140 | + if headers is not None: |
| 141 | + request_id = headers.get("x-request-id") or headers.get("request-id") |
| 142 | + if request_id is None and isinstance(document.get("id"), str): |
| 143 | + request_id = document["id"] |
| 144 | + return ProviderChatResult( |
| 145 | + content=content, |
| 146 | + configured_model=profile.configured_model, |
| 147 | + resolved_model=(document.get("model") if isinstance(document.get("model"), str) else None), |
| 148 | + provider_request_id=request_id, |
| 149 | + usage=dict(usage), |
| 150 | + attempts=attempts, |
| 151 | + transport_events=tuple(events), |
| 152 | + ) |
| 153 | + |
| 154 | + def complete_json( |
| 155 | + self, |
| 156 | + profile: ResolvedProviderProfile, |
| 157 | + *, |
| 158 | + system_prompt: str, |
| 159 | + context_json: str, |
| 160 | + ) -> ProviderChatResult: |
| 161 | + endpoint = self._endpoint(profile) |
| 162 | + payload = self._request_body(profile, system_prompt, context_json) |
| 163 | + max_request = int(profile.document["bounds"]["max_request_bytes"]) |
| 164 | + max_response = int(profile.document["bounds"]["max_response_bytes"]) |
| 165 | + if len(payload) > max_request: |
| 166 | + raise ModelOutputInvalid("provider_request_exceeds_profile_bound") |
| 167 | + maximum_attempts = int(profile.document["bounds"]["max_provider_attempts"]) |
| 168 | + timeout = float(profile.document["timeouts"]["request_seconds"]) |
| 169 | + events: list[dict[str, Any]] = [] |
| 170 | + |
| 171 | + for attempt in range(1, maximum_attempts + 1): |
| 172 | + request = urllib.request.Request( |
| 173 | + endpoint, |
| 174 | + data=payload, |
| 175 | + method="POST", |
| 176 | + headers={ |
| 177 | + "Authorization": "Bearer " + profile.api_key, |
| 178 | + "Content-Type": "application/json", |
| 179 | + "Accept": "application/json", |
| 180 | + }, |
| 181 | + ) |
| 182 | + started = self._monotonic() |
| 183 | + try: |
| 184 | + with self._opener.open(request, timeout=timeout) as response: |
| 185 | + body = response.read(max_response + 1) |
| 186 | + duration = max(0.0, self._monotonic() - started) |
| 187 | + events.append( |
| 188 | + { |
| 189 | + "attempt": attempt, |
| 190 | + "status": int(response.status), |
| 191 | + "duration_seconds": duration, |
| 192 | + "response_bytes": len(body), |
| 193 | + } |
| 194 | + ) |
| 195 | + if len(body) > max_response: |
| 196 | + raise ModelOutputInvalid( |
| 197 | + "provider_response_exceeds_profile_bound", attempts=attempt |
| 198 | + ) |
| 199 | + return self._decode_response( |
| 200 | + body, profile, attempt, events, response.headers |
| 201 | + ) |
| 202 | + except urllib.error.HTTPError as error: |
| 203 | + duration = max(0.0, self._monotonic() - started) |
| 204 | + status = int(error.code) |
| 205 | + error_headers = error.headers |
| 206 | + error.close() |
| 207 | + events.append( |
| 208 | + { |
| 209 | + "attempt": attempt, |
| 210 | + "status": status, |
| 211 | + "duration_seconds": duration, |
| 212 | + "response_bytes": 0, |
| 213 | + } |
| 214 | + ) |
| 215 | + if status in {301, 302, 303, 307, 308}: |
| 216 | + raise ProviderUnavailable("provider_redirect_refused", attempts=attempt) from None |
| 217 | + if status not in RETRYABLE_HTTP_STATUS or attempt >= maximum_attempts: |
| 218 | + reason = ( |
| 219 | + "provider_authentication_failed" |
| 220 | + if status in {401, 403} |
| 221 | + else "provider_http_failure" |
| 222 | + ) |
| 223 | + raise ProviderUnavailable(reason, attempts=attempt) from None |
| 224 | + self._sleep(self._retry_delay(error_headers, attempt)) |
| 225 | + except (TimeoutError, urllib.error.URLError, ConnectionError, OSError) as error: |
| 226 | + duration = max(0.0, self._monotonic() - started) |
| 227 | + events.append( |
| 228 | + { |
| 229 | + "attempt": attempt, |
| 230 | + "status": "transport_error", |
| 231 | + "duration_seconds": duration, |
| 232 | + "response_bytes": 0, |
| 233 | + "exception": type(error).__name__, |
| 234 | + } |
| 235 | + ) |
| 236 | + if attempt >= maximum_attempts: |
| 237 | + raise ProviderUnavailable( |
| 238 | + "provider_transport_failure", attempts=attempt |
| 239 | + ) from None |
| 240 | + self._sleep(self._retry_delay(None, attempt)) |
| 241 | + |
| 242 | + raise ProviderUnavailable("provider_attempt_budget_exhausted", attempts=maximum_attempts) |
0 commit comments