feat(server): resolve trusted Rails configurations

nemoguardrails/server/experimental/_application.py

@@ -21,41 +21,71 @@
 import httpx
 from fastapi import APIRouter, FastAPI
 
-from nemoguardrails.server.experimental._content_checker import ContentChecker
 from nemoguardrails.server.experimental._http_kernel import (
     DEFAULT_MAX_REQUEST_BODY_BYTES,
     DEFAULT_MAX_RESPONSE_BODY_BYTES,
 )
 from nemoguardrails.server.experimental._provider_routes import include_provider_routes
-from nemoguardrails.server.experimental._runtime import ProxyRuntime, _close_http_resource, get_runtime
+from nemoguardrails.server.experimental._runtime import (
+    DEFAULT_CONFIG_RESOLUTION_TIMEOUT,
+    ServerRuntime,
+    _close_http_resource,
+    _validate_config_resolution_timeout,
+    get_runtime,
+)
 from nemoguardrails.server.experimental.provider.integration import ProviderIntegration
+from nemoguardrails.server.experimental.rails import ConfigEntry, ConfigResolver
+from nemoguardrails.server.experimental.rails.config import validate_config_entries
 
 HttpTransportFactory = Callable[[], httpx.AsyncBaseTransport]
 
 
 @asynccontextmanager
 async def _lifespan(application: FastAPI):
+    """Publish only ready runtimes and retain unfinished Rails cleanup."""
+
+    pending = getattr(application.state, "proxy_cleanup_runtime", None)
+    if pending is not None:
+        await pending.close()
+        if pending.has_pending_cleanup:
+            raise RuntimeError("Previous Guardrails cleanup is incomplete.")
+        del application.state.proxy_cleanup_runtime
     factory = application.state.proxy_http_transport_factory
     transport = factory() if factory is not None else None
     if factory is not None and not isinstance(transport, httpx.AsyncBaseTransport):
         raise TypeError("The HTTP transport factory must return an httpx.AsyncBaseTransport.")
     try:
-        runtime = ProxyRuntime.create(http_transport=transport)
+        runtime = ServerRuntime.create(
+            application.state.proxy_rails_configs,
+            application.state.proxy_config_resolver,
+            http_transport=transport,
+            config_resolution_timeout=application.state.proxy_config_resolution_timeout,
+        )
     except BaseException:
         if transport is not None:
             await _close_http_resource(transport)
         raise
-    application.state.proxy_runtime = runtime
     try:
+        await runtime.startup()
+        application.state.proxy_runtime = runtime
         yield
     finally:
