"""Typed client for the fixed internal model-gateway trust boundary.""" from __future__ import annotations import json import math from collections.abc import AsyncIterator, Mapping, Sequence from typing import Any, Literal, Self from urllib.parse import urlsplit import httpx from app.adapters.bailian._base import ( extract_request_id, invalid_request, invalid_response, parse_usage, response_model, safe_identifier, sanitized_error, ) from app.core.config import Settings from app.core.secrets import read_secret_file from app.ports.model_providers import ( ChatCompletionResult, ChatMessage, ChatStreamEvent, EmbeddingResult, ModelProviderError, ProviderErrorKind, ProviderUsage, RankedItem, RerankResult, ) _GATEWAY_HOST = "model-gateway" _GATEWAY_PORT = 8000 _DIMENSION = 1024 _ALLOWED_ROLES = frozenset({"system", "user", "assistant"}) class ModelGatewayAdapter: """Expose internal gateway calls through the provider-neutral model ports.""" def __init__( self, *, token: str, caller: Literal["api", "worker"], base_url: str = "http://model-gateway:8000", embedding_model: str = "text-embedding-v4", rerank_model: str = "qwen3-rerank", chat_model: str = "deepseek-v4-flash", http_client: httpx.AsyncClient | None = None, timeout_seconds: float = 120.0, ) -> None: if not token or token != token.strip(): raise invalid_request("model_gateway.configuration", "invalid_token") if caller not in ("api", "worker"): raise invalid_request("model_gateway.configuration", "invalid_caller") self._base_url = self._validate_base_url(base_url) self._token = token self._caller = caller self._embedding_model = embedding_model self._rerank_model = rerank_model self._chat_model = chat_model self._owns_client = http_client is None self._client = http_client or httpx.AsyncClient( timeout=httpx.Timeout(timeout_seconds), follow_redirects=False, trust_env=False, ) @classmethod def from_settings( cls, settings: Settings, *, http_client: httpx.AsyncClient | None = None, ) -> Self: return cls( token=read_secret_file(settings.model_gateway_token_file), caller=settings.model_gateway_caller, base_url=settings.model_gateway_base_url, embedding_model=settings.embedding_model, rerank_model=settings.rerank_model, chat_model=settings.llm_model, http_client=http_client, timeout_seconds=settings.model_gateway_timeout_seconds, ) async def __aenter__(self) -> Self: return self async def __aexit__(self, *_: object) -> None: await self.aclose() async def aclose(self) -> None: if self._owns_client: await self._client.aclose() async def embed_documents(self, texts: Sequence[str]) -> EmbeddingResult: return await self._embed(texts, input_type="document") async def embed_query(self, text: str) -> EmbeddingResult: return await self._embed((text,), input_type="query") async def _embed( self, texts: Sequence[str], *, input_type: Literal["document", "query"], ) -> EmbeddingResult: operation = f"model_gateway.embedding.{input_type}" validated = self._texts(texts, operation=operation, maximum=10) if input_type == "query" and len(validated) != 1: raise invalid_request(operation, "query_requires_one_text") body = await self._post_json( operation=operation, path="embeddings", payload={"texts": list(validated), "input_type": input_type}, ) vectors = self._vectors(body, expected=len(validated), operation=operation) return EmbeddingResult( vectors=vectors, model=response_model(body, self._embedding_model, sensitive_values=validated), request_id=extract_request_id(body, sensitive_values=validated), usage=parse_usage(body.get("usage")), elapsed_ms=self._elapsed(body, operation=operation), ) async def rerank( self, query: str, documents: Sequence[str], *, top_n: int, instruct: str | None = None, ) -> RerankResult: operation = "model_gateway.rerank" if not isinstance(query, str) or not query: raise invalid_request(operation, "invalid_query") validated = self._texts(documents, operation=operation, maximum=500) if ( isinstance(top_n, bool) or not isinstance(top_n, int) or not 1 <= top_n <= len(validated) ): raise invalid_request(operation, "invalid_top_n") payload: dict[str, Any] = { "query": query, "documents": list(validated), "top_n": top_n, } if instruct is not None: if not isinstance(instruct, str) or not instruct: raise invalid_request(operation, "invalid_instruct") payload["instruct"] = instruct body = await self._post_json(operation=operation, path="rerank", payload=payload) raw_items = body.get("items") if not isinstance(raw_items, list) or len(raw_items) > top_n: raise invalid_response(operation, "invalid_items") items: list[RankedItem] = [] seen: set[int] = set() for raw_item in raw_items: if not isinstance(raw_item, Mapping): raise invalid_response(operation, "invalid_item") index = raw_item.get("index") score = raw_item.get("relevance_score") document = raw_item.get("document") if ( isinstance(index, bool) or not isinstance(index, int) or not 0 <= index < len(validated) or index in seen or isinstance(score, bool) or not isinstance(score, (int, float)) or not math.isfinite(float(score)) or document != validated[index] ): raise invalid_response(operation, "invalid_item") seen.add(index) items.append( RankedItem(index=index, relevance_score=float(score), document=validated[index]) ) return RerankResult( items=tuple(items), model=response_model(body, self._rerank_model, sensitive_values=(query, *validated)), request_id=extract_request_id(body, sensitive_values=(query, *validated)), usage=parse_usage(body.get("usage")), elapsed_ms=self._elapsed(body, operation=operation), ) async def complete( self, messages: Sequence[ChatMessage], *, max_tokens: int, ) -> ChatCompletionResult: operation = "model_gateway.chat.complete" validated = self._messages(messages, max_tokens=max_tokens, operation=operation) body = await self._post_json( operation=operation, path="chat/completions", payload={ "messages": [ {"role": message.role, "content": message.content} for message in validated ], "max_tokens": max_tokens, }, ) content = body.get("content") finish_reason = body.get("finish_reason") if not isinstance(content, str) or ( finish_reason is not None and not isinstance(finish_reason, str) ): raise invalid_response(operation, "invalid_completion") sensitive = tuple(message.content for message in validated) return ChatCompletionResult( content=content, finish_reason=finish_reason, model=response_model(body, self._chat_model, sensitive_values=sensitive), request_id=extract_request_id(body, sensitive_values=sensitive), usage=parse_usage(body.get("usage")), elapsed_ms=self._elapsed(body, operation=operation), ) async def stream( self, messages: Sequence[ChatMessage], *, max_tokens: int, ) -> AsyncIterator[ChatStreamEvent]: operation = "model_gateway.chat.stream" validated = self._messages(messages, max_tokens=max_tokens, operation=operation) payload = { "messages": [ {"role": message.role, "content": message.content} for message in validated ], "max_tokens": max_tokens, } try: async with self._client.stream( "POST", self._url("chat/stream"), headers=self._headers(), json=payload, ) as response: if response.status_code >= 400: await response.aread() self._raise_http_error(operation=operation, response=response) event_name: str | None = None complete_seen = False async for line in response.aiter_lines(): if not line: event_name = None continue if line.startswith(":"): continue if line.startswith("event:"): event_name = line[6:].strip() continue if not line.startswith("data:") or event_name is None: raise invalid_response(operation, "invalid_sse_event") body = self._json_object(line[5:].strip(), operation=operation) if event_name == "error": self._raise_stream_error(body, operation=operation) if event_name not in {"delta", "complete"}: raise invalid_response(operation, "unsupported_sse_event") if complete_seen: raise invalid_response(operation, "event_after_complete") terminal = event_name == "complete" if terminal: complete_seen = True yield self._stream_event( body, operation=operation, terminal=terminal, ) if not complete_seen: raise invalid_response(operation, "missing_complete_event") except ModelProviderError: raise except httpx.TimeoutException: raise sanitized_error( operation=operation, kind=ProviderErrorKind.TIMEOUT, provider_code="request_timeout", retryable=True, ) from None except httpx.HTTPError: raise sanitized_error( operation=operation, kind=ProviderErrorKind.TRANSPORT, provider_code="transport_error", retryable=True, ) from None def _headers(self) -> dict[str, str]: return { "Authorization": f"Bearer {self._token}", "Content-Type": "application/json", "X-RAG-Caller": self._caller, } def _url(self, path: str) -> str: return f"{self._base_url}/internal/v1/{path.lstrip('/')}" async def _post_json( self, *, operation: str, path: str, payload: Mapping[str, Any], ) -> Mapping[str, Any]: try: response = await self._client.post( self._url(path), headers=self._headers(), json=payload ) except httpx.TimeoutException: raise sanitized_error( operation=operation, kind=ProviderErrorKind.TIMEOUT, provider_code="request_timeout", retryable=True, ) from None except httpx.HTTPError: raise sanitized_error( operation=operation, kind=ProviderErrorKind.TRANSPORT, provider_code="transport_error", retryable=True, ) from None if response.status_code >= 400: self._raise_http_error(operation=operation, response=response) return self._json_object(response.text, operation=operation) def _raise_http_error(self, *, operation: str, response: httpx.Response) -> None: status = response.status_code if status == 400 or status == 422: kind, retryable = ProviderErrorKind.INVALID_REQUEST, False elif status == 401: kind, retryable = ProviderErrorKind.AUTHENTICATION, False elif status == 403: kind, retryable = ProviderErrorKind.PERMISSION_DENIED, False elif status == 404: kind, retryable = ProviderErrorKind.NOT_FOUND, False elif status == 408 or status == 504: kind, retryable = ProviderErrorKind.TIMEOUT, True elif status == 429: kind, retryable = ProviderErrorKind.RATE_LIMITED, True else: kind, retryable = ProviderErrorKind.UPSTREAM, status >= 500 # The trusted gateway exposes only a fixed provider-neutral error object. # Preserve that category for diagnostics while discarding all other body data. request_id: str | None = None try: decoded = response.json() except (ValueError, httpx.ResponseNotRead): decoded = None if isinstance(decoded, Mapping): raw_error = decoded.get("error") if isinstance(raw_error, Mapping): try: kind = ProviderErrorKind(str(raw_error.get("kind"))) except ValueError: pass retryable = raw_error.get("retryable") is True request_id = extract_request_id( raw_error, sensitive_values=(self._token,), ) raise sanitized_error( operation=operation, kind=kind, status_code=status, provider_code="model_gateway_rejected", request_id=request_id, retryable=retryable, ) def _raise_stream_error(self, body: Mapping[str, Any], *, operation: str) -> None: raw_kind = body.get("kind") try: kind = ProviderErrorKind(str(raw_kind)) except ValueError: kind = ProviderErrorKind.UPSTREAM retryable = body.get("retryable") is True request_id = body.get("request_id") raise sanitized_error( operation=operation, kind=kind, provider_code="model_gateway_stream_error", request_id=request_id if isinstance(request_id, str) else None, retryable=retryable, ) @staticmethod def _validate_base_url(value: str) -> str: normalized = value.rstrip("/") parsed = urlsplit(normalized) if ( parsed.scheme != "http" or parsed.hostname != _GATEWAY_HOST or parsed.port != _GATEWAY_PORT or parsed.path not in ("", "/") or parsed.username is not None or parsed.password is not None or parsed.query or parsed.fragment ): raise invalid_request("model_gateway.configuration", "invalid_base_url") return normalized @staticmethod def _texts( values: Sequence[str], *, operation: str, maximum: int, ) -> tuple[str, ...]: if isinstance(values, (str, bytes)) or not isinstance(values, Sequence): raise invalid_request(operation, "invalid_text_collection") validated = tuple(values) if not validated or len(validated) > maximum: raise invalid_request(operation, "invalid_text_count") if any(not isinstance(value, str) or not value for value in validated): raise invalid_request(operation, "invalid_text") return validated @staticmethod def _messages( messages: object, *, max_tokens: int, operation: str, ) -> tuple[ChatMessage, ...]: if isinstance(messages, (str, bytes)) or not isinstance(messages, Sequence): raise invalid_request(operation, "invalid_message_collection") validated = tuple(messages) if not validated or isinstance(max_tokens, bool) or not isinstance(max_tokens, int): raise invalid_request(operation, "invalid_messages") if not 1 <= max_tokens <= 8192: raise invalid_request(operation, "invalid_max_tokens") if any( not isinstance(message, ChatMessage) or message.role not in _ALLOWED_ROLES or not message.content for message in validated ): raise invalid_request(operation, "invalid_message") return validated @staticmethod def _json_object(value: str, *, operation: str) -> Mapping[str, Any]: try: body = json.loads(value) except (TypeError, ValueError): raise invalid_response(operation, "invalid_json") from None if not isinstance(body, Mapping): raise invalid_response(operation, "invalid_json_object") return body @staticmethod def _elapsed(body: Mapping[str, Any], *, operation: str) -> float: value = body.get("elapsed_ms") if isinstance(value, bool) or not isinstance(value, (int, float)): raise invalid_response(operation, "invalid_elapsed_ms") elapsed = float(value) if not math.isfinite(elapsed) or elapsed < 0: raise invalid_response(operation, "invalid_elapsed_ms") return elapsed @staticmethod def _vectors( body: Mapping[str, Any], *, expected: int, operation: str, ) -> tuple[tuple[float, ...], ...]: raw_vectors = body.get("vectors") if not isinstance(raw_vectors, list) or len(raw_vectors) != expected: raise invalid_response(operation, "invalid_embedding_count") vectors: list[tuple[float, ...]] = [] for raw_vector in raw_vectors: if not isinstance(raw_vector, list) or len(raw_vector) != _DIMENSION: raise invalid_response(operation, "invalid_embedding_dimensions") vector: list[float] = [] for component in raw_vector: if isinstance(component, bool) or not isinstance(component, (int, float)): raise invalid_response(operation, "invalid_embedding_component") number = float(component) if not math.isfinite(number): raise invalid_response(operation, "invalid_embedding_component") vector.append(number) if math.hypot(*vector) <= 0: raise invalid_response(operation, "invalid_embedding_norm") vectors.append(tuple(vector)) return tuple(vectors) def _stream_event( self, body: Mapping[str, Any], *, operation: str, terminal: bool, ) -> ChatStreamEvent: delta = body.get("delta", "") finish_reason = body.get("finish_reason") if not isinstance(delta, str) or ( finish_reason is not None and not isinstance(finish_reason, str) ): raise invalid_response(operation, "invalid_stream_event") raw_model = body.get("model") model = safe_identifier(raw_model, sensitive_values=(self._token,)) if model is None: raise invalid_response(operation, "invalid_stream_model") request_id = extract_request_id(body, sensitive_values=(self._token,)) if body.get("request_id") is not None and request_id is None: raise invalid_response(operation, "invalid_stream_request_id") if terminal and "elapsed_ms" not in body: raise invalid_response(operation, "missing_elapsed_ms") elapsed = body.get("elapsed_ms", 0.0) if isinstance(elapsed, bool) or not isinstance(elapsed, (int, float)): raise invalid_response(operation, "invalid_elapsed_ms") normalized_elapsed = float(elapsed) if not math.isfinite(normalized_elapsed) or normalized_elapsed < 0: raise invalid_response(operation, "invalid_elapsed_ms") if terminal and not isinstance(body.get("usage"), Mapping): raise invalid_response(operation, "missing_usage") return ChatStreamEvent( delta=delta, finish_reason=finish_reason, model=model, request_id=request_id, usage=parse_usage(body.get("usage")) if "usage" in body else ProviderUsage(), elapsed_ms=normalized_elapsed, )