diff --git a/nemoguardrails/server/experimental/_content_checker.py b/nemoguardrails/server/experimental/_content_checker.py index 101cc20ac6..425f2e0656 100644 --- a/nemoguardrails/server/experimental/_content_checker.py +++ b/nemoguardrails/server/experimental/_content_checker.py @@ -33,18 +33,40 @@ class UnsupportedContentModification(RuntimeError): """Report a checker modification that the proxy cannot safely apply.""" +@dataclass(frozen=True, slots=True) +class StreamBufferingPolicy: + """Declare how guarded stream deltas are checked before release.""" + + chunk_size: int + context_size: int + release_before_check: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.chunk_size, int) or self.chunk_size <= 0: + raise ValueError("A stream chunk size must be a positive integer.") + if not isinstance(self.context_size, int) or self.context_size < 0: + raise ValueError("A stream context size must be a non-negative integer.") + if not isinstance(self.release_before_check, bool): + raise TypeError("The stream release policy must be a boolean.") + + @dataclass(frozen=True, slots=True) class ContentInspectionPolicy: """Record whether a checker examines input and output content.""" inspect_input: bool inspect_output: bool + stream_buffering: StreamBufferingPolicy | None = None def __post_init__(self) -> None: """Reject non-boolean inspection flags.""" if not isinstance(self.inspect_input, bool) or not isinstance(self.inspect_output, bool): raise TypeError("Content inspection flags must be booleans.") + if self.stream_buffering is not None and not isinstance(self.stream_buffering, StreamBufferingPolicy): + raise TypeError("Stream buffering must be a StreamBufferingPolicy.") + if not self.inspect_output and self.stream_buffering is not None: + raise ValueError("Stream buffering requires output inspection.") @dataclass(frozen=True, slots=True) @@ -56,7 +78,14 @@ class InputContentCheck: @dataclass(frozen=True, slots=True) class OutputContentCheck: - """Carry generated text and its effective input context to one checker.""" + """Carry generated text and its effective input context to one checker. + + Buffered and streaming execution currently supply one projected user input + subject as context for assistant-output checks. This carrier does not + itself enforce the role. Supporting other roles or multiple input subjects + requires coordinated changes to bindings, checker semantics, and execution; + changing provider metadata alone does not extend that capability. + """ input_message: GuardedMessage output_content: str diff --git a/nemoguardrails/server/experimental/_guarded_stream.py b/nemoguardrails/server/experimental/_guarded_stream.py new file mode 100644 index 0000000000..bfd717d1b0 --- /dev/null +++ b/nemoguardrails/server/experimental/_guarded_stream.py @@ -0,0 +1,359 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Apply buffered output checks to provider-native event streams.""" + +import logging +from collections.abc import AsyncIterator, Callable +from dataclasses import dataclass + +from nemoguardrails.server.experimental._buffered_kernel import ( + InspectionStage, + OperationBlocked, + OperationCheckFailed, + OperationModificationUnsupported, + OperationProjectionFailed, +) +from nemoguardrails.server.experimental._content_checker import ( + ContentAllowed, + ContentBlocked, + ContentChecker, + ContentCheckFailed, + OutputContentCheck, + StreamBufferingPolicy, + UnsupportedContentModification, + validate_content_check_decision, +) +from nemoguardrails.server.experimental._guarded_operation import UnsupportedGuardedPayload +from nemoguardrails.server.experimental._http_kernel import BufferedHttpResponse +from nemoguardrails.server.experimental._stream_cleanup import StreamClose, close_stream +from nemoguardrails.server.experimental.provider.sse import ( + AmbiguousServerSentEvent, + ServerSentEvent, + ServerSentEventTooLarge, + TruncatedServerSentEvent, + iter_sse_events, +) +from nemoguardrails.server.experimental.provider.stream import ( + ProviderStreamAdapter, + StreamEventRole, + UnsupportedProviderStream, + validate_stream_event, +) +from nemoguardrails.server.experimental.provider.types import GuardedMessage + +log = logging.getLogger(__name__) + + +class UnsupportedStreamInspection(ValueError): + """Report an output policy that cannot safely release a guarded stream.""" + + +class PendingStreamBufferTooLarge(ValueError): + """Report unchecked provider events that exceed the buffering limit.""" + + +class StreamSnapshotHistoryTooLarge(ValueError): + """Report snapshot verification state that exceeds the stream limit.""" + + +@dataclass(frozen=True, slots=True) +class StreamUpstreamFailed: + """Stop a stream when its upstream byte iterator fails.""" + + failure: BaseException + + +@dataclass(frozen=True, slots=True) +class StreamProcessingFailed: + """Stop a stream when guarded stream processing fails unexpectedly.""" + + failure: BaseException + + +StreamOutcome = ( + OperationBlocked + | OperationCheckFailed + | OperationModificationUnsupported + | OperationProjectionFailed + | StreamUpstreamFailed + | StreamProcessingFailed +) +StreamOutcomeRenderer = Callable[[StreamOutcome], BufferedHttpResponse] + + +@dataclass(frozen=True, slots=True) +class _PendingStreamEvent: + event: ServerSentEvent + has_guarded_text: bool + + +@dataclass(frozen=True, slots=True) +class _TextWindow: + processing: tuple[str, ...] + release_count: int + + +class _UpstreamStreamFailure(Exception): + def __init__(self, failure: BaseException): + super().__init__(str(failure)) + self.failure = failure + + +def validate_streaming_policy(streaming_policy: StreamBufferingPolicy | None) -> None: + """Require output streams to buffer every guarded delta before release.""" + + if streaming_policy is None: + raise UnsupportedStreamInspection("Guarded output streaming requires a buffering policy.") + if streaming_policy.release_before_check: + raise UnsupportedStreamInspection("Guarded streams cannot release output before checking it.") + + +def _take_checked_events( + pending: list[_PendingStreamEvent], + text_count: int, +) -> tuple[ServerSentEvent, ...]: + checked_count = 0 + end = 0 + for end, event in enumerate(pending, start=1): + if event.has_guarded_text: + checked_count += 1 + if checked_count == text_count: + break + checked = pending[:end] + del pending[:end] + return tuple(event.event for event in checked) + + +async def _text_windows( + source: AsyncIterator[str | ServerSentEvent], + *, + chunk_size: int, + context_size: int, +) -> AsyncIterator[_TextWindow | ServerSentEvent]: + context: list[str] = [] + pending: list[str] = [] + async for text in source: + if isinstance(text, ServerSentEvent): + yield text + continue + pending.append(text) + if len(pending) < chunk_size: + continue + processing = context + pending + yield _TextWindow(tuple(processing), len(pending)) + context = processing[-context_size:] if context_size else [] + pending = [] + if pending: + yield _TextWindow(tuple(context + pending), len(pending)) + + +async def _guard_upstream(source: AsyncIterator[bytes]) -> AsyncIterator[bytes]: + try: + async for chunk in source: + if not isinstance(chunk, bytes): + raise TypeError("A provider stream must yield bytes.") + yield chunk + except Exception as failure: + raise _UpstreamStreamFailure(failure) from failure + + +def _log_outcome(outcome: StreamOutcome) -> None: + """Record why a guarded stream stopped, by type and stage only. + + Failure messages and tracebacks can carry provider or checked text, so + they are never logged. Blocks are normal checker decisions and are not + logged here. + """ + if isinstance(outcome, OperationBlocked): + return + failure = outcome.failure + if isinstance(failure, ContentCheckFailed): + failure = failure.cause + stage = "stream" if isinstance(outcome, StreamUpstreamFailed | StreamProcessingFailed) else outcome.stage.value + log.warning("Guarded stream stopped: %s at %s stage (%s)", type(outcome).__name__, stage, type(failure).__name__) + + +def _encoded_outcome( + outcome: StreamOutcome, + *, + adapter: ProviderStreamAdapter, + render_outcome: StreamOutcomeRenderer, +) -> tuple[bytes, ...]: + _log_outcome(outcome) + response = render_outcome(outcome) + return tuple(adapter.encode_error(response.body)) + + +async def guard_provider_stream( + source: AsyncIterator[bytes], + *, + checker: ContentChecker, + streaming_policy: StreamBufferingPolicy | None, + input_message: GuardedMessage, + adapter: ProviderStreamAdapter, + render_outcome: StreamOutcomeRenderer, + max_event_bytes: int = 1024 * 1024, + max_pending_bytes: int = 10 * 1024 * 1024, + close: StreamClose | None = None, +) -> AsyncIterator[bytes]: + """Release provider SSE events only after required output checks pass. + + The current execution model requires one projected user input subject as + context for assistant-output checks. This is a provider-neutral execution + invariant, not a restriction on roles in the provider's raw conversation. + Bindings supporting other input subjects require corresponding checker and + execution support. + + Change this invariant only when the framework implements a capability for + other input roles or multiple guarded input subjects, with matching bindings, + checker semantics, and buffered and streaming execution tests. Adding a + provider or changing a profile declaration alone does not relax it. + + Iteration attempts bounded, shielded cleanup of ``source`` and the optional + resource-owner ``close`` callback, including when arguments are rejected. + Iterator termination alone cannot recover interrupted generator cleanup; + callers owning external resources must supply their independent closer. + A stream never iterated does not run cleanup, so its caller retains ownership. + Validated data-less keepalive prefixes are + released promptly; they never overtake pending payloads or terminal events. + """ + + try: + if input_message.role != "user": + raise ValueError("A guarded stream requires a user input message.") + if any(type(limit) is not int or limit <= 0 for limit in (max_event_bytes, max_pending_bytes)): + raise ValueError("Stream byte limits must be positive.") + except ValueError: + await close_stream(source, close=close) + raise + + try: + if streaming_policy is None: + async for chunk in source: + yield chunk + return + validate_streaming_policy(streaming_policy) + + pending: list[_PendingStreamEvent] = [] + pending_bytes = 0 + observed_text: list[str] = [] + observed_text_bytes = 0 + snapshot_history_complete = True + + async def text_deltas() -> AsyncIterator[str | ServerSentEvent]: + nonlocal observed_text_bytes, pending_bytes, snapshot_history_complete + async for event in iter_sse_events(_guard_upstream(source), max_event_bytes=max_event_bytes): + classified = adapter.classify_event(event) + validate_stream_event(adapter.contract, classified) + if not pending and event.data is None and classified.role is StreamEventRole.OPAQUE_METADATA: + yield event + continue + if classified.role is StreamEventRole.TEXT_SNAPSHOT: + if not snapshot_history_complete: + raise StreamSnapshotHistoryTooLarge( + "Provider stream text exceeds the configured snapshot verification limit." + ) + if classified.text != "".join(observed_text): + raise UnsupportedProviderStream( + "A provider stream text snapshot does not match the guarded text deltas." + ) + text = classified.text if classified.role is StreamEventRole.GUARDED_TEXT else None + pending.append(_PendingStreamEvent(event, has_guarded_text=bool(text))) + pending_bytes += len(event.raw) + if pending_bytes > max_pending_bytes: + raise PendingStreamBufferTooLarge( + "Unchecked provider stream events exceed the configured buffering limit." + ) + if text: + if snapshot_history_complete: + # JSON escapes can produce lone surrogates, for example half of a + # pair split across deltas. Clients receive them as sent, so the + # proxy relays them too and only needs a byte count here. + observed_text_bytes += len(text.encode("utf-8", "surrogatepass")) + if observed_text_bytes > max_event_bytes: + observed_text.clear() + snapshot_history_complete = False + else: + observed_text.append(text) + yield text + adapter.validate_end_of_stream() + + async for window in _text_windows( + text_deltas(), + chunk_size=streaming_policy.chunk_size, + context_size=streaming_policy.context_size, + ): + if isinstance(window, ServerSentEvent): + yield window.raw + continue + checked_events = _take_checked_events(pending, window.release_count) + pending_bytes -= sum(len(event.raw) for event in checked_events) + try: + decision = validate_content_check_decision( + await checker.check_output(OutputContentCheck(input_message, "".join(window.processing))) + ) + except UnsupportedContentModification as failure: + outcome: StreamOutcome = OperationModificationUnsupported(InspectionStage.OUTPUT, failure) + except Exception as failure: + outcome = OperationCheckFailed( + InspectionStage.OUTPUT, + ContentCheckFailed("The output content check failed.", failure), + ) + else: + if isinstance(decision, ContentAllowed): + for event in checked_events: + yield event.raw + continue + if isinstance(decision, ContentBlocked): + outcome = OperationBlocked(InspectionStage.OUTPUT, decision) + else: + outcome = OperationCheckFailed(InspectionStage.OUTPUT, decision) + for chunk in _encoded_outcome(outcome, adapter=adapter, render_outcome=render_outcome): + yield chunk + return + + for event in pending: + yield event.event.raw + except ( + AmbiguousServerSentEvent, + PendingStreamBufferTooLarge, + ServerSentEventTooLarge, + TruncatedServerSentEvent, + StreamSnapshotHistoryTooLarge, + UnsupportedProviderStream, + ) as failure: + outcome = OperationProjectionFailed( + InspectionStage.OUTPUT, + UnsupportedGuardedPayload(str(failure) or "The provider stream does not match the guarded content schema."), + ) + for chunk in _encoded_outcome(outcome, adapter=adapter, render_outcome=render_outcome): + yield chunk + except _UpstreamStreamFailure as failure: + for chunk in _encoded_outcome( + StreamUpstreamFailed(failure.failure), + adapter=adapter, + render_outcome=render_outcome, + ): + yield chunk + except Exception as failure: + for chunk in _encoded_outcome( + StreamProcessingFailed(failure), + adapter=adapter, + render_outcome=render_outcome, + ): + yield chunk + finally: + await close_stream(source, close=close) diff --git a/nemoguardrails/server/experimental/_stream_cleanup.py b/nemoguardrails/server/experimental/_stream_cleanup.py new file mode 100644 index 0000000000..5785cc3aeb --- /dev/null +++ b/nemoguardrails/server/experimental/_stream_cleanup.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Best-effort, bounded cleanup independent of stream iteration state.""" + +import inspect +import logging +from collections.abc import Awaitable, Callable + +import anyio + +log = logging.getLogger(__name__) + +_STREAM_CLOSE_TIMEOUT_SECONDS = 5.0 +StreamClose = Callable[[], Awaitable[None]] + + +async def close_stream(source: object, *, close: StreamClose | None = None) -> None: + """Close an iterator and, independently, its resource owner. + + Each attempt has a shielded five-second budget. Timeouts and ordinary close + failures are logged without replacing the stream's outcome. Resource owners + must make closing idempotent and cancellation-cooperative. Iterator cleanup + alone cannot recover resources from an interrupted async-generator finally. + This does not forcibly stop blocking code or cancellation-suppressing code. + """ + + async def attempt(closer: Callable[[], object] | None) -> None: + """Bound one closer without swallowing cancellation from its caller.""" + if not callable(closer): + return + try: + with anyio.move_on_after(_STREAM_CLOSE_TIMEOUT_SECONDS, shield=True) as scope: + result = closer() + if inspect.isawaitable(result): + await result + if scope.cancelled_caught: + log.warning("Stream cleanup timed out; resources may remain open.") + except Exception: + log.warning("Stream cleanup failed; resources may remain open.") + + try: + await attempt(getattr(source, "aclose", None)) + finally: + await attempt(close) diff --git a/nemoguardrails/server/experimental/provider/sse.py b/nemoguardrails/server/experimental/provider/sse.py new file mode 100644 index 0000000000..e78bcc8a9d --- /dev/null +++ b/nemoguardrails/server/experimental/provider/sse.py @@ -0,0 +1,198 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Parse Server-Sent Events without provider-specific knowledge. + +The parser preserves complete upstream events byte-for-byte while exposing the +event type and joined data fields needed by provider adapters. +""" + +from collections.abc import AsyncIterator, Iterator +from dataclasses import dataclass + + +class ServerSentEventTooLarge(ValueError): + """Report an SSE event that exceeds the configured buffering limit.""" + + +class TruncatedServerSentEvent(ValueError): + """Report an unfinished SSE block containing fields at end of stream.""" + + +class AmbiguousServerSentEvent(ValueError): + """Report SSE bytes that common clients interpret differently.""" + + +@dataclass(frozen=True, slots=True) +class ServerSentEvent: + """Represent one complete SSE event. + + Attributes: + raw: Original event bytes, including its terminating blank line when + supplied by the upstream server. + lines: Event field and comment lines without line endings. + """ + + raw: bytes + lines: tuple[bytes, ...] + + @classmethod + def from_bytes(cls, raw: bytes) -> "ServerSentEvent": + """Create an event while retaining its original representation. + + Args: + raw: Exactly one complete SSE block, without stream-start BOM handling. + + Returns: + The parsed event representation. + """ + + framer = _SSEFramer(max_event_bytes=max(1, len(raw)), stream_start=False) + events = [*framer.feed(raw), *framer.finish()] + if len(events) != 1 or events[0].raw != raw: + raise ValueError("Expected exactly one complete SSE block.") + return events[0] + + @property + def data(self) -> bytes | None: + """Return joined SSE data fields, or ``None`` when none are present.""" + + values = [] + for line in self.lines: + name, separator, value = line.partition(b":") + if separator and name == b"data": + values.append(value[1:] if value.startswith(b" ") else value) + elif not separator and name == b"data": + values.append(b"") + return b"\n".join(values) if values else None + + @property + def event_type(self) -> bytes | None: + """Return the last SSE event field, or ``None`` when it is absent.""" + + value = None + for line in self.lines: + name, separator, field_value = line.partition(b":") + if name == b"event": + value = field_value[1:] if separator and field_value.startswith(b" ") else field_value + return value + + +class _SSEFramer: + def __init__(self, *, max_event_bytes: int, stream_start: bool = True): + self.max_event_bytes = max_event_bytes + self.raw = bytearray() + self.lines: list[bytes] = [] + self.line_start = 0 + self.pending_cr = False + self.stream_start = stream_start + + def _append(self, value: int) -> None: + if len(self.raw) >= self.max_event_bytes: + raise ServerSentEventTooLarge + self.raw.append(value) + + def _field_line(self, line: bytes) -> bytes: + if self.stream_start: + self.stream_start = False + line = line.removeprefix(b"\xef\xbb\xbf") + # The SSE standard ignores only one BOM at stream start. Some clients + # also strip a BOM at later line starts and read the line as a field, + # so such a line could carry data that this parser treats as unknown. + if line.startswith(b"\xef\xbb\xbf"): + raise AmbiguousServerSentEvent("A provider stream line starts with a byte order mark after stream start.") + return line + + def _end_line(self, ending_length: int) -> ServerSentEvent | None: + line = self._field_line(bytes(self.raw[self.line_start : -ending_length])) + if line: + self.lines.append(line) + self.line_start = len(self.raw) + return None + event = ServerSentEvent(bytes(self.raw), tuple(self.lines)) + self.raw.clear() + self.lines.clear() + self.line_start = 0 + return event + + def feed(self, chunk: bytes) -> Iterator[ServerSentEvent]: + for value in chunk: + if self.pending_cr: + self.pending_cr = False + if value == 10: + self._append(value) + event = self._end_line(2) + if event is not None: + yield event + continue + event = self._end_line(1) + if event is not None: + yield event + self._append(value) + if value == 13: + self.pending_cr = True + elif value == 10: + event = self._end_line(1) + if event is not None: + yield event + + def finish(self) -> Iterator[ServerSentEvent]: + if self.pending_cr: + self.pending_cr = False + event = self._end_line(1) + if event is not None: + yield event + if self.raw: + tail = self._field_line(bytes(self.raw[self.line_start :])) + if any(line and not line.startswith(b":") for line in (*self.lines, tail)): + raise TruncatedServerSentEvent("The provider stream ends with an unfinished SSE field block.") + self.raw.clear() + self.lines.clear() + self.line_start = 0 + + +async def iter_sse_events( + source: AsyncIterator[bytes], + *, + max_event_bytes: int = 1024 * 1024, +) -> AsyncIterator[ServerSentEvent]: + """Yield complete SSE blocks from arbitrary HTTP byte chunks. + + Preserve raw bytes, including an initial BOM ignored only for field parsing. + Drop unfinished comment tails; reject all other unfinished field blocks. + + Args: + source: Raw upstream response byte stream. + max_event_bytes: Maximum bytes buffered for one event. + + Yields: + Complete SSE events in upstream order. + + Raises: + ValueError: If ``max_event_bytes`` is not positive. + ServerSentEventTooLarge: If one event exceeds the limit. + TruncatedServerSentEvent: If the stream ends in an unfinished field block. + AmbiguousServerSentEvent: If a line after stream start begins with a BOM. + """ + + if type(max_event_bytes) is not int or max_event_bytes <= 0: + raise ValueError("max_event_bytes must be positive.") + + framer = _SSEFramer(max_event_bytes=max_event_bytes) + async for chunk in source: + for event in framer.feed(chunk): + yield event + for event in framer.finish(): + yield event diff --git a/nemoguardrails/server/experimental/provider/stream.py b/nemoguardrails/server/experimental/provider/stream.py new file mode 100644 index 0000000000..a175d6d604 --- /dev/null +++ b/nemoguardrails/server/experimental/provider/stream.py @@ -0,0 +1,219 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Define the stream boundary implemented by provider adapters.""" + +from collections.abc import Sequence +from dataclasses import dataclass +from enum import Enum +from typing import Callable, Protocol + +from nemoguardrails.server.experimental.provider.payload import ProjectionFieldCoverage +from nemoguardrails.server.experimental.provider.sse import ServerSentEvent + + +class UnsupportedProviderStream(ValueError): + """Report a provider event that the guarded stream cannot interpret.""" + + +class StreamEventRole(str, Enum): + """Describe how one provider event relates to guarded content.""" + + GUARDED_TEXT = "guarded_text" + TEXT_SNAPSHOT = "text_snapshot" + OPAQUE_METADATA = "opaque_metadata" + PROVIDER_ERROR = "provider_error" + + +class StreamCapabilityProfile(str, Enum): + """Identify a closed streaming capability understood by the framework.""" + + SINGLE_TEXT_DELTA_V1 = "single_text_delta.v1" + + +@dataclass(frozen=True, slots=True) +class StreamShapeCoverage: + """Classify every reviewed native event shape in a provider stream.""" + + guarded_shapes: frozenset[str] = frozenset() + snapshot_shapes: frozenset[str] = frozenset() + opaque_shapes: frozenset[str] = frozenset() + provider_error_shapes: frozenset[str] = frozenset() + + def __post_init__(self) -> None: + classifications = ( + self.guarded_shapes, + self.snapshot_shapes, + self.opaque_shapes, + self.provider_error_shapes, + ) + repeated = set() + for index, shapes in enumerate(classifications): + for other_shapes in classifications[index + 1 :]: + repeated.update(shapes & other_shapes) + if repeated: + raise ValueError(f"Stream shapes have multiple classifications: {sorted(repeated)}") + + @property + def reviewed_shapes(self) -> frozenset[str]: + """Return every shape deliberately covered by the stream contract.""" + + return self.guarded_shapes | self.snapshot_shapes | self.opaque_shapes | self.provider_error_shapes + + +@dataclass(frozen=True, slots=True) +class StreamProjectionContract: + """Describe the runtime semantics of one guarded stream projection.""" + + projection_id: str + profile: StreamCapabilityProfile + shapes: StreamShapeCoverage + fields: ProjectionFieldCoverage | None = None + + def __post_init__(self) -> None: + if not self.projection_id.strip(): + raise ValueError("A stream projection contract must have an identifier.") + if not self.shapes.reviewed_shapes: + raise ValueError("A stream projection contract must classify at least one shape.") + + +@dataclass(frozen=True, slots=True) +class GuardedStreamEvent: + """Describe one provider event at the Guardrails boundary.""" + + shape: str + role: StreamEventRole + text: str | None = None + + def __post_init__(self) -> None: + """Require text exactly for roles that carry generated content.""" + + if not self.shape: + raise ValueError("A guarded stream event must identify its provider shape.") + carries_text = self.role in {StreamEventRole.GUARDED_TEXT, StreamEventRole.TEXT_SNAPSHOT} + if carries_text != (self.text is not None): + raise ValueError("Guarded text and snapshot events must carry text; other event roles must not.") + + +def validate_stream_event(contract: StreamProjectionContract, event: GuardedStreamEvent) -> None: + """Require a classified event to match its declared shape semantics.""" + + roles = { + StreamEventRole.GUARDED_TEXT: contract.shapes.guarded_shapes, + StreamEventRole.TEXT_SNAPSHOT: contract.shapes.snapshot_shapes, + StreamEventRole.OPAQUE_METADATA: contract.shapes.opaque_shapes, + StreamEventRole.PROVIDER_ERROR: contract.shapes.provider_error_shapes, + } + declared_role = next((role for role, shapes in roles.items() if event.shape in shapes), None) + if declared_role is None: + raise UnsupportedProviderStream( + f"Provider stream shape {event.shape!r} is not declared by {contract.projection_id!r}." + ) + if event.role is not declared_role: + raise UnsupportedProviderStream( + f"Provider stream shape {event.shape!r} is classified as {event.role.value!r}; " + f"the contract declares {declared_role.value!r}." + ) + + +class ProviderStreamClassifier(Protocol): + """Classify native events using declarative projection semantics.""" + + contract: StreamProjectionContract + + def classify_event(self, event: ServerSentEvent) -> GuardedStreamEvent: + """Validate and classify one provider event.""" + + ... + + +class ProviderStreamHooks(Protocol): + """Implement stateful protocol checks and provider-native error framing.""" + + def observe_event(self, event: ServerSentEvent, classified: GuardedStreamEvent) -> None: + """Observe one classified event and update protocol state.""" + + ... + + def validate_end_of_stream(self) -> None: + """Validate provider protocol state when the byte stream ends.""" + + ... + + def encode_error(self, rendered_body: bytes) -> Sequence[bytes]: + """Frame a proxy error and its provider-native body as complete SSE blocks.""" + + ... + + +class ProviderStreamAdapter(Protocol): + """Expose the complete stream behavior consumed by the shared runtime.""" + + contract: StreamProjectionContract + + def classify_event(self, event: ServerSentEvent) -> GuardedStreamEvent: + """Validate, classify, and observe one provider event.""" + + ... + + def validate_end_of_stream(self) -> None: + """Validate provider protocol state when the byte stream ends.""" + + ... + + def encode_error(self, rendered_body: bytes) -> Sequence[bytes]: + """Frame a proxy error and its provider-native body as complete SSE blocks.""" + + ... + + +@dataclass(slots=True) +class ClassifiedStreamAdapter: + """Compose a generated classifier with handwritten protocol hooks.""" + + classifier: ProviderStreamClassifier + hooks: ProviderStreamHooks + + @property + def contract(self) -> StreamProjectionContract: + """Return the classifier's stream projection contract.""" + + return self.classifier.contract + + def classify_event(self, event: ServerSentEvent) -> GuardedStreamEvent: + """Classify an event and expose it to stateful protocol hooks.""" + + classified = self.classifier.classify_event(event) + self.hooks.observe_event(event, classified) + return classified + + def validate_end_of_stream(self) -> None: + """Delegate end-of-stream validation to provider hooks.""" + + self.hooks.validate_end_of_stream() + + def encode_error(self, rendered_body: bytes) -> Sequence[bytes]: + """Delegate provider-native error framing to provider hooks.""" + + return self.hooks.encode_error(rendered_body) + + +def create_classified_stream_adapter_factory( + classifier: ProviderStreamClassifier, + hooks_factory: Callable[[], ProviderStreamHooks], +) -> Callable[[], ProviderStreamAdapter]: + """Create a zero-argument adapter factory for an endpoint declaration.""" + + return lambda: ClassifiedStreamAdapter(classifier, hooks_factory()) diff --git a/nemoguardrails/server/experimental/provider/stream_classifier.py b/nemoguardrails/server/experimental/provider/stream_classifier.py new file mode 100644 index 0000000000..9cc54d75d5 --- /dev/null +++ b/nemoguardrails/server/experimental/provider/stream_classifier.py @@ -0,0 +1,262 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Validate declared event shapes without changing the inspected wire text.""" + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any, Literal + +from pydantic import BaseModel, ValidationError + +from nemoguardrails.server.experimental._json_payload import parse_json_object +from nemoguardrails.server.experimental.provider.payload import guarded_schema_error +from nemoguardrails.server.experimental.provider.sse import ServerSentEvent +from nemoguardrails.server.experimental.provider.stream import ( + GuardedStreamEvent, + StreamEventRole, + StreamProjectionContract, + UnsupportedProviderStream, + validate_stream_event, +) + +StreamMatchValue = str | int | float | bool | None + + +@dataclass(frozen=True, slots=True) +class StreamArraySelection: + """Select stream content by provider discriminator fields.""" + + match: tuple[tuple[str, StreamMatchValue], ...] + cardinality: Literal["at_most_one", "exactly_one"] + + def select(self, items: Sequence[object]) -> object | None: + matches = [ + item + for item in items + if all( + type(value := (item.get(name, object()) if isinstance(item, dict) else getattr(item, name, object()))) + is type(expected) + and value == expected + for name, expected in self.match + ) + ] + if len(matches) > 1 or (self.cardinality == "exactly_one" and not matches): + raise ValueError(f"Stream text selection requires {self.cardinality.replace('_', ' ')} matching item.") + return matches[0] if matches else None + + +StreamTextPath = tuple[str | int | StreamArraySelection, ...] + + +@dataclass(frozen=True, slots=True) +class StreamEventRule: + """Declare validation and classification for one provider event shape.""" + + shape: str + model: type[BaseModel] + role: StreamEventRole + match: tuple[tuple[str, StreamMatchValue], ...] = () + required_fields: frozenset[str] = frozenset() + text_path: StreamTextPath | None = None + missing_text_role: StreamEventRole | None = None + missing_text_shape: str | None = None + missing_text_value: str | None = None + + def matches(self, payload: dict[str, Any]) -> bool: + """Return whether a provider payload selects this event rule.""" + + missing = object() + + def resolve(path: str) -> object: + value: object = payload + for segment in path.split("."): + if not isinstance(value, dict): + return missing + value = value.get(segment, missing) + return value + + return self.required_fields <= payload.keys() and all( + type(value := resolve(name)) is type(expected) and value == expected for name, expected in self.match + ) + + +@dataclass(frozen=True, slots=True) +class StreamClassifierDefinition: + """Describe the event-shape semantics of one guarded provider stream.""" + + subject: str + contract: StreamProjectionContract + rules: tuple[StreamEventRule, ...] + sentinels: tuple[tuple[bytes, str], ...] = () + non_data_shape: str | None = None + event_type_field: str | None = None + require_event_type: bool = False + + +@dataclass(frozen=True, slots=True) +class StreamClassifier: + """Classify provider events using validated stream rules.""" + + subject: str + contract: StreamProjectionContract + rules: tuple[StreamEventRule, ...] + sentinels: tuple[tuple[bytes, str], ...] + non_data_shape: str | None + event_type_field: str | None + require_event_type: bool + + def classify_event(self, event: ServerSentEvent) -> GuardedStreamEvent: + """Validate and classify one native provider event.""" + + data = event.data + if data is None: + if self.non_data_shape is None: + raise UnsupportedProviderStream(f"{self.subject} events must contain data.") + return GuardedStreamEvent(self.non_data_shape, StreamEventRole.OPAQUE_METADATA) + stripped = data.strip() + for sentinel, shape in self.sentinels: + if stripped == sentinel: + return GuardedStreamEvent(shape, StreamEventRole.OPAQUE_METADATA) + try: + payload = parse_json_object(data) + if self.event_type_field is not None: + event_type = payload.get(self.event_type_field) + if not isinstance(event_type, str): + raise UnsupportedProviderStream( + f"{self.subject} events must declare string field {self.event_type_field!r}." + ) + if self.require_event_type and event.event_type is None: + raise UnsupportedProviderStream(f"{self.subject} SSE event field is required.") + if event.event_type is not None and event.event_type != event_type.encode("ascii"): + raise UnsupportedProviderStream( + f"{self.subject} SSE event field must match data field {self.event_type_field!r}." + ) + matching_rules = [candidate for candidate in self.rules if candidate.matches(payload)] + if not matching_rules: + raise UnsupportedProviderStream(f"{self.subject} event is not covered by the guarded profile.") + if len(matching_rules) > 1: + shapes = ", ".join(repr(rule.shape) for rule in matching_rules) + raise UnsupportedProviderStream(f"{self.subject} event ambiguously matches {shapes}.") + rule = matching_rules[0] + text = _resolve_text(payload, rule.text_path) + projection = rule.model.model_validate(payload) + if _resolve_text(projection, rule.text_path) != text: + raise UnsupportedProviderStream(f"{self.subject} validation changed the original guarded text.") + if text is None and rule.missing_text_role is not None: + return GuardedStreamEvent( + rule.missing_text_shape or rule.shape, + rule.missing_text_role, + rule.missing_text_value, + ) + if text is None and rule.role in {StreamEventRole.GUARDED_TEXT, StreamEventRole.TEXT_SNAPSHOT}: + raise UnsupportedProviderStream(f"{self.subject} event {rule.shape!r} has no guarded text value.") + return GuardedStreamEvent(rule.shape, rule.role, text) + except UnsupportedProviderStream: + raise + except ValidationError as error: + raise UnsupportedProviderStream(guarded_schema_error(error, self.subject)) from error + except (IndexError, TypeError, UnicodeError, ValueError) as error: + raise UnsupportedProviderStream(str(error)) from error + + +def _resolve_text(projection: BaseModel | dict[str, Any], path: StreamTextPath | None) -> str | None: + if path is None: + return None + value: object = projection + for segment in path: + if isinstance(segment, str) and isinstance(value, BaseModel): + value = getattr(value, segment) + elif isinstance(segment, str) and isinstance(value, dict): + value = value.get(segment) + elif isinstance(segment, int) and isinstance(value, list): + if segment < 0: + raise ValueError("Stream event text paths cannot use negative indexes.") + if segment >= len(value): + return None + value = value[segment] + elif isinstance(segment, StreamArraySelection) and isinstance(value, list): + value = segment.select(value) + else: + raise TypeError("Stream event text path does not match its projection.") + if value is None: + return None + if not isinstance(value, str): + raise TypeError("Stream event text must resolve to a string.") + return value + + +def build_stream_classifier(definition: StreamClassifierDefinition) -> StreamClassifier: + """Validate declared event rules and return their runtime classifier.""" + + if not definition.subject: + raise ValueError("A stream classifier definition must identify its provider event subject.") + if not definition.rules: + raise ValueError("A stream classifier definition must declare at least one event rule.") + if definition.subject != definition.contract.projection_id: + raise ValueError("A stream classifier subject must match its projection contract identifier.") + if definition.require_event_type and definition.event_type_field is None: + raise ValueError("A required SSE event field needs a data discriminator.") + signatures: set[tuple[frozenset[tuple[str, type, StreamMatchValue]], frozenset[str]]] = set() + + def validate_declared_event(shape: str, role: StreamEventRole, text: str | None = None) -> None: + try: + validate_stream_event(definition.contract, GuardedStreamEvent(shape, role, text)) + except UnsupportedProviderStream as error: + raise ValueError(str(error)) from error + + if definition.non_data_shape is not None: + validate_declared_event(definition.non_data_shape, StreamEventRole.OPAQUE_METADATA) + for _, shape in definition.sentinels: + validate_declared_event(shape, StreamEventRole.OPAQUE_METADATA) + for rule in definition.rules: + if not rule.shape: + raise ValueError("A stream event rule must identify its provider shape.") + if not rule.match and not rule.required_fields: + raise ValueError(f"Stream event rule {rule.shape!r} must declare how it is selected.") + signature = (frozenset((name, type(value), value) for name, value in rule.match), rule.required_fields) + if signature in signatures: + raise ValueError(f"Stream event rule {rule.shape!r} duplicates another selection.") + signatures.add(signature) + carries_text = rule.role in {StreamEventRole.GUARDED_TEXT, StreamEventRole.TEXT_SNAPSHOT} + if carries_text != (rule.text_path is not None): + raise ValueError(f"Stream event rule {rule.shape!r} has an inconsistent text role and path.") + if rule.missing_text_role is not None and rule.text_path is None: + raise ValueError(f"Stream event rule {rule.shape!r} cannot classify missing text without a text path.") + if rule.missing_text_shape is not None and rule.missing_text_role is None: + raise ValueError(f"Stream event rule {rule.shape!r} cannot rename missing text without classifying it.") + if rule.missing_text_value is not None and rule.missing_text_role is None: + raise ValueError(f"Stream event rule {rule.shape!r} cannot supply missing text without classifying it.") + if ( + rule.missing_text_role in {StreamEventRole.GUARDED_TEXT, StreamEventRole.TEXT_SNAPSHOT} + and rule.missing_text_value is None + ): + raise ValueError(f"Stream event rule {rule.shape!r} cannot classify missing text as content.") + validate_declared_event(rule.shape, rule.role, "" if carries_text else None) + if rule.missing_text_role is not None: + validate_declared_event( + rule.missing_text_shape or rule.shape, + rule.missing_text_role, + rule.missing_text_value, + ) + return StreamClassifier( + subject=definition.subject, + contract=definition.contract, + rules=definition.rules, + sentinels=definition.sentinels, + non_data_shape=definition.non_data_shape, + event_type_field=definition.event_type_field, + require_event_type=definition.require_event_type, + ) diff --git a/tests/server/experimental/test_content_checker.py b/tests/server/experimental/test_content_checker.py index d4a8a8b67d..f52713d515 100644 --- a/tests/server/experimental/test_content_checker.py +++ b/tests/server/experimental/test_content_checker.py @@ -25,6 +25,7 @@ InputContentCheck, InvalidContentChecker, OutputContentCheck, + StreamBufferingPolicy, UnsupportedContentCheckerConfiguration, UnsupportedContentModification, validate_content_check_decision, @@ -157,3 +158,32 @@ def test_input_check_carries_the_guarded_message(): message = GuardedMessage("user", "question") assert InputContentCheck(message).message is message + + +def test_stream_buffering_policy_is_part_of_output_inspection(): + buffering = StreamBufferingPolicy(chunk_size=2, context_size=1) + + assert ContentInspectionPolicy(True, True, buffering).stream_buffering is buffering + + +@pytest.mark.parametrize( + "buffering", + [ + pytest.param(lambda: StreamBufferingPolicy(0, 0), id="zero-chunk"), + pytest.param(lambda: StreamBufferingPolicy(1, -1), id="negative-context"), + pytest.param(lambda: StreamBufferingPolicy(1, 0, 1), id="non-boolean-release"), + ], +) +def test_stream_buffering_policy_rejects_invalid_values(buffering): + with pytest.raises((TypeError, ValueError)): + buffering() + + +def test_stream_buffering_requires_output_inspection(): + with pytest.raises(ValueError, match="requires output inspection"): + ContentInspectionPolicy(True, False, StreamBufferingPolicy(1, 0)) + + +def test_stream_buffering_must_be_a_buffering_policy(): + with pytest.raises(TypeError, match="StreamBufferingPolicy"): + ContentInspectionPolicy(True, True, {"chunk_size": 1}) diff --git a/tests/server/experimental/test_guarded_stream.py b/tests/server/experimental/test_guarded_stream.py new file mode 100644 index 0000000000..2bdb787782 --- /dev/null +++ b/tests/server/experimental/test_guarded_stream.py @@ -0,0 +1,664 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import json + +import anyio +import pytest + +from nemoguardrails.server.experimental import _stream_cleanup +from nemoguardrails.server.experimental._buffered_kernel import ( + OperationBlocked, + OperationCheckFailed, + OperationModificationUnsupported, + OperationProjectionFailed, +) +from nemoguardrails.server.experimental._content_checker import ( + ContentAllowed, + ContentBlocked, + ContentCheckFailed, + StreamBufferingPolicy, +) +from nemoguardrails.server.experimental._guarded_stream import ( + StreamProcessingFailed, + StreamUpstreamFailed, + UnsupportedStreamInspection, + guard_provider_stream, + validate_streaming_policy, +) +from nemoguardrails.server.experimental._http_kernel import BufferedHttpResponse +from nemoguardrails.server.experimental._json_payload import parse_json_object +from nemoguardrails.server.experimental.provider.stream import ( + ClassifiedStreamAdapter, + GuardedStreamEvent, + StreamCapabilityProfile, + StreamEventRole, + StreamProjectionContract, + StreamShapeCoverage, + UnsupportedProviderStream, +) +from nemoguardrails.server.experimental.provider.types import GuardedMessage + +_CONTRACT = StreamProjectionContract( + projection_id="test.stream", + profile=StreamCapabilityProfile.SINGLE_TEXT_DELTA_V1, + shapes=StreamShapeCoverage( + guarded_shapes=frozenset({"text"}), + snapshot_shapes=frozenset({"snapshot"}), + opaque_shapes=frozenset({"metadata", "end"}), + provider_error_shapes=frozenset({"error"}), + ), +) + + +class StaticChecker: + def __init__(self, decisions=()): + self.decisions = list(decisions) + self.calls = [] + + async def check_output(self, check): + self.calls.append(check) + decision = self.decisions.pop(0) if self.decisions else ContentAllowed() + if isinstance(decision, Exception): + raise decision + return decision + + +class Classifier: + contract = _CONTRACT + + def classify_event(self, event): + if event.data == b"[END]": + return GuardedStreamEvent("end", StreamEventRole.OPAQUE_METADATA) + try: + payload = parse_json_object(event.data or b"") + except ValueError as error: + raise UnsupportedProviderStream("unsupported fake stream event") from error + if "error" in payload: + return GuardedStreamEvent("error", StreamEventRole.PROVIDER_ERROR) + if "snapshot" in payload: + return GuardedStreamEvent("snapshot", StreamEventRole.TEXT_SNAPSHOT, payload["snapshot"]) + if "text" in payload: + return GuardedStreamEvent("text", StreamEventRole.GUARDED_TEXT, payload["text"]) + return GuardedStreamEvent("metadata", StreamEventRole.OPAQUE_METADATA) + + +class Hooks: + def __init__(self): + self.observed = [] + self.finished = False + self.encoded = [] + + def observe_event(self, event, classified): + self.observed.append((event, classified)) + + def validate_end_of_stream(self): + self.finished = True + + def encode_error(self, rendered_body): + self.encoded.append(rendered_body) + return (b"data: " + rendered_body + b"\n\n", b"data: [END]\n\n") + + +class Source: + def __init__(self, chunks, failure=None): + self.chunks = iter(chunks) + self.failure = failure + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self.chunks) + except StopIteration: + if self.failure is not None: + failure, self.failure = self.failure, None + raise failure + raise StopAsyncIteration + + async def aclose(self): + self.closed = True + + +def render_outcome(outcome): + if isinstance(outcome, OperationBlocked): + code = "blocked" + elif isinstance(outcome, OperationCheckFailed): + code = "check_failed" + elif isinstance(outcome, OperationModificationUnsupported): + code = "modification_unsupported" + elif isinstance(outcome, OperationProjectionFailed): + code = "unsupported_stream" + elif isinstance(outcome, StreamUpstreamFailed): + code = "upstream_failed" + elif isinstance(outcome, StreamProcessingFailed): + code = "processing_failed" + else: + raise AssertionError("unknown outcome") + return BufferedHttpResponse(200, (), json.dumps({"code": code}).encode()) + + +def stream_event(payload): + return b"data: " + json.dumps(payload, separators=(",", ":")).encode() + b"\n\n" + + +async def collect(source, checker, hooks, policy=StreamBufferingPolicy(1, 0), **limits): + return b"".join( + [ + chunk + async for chunk in guard_provider_stream( + source, + checker=checker, + streaming_policy=policy, + input_message=GuardedMessage("user", "question"), + adapter=ClassifiedStreamAdapter(Classifier(), hooks), + render_outcome=render_outcome, + **limits, + ) + ] + ) + + +@pytest.mark.asyncio +async def test_guarded_stream_checks_windows_before_releasing_original_events(): + events = [ + stream_event({"meta": 1}), + stream_event({"text": "one "}), + stream_event({"text": "two "}), + stream_event({"text": "three"}), + b"data: [END]\n\n", + ] + source = Source([b"".join(events)[:23], b"".join(events)[23:]]) + checker = StaticChecker() + hooks = Hooks() + + result = await collect(source, checker, hooks, StreamBufferingPolicy(chunk_size=2, context_size=1)) + + assert result == b"".join(events) + assert [call.output_content for call in checker.calls] == ["one two ", "two three"] + assert hooks.finished is True + assert source.closed is True + + +@pytest.mark.asyncio +async def test_zero_context_does_not_retain_the_previous_window(): + events = [stream_event({"text": text}) for text in ("one ", "two ", "three ", "four")] + checker = StaticChecker() + + result = await collect( + Source([b"".join(events)]), + checker, + Hooks(), + StreamBufferingPolicy(chunk_size=2, context_size=0), + ) + + assert result == b"".join(events) + assert [call.output_content for call in checker.calls] == ["one two ", "three four"] + + +class KeepaliveClassifier(Classifier): + """Recognize data-less transport events without inventing text.""" + + def classify_event(self, event): + if event.data is None: + return GuardedStreamEvent("metadata", StreamEventRole.OPAQUE_METADATA) + return super().classify_event(event) + + +@pytest.mark.asyncio +async def test_keepalive_prefix_is_released_before_reading_the_next_event(): + """A live idle stream must make progress without waiting for text or EOF.""" + checker = StaticChecker() + closed = [] + + async def source(): + try: + yield b": keepalive\n\n" + await asyncio.Event().wait() + finally: + closed.append(True) + + stream = guard_provider_stream( + source(), + checker=checker, + streaming_policy=StreamBufferingPolicy(2, 0), + input_message=GuardedMessage("user", "question"), + adapter=ClassifiedStreamAdapter(KeepaliveClassifier(), Hooks()), + render_outcome=render_outcome, + ) + try: + assert await asyncio.wait_for(anext(stream), 1) == b": keepalive\n\n" + assert checker.calls == [] + finally: + await stream.aclose() + assert closed == [True] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("prefix", [b"", stream_event({"text": "pending"}), b"data: [END]\n\n"]) +async def test_keepalives_preserve_pending_payload_order_and_bounds(prefix): + """Only a transport-only prefix can bypass the pending payload window.""" + keepalives = b": alive\n\n" * 30 + source = Source([prefix + keepalives]) + checker = StaticChecker() + result = b"".join( + [ + part + async for part in guard_provider_stream( + source, + checker=checker, + streaming_policy=StreamBufferingPolicy(2, 0), + input_message=GuardedMessage("user", "question"), + adapter=ClassifiedStreamAdapter(KeepaliveClassifier(), Hooks()), + render_outcome=render_outcome, + max_event_bytes=128, + max_pending_bytes=128, + ) + ] + ) + if prefix: + assert b"unsupported_stream" in result + assert b": alive" not in result + assert b"pending" not in result + else: + assert result == keepalives + assert checker.calls == [] + assert source.closed is True + + +@pytest.mark.asyncio +async def test_stream_block_suppresses_the_complete_pending_window(): + first = stream_event({"text": "unsafe "}) + second = stream_event({"text": "content"}) + checker = StaticChecker([ContentBlocked("blocked", "rule")]) + + result = await collect( + Source([first + second + b"data: [END]\n\n"]), + checker, + Hooks(), + StreamBufferingPolicy(2, 0), + ) + + assert first not in result + assert second not in result + assert b'"code": "blocked"' in result + assert result.endswith(b"data: [END]\n\n") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("decision", [ContentCheckFailed("failed"), RuntimeError("checker failed")]) +async def test_stream_check_failure_is_distinct_from_provider_drift(decision): + result = await collect(Source([stream_event({"text": "answer"})]), StaticChecker([decision]), Hooks()) + + assert b'"code": "check_failed"' in result + assert b"unsupported_stream" not in result + + +@pytest.mark.asyncio +async def test_stream_replacement_is_rejected_without_releasing_pending_text(): + event = stream_event({"text": "answer"}) + + result = await collect(Source([event]), StaticChecker([ContentAllowed(replacement="changed")]), Hooks()) + + assert event not in result + assert b'"code": "modification_unsupported"' in result + + +@pytest.mark.asyncio +async def test_provider_error_events_are_preserved_without_being_checked_as_text(): + error = stream_event({"error": {"message": "upstream failed"}}) + terminal = b"data: [END]\n\n" + checker = StaticChecker() + + result = await collect(Source([error + terminal]), checker, Hooks(), StreamBufferingPolicy(2, 0)) + + assert result == error + terminal + assert checker.calls == [] + + +@pytest.mark.asyncio +async def test_snapshot_must_equal_the_guarded_delta_history(): + source = Source([stream_event({"text": "answer"}) + stream_event({"snapshot": "different"})]) + + result = await collect(source, StaticChecker(), Hooks()) + + assert b'"code": "unsupported_stream"' in result + + +@pytest.mark.asyncio +async def test_unsupported_event_and_pending_limit_fail_closed(): + unsupported = await collect(Source([b"data: not-json\n\n"]), StaticChecker(), Hooks()) + oversized = await collect( + Source([stream_event({"meta": "long"})]), + StaticChecker(), + Hooks(), + max_pending_bytes=5, + ) + + assert b'"code": "unsupported_stream"' in unsupported + assert b'"code": "unsupported_stream"' in oversized + + +@pytest.mark.asyncio +async def test_manual_adapter_contract_mismatch_fails_closed(): + class MismatchedClassifier(Classifier): + def classify_event(self, _event): + return GuardedStreamEvent("text", StreamEventRole.OPAQUE_METADATA) + + hooks = Hooks() + result = b"".join( + [ + chunk + async for chunk in guard_provider_stream( + Source([stream_event({"text": "answer"})]), + checker=StaticChecker(), + streaming_policy=StreamBufferingPolicy(1, 0), + input_message=GuardedMessage("user", "question"), + adapter=ClassifiedStreamAdapter(MismatchedClassifier(), hooks), + render_outcome=render_outcome, + ) + ] + ) + + assert b'"code": "unsupported_stream"' in result + assert hooks.finished is False + + +@pytest.mark.asyncio +async def test_uninspected_stream_is_relayed_by_chunk_without_parsing(): + chunks = [b"not sse", b" and still opaque"] + source = Source(chunks) + + result = await collect(source, StaticChecker(), Hooks(), policy=None) + + assert result == b"".join(chunks) + assert source.closed is True + + +@pytest.mark.asyncio +async def test_cancelling_downstream_iteration_closes_the_upstream_source(): + source = Source([b"first", b"second"]) + stream = guard_provider_stream( + source, + checker=StaticChecker(), + streaming_policy=None, + input_message=GuardedMessage("user", "question"), + adapter=ClassifiedStreamAdapter(Classifier(), Hooks()), + render_outcome=render_outcome, + ) + + assert await anext(stream) == b"first" + await stream.aclose() + + assert source.closed is True + + +@pytest.mark.asyncio +async def test_upstream_failure_is_distinct_from_processing_failure(): + source = Source([], RuntimeError("upstream disconnected")) + + result = await collect(source, StaticChecker(), Hooks()) + + assert b'"code": "upstream_failed"' in result + assert source.closed is True + + +def test_stream_policy_must_buffer_before_release(): + with pytest.raises(UnsupportedStreamInspection, match="requires a buffering policy"): + validate_streaming_policy(None) + + with pytest.raises(UnsupportedStreamInspection, match="cannot release"): + validate_streaming_policy(StreamBufferingPolicy(1, 0, True)) + + +@pytest.mark.asyncio +async def test_truncation_withholds_pending_bytes_and_error_starts_on_complete_boundary(): + approved = stream_event({"text": "approved"}) + incomplete = stream_event({"text": "hidden"})[:-2] + source = Source([approved, incomplete]) + hooks = Hooks() + + result = await collect(source, StaticChecker(), hooks) + + assert result.startswith(approved + b"data: ") + assert b"hidden" not in result + assert b'"code": "unsupported_stream"' in result + assert hooks.finished is False + assert source.closed is True + + +@pytest.mark.asyncio +async def test_end_validation_failure_withholds_pending_window(): + class RejectingHooks(Hooks): + def validate_end_of_stream(self): + raise UnsupportedProviderStream("truncated provider lifecycle") + + checker = StaticChecker() + source = Source([stream_event({"text": "hidden"})]) + + result = await collect(source, checker, RejectingHooks(), policy=StreamBufferingPolicy(2, 0)) + + assert b"hidden" not in result + assert b'"code": "unsupported_stream"' in result + assert checker.calls == [] + assert source.closed is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("role", "limits"), + [("assistant", {}), ("user", {"max_event_bytes": 0}), ("user", {"max_pending_bytes": -1})], +) +async def test_rejected_arguments_raise_and_close_source(role, limits): + source = Source([stream_event({"text": "answer"})]) + stream = guard_provider_stream( + source, + checker=StaticChecker(), + streaming_policy=StreamBufferingPolicy(1, 0), + input_message=GuardedMessage(role, "question"), + adapter=ClassifiedStreamAdapter(Classifier(), Hooks()), + render_outcome=render_outcome, + **limits, + ) + + with pytest.raises(ValueError): + await anext(stream) + + assert source.closed is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("split", [0, 1, 2, 3]) +async def test_line_leading_bom_mid_stream_fails_closed_without_release(split): + """An ambiguous line is never forwarded, even when its BOM bytes are fragmented.""" + approved = stream_event({"text": "approved"}) + ambiguous = b"\xef\xbb\xbf" + stream_event({"text": "EVIL"}) + source = Source([approved, ambiguous[:split], ambiguous[split:]]) + + result = await collect(source, StaticChecker(), Hooks()) + + assert result.startswith(approved) + assert b"EVIL" not in result + assert b'"code": "unsupported_stream"' in result + assert source.closed is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("events", "text"), + [ + ([stream_event({"text": "a\ud800b"})], "a\ud800b"), + ([b'data: {"text":"a\xed\xa0\x80b"}\n\n'], "a\ud800b"), + ([stream_event({"text": "\ud83d"}), stream_event({"text": "\ude00"})], "\ud83d\ude00"), + ], + ids=["escape", "bytes", "split-pair"], +) +async def test_lone_surrogate_text_is_checked_and_relayed_unchanged(events, text): + """Lone surrogates reach clients as sent, so the proxy checks and relays them too.""" + checker = StaticChecker() + result = await collect(Source(events), checker, Hooks(), policy=StreamBufferingPolicy(len(events), 0)) + + assert result == b"".join(events) + assert [call.output_content for call in checker.calls] == [text] + + +@pytest.mark.asyncio +async def test_checker_that_cannot_encode_lone_surrogates_fails_closed(): + event = stream_event({"text": "a\ud800b"}) + checker = StaticChecker([UnicodeEncodeError("utf-8", "a\ud800b", 1, 2, "surrogates not allowed")]) + + result = await collect(Source([event]), checker, Hooks()) + + assert b'"code": "check_failed"' in result + assert b"ud800" not in result + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["stall", "raise"]) +@pytest.mark.parametrize("decision", [ContentAllowed(), ContentBlocked("blocked")]) +async def test_cleanup_is_bounded_without_replacing_the_stream_outcome(monkeypatch, caplog, failure, decision): + monkeypatch.setattr(_stream_cleanup, "_STREAM_CLOSE_TIMEOUT_SECONDS", 0.01) + + class FailingClose(Source): + async def aclose(self): + if failure == "stall": + await anyio.sleep_forever() + raise RuntimeError("private upstream detail") + + event = stream_event({"text": "answer"}) + with anyio.fail_after(1): + result = await collect(FailingClose([event]), StaticChecker([decision]), Hooks()) + + assert result == event if isinstance(decision, ContentAllowed) else b'"code": "blocked"' in result + assert "Stream cleanup" in caplog.text + assert "private upstream detail" not in caplog.text + + +@pytest.mark.asyncio +async def test_explicit_resource_owner_closes_after_iterator_cleanup_fails(): + closed = [] + + class FailingClose(Source): + async def aclose(self): + raise RuntimeError("iterator cleanup failed") + + async def close_owner(): + await anyio.sleep(0) + closed.append(True) + + await collect(FailingClose([]), StaticChecker(), Hooks(), close=close_owner) + assert closed == [True] + + +@pytest.mark.asyncio +async def test_iterator_close_cancellation_does_not_skip_the_resource_owner(): + closed = [] + + class CancelledClose(Source): + async def aclose(self): + raise asyncio.CancelledError + + async def close_owner(): + await anyio.sleep(0) + closed.append(True) + + with pytest.raises(asyncio.CancelledError): + await collect(CancelledClose([]), StaticChecker(), Hooks(), close=close_owner) + assert closed == [True] + + +class SecretRejectingHooks(Hooks): + def observe_event(self, event, classified): + raise UnsupportedProviderStream("Duplicate JSON member 'secret' is not supported.") + + +class FailingEndHooks(Hooks): + def validate_end_of_stream(self): + raise RuntimeError("secret") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("source", "checker", "hooks", "expected"), + [ + ( + lambda: Source([stream_event({"text": "hi"})]), + lambda: StaticChecker(), + SecretRejectingHooks, + "OperationProjectionFailed at output stage (UnsupportedGuardedPayload)", + ), + ( + lambda: Source([stream_event({"text": "hi"})]), + lambda: StaticChecker([RuntimeError("secret")]), + Hooks, + "OperationCheckFailed at output stage (RuntimeError)", + ), + ( + lambda: Source([stream_event({"text": "hi"})], failure=ConnectionError("secret")), + lambda: StaticChecker(), + Hooks, + "StreamUpstreamFailed at stream stage (ConnectionError)", + ), + ( + lambda: Source([stream_event({"text": "hi"})]), + lambda: StaticChecker(), + FailingEndHooks, + "StreamProcessingFailed at stream stage (RuntimeError)", + ), + ], + ids=["projection", "check", "upstream", "processing"], +) +async def test_stream_failures_are_logged_by_type_only(caplog, source, checker, hooks, expected): + """Failure messages and tracebacks can carry private text, so only types reach the log.""" + caplog.set_level("DEBUG", logger="nemoguardrails.server.experimental._guarded_stream") + + await collect(source(), checker(), hooks(), policy=StreamBufferingPolicy(2, 0)) + + records = [record for record in caplog.records if record.name.endswith("_guarded_stream")] + assert [record.getMessage() for record in records] == [f"Guarded stream stopped: {expected}"] + assert records[0].exc_info is None + assert "secret" not in caplog.text + + +@pytest.mark.asyncio +async def test_blocked_stream_is_not_logged_as_a_failure(caplog): + caplog.set_level("DEBUG", logger="nemoguardrails.server.experimental._guarded_stream") + + await collect(Source([stream_event({"text": "secret"})]), StaticChecker([ContentBlocked("blocked")]), Hooks()) + + assert not [record for record in caplog.records if record.name.endswith("_guarded_stream")] + + +@pytest.mark.asyncio +async def test_non_byte_provider_chunks_are_upstream_failures(): + result = await collect(Source(["data: text\n\n"]), StaticChecker(), Hooks()) + + assert b'"code": "upstream_failed"' in result + + +@pytest.mark.asyncio +async def test_snapshot_after_unverifiable_history_fails_closed(): + """Once delta history exceeds the limit, a later snapshot cannot be verified and is rejected.""" + deltas = [stream_event({"text": "x" * 10}) for _ in range(5)] + checker = StaticChecker() + + result = await collect(Source([*deltas, stream_event({"snapshot": "x"})]), checker, Hooks(), max_event_bytes=40) + + assert result.startswith(b"".join(deltas)) + assert b'"code": "unsupported_stream"' in result + assert b'"snapshot"' not in result diff --git a/tests/server/experimental/test_sse.py b/tests/server/experimental/test_sse.py new file mode 100644 index 0000000000..70277bccd2 --- /dev/null +++ b/tests/server/experimental/test_sse.py @@ -0,0 +1,268 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio + +import pytest +from openai._streaming import SSEDecoder + +from nemoguardrails.server.experimental.provider.sse import ( + AmbiguousServerSentEvent, + ServerSentEvent, + ServerSentEventTooLarge, + TruncatedServerSentEvent, + iter_sse_events, +) + + +async def _chunks(*values): + for value in values: + yield value + + +async def _byte_chunks(raw): + for byte in raw: + yield bytes([byte]) + + +def test_sse_event_exposes_joined_data_and_last_event_type_without_changing_bytes(): + raw = b"event: first\r\nevent: final\r\ndata: one\r\ndata:two\r\n\r\n" + event = ServerSentEvent.from_bytes(raw) + + assert event.raw == raw + assert event.event_type == b"final" + assert event.data == b"one\ntwo" + + +@pytest.mark.asyncio +async def test_sse_parser_preserves_events_split_across_arbitrary_chunks(): + raw = b": metadata\n\ndata: first\n\ndata: second\n\n" + + events = [event async for event in iter_sse_events(_chunks(raw[:7], raw[7:19], raw[19:]), max_event_bytes=30)] + + assert [event.raw for event in events] == [b": metadata\n\n", b"data: first\n\n", b"data: second\n\n"] + + +@pytest.mark.asyncio +async def test_sse_parser_bounds_complete_and_incomplete_events(): + with pytest.raises(ServerSentEventTooLarge): + _ = [event async for event in iter_sse_events(_chunks(b"data: oversized\n\n"), max_event_bytes=8)] + with pytest.raises(ServerSentEventTooLarge): + _ = [event async for event in iter_sse_events(_chunks(b"data: unfinished"), max_event_bytes=8)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ending", [b"\n\n", b"\n\r", b"\n\r\n", b"\r\r", b"\r\r\n", b"\r\n\n", b"\r\n\r", b"\r\n\r\n"]) +async def test_all_event_line_endings_at_every_chunk_split(ending): + blocks = [b"data: one" + ending, b"data: two" + ending] + raw = b"".join(blocks) + + for split in range(len(raw) + 1): + events = [event async for event in iter_sse_events(_chunks(raw[:split], b"", raw[split:]))] + assert [event.raw for event in events] == blocks + assert [event.data for event in events] == [b"one", b"two"] + assert ServerSentEvent.from_bytes(blocks[0]) == events[0] + + +@pytest.mark.asyncio +async def test_bom_is_ignored_only_at_stream_start_and_preserved_in_raw(): + blocks = [b"\xef\xbb\xbfdata: one\n\n", b"data: \xef\xbb\xbfvalue\n\n"] + raw = b"".join(blocks) + + for split in range(len(raw) + 1): + events = [event async for event in iter_sse_events(_chunks(raw[:split], raw[split:]))] + assert [event.raw for event in events] == blocks + assert [event.data for event in events] == [b"one", b"\xef\xbb\xbfvalue"] + with pytest.raises(AmbiguousServerSentEvent): + ServerSentEvent.from_bytes(blocks[0]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "raw", + [ + b"data: one\n\n\xef\xbb\xbfdata: two\n\n", + b"data: one\n\xef\xbb\xbf\ndata: two\n\n", + b"\xef\xbb\xbf\xef\xbb\xbfdata: one\n\n", + b"data: one\n\n\xef\xbb\xbf: comment\n\n", + b"data: one\r\r\xef\xbb\xbfdata: two\r\r", + b"data: one\r\n\r\n\xef\xbb\xbfdata: two\r\n\r\n", + b"data: one\n\n\xef\xbb\xbfdata: two", + ], +) +async def test_line_leading_bom_after_stream_start_is_rejected_at_every_split(raw): + """Clients disagree on such lines, so the parser never interprets them.""" + for split in range(len(raw) + 1): + with pytest.raises(AmbiguousServerSentEvent): + _ = [event async for event in iter_sse_events(_chunks(raw[:split], raw[split:]))] + + +@pytest.mark.asyncio +async def test_comments_empty_blocks_and_unknown_fields_preserve_exact_bytes(): + blocks = [b"\n", b": keepalive\r\n\r", b"unknown: value\ndata\n\n"] + raw = b"".join(blocks) + + events = [event async for event in iter_sse_events(_chunks(*(bytes([byte]) for byte in raw)))] + + assert [event.raw for event in events] == blocks + assert [event.data for event in events] == [None, None, b""] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tail", + [ + b"data: partial", + b"data: partial\n", + b"data: partial\r", + b"event: future", + b"id: 1", + b"retry: 1", + b"unknown", + b" ", + b": comment\ndata: partial", + ], +) +async def test_unfinished_field_blocks_are_rejected(tail): + with pytest.raises(TruncatedServerSentEvent): + _ = [event async for event in iter_sse_events(_chunks(tail))] + with pytest.raises(ValueError): + ServerSentEvent.from_bytes(tail) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tail", [b"", b": comment", b": comment\n", b": comment\r", b": one\n: two", b"\xef\xbb\xbf: comment"] +) +async def test_unfinished_comment_tails_are_discarded(tail): + events = [event async for event in iter_sse_events(_chunks(tail))] + + assert events == [] + + +@pytest.mark.asyncio +async def test_event_limit_counts_terminators_bom_and_unfinished_lines(): + block = b"\xef\xbb\xbfdata: one\r\n\r\n" + events = [event async for event in iter_sse_events(_chunks(block + block[3:] * 2), max_event_bytes=len(block))] + assert len(events) == 3 + for raw in (block, b"x" * len(block)): + with pytest.raises(ServerSentEventTooLarge): + _ = [ + event + async for event in iter_sse_events( + _chunks(*(bytes([byte]) for byte in raw)), max_event_bytes=len(block) - 1 + ) + ] + + +def test_from_bytes_rejects_multiple_blocks(): + with pytest.raises(ValueError, match="exactly one"): + ServerSentEvent.from_bytes(b"data: one\n\ndata: two\n\n") + + +@pytest.mark.asyncio +async def test_many_small_events_in_one_large_chunk(): + block = b"data: x\n\n" + count = 1024 * 1024 // len(block) + observed = 0 + + async for event in iter_sse_events(_chunks(block * count), max_event_bytes=len(block)): + assert event.raw == block + observed += 1 + + assert observed == count + + +@pytest.mark.asyncio +async def test_large_event_arriving_one_byte_per_chunk(): + raw = b"data: " + b"x" * (1024 * 1024 - 8) + b"\n\n" + + events = [event async for event in iter_sse_events(_byte_chunks(raw), max_event_bytes=len(raw))] + + assert len(events) == 1 + assert events[0].raw == raw + + +@pytest.mark.parametrize( + "raw", + [ + b"\xef\xbb\xbfdata: one\r\n\r\ndata: two\n\n", + b"data: a\ndata: b\r\rdata: c\r\n\r\n", + b": comment\n\nid: 1\nevent: delta\ndata: x\n\nretry: 5\n\n", + b"data:no-space\n\ndata: \xef\xbb\xbfvalue\n\nunknown: field\ndata\n\n", + b"Data: wrong-case\n\ndata: right\n\n", + ], +) +def test_python_sdk_never_sees_data_this_parser_did_not_parse(raw): + """Compare Python SDK parsing, separately from the Node BOM regression.""" + + async def parse(): + return [event.data for event in [event async for event in iter_sse_events(_chunks(raw))]] + + parsed = [data for data in asyncio.run(parse()) if data is not None] + client = [event.data.encode() for event in SSEDecoder().iter_bytes(iter([raw])) if event.data] + remaining = iter(parsed) + assert all(any(data == candidate for candidate in remaining) for data in client) + + +def _node_like_events(raw: bytes) -> list[bytes]: + """Model openai-node 7.31.0, which strips a BOM at the start of every line. + + The preserved review evidence shows Node reading a later BOM-prefixed data + line that a stream-start-only parser ignores. This is a model, not Node. + """ + events, data = [], [] + for line in raw.replace(b"\r\n", b"\n").replace(b"\r", b"\n").split(b"\n"): + line = line.removeprefix(b"\xef\xbb\xbf") + if not line: + if data: + events.append(b"\n".join(data)) + data = [] + continue + name, _, value = line.partition(b":") + if name == b"data": + data.append(value.removeprefix(b" ")) + return events + + +@pytest.mark.parametrize( + "raw", + [ + b"data: safe\n\n\xef\xbb\xbfdata: EVIL\n\n", + b"data: safe\r\r\xef\xbb\xbfdata: EVIL\r\r", + b"data: safe\n\xef\xbb\xbfdata: EVIL\n\n", + b": note\n\xef\xbb\xbfdata: EVIL\n\n", + b"\xef\xbb\xbfdata: one\n\ndata: two\n\n", + ], +) +@pytest.mark.parametrize("split", range(4)) +def test_node_like_client_never_sees_data_this_parser_did_not_parse(raw, split): + """Guard the Node BOM differential: accepted streams expose no extra data.""" + + async def parse(): + return [event.data async for event in iter_sse_events(_chunks(raw[:split], raw[split:]))] + + try: + parsed = [data for data in asyncio.run(parse()) if data is not None] + except AmbiguousServerSentEvent: + return + remaining = iter(parsed) + assert all(any(data == candidate for candidate in remaining) for data in _node_like_events(raw)) + + +@pytest.mark.parametrize("limit", [0, -1, True, 1.5]) +def test_event_limit_must_be_a_positive_integer(limit): + with pytest.raises(ValueError, match="max_event_bytes must be positive"): + asyncio.run(anext(iter_sse_events(_chunks(b"data: x\n\n"), max_event_bytes=limit))) diff --git a/tests/server/experimental/test_stream_classifier.py b/tests/server/experimental/test_stream_classifier.py new file mode 100644 index 0000000000..aa82dfa7c5 --- /dev/null +++ b/tests/server/experimental/test_stream_classifier.py @@ -0,0 +1,411 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +from dataclasses import replace +from typing import Literal + +import pytest +from pydantic import BaseModel, model_validator + +from nemoguardrails.server.experimental.provider.sse import ServerSentEvent +from nemoguardrails.server.experimental.provider.stream import ( + GuardedStreamEvent, + StreamCapabilityProfile, + StreamEventRole, + StreamProjectionContract, + StreamShapeCoverage, + UnsupportedProviderStream, + validate_stream_event, +) +from nemoguardrails.server.experimental.provider.stream_classifier import ( + StreamArraySelection, + StreamClassifierDefinition, + StreamEventRule, + build_stream_classifier, +) + + +class Delta(BaseModel): + text: str | None = None + + +class Chunk(BaseModel): + kind: Literal["chunk"] + deltas: list[Delta] + + +class Error(BaseModel): + error: str + + +class SnapshotItem(BaseModel): + kind: str + text: str + + +class Snapshot(BaseModel): + kind: Literal["snapshot"] + items: list[SnapshotItem] + + +CONTRACT = StreamProjectionContract( + projection_id="test.stream", + profile=StreamCapabilityProfile.SINGLE_TEXT_DELTA_V1, + shapes=StreamShapeCoverage( + guarded_shapes=frozenset({"chunk:text"}), + snapshot_shapes=frozenset({"snapshot"}), + opaque_shapes=frozenset({"done", "chunk:metadata"}), + provider_error_shapes=frozenset({"error"}), + ), +) + + +def classifier(): + return build_stream_classifier( + StreamClassifierDefinition( + subject="test.stream", + contract=CONTRACT, + sentinels=((b"[DONE]", "done"),), + rules=( + StreamEventRule( + shape="error", + model=Error, + role=StreamEventRole.PROVIDER_ERROR, + required_fields=frozenset({"error"}), + ), + StreamEventRule( + shape="chunk:text", + model=Chunk, + role=StreamEventRole.GUARDED_TEXT, + match=(("kind", "chunk"),), + text_path=("deltas", 0, "text"), + missing_text_role=StreamEventRole.OPAQUE_METADATA, + missing_text_shape="chunk:metadata", + ), + ), + ) + ) + + +@pytest.mark.parametrize( + ("raw", "shape", "role", "text"), + [ + (b"data: [DONE]\n\n", "done", StreamEventRole.OPAQUE_METADATA, None), + ( + b'data: {"kind":"chunk","deltas":[{"text":"hello"}]}\n\n', + "chunk:text", + StreamEventRole.GUARDED_TEXT, + "hello", + ), + ( + b'data: {"kind":"chunk","deltas":[]}\n\n', + "chunk:metadata", + StreamEventRole.OPAQUE_METADATA, + None, + ), + (b'data: {"error":"failed"}\n\n', "error", StreamEventRole.PROVIDER_ERROR, None), + ], +) +def test_classifier_validates_and_classifies_reviewed_shapes(raw, shape, role, text): + event = classifier().classify_event(ServerSentEvent.from_bytes(raw)) + + assert (event.shape, event.role, event.text) == (shape, role, text) + + +def test_classifier_rejects_unknown_and_ambiguous_events(): + with pytest.raises(UnsupportedProviderStream, match="not covered"): + classifier().classify_event(ServerSentEvent.from_bytes(b'data: {"kind":"future"}\n\n')) + + with pytest.raises(UnsupportedProviderStream, match="ambiguously matches"): + classifier().classify_event( + ServerSentEvent.from_bytes(b'data: {"kind":"chunk","error":"failed","deltas":[]}\n\n') + ) + + +def test_classifier_can_select_one_nested_text_item(): + snapshot_classifier = build_stream_classifier( + StreamClassifierDefinition( + subject="test.stream", + contract=CONTRACT, + rules=( + StreamEventRule( + shape="snapshot", + model=Snapshot, + role=StreamEventRole.TEXT_SNAPSHOT, + match=(("kind", "snapshot"),), + text_path=( + "items", + StreamArraySelection(match=(("kind", "message"),), cardinality="at_most_one"), + "text", + ), + missing_text_role=StreamEventRole.TEXT_SNAPSHOT, + missing_text_value="", + ), + ), + ) + ) + + empty = snapshot_classifier.classify_event(ServerSentEvent.from_bytes(b'data: {"kind":"snapshot","items":[]}\n\n')) + populated = snapshot_classifier.classify_event( + ServerSentEvent.from_bytes(b'data: {"kind":"snapshot","items":[{"kind":"message","text":"hello"}]}\n\n') + ) + + assert empty.text == "" + assert populated.text == "hello" + + +def test_stream_contract_rejects_unknown_or_misclassified_events(): + with pytest.raises(UnsupportedProviderStream, match="contract declares"): + validate_stream_event(CONTRACT, GuardedStreamEvent("chunk:text", StreamEventRole.OPAQUE_METADATA)) + + +def test_classifier_definition_must_agree_with_the_stream_contract(): + definition = StreamClassifierDefinition( + subject="test.stream", + contract=CONTRACT, + rules=( + StreamEventRule( + shape="chunk:text", + model=Chunk, + role=StreamEventRole.OPAQUE_METADATA, + match=(("kind", "chunk"),), + ), + ), + ) + + with pytest.raises(ValueError, match="contract declares 'guarded_text'"): + build_stream_classifier(definition) + + +def test_stream_shape_coverage_rejects_multiple_classifications(): + with pytest.raises(ValueError, match="multiple classifications"): + StreamShapeCoverage( + guarded_shapes=frozenset({"chunk"}), + opaque_shapes=frozenset({"chunk"}), + ) + + +@pytest.mark.parametrize("expected,other", [(False, 0), (True, 1), (0, False), (1, True), (1, 1.0)]) +def test_stream_discriminators_require_matching_types(expected, other): + correct = {"kind": expected} + incorrect = {"kind": other} + selector = StreamArraySelection(match=(("kind", expected),), cardinality="exactly_one") + rule = StreamEventRule( + shape="chunk:text", model=Chunk, role=StreamEventRole.GUARDED_TEXT, match=(("kind", expected),) + ) + + assert selector.select([incorrect, correct]) is correct + assert rule.matches(correct) is True + assert rule.matches(incorrect) is False + + +@pytest.mark.parametrize("mutation", ["text", "drop", "reorder"]) +def test_classifier_rejects_validation_that_changes_the_original_subject(mutation): + class TransformingChunk(Chunk): + @model_validator(mode="before") + @classmethod + def transform(cls, value): + if mutation == "text": + value["deltas"][0]["text"] = "safe" + elif mutation == "drop": + value["deltas"] = [] + else: + value["deltas"].reverse() + return value + + original = classifier() + modified = replace(original, rules=(replace(original.rules[1], model=TransformingChunk),)) + event = ServerSentEvent.from_bytes(b'data: {"kind":"chunk","deltas":[{"text":"unsafe"},{"text":"safe"}]}\n\n') + + with pytest.raises(UnsupportedProviderStream, match="original guarded text"): + modified.classify_event(event) + + +@pytest.mark.parametrize("value", [None, "", "hello"]) +def test_classifier_preserves_nullable_and_empty_original_text(value): + event = ServerSentEvent.from_bytes( + b"data: " + json.dumps({"kind": "chunk", "deltas": [{"text": value}]}).encode() + b"\n\n" + ) + assert classifier().classify_event(event).text == value + + +def test_classifier_distinguishes_boolean_integer_and_float_selections(): + rules = tuple( + StreamEventRule( + shape="chunk:metadata", model=BaseModel, role=StreamEventRole.OPAQUE_METADATA, match=(("kind", value),) + ) + for value in (True, 1, 1.0) + ) + result = build_stream_classifier(StreamClassifierDefinition(subject="test.stream", contract=CONTRACT, rules=rules)) + for value in (True, 1, 1.0): + assert sum(rule.matches({"kind": value}) for rule in result.rules) == 1 + + +def test_array_selection_distinguishes_missing_from_null(): + selector = StreamArraySelection(match=(("kind", None),), cardinality="exactly_one") + explicit_null = {"kind": None} + assert selector.select([{}, explicit_null]) is explicit_null + + +def chunk_rule(**changes): + rule = StreamEventRule( + shape="chunk:text", + model=Chunk, + role=StreamEventRole.GUARDED_TEXT, + match=(("kind", "chunk"),), + text_path=("deltas", 0, "text"), + ) + return replace(rule, **changes) + + +def definition(**changes): + return replace( + StreamClassifierDefinition(subject="test.stream", contract=CONTRACT, rules=(chunk_rule(),)), **changes + ) + + +def event(raw): + return ServerSentEvent.from_bytes(raw) + + +@pytest.mark.parametrize( + ("changes", "message"), + [ + ({"subject": ""}, "identify its provider event subject"), + ({"rules": ()}, "at least one event rule"), + ({"subject": "other.stream"}, "match its projection contract identifier"), + ({"require_event_type": True}, "needs a data discriminator"), + ({"non_data_shape": "chunk:text"}, "contract declares"), + ({"rules": (chunk_rule(shape=""),)}, "identify its provider shape"), + ({"rules": (chunk_rule(match=()),)}, "must declare how it is selected"), + ({"rules": (chunk_rule(), chunk_rule())}, "duplicates another selection"), + ({"rules": (chunk_rule(text_path=None),)}, "inconsistent text role and path"), + ( + { + "rules": ( + chunk_rule( + shape="chunk:metadata", + role=StreamEventRole.OPAQUE_METADATA, + text_path=None, + missing_text_role=StreamEventRole.OPAQUE_METADATA, + ), + ) + }, + "without a text path", + ), + ({"rules": (chunk_rule(missing_text_shape="chunk:metadata"),)}, "rename missing text without classifying"), + ({"rules": (chunk_rule(missing_text_value=""),)}, "supply missing text without classifying"), + ({"rules": (chunk_rule(missing_text_role=StreamEventRole.GUARDED_TEXT),)}, "classify missing text as content"), + ], +) +def test_classifier_definition_rejects_inconsistent_declarations(changes, message): + with pytest.raises(ValueError, match=message): + build_stream_classifier(definition(**changes)) + + +def test_data_less_events_need_a_declared_shape(): + data_less = event(b"event: ping\n\n") + with pytest.raises(UnsupportedProviderStream, match="must contain data"): + build_stream_classifier(definition()).classify_event(data_less) + + classified = build_stream_classifier(definition(non_data_shape="done")).classify_event(data_less) + assert classified == GuardedStreamEvent("done", StreamEventRole.OPAQUE_METADATA) + + +def test_event_type_field_must_agree_with_the_sse_event_name(): + typed = build_stream_classifier(definition(event_type_field="kind", require_event_type=True)) + chunk = b'data: {"kind":"chunk","deltas":[{"text":"hi"}]}\n\n' + + assert typed.classify_event(event(b"event: chunk\n" + chunk)).text == "hi" + with pytest.raises(UnsupportedProviderStream, match="SSE event field is required"): + typed.classify_event(event(chunk)) + with pytest.raises(UnsupportedProviderStream, match="must match data field 'kind'"): + typed.classify_event(event(b"event: other\n" + chunk)) + with pytest.raises(UnsupportedProviderStream, match="must declare string field 'kind'"): + typed.classify_event(event(b'event: chunk\ndata: {"kind":1,"deltas":[]}\n\n')) + + +@pytest.mark.parametrize( + ("raw", "message"), + [ + (b'data: {"kind":"chunk","deltas":[{"text":null}]}\n\n', "has no guarded text value"), + (b'data: {"kind":"chunk","deltas":[{"text":"hi"},5]}\n\n', "test.stream"), + (b"data: {not json}\n\n", "JSON"), + ], + ids=["missing-text", "schema", "json"], +) +def test_classifier_rejects_unusable_event_payloads(raw, message): + with pytest.raises(UnsupportedProviderStream, match=message): + build_stream_classifier(definition()).classify_event(event(raw)) + + +@pytest.mark.parametrize( + ("text_path", "message"), + [ + (("deltas", -1, "text"), "negative indexes"), + (("deltas", "text"), "does not match its projection"), + (("deltas",), "must resolve to a string"), + ], +) +def test_classifier_rejects_text_paths_that_do_not_reach_a_string(text_path, message): + classifier = build_stream_classifier(definition(rules=(chunk_rule(text_path=text_path),))) + with pytest.raises(UnsupportedProviderStream, match=message): + classifier.classify_event(event(b'data: {"kind":"chunk","deltas":[{"text":"hi"}]}\n\n')) + + +def test_exactly_one_array_selection_rejects_missing_or_repeated_items(): + selection = StreamArraySelection(match=(("kind", "message"),), cardinality="exactly_one") + with pytest.raises(ValueError, match="exactly one matching item"): + selection.select([{"kind": "other"}]) + with pytest.raises(ValueError, match="exactly one matching item"): + selection.select([{"kind": "message"}, {"kind": "message"}]) + + +def test_dotted_match_keys_do_not_match_through_scalars(): + rule = chunk_rule(match=(("kind.name", "chunk"),)) + assert rule.matches({"kind": {"name": "chunk"}}) is True + assert rule.matches({"kind": "chunk"}) is False + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"projection_id": " "}, "must have an identifier"), + ({"shapes": StreamShapeCoverage()}, "at least one shape"), + ], +) +def test_stream_projection_contract_requires_identity_and_shapes(kwargs, message): + values = {"projection_id": "p", "profile": StreamCapabilityProfile.SINGLE_TEXT_DELTA_V1, "shapes": CONTRACT.shapes} + with pytest.raises(ValueError, match=message): + StreamProjectionContract(**{**values, **kwargs}) + + +@pytest.mark.parametrize( + ("args", "message"), + [ + (("", StreamEventRole.OPAQUE_METADATA), "identify its provider shape"), + (("chunk:text", StreamEventRole.GUARDED_TEXT), "must carry text"), + (("done", StreamEventRole.OPAQUE_METADATA, "text"), "must carry text"), + ], +) +def test_guarded_stream_event_requires_shape_and_role_consistent_text(args, message): + with pytest.raises(ValueError, match=message): + GuardedStreamEvent(*args) + + +def test_stream_contract_rejects_undeclared_shapes(): + with pytest.raises(UnsupportedProviderStream, match="is not declared"): + validate_stream_event(CONTRACT, GuardedStreamEvent("future", StreamEventRole.OPAQUE_METADATA))