model-library-client 0.1.35__tar.gz

model_library_client-0.1.35/PKG-INFO

@@ -0,0 +1,21 @@
+Metadata-Version: 2.4
+Name: model-library-client
+Version: 0.1.35
+Summary: SDK-free Gateway client for model-library
+Author-email: "Vals AI, Inc." <contact@vals.ai>
+License: MIT
+Requires-Python: >=3.11
+Description-Content-Type: text/plain
+Requires-Dist: typing-extensions<5.0,>=4.15.0
+Requires-Dist: pydantic<3.0,>=2.12.5
+Requires-Dist: pyyaml<7.0,>=6.0.3
+Requires-Dist: rich<15.0,>=14.3.3
+Requires-Dist: httpx[http2]<1.0,>=0.28.0
+Requires-Dist: backoff<3.0,>=2.2.1
+Requires-Dist: redis<7.0,>=6.4.0
+Requires-Dist: tiktoken<1.0,>=0.12.0
+Requires-Dist: pillow<13.0,>=12.1.1
+Requires-Dist: python-dotenv<2.0,>=1.0.0
+Requires-Dist: aiohttp<4.0,>=3.13.5
+
+Install model-library-client to use the model_library Gateway client without provider SDK dependencies. Install model-library for the full provider distribution.

model_library_client-0.1.35/model_gateway/app.py

