feat(server): integrate OpenAI Chat streaming

nemoguardrails/server/experimental/_buffered_kernel.py

@@ -91,6 +91,15 @@ class OperationProjectionFailed:
     failure: UnsupportedGuardedPayload
 
 
+@dataclass(frozen=True, slots=True)
+class PreparedOperationInput(Generic[RequestT]):
+    """Carry one projected and inspected input into its response mode."""
+
+    request: RequestT
+    input_message: GuardedMessage
+    resolved: _ResolvedContentChecker
+
+
 def _project_message(
     projection: GuardedMessageProjection[PayloadT],
     payload: PayloadT,
@@ -146,19 +155,18 @@ async def _run_check(
     return _stopped_operation(stage, decision)
 
 
-async def execute_buffered_operation(
+async def prepare_operation_input(
     operation: BufferedGuardedOperation[RequestT, ResponseT],
     checker: ContentChecker | _ResolvedContentChecker,
     request: RequestT,
-    dispatch: Callable[[RequestT], Awaitable[ResponseT]],
 ) -> (
-    OperationCompleted[ResponseT]
+    PreparedOperationInput[RequestT]
     | OperationBlocked
     | OperationCheckFailed
     | OperationModificationUnsupported
     | OperationProjectionFailed
 ):
-    """Execute one buffered operation with exactly one statically bound checker."""
+    """Project and inspect one operation input exactly once."""
 
     validated = checker if isinstance(checker, _ResolvedContentChecker) else validate_content_checker(checker)
     validated_checker = validated.checker
@@ -178,9 +186,25 @@ async def execute_buffered_operation(
         if stopped is not None:
             return stopped
 
-    response = await dispatch(request)
+    return PreparedOperationInput(request, input_message, validated)
+
+
+async def execute_prepared_buffered_operation(
+    operation: BufferedGuardedOperation[RequestT, ResponseT],
+    prepared: PreparedOperationInput[RequestT],
+    dispatch: Callable[[RequestT], Awaitable[ResponseT]],
+) -> (
+    OperationCompleted[ResponseT]
+    | OperationBlocked
+    | OperationCheckFailed
+    | OperationModificationUnsupported
+    | OperationProjectionFailed
+):
+    """Dispatch and inspect a buffered response for one prepared input."""
+
+    response = await dispatch(prepared.request)
 
-    if policy.inspect_output:
+    if prepared.resolved.policy.inspect_output:
         try:
             output_message = _project_message(operation.output_projection, response, "assistant")
         except UnsupportedGuardedPayload as failure:
@@ -189,9 +213,31 @@ async def execute_buffered_operation(
             return OperationCompleted(response)
         stopped = await _run_check(
             InspectionStage.OUTPUT,
-            lambda: validated_checker.check_output(OutputContentCheck(input_message, output_message.content)),
+            lambda: prepared.resolved.checker.check_output(
+                OutputContentCheck(prepared.input_message, output_message.content)
+            ),
         )
         if stopped is not None:
             return stopped
 
     return OperationCompleted(response)
+
+
+async def execute_buffered_operation(
+    operation: BufferedGuardedOperation[RequestT, ResponseT],
+    checker: ContentChecker | _ResolvedContentChecker,
+    request: RequestT,
+    dispatch: Callable[[RequestT], Awaitable[ResponseT]],
+) -> (
+    OperationCompleted[ResponseT]
+    | OperationBlocked
+    | OperationCheckFailed
+    | OperationModificationUnsupported
+    | OperationProjectionFailed
+):
+    """Execute one buffered operation with exactly one statically bound checker."""
+
+    prepared = await prepare_operation_input(operation, checker, request)
+    if not isinstance(prepared, PreparedOperationInput):
+        return prepared
+    return await execute_prepared_buffered_operation(operation, prepared, dispatch)

nemoguardrails/server/experimental/_guarded_proxy.py

@@ -13,19 +13,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Compose provider declarations with the guarded buffered HTTP pipeline."""
+"""Compose provider declarations with the guarded HTTP pipeline."""
 
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 
 from pydantic import ValidationError
 
+from nemoguardrails.server.experimental._buffered_kernel import (
+    InspectionStage,
+    OperationProjectionFailed,
+    PreparedOperationInput,
+)
 from nemoguardrails.server.experimental._guarded_operation import (
     BufferedGuardedOperation,
     ContentInspectionNotApplicable,
     InvalidGuardedPayload,
     UnsupportedGuardedPayload,
     UnsupportedGuardedRepresentation,
 )
+from nemoguardrails.server.experimental._guarded_stream import (
+    StreamInspectionUnsupported,
+    StreamOutcome,
+    UnsupportedStreamInspection,
+)
 from nemoguardrails.server.experimental._http_kernel import (
     BufferedHttpRequest,
     BufferedHttpResponse,
@@ -35,6 +45,7 @@
     _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.provider.endpoint import GuardedJsonEndpoint
 from nemoguardrails.server.experimental.provider.errors import ProviderErrorMapping
 from nemoguardrails.server.experimental.provider.payload import GuardedMessageTarget, guarded_schema_error
@@ -49,6 +60,7 @@ class GuardableProviderRequest:
 
     request: BufferedHttpRequest
     target: GuardedMessageTarget
+    streaming: bool = False
 
 
 def _require_json_representation(request: BufferedHttpRequest) -> None:
@@ -89,21 +101,31 @@ def _prepare_guardable_request(
         raise UnsupportedGuardedPayload(
             guarded_schema_error(error, "request"), endpoint.unsupported_request_code
         ) from error
-    if projection.streams_response:
+    streaming = projection.streams_response
+    if streaming and endpoint.stream is None:
         raise UnsupportedGuardedPayload(
-            "The guarded endpoint does not support streaming responses yet.", endpoint.unsupported_request_code
+            "The guarded endpoint does not support streaming responses.", endpoint.unsupported_request_code
         )
     return GuardableProviderRequest(
         request=request,
         target=target,
+        streaming=streaming,
     )
 
 
-def create_buffered_guarded_http_operation(
+def create_guarded_http_operation(
     endpoint: GuardedJsonEndpoint,
     errors: ProviderErrorMapping,
+    *,
+    stream_dispatch: StreamingHttpDispatch | None = None,
+    max_stream_event_bytes: int,
+    max_pending_stream_bytes: int,
 ) -> GuardedHttpOperation:
-    """Create one buffered HTTP operation from the provider endpoint contract."""
+    """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:
+        raise ValueError("A streaming dispatcher requires an endpoint stream binding.")
 
     def prepare_request(request: BufferedHttpRequest) -> GuardableProviderRequest:
         """Prepare one buffered HTTP request for guarded execution."""
@@ -135,8 +157,9 @@ def project_response(
             document = parse_json_object(payload.body)
             projection = endpoint.guarded_response_model.validate_payload(document)
             target = projection.locate_guarded_message(document)
-            # Validation has already rejected content in any other reviewed
-            # field, so a response without text has nothing to inspect.
+            # Validation has already rejected generated content in other
+            # reviewed fields, and opaque metadata is not inspected, so a
+            # response without text has no guarded text to check.
             if not target.has_text:
                 return ContentInspectionNotApplicable()
             return target.message
@@ -153,6 +176,62 @@ def project_response(
         if issubclass(request_model, PolicyModel)
         else request_model.model_json_schema()
     )
+
+    def render_stream_outcome(outcome: StreamOutcome) -> BufferedHttpResponse:
+        """Bind generic stream shape failures to the endpoint's declared codes."""
+        if (
+            isinstance(outcome, OperationProjectionFailed)
+            and type(outcome.failure) is UnsupportedGuardedPayload
+            and outcome.failure.code is None
+        ):
+            code = (
+                endpoint.unsupported_request_code
+                if outcome.stage is InspectionStage.INPUT
+                else endpoint.unsupported_response_code
+            )
+            outcome = replace(outcome, failure=UnsupportedGuardedPayload(str(outcome.failure), code))
+        return errors.renderer(outcome)
+
+    async def handle_prepared_request(
+        prepared_input: PreparedOperationInput[GuardableProviderRequest],
+    ):
+        guardable = prepared_input.request
+        if not guardable.streaming:
+            return None
+        if stream_dispatch is None or stream_binding is None:
+            return render_stream_outcome(
+                OperationProjectionFailed(
+                    InspectionStage.INPUT,
+                    UnsupportedGuardedPayload("The guarded endpoint does not have streaming dispatch configured."),
+                )
+            )
+        streaming_policy = prepared_input.resolved.policy.stream_buffering
+        if prepared_input.resolved.policy.inspect_output and streaming_policy is None:
+            return errors.renderer(
+                StreamInspectionUnsupported(
+                    UnsupportedStreamInspection("Guarded output streaming requires a buffering policy.")
+                )
+            )
+        return await execute_streaming_http(
+            guardable.request,
+            dispatch=stream_dispatch,
+            checker=prepared_input.resolved.checker,
+            streaming_policy=streaming_policy,
+            input_message=prepared_input.input_message,
+            adapter=stream_binding.create_adapter(),
+            render_outcome=render_stream_outcome,
+            max_event_bytes=max_stream_event_bytes,
+            max_pending_bytes=max_pending_stream_bytes,
+        )
+
+    success_content = {"application/json": {"schema": {}}}
+    if endpoint.stream is not None and stream_dispatch is not None:
+        success_content["text/event-stream"] = {"schema": {"type": "string"}}
+    documented_responses = errors.documented_responses
+    documented_responses[200] = {
+        "description": "Provider-native successful response.",
+        "content": success_content,
+    }
     openapi_extra = {
         "requestBody": {
             "required": True,
@@ -177,7 +256,8 @@ def project_response(
         operation=operation,
         prepare_request=prepare_request,
         forward_request=forward_request,
-        documented_responses=errors.documented_responses,
+        prepared_request_handler=handle_prepared_request,
+        documented_responses=documented_responses,
         guarded_operation_paths=endpoint.operation_paths,
         openapi_extra=openapi_extra,
     )

nemoguardrails/server/experimental/_guarded_stream.py

@@ -83,11 +83,19 @@ class StreamProcessingFailed:
     failure: BaseException
 
 
+@dataclass(frozen=True, slots=True)
+class StreamInspectionUnsupported:
+    """Stop before dispatch when output inspection cannot guard a stream."""
+
+    failure: UnsupportedStreamInspection
+
+
 StreamOutcome = (
     OperationBlocked
     | OperationCheckFailed
     | OperationModificationUnsupported
     | OperationProjectionFailed
+    | StreamInspectionUnsupported
     | StreamUpstreamFailed
     | StreamProcessingFailed
 )
@@ -182,7 +190,8 @@ def _log_outcome(outcome: StreamOutcome) -> None:
     failure = outcome.failure
     if isinstance(failure, ContentCheckFailed):
         failure = failure.cause
-    stage = "stream" if isinstance(outcome, StreamUpstreamFailed | StreamProcessingFailed) else outcome.stage.value
+    staged = OperationCheckFailed | OperationModificationUnsupported | OperationProjectionFailed
+    stage = outcome.stage.value if isinstance(outcome, staged) else "stream"
     log.warning("Guarded stream stopped: %s at %s stage (%s)", type(outcome).__name__, stage, type(failure).__name__)
 
 

nemoguardrails/server/experimental/_http_kernel.py

@@ -34,7 +34,9 @@
     OperationCompleted,
     OperationModificationUnsupported,
     OperationProjectionFailed,
-    execute_buffered_operation,
+    PreparedOperationInput,
+    execute_prepared_buffered_operation,
+    prepare_operation_input,
 )
 from nemoguardrails.server.experimental._content_checker import (
     ContentChecker,
@@ -148,12 +150,20 @@ def __post_init__(self) -> None:
 
 @dataclass(frozen=True, slots=True)
 class GuardedHttpOperation:
-    """Bind one buffered guarded operation to an owned HTTP path."""
+    """Bind one guarded operation to an owned HTTP path.
+
+    The prepared-request handler is the private response-mode seam used while
+    buffered and streaming transports remain separately injected. It runs only
+    after the shared input projection and check have completed.
+    """
 
     operation_path: GuardedOperationPath
     operation: BufferedGuardedOperation[Any, BufferedHttpResponse]
     prepare_request: Callable[[BufferedHttpRequest], Any] | None = None
     forward_request: Callable[[Any], BufferedHttpRequest] | None = None
+    prepared_request_handler: (
+        Callable[[PreparedOperationInput[Any]], Awaitable[BufferedHttpResponse | Response | None]] | None
+    ) = None
     documented_responses: dict[int | str, dict[str, Any]] | None = None
     guarded_operation_paths: tuple[GuardedOperationPath, ...] = ()
     openapi_extra: dict[str, Any] | None = None
@@ -510,11 +520,24 @@ async def handle(request: Request) -> Response:
                 return _render_failure(
                     OperationProjectionFailed(InspectionStage.INPUT, failure),
                     render_outcome,
+                    request.method,
                 )
-            outcome = await execute_buffered_operation(
+            prepared = await prepare_operation_input(
                 declaration.operation,
                 checker,
                 operation_request,
+            )
+            if not isinstance(prepared, PreparedOperationInput):
+                return _render_failure(prepared, render_outcome, request.method)
+            if declaration.prepared_request_handler is not None:
+                response = await declaration.prepared_request_handler(prepared)
+                if isinstance(response, BufferedHttpResponse):
+                    return _render_response(response, request.method)
+                if response is not None:
+                    return response
+            outcome = await execute_prepared_buffered_operation(
+                declaration.operation,
+                prepared,
                 bounded_dispatch,
             )
         except RequestBodyTooLarge as failure:

nemoguardrails/server/experimental/_streaming_http.py

@@ -27,9 +27,11 @@
 from nemoguardrails.server.experimental._content_checker import ContentChecker, StreamBufferingPolicy
 from nemoguardrails.server.experimental._guarded_operation import UnsupportedGuardedPayload
 from nemoguardrails.server.experimental._guarded_stream import (
+    StreamInspectionUnsupported,
     StreamOutcomeRenderer,
     StreamProcessingFailed,
     StreamUpstreamFailed,
+    UnsupportedStreamInspection,
     guard_provider_stream,
     validate_streaming_policy,
 )
@@ -180,7 +182,10 @@ async def execute_streaming_http(
     if input_message.role != "user":
         raise ValueError("A guarded stream requires a user input message.")
     if streaming_policy is not None:
-        validate_streaming_policy(streaming_policy)
+        try:
+            validate_streaming_policy(streaming_policy)
+        except UnsupportedStreamInspection as failure:
+            return _render_response(render_outcome(StreamInspectionUnsupported(failure)), request.method)
         # An encoded stream cannot be inspected; uninspected relays keep the client's preference.
         request = replace(request, headers=_request_identity_encoding(request.headers))