-        await runtime.close()
+        application.state.proxy_cleanup_runtime = runtime
+        try:
+            await runtime.close()
+        finally:
+            if not runtime.has_pending_cleanup:
+                del application.state.proxy_cleanup_runtime
+            if getattr(application.state, "proxy_runtime", None) is runtime:
+                del application.state.proxy_runtime
 
 
 def create_proxy_application(
     *,
     providers: Sequence[ProviderIntegration],
-    checker: ContentChecker,
+    rails_configs: Sequence[ConfigEntry] = (),
+    config_resolver: ConfigResolver | None = None,
+    config_resolution_timeout: float = DEFAULT_CONFIG_RESOLUTION_TIMEOUT,
     local_routers: Sequence[APIRouter] = (),
     path_prefix: str = "",
     http_transport_factory: HttpTransportFactory | None = None,
@@ -67,14 +97,41 @@ def create_proxy_application(
 
     if http_transport_factory is not None and not callable(http_transport_factory):
         raise TypeError("The HTTP transport factory must be callable.")
+    config_resolution_timeout = _validate_config_resolution_timeout(config_resolution_timeout)
+    resolved_configs = validate_config_entries(rails_configs)
+    if config_resolver is None and len(resolved_configs) > 1:
+        raise ValueError("Multiple Rails configs require a resolver.")
+    if config_resolver is not None and not resolved_configs:
+        raise ValueError("A config resolver requires at least one config entry.")
+    resolved_resolver = config_resolver
+    if resolved_resolver is None and len(resolved_configs) == 1:
+        config_id = resolved_configs[0].id
+
+        def resolve_static_config(_context):
+            return config_id
+
+        resolved_resolver = resolve_static_config
+
     application = FastAPI(lifespan=_lifespan)
     application.state.proxy_http_transport_factory = http_transport_factory
+    application.state.proxy_rails_configs = resolved_configs
+    application.state.proxy_config_resolver = resolved_resolver
+    application.state.proxy_config_resolution_timeout = config_resolution_timeout
+
+    async def resolve_checker(request, provider_type, provider_instance_id, operation_name):
+        return await get_runtime(application).resolve_content_checker(
+            request,
+            provider_type,
+            provider_instance_id,
+            operation_name,
+        )
+
     for router in local_routers:
         application.include_router(router)
     include_provider_routes(
         application,
         providers,
-        checker=checker,
+        resolve_checker=resolve_checker,
         http_client=lambda: get_runtime(application).http_client,
         path_prefix=path_prefix,
         max_request_body_bytes=max_request_body_bytes,

nemoguardrails/server/experimental/_content_checker.py

@@ -15,9 +15,12 @@
 
 """Define private content-checking declarations for guarded operations."""
 
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass, field
 from typing import Protocol, TypeAlias, cast
 
+from fastapi import Request
+
 from nemoguardrails.server.experimental.provider.types import GuardedMessage
 
 
@@ -33,6 +36,20 @@ class UnsupportedContentModification(RuntimeError):
     """Report a checker modification that the proxy cannot safely apply."""
 
 
+@dataclass(frozen=True, slots=True)
+class ContentCheckerResolutionFailed:
+    """Stop an operation when its trusted checker cannot be resolved."""
+
+    failure: BaseException = field(compare=False, repr=False)
+
+
+@dataclass(frozen=True, slots=True)
+class ContentCheckerConfigurationUnsupported:
+    """Stop an operation when its resolved checker uses unsupported capabilities."""
+
+    failure: UnsupportedContentCheckerConfiguration
+
+
 @dataclass(frozen=True, slots=True)
 class StreamBufferingPolicy:
     """Declare how guarded stream deltas are checked before release."""
@@ -158,6 +175,12 @@ async def check_output(self, check: OutputContentCheck) -> ContentCheckDecision:
         ...
 
 
+ContentCheckerResolver: TypeAlias = Callable[
+    [Request, str, str, str],
+    ContentChecker | Awaitable[ContentChecker],
+]
+
+
 @dataclass(frozen=True, slots=True)
 class _ResolvedContentChecker:
     """Bind one checker to the policy validated for its execution."""

nemoguardrails/server/experimental/_guarded_proxy.py

@@ -15,6 +15,7 @@
 
 """Compose provider declarations with the guarded HTTP pipeline."""
 
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass, replace
 
 from pydantic import ValidationError
@@ -45,14 +46,23 @@
     _has_media_type,
 )
 from nemoguardrails.server.experimental._json_payload import InvalidJson, UnsupportedJsonShape, parse_json_object
-from nemoguardrails.server.experimental._streaming_http import StreamingHttpDispatch, execute_streaming_http
+from nemoguardrails.server.experimental._streaming_http import (
+    StreamingHttpDispatch,
+    StreamingHttpResponse,
+    execute_streaming_http,
+)
 from nemoguardrails.server.experimental.provider.endpoint import GuardedJsonEndpoint
 from nemoguardrails.server.experimental.provider.errors import ProviderErrorMapping
 from nemoguardrails.server.experimental.provider.payload import GuardedMessageTarget, guarded_schema_error
 from nemoguardrails.server.experimental.provider.projection_policy import PolicyModel, policy_json_schema
 from nemoguardrails.server.experimental.provider.transport import ProviderBindingViolation
 from nemoguardrails.server.experimental.provider.types import GuardedMessage
 
+PolicyStreamingHttpDispatch = Callable[
+    [BufferedHttpRequest, bool],
+    Awaitable[StreamingHttpResponse | BufferedHttpResponse],
+]
+
 
 @dataclass(slots=True)
 class GuardableProviderRequest:
@@ -118,13 +128,14 @@ def create_guarded_http_operation(
     errors: ProviderErrorMapping,
     *,
     stream_dispatch: StreamingHttpDispatch | None = None,
+    policy_stream_dispatch: PolicyStreamingHttpDispatch | None = None,
     max_stream_event_bytes: int,
     max_pending_stream_bytes: int,
 ) -> GuardedHttpOperation:
     """Create one buffered-and-streaming operation from an endpoint contract."""
 
     stream_binding = endpoint.stream
-    if stream_dispatch is not None and stream_binding is None:
+    if (stream_dispatch is not None or policy_stream_dispatch is not None) and stream_binding is None:
         raise ValueError("A streaming dispatcher requires an endpoint stream binding.")
 
     def prepare_request(request: BufferedHttpRequest) -> GuardableProviderRequest:
@@ -198,7 +209,7 @@ async def handle_prepared_request(
         guardable = prepared_input.request
         if not guardable.streaming:
             return None
-        if stream_dispatch is None or stream_binding is None:
+        if (stream_dispatch is None and policy_stream_dispatch is None) or stream_binding is None:
             return render_stream_outcome(
                 OperationProjectionFailed(
                     InspectionStage.INPUT,
@@ -212,9 +223,17 @@ async def handle_prepared_request(
                     UnsupportedStreamInspection("Guarded output streaming requires a buffering policy.")
                 )
             )
+
+        async def dispatch(request: BufferedHttpRequest):
+            if policy_stream_dispatch is not None:
+                return await policy_stream_dispatch(request, prepared_input.resolved.policy.inspect_output)
+            if stream_dispatch is None:
+                raise AssertionError("A configured streaming request requires streaming dispatch.")
+            return await stream_dispatch(request)
+
         return await execute_streaming_http(
             guardable.request,
-            dispatch=stream_dispatch,
+            dispatch=dispatch,
             checker=prepared_input.resolved.checker,
             streaming_policy=streaming_policy,
             input_message=prepared_input.input_message,
@@ -225,7 +244,7 @@ async def handle_prepared_request(
         )
 
     success_content = {"application/json": {"schema": {}}}
-    if endpoint.stream is not None and stream_dispatch is not None:
+    if endpoint.stream is not None and (stream_dispatch is not None or policy_stream_dispatch is not None):
         success_content["text/event-stream"] = {"schema": {"type": "string"}}
     documented_responses = errors.documented_responses
     documented_responses[200] = {