"""Credential-isolated internal gateway for all cloud model capabilities.""" import asyncio import json import os import secrets from collections.abc import AsyncIterator, Callable, Mapping from contextlib import asynccontextmanager from pathlib import Path from typing import Annotated, Literal, Protocol, Self, runtime_checkable from fastapi import Depends, FastAPI, Request, status from fastapi.exceptions import RequestValidationError from fastapi.responses import JSONResponse, StreamingResponse from pydantic import BaseModel, ConfigDict, Field, model_validator from starlette.types import Receive, Scope, Send from app import __version__ from app.adapters.bailian import ( BailianChatAdapter, BailianEmbeddingAdapter, BailianRerankerAdapter, ) from app.core.config import Settings from app.core.secrets import SecretFileError, read_secret_file from app.ports.model_providers import ( ChatMessage, ChatProvider, ChatStreamEvent, EmbeddingProvider, ModelProviderError, ProviderErrorKind, ProviderUsage, Reranker, ) Caller = Literal["api", "worker"] InputType = Literal["document", "query"] Role = Literal["system", "user", "assistant"] SettingsFactory = Callable[[], Settings] CALLERS: tuple[Caller, Caller] = ("api", "worker") DEFAULT_ALLOWED_TOKEN_FILES = ( "/run/secrets/model_gateway_api_token,/run/secrets/model_gateway_worker_token" # noqa: S105 ) MAX_CHAT_MESSAGES = 100 MAX_CHAT_CONTENT_CHARS = 100_000 MAX_CHAT_OUTPUT_TOKENS = 8_192 ALLOWED_FINISH_REASONS = frozenset( {"stop", "length", "content_filter", "tool_calls", "function_call"} ) class _StrictModel(BaseModel): model_config = ConfigDict(extra="forbid") class EmbeddingRequest(_StrictModel): texts: list[Annotated[str, Field(min_length=1, max_length=8_192)]] = Field( min_length=1, max_length=10, ) input_type: InputType @model_validator(mode="after") def validate_query_count(self) -> Self: if self.input_type == "query" and len(self.texts) != 1: raise ValueError("query embedding accepts exactly one text") return self class RerankRequest(_StrictModel): query: str = Field(min_length=1, max_length=4_000) documents: list[Annotated[str, Field(min_length=1, max_length=4_000)]] = Field( min_length=1, max_length=500, ) top_n: int = Field(ge=1, le=500) instruct: str | None = Field(default=None, min_length=1, max_length=4_000) @model_validator(mode="after") def validate_top_n(self) -> Self: if self.top_n > len(self.documents): raise ValueError("top_n must not exceed document count") return self class ChatMessageRequest(_StrictModel): role: Role content: str = Field(min_length=1, max_length=MAX_CHAT_CONTENT_CHARS) class ChatRequest(_StrictModel): messages: list[ChatMessageRequest] = Field(min_length=1, max_length=MAX_CHAT_MESSAGES) max_tokens: int = Field(default=1_024, ge=1, le=MAX_CHAT_OUTPUT_TOKENS) class UsageResponse(_StrictModel): input_tokens: int | None output_tokens: int | None total_tokens: int | None class EmbeddingResponse(_StrictModel): vectors: list[list[float]] model: str request_id: str | None usage: UsageResponse elapsed_ms: float class RankedItemResponse(_StrictModel): index: int relevance_score: float document: str class RerankResponse(_StrictModel): items: list[RankedItemResponse] model: str request_id: str | None usage: UsageResponse elapsed_ms: float class ChatResponse(_StrictModel): content: str finish_reason: str | None model: str request_id: str | None usage: UsageResponse elapsed_ms: float class _UnauthorizedError(RuntimeError): pass class _ForbiddenError(RuntimeError): pass class _UnavailableError(RuntimeError): pass class _RestartRequiredError(RuntimeError): pass @runtime_checkable class _SupportsAclose(Protocol): async def aclose(self) -> None: ... class _ClosingStreamingResponse(StreamingResponse): """Close the response iterator on completion, cancellation, or send failure.""" async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: try: await super().__call__(scope, receive, send) finally: await _close_stream(self.body_iterator) class _Runtime: def __init__(self) -> None: self.embedding: EmbeddingProvider | None = None self.reranker: Reranker | None = None self.chat: ChatProvider | None = None self.semaphore: asyncio.Semaphore | None = None self.allowed_tokens: dict[Caller, str] = {} self.local_configuration_check: Callable[[], bool] = lambda: False self.restart_required = False @property def available(self) -> bool: return all( provider is not None for provider in (self.embedding, self.reranker, self.chat, self.semaphore) ) and set(self.allowed_tokens) == {"api", "worker"} def invalidate_for_restart(self) -> None: """Atomically stop accepting work after mounted configuration changes.""" self.restart_required = True self.embedding = None self.reranker = None self.chat = None self.semaphore = None self.allowed_tokens = {} def _usage_response(usage: ProviderUsage) -> UsageResponse: return UsageResponse( input_tokens=usage.input_tokens, output_tokens=usage.output_tokens, total_tokens=usage.total_tokens, ) def _merge_usage(current: ProviderUsage, update: ProviderUsage) -> ProviderUsage: return ProviderUsage( input_tokens=( update.input_tokens if update.input_tokens is not None else current.input_tokens ), output_tokens=( update.output_tokens if update.output_tokens is not None else current.output_tokens ), total_tokens=( update.total_tokens if update.total_tokens is not None else current.total_tokens ), ) def _boundary_error(operation: str) -> ModelProviderError: return ModelProviderError( operation=operation, kind=ProviderErrorKind.INVALID_RESPONSE, provider_code="gateway_boundary_error", ) def _normalize_allowed_tokens(allowed_tokens: Mapping[str, str]) -> dict[Caller, str]: if set(allowed_tokens) != {"api", "worker"}: raise ValueError("allowed_tokens must define api and worker identities") normalized: dict[Caller, str] = {} for caller in CALLERS: token = allowed_tokens.get(caller) if ( not isinstance(token, str) or not token or token != token.strip() or "\n" in token or "\r" in token or len(token) > 4_096 ): raise ValueError("allowed token is invalid") normalized[caller] = token if secrets.compare_digest(normalized["api"], normalized["worker"]): raise ValueError("api and worker tokens must be different") return normalized def _load_allowed_tokens_from_files() -> dict[Caller, str]: raw = os.environ.get("MODEL_GATEWAY_ALLOWED_TOKEN_FILES", DEFAULT_ALLOWED_TOKEN_FILES) entries = raw.split(",") if len(entries) != 2 or any(not entry.strip() for entry in entries): raise ValueError("MODEL_GATEWAY_ALLOWED_TOKEN_FILES must define two files") paths: dict[str, str] = {} if all("=" not in entry for entry in entries): paths = {caller: entry.strip() for caller, entry in zip(CALLERS, entries, strict=True)} else: for entry in entries: caller, separator, path = entry.partition("=") if not separator or caller.strip() in paths: raise ValueError("invalid model gateway token file mapping") paths[caller.strip()] = path.strip() if set(paths) != {"api", "worker"}: raise ValueError("model gateway token files must map api and worker") return _normalize_allowed_tokens( { "api": read_secret_file(Path(paths["api"])), "worker": read_secret_file(Path(paths["worker"])), } ) def _provider_status(error: ModelProviderError) -> int: return { ProviderErrorKind.INVALID_REQUEST: status.HTTP_422_UNPROCESSABLE_CONTENT, ProviderErrorKind.RATE_LIMITED: status.HTTP_429_TOO_MANY_REQUESTS, ProviderErrorKind.TIMEOUT: status.HTTP_504_GATEWAY_TIMEOUT, ProviderErrorKind.TRANSPORT: status.HTTP_503_SERVICE_UNAVAILABLE, }.get(error.kind, status.HTTP_502_BAD_GATEWAY) def _error_payload( kind: str, *, retryable: bool = False, request_id: str | None = None, ) -> dict[str, dict[str, str | bool | None]]: return { "error": { "kind": kind, "retryable": retryable, "request_id": request_id, } } def _sse(event: str, payload: Mapping[str, object]) -> bytes: data = json.dumps(payload, ensure_ascii=False, separators=(",", ":"), sort_keys=True) return f"event: {event}\ndata: {data}\n\n".encode() async def _close_stream(stream: object | None) -> None: if isinstance(stream, _SupportsAclose): try: await stream.aclose() except Exception: # Closing is best-effort and must never expose provider exception text. return def create_model_gateway_app( *, embedding_provider: EmbeddingProvider | None = None, reranker: Reranker | None = None, chat_provider: ChatProvider | None = None, allowed_tokens: Mapping[str, str] | None = None, max_concurrency: int | None = None, settings_factory: SettingsFactory = Settings, ) -> FastAPI: """Create the internal model gateway with injectable hermetic providers.""" injected = (embedding_provider, reranker, chat_provider) if any(provider is not None for provider in injected) != all( provider is not None for provider in injected ): raise ValueError("all three providers must be injected together") providers_are_injected = all(provider is not None for provider in injected) if providers_are_injected != (allowed_tokens is not None): raise ValueError("injected providers and allowed_tokens must be supplied together") if max_concurrency is not None and ( isinstance(max_concurrency, bool) or max_concurrency < 1 or max_concurrency > 100 ): raise ValueError("max_concurrency must be between 1 and 100") runtime = _Runtime() owned_adapters: list[BailianEmbeddingAdapter | BailianRerankerAdapter | BailianChatAdapter] = [] @asynccontextmanager async def lifespan(_: FastAPI) -> AsyncIterator[None]: if providers_are_injected: assert embedding_provider is not None assert reranker is not None assert chat_provider is not None assert allowed_tokens is not None runtime.embedding = embedding_provider runtime.reranker = reranker runtime.chat = chat_provider runtime.allowed_tokens = _normalize_allowed_tokens(allowed_tokens) runtime.semaphore = asyncio.Semaphore(max_concurrency or 4) runtime.local_configuration_check = lambda: True runtime.restart_required = False else: try: settings = settings_factory() api_key = settings.bailian_api_key() loaded_tokens = _load_allowed_tokens_from_files() embedding_adapter = BailianEmbeddingAdapter( api_key=api_key, base_url=settings.bailian_openai_base_url, model=settings.embedding_model, dimensions=settings.embedding_dimension, timeout_seconds=settings.model_timeout_seconds, max_retries=settings.model_max_retries, ) owned_adapters.append(embedding_adapter) rerank_adapter = BailianRerankerAdapter( api_key=api_key, base_url=settings.bailian_rerank_base_url, model=settings.rerank_model, timeout_seconds=settings.model_timeout_seconds, max_retries=settings.model_max_retries, ) owned_adapters.append(rerank_adapter) chat_adapter = BailianChatAdapter( api_key=api_key, base_url=settings.bailian_openai_base_url, model=settings.llm_model, timeout_seconds=settings.model_timeout_seconds, max_retries=settings.model_max_retries, ) owned_adapters.append(chat_adapter) runtime.embedding = embedding_adapter runtime.reranker = rerank_adapter runtime.chat = chat_adapter runtime.allowed_tokens = loaded_tokens runtime.semaphore = asyncio.Semaphore( max_concurrency or settings.model_max_concurrency ) runtime.restart_required = False def check_local_configuration() -> bool: try: current_settings = settings_factory() current_api_key = current_settings.bailian_api_key() current_tokens = _load_allowed_tokens_from_files() same_key = secrets.compare_digest(current_api_key, api_key) same_tokens = all( secrets.compare_digest(current_tokens[caller], loaded_tokens[caller]) for caller in CALLERS ) same_provider_contract = ( current_settings.bailian_openai_base_url == settings.bailian_openai_base_url and current_settings.bailian_rerank_base_url == settings.bailian_rerank_base_url and current_settings.embedding_model == settings.embedding_model and current_settings.embedding_dimension == settings.embedding_dimension and current_settings.rerank_model == settings.rerank_model and current_settings.llm_model == settings.llm_model ) return same_key and same_tokens and same_provider_contract except (OSError, SecretFileError, ValueError, ModelProviderError): return False runtime.local_configuration_check = check_local_configuration except (OSError, SecretFileError, ValueError, ModelProviderError): runtime.local_configuration_check = lambda: False try: yield finally: runtime.embedding = None runtime.reranker = None runtime.chat = None runtime.semaphore = None runtime.allowed_tokens = {} runtime.local_configuration_check = lambda: False runtime.restart_required = False for adapter in reversed(owned_adapters): await adapter.aclose() owned_adapters.clear() gateway = FastAPI( title="Geological RAG Model Gateway", version=__version__, docs_url=None, redoc_url=None, openapi_url=None, lifespan=lifespan, ) @gateway.exception_handler(RequestValidationError) async def request_validation_error( _: Request, __: RequestValidationError, ) -> JSONResponse: return JSONResponse( status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, content=_error_payload("invalid_request"), ) @gateway.exception_handler(ModelProviderError) async def model_provider_error(_: Request, error: ModelProviderError) -> JSONResponse: return JSONResponse( status_code=_provider_status(error), content=_error_payload( error.kind.value, retryable=error.retryable, request_id=error.request_id, ), ) @gateway.exception_handler(_UnauthorizedError) async def unauthorized_error(_: Request, __: _UnauthorizedError) -> JSONResponse: return JSONResponse( status_code=status.HTTP_401_UNAUTHORIZED, content=_error_payload("unauthorized"), headers={"WWW-Authenticate": "Bearer"}, ) @gateway.exception_handler(_ForbiddenError) async def forbidden_error(_: Request, __: _ForbiddenError) -> JSONResponse: return JSONResponse( status_code=status.HTTP_403_FORBIDDEN, content=_error_payload("forbidden"), ) @gateway.exception_handler(_UnavailableError) async def unavailable_error(_: Request, __: _UnavailableError) -> JSONResponse: return JSONResponse( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, content=_error_payload("unavailable", retryable=True), ) @gateway.exception_handler(_RestartRequiredError) async def restart_required_error(_: Request, __: _RestartRequiredError) -> JSONResponse: return JSONResponse( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, content=_error_payload("restart_required"), ) @gateway.exception_handler(Exception) async def unexpected_error(_: Request, __: Exception) -> JSONResponse: return JSONResponse( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, content=_error_payload("internal_error"), ) def ensure_current_configuration() -> None: if runtime.restart_required: raise _RestartRequiredError if not runtime.available: raise _UnavailableError if not runtime.local_configuration_check(): runtime.invalidate_for_restart() raise _RestartRequiredError def require_runtime() -> tuple[ EmbeddingProvider, Reranker, ChatProvider, asyncio.Semaphore, ]: ensure_current_configuration() assert runtime.embedding is not None assert runtime.reranker is not None assert runtime.chat is not None assert runtime.semaphore is not None return runtime.embedding, runtime.reranker, runtime.chat, runtime.semaphore async def authorize(request: Request) -> Caller: ensure_current_configuration() authorization = request.headers.get("authorization", "") scheme, separator, credential = authorization.partition(" ") caller_value = request.headers.get("x-rag-caller", "") if ( not separator or scheme.lower() != "bearer" or not credential or len(credential) > 4_096 or caller_value not in {"api", "worker"} ): raise _UnauthorizedError matched_identity: Caller | None = None for identity, allowed_token in runtime.allowed_tokens.items(): if secrets.compare_digest(credential, allowed_token): matched_identity = identity if matched_identity is None or matched_identity != caller_value: raise _UnauthorizedError return matched_identity @gateway.get("/health/live", include_in_schema=False) async def live() -> dict[str, str]: return {"status": "ok", "version": __version__} @gateway.get("/health/ready", include_in_schema=False) async def ready() -> JSONResponse: if runtime.restart_required: configuration_status = "restart_required" elif runtime.available and runtime.local_configuration_check(): return JSONResponse( status_code=status.HTTP_200_OK, content={"status": "ready", "checks": {"configuration": "ok"}}, ) elif runtime.available: runtime.invalidate_for_restart() configuration_status = "restart_required" else: configuration_status = "unavailable" return JSONResponse( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, content={ "status": "not_ready", "checks": {"configuration": configuration_status}, }, ) @gateway.post("/internal/v1/embeddings", response_model=EmbeddingResponse) async def embeddings( payload: EmbeddingRequest, caller: Annotated[Caller, Depends(authorize)], ) -> EmbeddingResponse: embedding, _, _, semaphore = require_runtime() if payload.input_type == "document" and caller != "worker": raise _ForbiddenError try: async with semaphore: if payload.input_type == "query": result = await embedding.embed_query(payload.texts[0]) else: result = await embedding.embed_documents(payload.texts) return EmbeddingResponse( vectors=[list(vector) for vector in result.vectors], model=result.model, request_id=result.request_id, usage=_usage_response(result.usage), elapsed_ms=result.elapsed_ms, ) except ModelProviderError: raise except Exception: raise _boundary_error("model_gateway.embedding") from None @gateway.post("/internal/v1/rerank", response_model=RerankResponse) async def rerank( payload: RerankRequest, _: Annotated[Caller, Depends(authorize)], ) -> RerankResponse: embedding_provider_unused, reranker_provider, chat_provider_unused, semaphore = ( require_runtime() ) del embedding_provider_unused, chat_provider_unused try: async with semaphore: result = await reranker_provider.rerank( payload.query, payload.documents, top_n=payload.top_n, instruct=payload.instruct, ) return RerankResponse( items=[ RankedItemResponse( index=item.index, relevance_score=item.relevance_score, document=item.document, ) for item in result.items ], model=result.model, request_id=result.request_id, usage=_usage_response(result.usage), elapsed_ms=result.elapsed_ms, ) except ModelProviderError: raise except Exception: raise _boundary_error("model_gateway.rerank") from None def chat_messages(payload: ChatRequest) -> list[ChatMessage]: return [ ChatMessage(role=message.role, content=message.content) for message in payload.messages ] @gateway.post("/internal/v1/chat/completions", response_model=ChatResponse) async def chat_completion( payload: ChatRequest, _: Annotated[Caller, Depends(authorize)], ) -> ChatResponse: embedding_provider_unused, reranker_unused, chat, semaphore = require_runtime() del embedding_provider_unused, reranker_unused try: async with semaphore: result = await chat.complete(chat_messages(payload), max_tokens=payload.max_tokens) return ChatResponse( content=result.content, finish_reason=result.finish_reason, model=result.model, request_id=result.request_id, usage=_usage_response(result.usage), elapsed_ms=result.elapsed_ms, ) except ModelProviderError: raise except Exception: raise _boundary_error("model_gateway.chat_complete") from None @gateway.post("/internal/v1/chat/stream") async def chat_stream( payload: ChatRequest, _: Annotated[Caller, Depends(authorize)], ) -> StreamingResponse: embedding_provider_unused, reranker_unused, chat, semaphore = require_runtime() del embedding_provider_unused, reranker_unused async def event_stream() -> AsyncIterator[bytes]: events: AsyncIterator[ChatStreamEvent] | None = None finish_reason: str | None = None model: str | None = None request_id: str | None = None usage = ProviderUsage() elapsed_ms = 0.0 try: events = chat.stream(chat_messages(payload), max_tokens=payload.max_tokens) async with semaphore: async for event in events: if event.finish_reason is not None: if event.finish_reason not in ALLOWED_FINISH_REASONS: raise ModelProviderError( operation="chat.stream", kind=ProviderErrorKind.INVALID_RESPONSE, provider_code="invalid_finish_reason", ) finish_reason = event.finish_reason if event.model: model = event.model if event.request_id: request_id = event.request_id usage = _merge_usage(usage, event.usage) elapsed_ms = max(elapsed_ms, event.elapsed_ms) if event.delta: yield _sse( "delta", { "delta": event.delta, "finish_reason": ( event.finish_reason if event.finish_reason in ALLOWED_FINISH_REASONS else None ), "model": event.model or model, "request_id": event.request_id or request_id, }, ) if finish_reason is None or model is None: raise ModelProviderError( operation="chat.stream", kind=ProviderErrorKind.INVALID_RESPONSE, provider_code="missing_terminal_event", ) yield _sse( "complete", { "finish_reason": finish_reason, "model": model, "request_id": request_id, "usage": _usage_response(usage).model_dump(), "elapsed_ms": elapsed_ms, }, ) except ModelProviderError as error: yield _sse( "error", { "kind": error.kind.value, "retryable": error.retryable, "request_id": error.request_id, }, ) except Exception: boundary_failure = _boundary_error("model_gateway.chat_stream") yield _sse( "error", { "kind": boundary_failure.kind.value, "retryable": boundary_failure.retryable, "request_id": boundary_failure.request_id, }, ) finally: await _close_stream(events) return _ClosingStreamingResponse( event_stream(), media_type="text/event-stream", headers={ "Cache-Control": "no-store", "X-Accel-Buffering": "no", }, ) return gateway app = create_model_gateway_app()