@@ -0,0 +1,335 @@
+"""FastAPI app factory for the model gateway."""
+
+import asyncio
+import json
+import logging
+import os
+from collections.abc import Awaitable, Callable
+from contextlib import asynccontextmanager, suppress
+from functools import partial
+from typing import Literal, cast
+
+import redis.asyncio as async_redis
+from dotenv import load_dotenv
+from redis.exceptions import TimeoutError as RedisTimeoutError
+from fastapi import FastAPI, Request
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import JSONResponse
+from starlette.responses import Response
+
+import model_library.telemetry as telemetry
+from model_library import model_library_settings
+from model_library.registry_utils import get_model_names
+from model_library.retriers.token import set_redis_client
+
+from model_gateway.asgi_observability import GatewayObservabilityMiddleware
+from model_gateway.auth import create_auth_middleware
+from model_gateway.cache import ModelCache
+from model_gateway.capacity import GatewayCapacityLimiter, create_capacity_middleware
+from model_gateway.errors import ErrorBody, ErrorResponse, RunTokenAuthorizationError
+from model_gateway.metrics import (
+    MetricPublisher,
+    create_metrics_middleware,
+    install_telemetry_delivery_metric_handler,
+    publish_metrics_periodically,
+    record_runtime,
+    remove_telemetry_delivery_metric_handler,
+)
+from model_gateway.observability import (
+    install_loop_exception_handler,
+    log_process_lifecycle,
+    runtime_snapshot,
+)
+from model_gateway.rate_limit_monitor import RateLimitMonitor
+from model_gateway.rate_limit_monitor.routes import register_rate_limit_monitor_routes
+from model_gateway.rate_limit_monitor.state import MonitorRedis, RateLimitMonitorStore
+from model_gateway.routes.benchmark_admission import register_benchmark_admission_routes
+from model_gateway.routes.health import register_health_routes
+from model_gateway.routes.models import register_model_routes
+from model_gateway.routes.provider_ops import register_provider_ops_routes
+from model_gateway.routes.rate_limit import (
+    RateLimitProbeService,
+    register_rate_limit_route,
+)
+from model_gateway.routes.query import register_query_routes
+from model_gateway.routes.run_tokens import register_run_token_routes
+from model_gateway.routes.token_retry import register_token_retry_routes
+from model_gateway.runtime_config import (
+    RuntimeConfigState,
+    runtime_config_manager_from_env,
+)
+from model_gateway.startup_canary import run_startup_canary, startup_canary_state
+from model_gateway.telemetry_helpers import error_telemetry_attributes
+from model_gateway.usage_ledger.store import (
+    NoopUsageLedger,
+    create_usage_ledger_from_env,
+)
+
+logger = logging.getLogger("model_proxy_server")
+
+load_dotenv()
+
+# The server must never proxy to itself — clear the client-side gateway URL.
+model_library_settings.unset("MODEL_GATEWAY_URL")
+
+
+GatewayRuntimeRole = Literal["combined", "query", "control"]
+
+
+def env_flag(name: str, *, default: bool = False) -> bool:
+    value = os.environ.get(name)
+    if value is None:
+        return default
+    return value.strip().lower() in {"1", "true", "yes", "on"}
+
+
+def gateway_runtime_role() -> GatewayRuntimeRole:
+    value = os.environ.get("GATEWAY_RUNTIME_ROLE", "combined").strip().lower()
+    if value not in {"combined", "query", "control"}:
+        raise ValueError("GATEWAY_RUNTIME_ROLE must be combined, query, or control")
+    return cast(GatewayRuntimeRole, value)
+
+
+async def _record_runtime_current(loop: asyncio.AbstractEventLoop) -> None:
+    lag_start = loop.time()
+    await asyncio.sleep(0)
+    event_loop_lag_ms = int((loop.time() - lag_start) * 1000)
+    snapshot = runtime_snapshot()
+    snapshot["event_loop_lag_ms"] = event_loop_lag_ms
+    record_runtime(snapshot)
+
+
+def create_app() -> FastAPI:
+    runtime_role = gateway_runtime_role()
+    query_enabled = runtime_role in {"combined", "query"}
+    control_enabled = runtime_role in {"combined", "control"}
+    api_keys = model_library_settings.get("MODEL_GATEWAY_API_KEYS", None)
+    if not isinstance(api_keys, str) or not api_keys:
+        raise RuntimeError("MODEL_GATEWAY_API_KEYS must be set")
+    api_keys_by_name = cast(dict[str, str], json.loads(api_keys))
+    valid_keys = set(api_keys_by_name.values())
+
+    hmac_secret_value = model_library_settings.get("MODEL_GATEWAY_HMAC_SECRET", None)
+    if not isinstance(hmac_secret_value, str) or not hmac_secret_value:
+        raise RuntimeError("MODEL_GATEWAY_HMAC_SECRET must be set")
+    hmac_secret = hmac_secret_value.encode()
+    cache = ModelCache()
+    runtime_config_state = RuntimeConfigState.bootstrap()
+    runtime_config_manager = runtime_config_manager_from_env(runtime_config_state)
+    capacity_limiter = GatewayCapacityLimiter()
+    startup_canary_enabled = query_enabled and env_flag(
+        "GATEWAY_STARTUP_CANARY_ENABLED"
+    )
+    usage_ledger = (
+        create_usage_ledger_from_env() if query_enabled else NoopUsageLedger()
+    )
+    rate_limit_probe_service: RateLimitProbeService | None = None
+
+    @asynccontextmanager
+    async def lifespan(app: FastAPI):
+        redis_url = os.environ.get("REDIS_URL", "")
+        if control_enabled and not redis_url:
+            raise RuntimeError("REDIS_URL must be set for a control-enabled runtime")
+
+        loop = asyncio.get_running_loop()
+        previous_exception_handler = install_loop_exception_handler(loop)
+        redis_client: async_redis.Redis | None = None
+        if redis_url:
+            redis_client = async_redis.from_url(  # pyright: ignore[reportUnknownMemberType]
+                redis_url,
+                decode_responses=True,
+                socket_timeout=30,
+                socket_connect_timeout=10,
+                socket_keepalive=True,
+                health_check_interval=30,
+                retry_on_error=[RedisTimeoutError],
+                max_connections=200,
+            )
+            set_redis_client(redis_client)
+
+        metrics_stop = asyncio.Event()
+        metrics_publishers: list[MetricPublisher] = [
+            partial(_record_runtime_current, loop)
+        ]
+        if query_enabled:
+            metrics_publishers.insert(0, capacity_limiter.record_current)
+        metrics_task: asyncio.Task[None] | None = None
+        usage_ledger_started = False
+        canary_task: asyncio.Task[None] | None = None
+        rate_limit_monitor: RateLimitMonitor | None = None
+        runtime_refresh_stop: asyncio.Event | None = None
+        runtime_refresh_task: asyncio.Task[None] | None = None
+        delivery_metric_handler = install_telemetry_delivery_metric_handler()
+        try:
+            if runtime_config_manager is not None:
+                initial_refresh = await runtime_config_manager.refresh_once()
+                if initial_refresh.failed or initial_refresh.rejected:
+                    raise RuntimeError("Initial runtime configuration refresh failed")
+                runtime_refresh_stop = asyncio.Event()
+                runtime_refresh_task = asyncio.create_task(
+                    runtime_config_manager.run(runtime_refresh_stop)
+                )
+            with runtime_config_state.capture():
+                model_count = len(get_model_names())
+            logger.info("Loaded gateway model registry with %s models", model_count)
+            log_process_lifecycle("gateway.process.startup")
+            telemetry.configure_telemetry(app)
+            if control_enabled:
+                assert redis_client is not None
+                rate_limit_monitor = RateLimitMonitor(
+                    RateLimitMonitorStore(
+                        cast(MonitorRedis, redis_client),
+                        stage=os.environ.get("GATEWAY_STAGE", "local"),
+                    ),
+                    runtime_capture=runtime_config_state.capture,
+                )
+                app.state.rate_limit_monitor = rate_limit_monitor
+                rate_limit_monitor.start()
+            metrics_task = asyncio.create_task(
+                publish_metrics_periodically(
+                    metrics_stop,
+                    publishers=metrics_publishers,
+                )
+            )
+            if startup_canary_enabled:
+                canary_task = asyncio.create_task(
+                    run_startup_canary(app, next(iter(api_keys_by_name.values())))
+                )
+            await usage_ledger.start()
+            usage_ledger_started = True
+            yield
+        finally:
+            log_process_lifecycle("gateway.process.shutdown_start")
+            if runtime_refresh_stop is not None:
+                runtime_refresh_stop.set()
+            if runtime_refresh_task is not None:
+                await runtime_refresh_task
+            if runtime_config_manager is not None:
+                await runtime_config_manager.aclose()
+            if canary_task is not None and not canary_task.done():
+                canary_task.cancel()
+                with suppress(asyncio.CancelledError):
+                    await canary_task
+            if rate_limit_probe_service is not None:
+                try:
+                    await rate_limit_probe_service.close()
+                except Exception:
+                    logger.exception(
+                        "Gateway rate-limit probe service close failed during shutdown"
+                    )
+            if usage_ledger_started:
+                try:
+                    await usage_ledger.close()
+                except Exception:
+                    logger.exception(
+                        "Gateway usage ledger close failed during shutdown"
+                    )
+            if rate_limit_monitor is not None:
+                try:
+                    await rate_limit_monitor.close()
+                except Exception:
+                    logger.exception(
+                        "Gateway rate-limit monitor close failed during shutdown"
+                    )
+                finally:
+                    app.state.rate_limit_monitor = None
+            if redis_client is not None:
+                try:
+                    await redis_client.aclose()
+                except Exception:
+                    logger.exception("Gateway Redis close failed during shutdown")
+            try:
+                telemetry.shutdown_telemetry()
+            finally:
+                try:
+                    remove_telemetry_delivery_metric_handler(delivery_metric_handler)
+                finally:
+                    metrics_stop.set()
+                    try:
+                        if metrics_task is not None:
+                            await metrics_task
+                    finally:
+                        loop.set_exception_handler(previous_exception_handler)
+                        log_process_lifecycle("gateway.process.shutdown_done")
+
+    app = FastAPI(
+        title="Model Proxy",
+        lifespan=lifespan,
+        redirect_slashes=False,
+        docs_url=None if runtime_role == "control" else "/docs",
+        redoc_url=None if runtime_role == "control" else "/redoc",
+        openapi_url=None if runtime_role == "control" else "/openapi.json",
+    )
+    app.state.cache = cache
+    app.state.runtime_config_state = runtime_config_state
+    app.state.hmac_secret = hmac_secret
+    app.state.runtime_role = runtime_role
+    app.state.capacity_limiter = capacity_limiter
+    app.state.usage_ledger = usage_ledger
+    app.state.rate_limit_monitor = None
+    app.state.startup_canary = startup_canary_state(startup_canary_enabled)
+
+    @app.middleware("http")
+    async def capture_runtime_configuration(
+        request: Request,
+        call_next: Callable[[Request], Awaitable[Response]],
+    ) -> Response:
+        with runtime_config_state.capture():
+            return await call_next(request)
+
+    if query_enabled:
+        app.middleware("http")(create_capacity_middleware())
+    app.middleware("http")(create_metrics_middleware())
+    app.middleware("http")(create_auth_middleware(api_keys_by_name))
+    app.add_middleware(GatewayObservabilityMiddleware)
+
+    @app.exception_handler(RunTokenAuthorizationError)
+    async def run_token_authorization_handler(
+        request: Request, exc: RunTokenAuthorizationError
+    ) -> JSONResponse:
+        claims = request.state.run_token_claims
+        telemetry.add_event(
+            "gateway.run_token.violation",
+            {
+                "gateway.run_token.run_id": claims.run_id,
+                "gateway.run_token.task_id": claims.task_id,
+                "gateway.error.message": str(exc),
+            },
+        )
+        return JSONResponse(
+            status_code=403,
+            content={"code": "model_not_authorized", "message": str(exc)},
+        )
+
+    @app.exception_handler(RequestValidationError)
+    async def validation_exception_handler(
+        request: Request, exc: RequestValidationError
+    ) -> JSONResponse:
+        err = ErrorResponse(
+            400,
+            ErrorBody(code="invalid_request", message=str(exc)),
+        )
+        error_attrs = error_telemetry_attributes(err, phase="request_validation")
+        telemetry.set_attributes(error_attrs)
+        telemetry.set_status_error(err.body.code)
+        telemetry.add_event("gateway.request_validation.error", error_attrs)
+        return JSONResponse(status_code=err.status_code, content=err.body.model_dump())
+
+    register_health_routes(
+        app,
+        valid_keys=valid_keys,
+        hmac_secret=hmac_secret,
+        control_enabled=control_enabled,
+    )
+    if control_enabled:
+        register_benchmark_admission_routes(app, cache=cache)
+        register_run_token_routes(app)
+        register_rate_limit_monitor_routes(app)
+    if control_enabled:
+        register_model_routes(app)
+    if query_enabled:
+        register_token_retry_routes(app)
+        register_query_routes(app, cache=cache)
+        rate_limit_probe_service = register_rate_limit_route(app, cache=cache)
+        register_provider_ops_routes(app, cache=cache)
+    return app

model_library_client-0.1.35/model_gateway/asgi_observability.py

@@ -0,0 +1,145 @@
+"""ASGI-level gateway request/response boundary diagnostics."""
+
+from __future__ import annotations
+
+import logging
+import time
+from typing import cast
+
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+from model_gateway.observability import (
+    log_gateway_event,
+    request_log_fields_from_scope,
+    runtime_snapshot,
+)
+
+
+class GatewayObservabilityMiddleware:
+    """Capture request and ASGI response-boundary events without payload data."""
+
+    def __init__(self, app: ASGIApp) -> None:
+        self.app = app
+
+    async def __call__(
+        self,
+        scope: Scope,
+        receive: Receive,
+        send: Send,
+    ) -> None:
+        if scope.get("type") != "http":
+            await self.app(scope, receive, send)
+            return
+
+        path = cast(str, scope["path"])
+        method = cast(str, scope["method"])
+        request_fields = request_log_fields_from_scope(scope)
+        state = _RequestBoundaryState(start=time.perf_counter())
+
+        def log_anomaly(
+            event: str, *, level: int = logging.INFO, **fields: object
+        ) -> None:
+            log_gateway_event(
+                event,
+                level=level,
+                route=path,
+                method=method,
+                runtime=_runtime_snapshot_for_anomaly(state),
+                **request_fields,
+                **fields,
+            )
+
+        async def wrapped_receive() -> Message:
+            message = await receive()
+            if message["type"] == "http.disconnect" and not state.client_disconnected:
+                state.client_disconnected = True
+                log_anomaly(
+                    "gateway.request.client_disconnect",
+                    response_started=state.response_started,
+                )
+            return message
+
+        async def wrapped_send(message: Message) -> None:
+            message_type = message["type"]
+            try:
+                await send(message)
+            except Exception as exc:
+                state.send_error = True
+                log_anomaly(
+                    "gateway.response.send_error",
+                    level=logging.WARNING,
+                    exception_type=type(exc).__name__,
+                    response_started=state.response_started,
+                    first_body_sent=state.first_body_sent,
+                    status_code=state.status_code,
+                )
+                raise
+            if message_type == "http.response.start":
+                state.response_started = True
+                state.status_code = cast(int, message["status"])
+            elif message_type == "http.response.body":
+                if not state.first_body_sent:
+                    state.first_body_sent = True
+                more_body = bool(message.get("more_body"))
+                if not more_body:
+                    state.response_done = True
+
+        try:
+            await self.app(scope, wrapped_receive, wrapped_send)
+        except Exception as exc:
+            if not state.send_error:
+                log_anomaly(
+                    "gateway.request.error",
+                    level=logging.WARNING,
+                    exception_type=type(exc).__name__,
+                    response_started=state.response_started,
+                    first_body_sent=state.first_body_sent,
+                    status_code=state.status_code,
+                )
+            raise
+        finally:
+            duration_ms = (time.perf_counter() - state.start) * 1000
+            if not state.response_started:
+                log_anomaly(
+                    "gateway.response.missing_start",
+                    level=logging.WARNING,
+                    client_disconnected=state.client_disconnected,
+                )
+            if _is_anomaly(state):
+                log_anomaly(
+                    "gateway.request.done",
+                    status_code=state.status_code,
+                    duration_ms=duration_ms,
+                    response_started=state.response_started,
+                    first_body_sent=state.first_body_sent,
+                    response_done=state.response_done,
+                    send_error=state.send_error,
+                    client_disconnected=state.client_disconnected,
+                )
+
+
+def _runtime_snapshot_for_anomaly(state: "_RequestBoundaryState") -> dict[str, int]:
+    if state.anomaly_runtime is None:
+        state.anomaly_runtime = runtime_snapshot()
+    return state.anomaly_runtime
+
+
+def _is_anomaly(state: "_RequestBoundaryState") -> bool:
+    return (
+        state.send_error
+        or state.client_disconnected
+        or not state.response_done
+        or (state.status_code is not None and not 200 <= state.status_code < 300)
+    )
+
+
+class _RequestBoundaryState:
+    def __init__(self, *, start: float):
+        self.start = start
+        self.response_started = False
+        self.first_body_sent = False
+        self.response_done = False
+        self.send_error = False
+        self.client_disconnected = False
+        self.status_code: int | None = None
+        self.anomaly_runtime: dict[str, int] | None = None

model_library_client-0.1.35/model_gateway/auth.py

@@ -0,0 +1,113 @@
+"""Bearer token auth middleware."""
+
+import hashlib
+import hmac
+from collections.abc import Mapping
+
+from fastapi import Request
+from fastapi.responses import JSONResponse
+
+import model_library.telemetry as telemetry
+from starlette.middleware.base import RequestResponseEndpoint
+from starlette.responses import Response
+
+from model_gateway.run_token_types import RUN_TOKEN_OPERATION_PATHS
+from model_gateway.run_tokens import RUN_TOKEN_PREFIX, get_run_token, sha256_hex
+
+EXEMPT_PATHS = {"/health/live", "/health/ready"}
+TRACE_AUTH_FAILURE_PATHS = telemetry.HTTP_TRACE_ALLOWED_ROUTES
+
+RUN_TOKEN_ALLOWED_PATHS = frozenset({*RUN_TOKEN_OPERATION_PATHS.values(), "/registry"})
+
+
+def create_auth_middleware(api_keys_by_name: Mapping[str, str]):
+    # Pre-hash keys for constant-time comparison (prevents timing side-channel)
+    hashed_keys = tuple(
+        (name, hashlib.sha256(api_key.encode()).digest())
+        for name, api_key in api_keys_by_name.items()
+    )
+    accepted_digests = {digest.hex() for _, digest in hashed_keys}
+
+    async def auth_middleware(
+        request: Request, call_next: RequestResponseEndpoint
+    ) -> Response:
+        if request.url.path in EXEMPT_PATHS:
+            return await call_next(request)
+
+        auth_header = request.headers.get("Authorization", "")
+        if not auth_header.startswith("Bearer "):
+            return _unauthorized_response(
+                request,
+                "Missing or malformed Authorization header",
+            )
+
+        presented = auth_header[7:]
+        request.state.run_token_claims = None
+
+        if presented.startswith(RUN_TOKEN_PREFIX):
+            # Before the lookup, so a denied path cannot leak token validity.
+            if request.url.path not in RUN_TOKEN_ALLOWED_PATHS:
+                return _auth_error_response(
+                    request, 403, "forbidden", "Run tokens cannot call this route"
+                )
+            claims = await get_run_token(sha256_hex(presented))
+            # Dev and prod share Redis but not keys.
+            if claims is None or claims.minted_with not in accepted_digests:
+                return _unauthorized_response(request, "Invalid or expired run token")
+            if claims.allowed_operations is not None and request.url.path not in {
+                "/registry",
+                *(RUN_TOKEN_OPERATION_PATHS[op] for op in claims.allowed_operations),
+            }:
+                return _auth_error_response(
+                    request, 403, "forbidden", "Run token does not allow this operation"
+                )
+            request.state.gateway_api_key_name = f"run-token:{claims.minted_by}"
+            request.state.gateway_api_key_digest = claims.minted_with
+            request.state.run_token_claims = claims
+            return await call_next(request)
+
+        token_hash = hashlib.sha256(presented.encode()).digest()
+        matched_name: str | None = None
+        for name, expected_hash in hashed_keys:
+            if hmac.compare_digest(token_hash, expected_hash):
+                matched_name = name
+        if matched_name is None:
+            return _unauthorized_response(request, "Invalid API key")
+
+        request.state.gateway_api_key_name = matched_name
+        request.state.gateway_api_key_digest = token_hash.hex()
+        return await call_next(request)
+
+    return auth_middleware
+
+
+def _unauthorized_response(request: Request, message: str) -> JSONResponse:
+    return _auth_error_response(request, 401, "unauthorized", message)
+
+
+def _auth_error_response(
+    request: Request, status_code: int, code: str, message: str
+) -> JSONResponse:
+    attrs = {
+        "gateway.route": request.url.path,
+        "gateway.operation": "access_check",
+        "gateway.error.code": "access_denied",
+        "gateway.error.phase": "access_control",
+        "gateway.status_code": status_code,
+        "http.request.method": request.method,
+        "http.status_code": status_code,
+        "http.response.status_code": status_code,
+    }
+    if request.url.path in TRACE_AUTH_FAILURE_PATHS:
+        with telemetry.start_span(
+            telemetry.http_server_span_name(request.method, request.url.path),
+            attrs,
+            kind="server",
+        ):
+            telemetry.set_attributes(attrs)
+            telemetry.set_status_error("access_denied")
+            telemetry.add_event("gateway.auth.error", attrs)
+    return JSONResponse(
+        status_code=status_code,
+        content={"code": code, "message": message},
+    )