diff --git a/.chronus/changes/structured-streaming-2026-0-0.md b/.chronus/changes/structured-streaming-2026-0-0.md new file mode 100644 index 00000000000..0c3ec139d78 --- /dev/null +++ b/.chronus/changes/structured-streaming-2026-0-0.md @@ -0,0 +1,17 @@ +--- +changeKind: feature +packages: + - "@typespec/http-client-python" +--- + +Generate structured streaming client methods: operations whose HTTP response is a JSONL (`application/jsonl`) or SSE (`text/event-stream`) stream now return `Stream[T]` / `AsyncStream[T]`, yielding deserialized model instances instead of raw bytes. + +The `Stream` / `AsyncStream` runtime (plus the JSONL / SSE decoders) is vendored at `_utils/streaming_base.py` and depends only on the released core runtime for the flavor — `azure.core.rest` for the Azure flavor and `corehttp.rest` for the unbranded flavor. These types are an internal implementation detail and are not part of the package's public API. + +```python +stream = client.receive() +for thing in stream: + ... +``` + +For SSE streams, the most recently received event `id` and `retry` value (if provided by the server) are exposed via `stream.last_event_id` / `stream.retry`. Event envelopes yield the `@Events.data` payload using its payload type and content type. JSON payload media types (`application/json` and `+json`) are decoded as JSON, while other SSE payload media types remain UTF-8 text for type-specific deserialization. Pass `last_event_id=` to an SSE operation to send `Last-Event-ID` when manually resuming a stream. diff --git a/cspell.yaml b/cspell.yaml index af8ebc2a006..9c26d9f1de2 100644 --- a/cspell.yaml +++ b/cspell.yaml @@ -5,13 +5,19 @@ dictionaries: - node - typescript words: + - aclose + - aclosing - Ablack - Adoptium + - aenter + - aexit - agentic - agentics - aiohttp + - aiter - alzimmer - amqp + - anext - AQID - Arize - arizeaiobservabilityeval @@ -121,6 +127,7 @@ words: - intrinsics - ints - IOHTTP + - isascii - isdigit - isinstance - issecret diff --git a/packages/http-client-python/emitter/src/http.ts b/packages/http-client-python/emitter/src/http.ts index 31d970d4b65..e1fe7efb351 100644 --- a/packages/http-client-python/emitter/src/http.ts +++ b/packages/http-client-python/emitter/src/http.ts @@ -18,6 +18,7 @@ import type { SdkQueryParameter, SdkServiceMethod, SdkServiceResponseHeader, + SdkSseEventMetadata, SdkType, } from "@azure-tools/typespec-client-generator-core"; import { getHttpOperationParameter, UsageFlags } from "@azure-tools/typespec-client-generator-core"; @@ -42,6 +43,158 @@ export enum ReferredByOperationTypes { NonPagingOnly = 2, } +type StructuredStreamKind = "jsonl" | "sse"; +type EmittedType = ReturnType; + +interface StructuredStreamEvent { + eventType: string | undefined; + /** + * Payload type for this one SSE event. For an event envelope, this is the type of the property + * marked `@Events.data`. Together with {@link eventType} these form the + * runtime dispatch table (wire event name -> model to deserialize) inside the generated + * `_callback`. This is a narrower type than {@link StructuredStreamingInfo.itemType}. + */ + payloadType: EmittedType; + /** + * True when this event is a `@terminalEvent` that carries a payload (a named / model event, + * not a bare string-constant sentinel). Such events are deserialized and yielded like any + * other event, and iteration stops immediately after one is yielded. Contrast with + * {@link StructuredStreamingInfo.terminalEvent}, the sentinel that stops without yielding. + */ + isTerminal?: boolean; + /** Content type of the payload, not the enclosing event. */ + payloadContentType?: string; +} + +interface StructuredStreamingInfo { + kind: StructuredStreamKind; + /** + * The aggregate stream element type used for the `Stream[T]` / `AsyncStream[T]` return + * annotation (a single type expression). For homogeneous JSONL this is the one model; for + * heterogeneous SSE this is the union of every event payload. + * + * Note the deliberate overlap with the per-event {@link StructuredStreamEvent.payloadType}: for + * heterogeneous SSE this union is exactly the sum of the `events[]` payload types. Both are + * carried because the union alone cannot recover the wire-name -> member mapping needed for + * dispatch, and the events list alone is not a single valid type expression for the annotation. + */ + itemType: EmittedType; + events?: StructuredStreamEvent[]; + /** + * A bare string-constant `@terminalEvent` with no event name (e.g. `"[DONE]"`). Iteration + * stops when an event's `data` equals this value, and the sentinel is NOT yielded. Named / + * model terminal events are carried in {@link events} with `isTerminal: true` instead. + */ + terminalEvent?: string; +} + +/** Whether pygen can deserialize the stream item type. */ +export function isStructuredStreamType(type: SdkType): boolean { + switch (type.kind) { + case "model": + case "union": + return true; + case "nullable": + return isStructuredStreamType(type.type); + default: + return false; + } +} + +export function getStructuredStreamKind( + response: SdkHttpResponse | SdkHttpErrorResponse, +): StructuredStreamKind | undefined { + if (response.sseMetadata) return "sse"; + + const contentTypes = response.streamMetadata?.contentTypes ?? response.contentTypes ?? []; + for (const contentType of contentTypes) { + const mediaType = contentType.split(";", 1)[0].trim().toLowerCase(); + if (mediaType === "text/event-stream") return "sse"; + if (mediaType === "application/jsonl") return "jsonl"; + } + return undefined; +} + +function getStringConstantValue(type: SdkType): string | undefined { + if (type.kind === "nullable") return getStringConstantValue(type.type); + return type.kind === "constant" && typeof type.value === "string" ? type.value : undefined; +} + +/** + * Split the SSE events into the runtime dispatch table and a bare string-constant sentinel. + * + * A `@terminalEvent` comes in two shapes: + * * a nameless string constant (e.g. `"[DONE]"`) -> a pure sentinel: iteration stops when an + * event's `data` equals this value and the event is NOT yielded. Returned as `terminalEvent`. + * * a named / model event (e.g. `error`, `response.completed`) -> carries a payload the consumer + * needs, so it is deserialized and yielded like any other event, then iteration stops. Returned + * in `events` with `isTerminal: true`. + * + * `toPayloadType` maps event payloads to emitted types; it is injected so this partitioning stays + * a pure function that can be unit-tested without a full emitter context. + */ +export function partitionSSEEvents( + events: readonly SdkSseEventMetadata[], + toPayloadType: (type: SdkType) => EmittedType, +): { events: StructuredStreamEvent[]; terminalEvent?: string } { + const dispatch: StructuredStreamEvent[] = []; + let terminalEvent: string | undefined; + for (const event of events) { + if (event.isTerminalEvent) { + const sentinelValue = + event.eventType === undefined + ? (getStringConstantValue(event.payloadType) ?? getStringConstantValue(event.type)) + : undefined; + if (sentinelValue !== undefined) { + // Keep the first sentinel; no current spec defines more than one. + terminalEvent ??= sentinelValue; + continue; + } + dispatch.push({ + eventType: event.eventType, + payloadType: toPayloadType(event.payloadType), + isTerminal: true, + payloadContentType: event.payloadContentType, + }); + continue; + } + dispatch.push({ + eventType: event.eventType, + payloadType: toPayloadType(event.payloadType), + payloadContentType: event.payloadContentType, + }); + } + return terminalEvent !== undefined ? { events: dispatch, terminalEvent } : { events: dispatch }; +} + +function emitStructuredStreamingInfo( + context: PythonSdkContext, + response: SdkHttpResponse | SdkHttpErrorResponse, +): StructuredStreamingInfo | undefined { + const streamMetadata = response.streamMetadata; + if (!streamMetadata || !isStructuredStreamType(streamMetadata.streamType)) return undefined; + + const kind = getStructuredStreamKind(response); + if (!kind) return undefined; + + const streaming: StructuredStreamingInfo = { + kind, + itemType: getType(context, streamMetadata.streamType), + }; + if (kind !== "sse") return streaming; + + const sseMetadata = response.sseMetadata; + if (!sseMetadata || sseMetadata.events.length === 0) return undefined; + + const { events, terminalEvent } = partitionSSEEvents(sseMetadata.events, (type) => + getType(context, type), + ); + if (events.length > 0) streaming.events = events; + if (terminalEvent !== undefined) streaming.terminalEvent = terminalEvent; + + return streaming; +} + function isEtagType(type: SdkType): boolean { if (type.kind === "nullable") return isEtagType(type.type); const raw = type.__raw; @@ -682,6 +835,7 @@ function emitHttpResponse( "invalid-lro-result", method, ), + streaming: isException ? undefined : emitStructuredStreamingInfo(context, response), }; } diff --git a/packages/http-client-python/emitter/test/streaming.test.ts b/packages/http-client-python/emitter/test/streaming.test.ts new file mode 100644 index 00000000000..b6cae795d2f --- /dev/null +++ b/packages/http-client-python/emitter/test/streaming.test.ts @@ -0,0 +1,155 @@ +import { strictEqual } from "assert"; +import { describe, it } from "vitest"; +import { + getStructuredStreamKind, + isStructuredStreamType, + partitionSSEEvents, +} from "../src/http.js"; + +describe("typespec-python: structured streaming", () => { + it("treats model and union payloads as structured", () => { + strictEqual(isStructuredStreamType({ kind: "model" } as any), true); + strictEqual(isStructuredStreamType({ kind: "union" } as any), true); + }); + + it("unwraps nullable payloads", () => { + strictEqual(isStructuredStreamType({ kind: "nullable", type: { kind: "model" } } as any), true); + strictEqual( + isStructuredStreamType({ kind: "nullable", type: { kind: "bytes" } } as any), + false, + ); + }); + + it("treats bare byte/string payloads as unstructured", () => { + strictEqual(isStructuredStreamType({ kind: "bytes" } as any), false); + strictEqual(isStructuredStreamType({ kind: "string" } as any), false); + }); + + it("detects the stream protocol explicitly", () => { + strictEqual(getStructuredStreamKind({ sseMetadata: { events: [] } } as any), "sse"); + strictEqual( + getStructuredStreamKind({ + streamMetadata: { contentTypes: ["text/event-stream; charset=utf-8"] }, + } as any), + "sse", + ); + strictEqual( + getStructuredStreamKind({ + streamMetadata: { contentTypes: ["application/jsonl"] }, + } as any), + "jsonl", + ); + strictEqual( + getStructuredStreamKind({ + streamMetadata: { contentTypes: ["application/json"] }, + } as any), + undefined, + ); + }); + + describe("terminal-event partitioning", () => { + const identity = (payloadType: any) => payloadType; + const model = (name: string) => ({ kind: "model", name }); + const constant = (value: string) => ({ kind: "constant", value }); + + it("keeps a nameless string-constant `[DONE]` as a drop-and-stop sentinel", () => { + const created = model("ResponseCreated"); + const done = constant("[DONE]"); + const { events, terminalEvent } = partitionSSEEvents( + [ + { + eventType: "response.created", + isTerminalEvent: false, + type: created, + payloadType: created, + }, + { eventType: undefined, isTerminalEvent: true, type: done, payloadType: done }, + ] as any, + identity, + ); + // The sentinel is NOT a dispatch event; it only sets `terminalEvent`. + strictEqual(terminalEvent, "[DONE]"); + strictEqual(events.length, 1); + strictEqual(events[0].eventType, "response.created"); + strictEqual(events[0].isTerminal, undefined); + }); + + it("keeps named / model `@terminalEvent`s in the dispatch table as yield-and-stop events", () => { + const created = model("ResponseCreated"); + const completed = model("ResponseCompleted"); + const errored = model("StreamError"); + const { events, terminalEvent } = partitionSSEEvents( + [ + { + eventType: "response.created", + isTerminalEvent: false, + type: created, + payloadType: created, + }, + { + eventType: "response.completed", + isTerminalEvent: true, + type: completed, + payloadType: completed, + }, + { eventType: "error", isTerminalEvent: true, type: errored, payloadType: errored }, + ] as any, + identity, + ); + // No bare sentinel: the two terminals carry payloads, so they stay in `events`. + strictEqual(terminalEvent, undefined); + strictEqual(events.length, 3); + strictEqual(events[0].isTerminal, undefined); + strictEqual(events[1].eventType, "response.completed"); + strictEqual(events[1].isTerminal, true); + strictEqual(events[1].payloadType, completed); + strictEqual(events[2].eventType, "error"); + strictEqual(events[2].isTerminal, true); + strictEqual(events[2].payloadType, errored); + }); + + it("supports a sentinel and named terminals together", () => { + const delta = model("ResponseDelta"); + const completed = model("ResponseCompleted"); + const done = constant("[DONE]"); + const { events, terminalEvent } = partitionSSEEvents( + [ + { eventType: "response.delta", isTerminalEvent: false, type: delta, payloadType: delta }, + { + eventType: "response.completed", + isTerminalEvent: true, + type: completed, + payloadType: completed, + }, + { eventType: undefined, isTerminalEvent: true, type: done, payloadType: done }, + ] as any, + identity, + ); + strictEqual(terminalEvent, "[DONE]"); + strictEqual(events.length, 2); + strictEqual(events[0].isTerminal, undefined); + strictEqual(events[1].eventType, "response.completed"); + strictEqual(events[1].isTerminal, true); + }); + + it("emits the payload metadata for event envelopes", () => { + const envelope = model("Envelope"); + const payload = { kind: "string" }; + const { events } = partitionSSEEvents( + [ + { + eventType: "withEnvelope", + isTerminalEvent: false, + type: envelope, + payloadType: payload, + isEventEnvelope: true, + payloadContentType: "text/plain", + }, + ] as any, + identity, + ); + strictEqual(events[0].payloadType, payload); + strictEqual(events[0].payloadContentType, "text/plain"); + }); + }); +}); diff --git a/packages/http-client-python/generator/pygen/codegen/models/code_model.py b/packages/http-client-python/generator/pygen/codegen/models/code_model.py index 73cd410eb84..5618f156055 100644 --- a/packages/http-client-python/generator/pygen/codegen/models/code_model.py +++ b/packages/http-client-python/generator/pygen/codegen/models/code_model.py @@ -279,12 +279,22 @@ def need_utils_folder(self, async_mode: bool, client_namespace: str) -> bool: self.need_utils_utils(async_mode, client_namespace) or self.need_utils_serialization or self.options["models-mode"] == "dpg" + or self.has_structured_stream ) @property def need_utils_serialization(self) -> bool: return not self.options["client-side-validation"] + @property + def has_structured_stream(self) -> bool: + return any( + op.has_structured_stream_response + for client in self.clients + for og in client.operation_groups + for op in og.operations + ) + def need_utils_utils(self, async_mode: bool, client_namespace: str) -> bool: return ( self.need_utils_form_data(async_mode, client_namespace) @@ -395,22 +405,26 @@ def core_library(self) -> Literal["azure.core", "corehttp"]: def _sort_model_types_helper( self, current: ModelType, - seen_schema_names: set[str], + seen_schema_keys: set[tuple[str, str]], seen_schema_yaml_ids: set[int], ): if current.id in seen_schema_yaml_ids: return [] - if current.name in seen_schema_names: - raise ValueError(f"We have already generated a schema with name {current.name}") + current_schema_key = (current.client_namespace, current.name) + if current_schema_key in seen_schema_keys: + raise ValueError( + f"We have already generated a schema with name {current.name} " + f"in namespace {current.client_namespace}" + ) ancestors = [current] if current.parents: for parent in current.parents: if parent.id in seen_schema_yaml_ids: continue - seen_schema_names.add(current.name) + seen_schema_keys.add(current_schema_key) seen_schema_yaml_ids.add(current.id) - ancestors = self._sort_model_types_helper(parent, seen_schema_names, seen_schema_yaml_ids) + ancestors - seen_schema_names.add(current.name) + ancestors = self._sort_model_types_helper(parent, seen_schema_keys, seen_schema_yaml_ids) + ancestors + seen_schema_keys.add(current_schema_key) seen_schema_yaml_ids.add(current.id) return ancestors @@ -420,11 +434,11 @@ def sort_model_types(self) -> None: :return: None :rtype: None """ - seen_schema_names: set[str] = set() + seen_schema_keys: set[tuple[str, str]] = set() seen_schema_yaml_ids: set[int] = set() sorted_object_schemas: list[ModelType] = [] for schema in sorted(self.model_types, key=lambda x: x.name.lower()): - sorted_object_schemas.extend(self._sort_model_types_helper(schema, seen_schema_names, seen_schema_yaml_ids)) + sorted_object_schemas.extend(self._sort_model_types_helper(schema, seen_schema_keys, seen_schema_yaml_ids)) self.model_types = sorted_object_schemas @property diff --git a/packages/http-client-python/generator/pygen/codegen/models/operation.py b/packages/http-client-python/generator/pygen/codegen/models/operation.py index 3c6d525f122..d1cac714727 100644 --- a/packages/http-client-python/generator/pygen/codegen/models/operation.py +++ b/packages/http-client-python/generator/pygen/codegen/models/operation.py @@ -104,6 +104,11 @@ def stream_value(self) -> Union[str, bool]: else self.has_stream_response ) + @property + def has_structured_stream_response(self) -> bool: + """Whether any success response is a structured (JSONL / SSE) stream returning Stream[T].""" + return any(getattr(r, "is_structured_stream", False) for r in self.responses) + @property def has_form_data_body(self): return self.parameters.has_form_data_body @@ -506,7 +511,7 @@ def filename(self) -> str: @property def has_stream_response(self) -> bool: - return any(r.is_stream_response for r in self.responses) + return any(r.is_stream_response or r.is_structured_stream for r in self.responses) @classmethod def get_request_builder(cls, yaml_data: dict[str, Any], client: "Client"): diff --git a/packages/http-client-python/generator/pygen/codegen/models/response.py b/packages/http-client-python/generator/pygen/codegen/models/response.py index 99a90481319..7f08d0251b5 100644 --- a/packages/http-client-python/generator/pygen/codegen/models/response.py +++ b/packages/http-client-python/generator/pygen/codegen/models/response.py @@ -3,6 +3,7 @@ # Licensed under the MIT License. See License.txt in the project root for # license information. # -------------------------------------------------------------------------- +from dataclasses import dataclass from typing import Optional, Any, TYPE_CHECKING, Union from .base import BaseModel @@ -17,6 +18,53 @@ if TYPE_CHECKING: from .code_model import CodeModel +DEFAULT_SSE_EVENT_TYPE = "message" + + +@dataclass(frozen=True) +class StreamingEvent: + event_type: Optional[str] + payload_type: BaseType + payload_content_type: Optional[str] + is_terminal: bool + + +def get_streaming_event_discriminator( + events: list[StreamingEvent], +) -> Optional[tuple[str, list[tuple[str, StreamingEvent]]]]: + discriminator_name: Optional[str] = None + variants: list[tuple[str, StreamingEvent]] = [] + values: set[str] = set() + for event in events: + discriminator_property = getattr(event.payload_type, "discriminator_property", None) + discriminator_value = getattr(event.payload_type, "discriminator_value", None) + if discriminator_property is None or discriminator_value is None: + return None + if discriminator_name is None: + discriminator_name = discriminator_property.wire_name + elif discriminator_property.wire_name != discriminator_name: + return None + if discriminator_value in values: + return None + values.add(discriminator_value) + variants.append((discriminator_value, event)) + if discriminator_name is None: + return None + return discriminator_name, variants + + +def _get_terminal_event_names(events: list[StreamingEvent]) -> list[str]: + unnamed_events = [event for event in events if event.event_type is None] + terminal_event_names: list[str] = [] + for event in events: + if not event.is_terminal: + continue + if event.event_type is not None: + terminal_event_names.append(event.event_type) + elif len(unnamed_events) == 1: + terminal_event_names.append(DEFAULT_SSE_EVENT_TYPE) + return list(dict.fromkeys(terminal_event_names)) + class ResponseHeader(BaseModel): def __init__( @@ -58,6 +106,25 @@ def __init__( self.type = type self.nullable = yaml_data.get("nullable") self.default_content_type = yaml_data.get("defaultContentType") + streaming = yaml_data.get("streaming") + self.streaming_kind: Optional[str] = streaming["kind"] if streaming else None + self.streaming_events: list[StreamingEvent] = [] + self.terminal_event: Optional[str] = streaming.get("terminalEvent") if streaming else None + # Named / model ``@terminalEvent`` events: deserialized and yielded like any other event, + # then iteration stops. The bare string-constant sentinel (``terminal_event``) + # stops WITHOUT yielding and is matched on event ``data`` instead of the event name. + self.terminal_event_names: list[str] = [] + if streaming: + self.streaming_events = [ + StreamingEvent( + event_type=event.get("eventType"), + payload_type=self.code_model.lookup_type(id(event["payloadType"])), + payload_content_type=event.get("payloadContentType"), + is_terminal=event.get("isTerminal", False), + ) + for event in streaming.get("events", []) + ] + self.terminal_event_names = _get_terminal_event_names(self.streaming_events) @property def result_property(self) -> str: @@ -92,12 +159,48 @@ def is_stream_response(self) -> bool: ) return retval + @property + def is_structured_stream(self) -> bool: + """Is the response a structured (JSONL / SSE) stream rendered as Stream[T] / AsyncStream[T].""" + return self.streaming_kind is not None + + def stream_class_name(self, async_mode: bool) -> str: + return "AsyncStream" if async_mode else "Stream" + + @property + def stream_item_type(self) -> Optional[BaseType]: + event_item_types = list(dict.fromkeys(event.payload_type for event in self.streaming_events)) + if len(event_item_types) == 1: + return event_item_types[0] + if event_item_types: + return CombinedType({"type": "combined"}, self.code_model, event_item_types) + return self.type + def serialization_type(self, **kwargs: Any) -> str: if self.type: return self.type.serialization_type(**kwargs) return "None" + def stream_item_annotation(self, **kwargs: Any) -> str: + """Valid type expression for a structured stream's item type. + + A named ``CombinedType`` (``@events`` union) renders its ``type_annotation`` as the + ``_unions.`` alias, which is a module-level variable and therefore rejected by + pyright/mypy inside ``Stream[...]`` ("Variable not allowed in type expression"). Expand + the union inline (``Union[Model, ...]`` / the single member) so the annotation is a + valid type expression. + """ + item_type = self.stream_item_type + if isinstance(item_type, CombinedType): + return item_type.type_definition(**kwargs) + return item_type.type_annotation(**kwargs) if item_type else "None" + def type_annotation(self, **kwargs: Any) -> str: + if self.is_structured_stream and self.type: + kwargs["is_operation_file"] = True + kwargs["is_response"] = True + stream_class = self.stream_class_name(kwargs.get("async_mode", False)) + return f"{stream_class}[{self.stream_item_annotation(**kwargs)}]" if self.type: kwargs["is_operation_file"] = True kwargs["is_response"] = True @@ -109,30 +212,63 @@ def type_annotation(self, **kwargs: Any) -> str: def docstring_text(self, **kwargs: Any) -> str: kwargs["is_response"] = True + if self.is_structured_stream and self.type: + stream_class = self.stream_class_name(kwargs.get("async_mode", False)) + item_type = self.stream_item_type or self.type + return f"An instance of {stream_class} that iterates over {item_type.docstring_text(**kwargs)}" if self.nullable and self.type: return f"{self.type.docstring_text(**kwargs)} or None" return self.type.docstring_text(**kwargs) if self.type else "None" def docstring_type(self, **kwargs: Any) -> str: kwargs["is_response"] = True + if self.is_structured_stream and self.type: + stream_class = self.stream_class_name(kwargs.get("async_mode", False)) + item_type = (self.stream_item_type or self.type).docstring_type(**kwargs) + return f"~{self.code_model.namespace}.{stream_class}[{item_type}]" if self.nullable and self.type: return f"{self.type.docstring_type(**kwargs)} or None" return self.type.docstring_type(**kwargs) if self.type else "None" def imports(self, **kwargs: Any) -> FileImport: file_import = FileImport(self.code_model) - if self.type: - file_import.merge(self.type.imports(**kwargs)) + item_type = self.stream_item_type if self.is_structured_stream else self.type + # For a structured stream whose item type is a named ``@events`` union, the annotation + # is expanded inline (see ``stream_item_annotation``), so import the union member types + # rather than the ``_unions`` alias. + if self.is_structured_stream and isinstance(item_type, CombinedType): + for member in item_type.types: + file_import.merge(member.imports(**kwargs)) + # ``Union`` is only needed when the inline expansion actually yields a union of + # 2+ distinct member types (a single member collapses to that member; a union of + # only literals collapses to a single ``Literal[...]``). + distinct = list(dict.fromkeys(m.type_annotation(**kwargs) for m in item_type.types)) + all_constant = all(t.type == "constant" for t in item_type.types) + if len(distinct) > 1 and not all_constant: + file_import.add_submodule_import("typing", "Union", ImportType.STDLIB) + elif item_type: + file_import.merge(item_type.imports(**kwargs)) + if not self.is_structured_stream and isinstance(item_type, CombinedType) and item_type.name: + serialize_namespace = kwargs.get("serialize_namespace", self.code_model.namespace) + file_import.add_submodule_import( + self.code_model.get_relative_import_path(serialize_namespace), + "_unions", + ImportType.LOCAL, + TypingSection.TYPING, + ) if self.nullable: file_import.add_submodule_import("typing", "Optional", ImportType.STDLIB) - if isinstance(self.type, CombinedType) and self.type.name: + if self.is_structured_stream: + stream_class = self.stream_class_name(kwargs.get("async_mode", False)) serialize_namespace = kwargs.get("serialize_namespace", self.code_model.namespace) - file_import.add_submodule_import( - self.code_model.get_relative_import_path(serialize_namespace), - "_unions", - ImportType.LOCAL, - TypingSection.TYPING, + relative_path = self.code_model.get_relative_import_path( + serialize_namespace, module_name="_utils.streaming_base" ) + file_import.add_submodule_import(relative_path, stream_class, ImportType.LOCAL) + if self.streaming_kind == "sse": + file_import.add_import("json", ImportType.STDLIB) + for event in self.streaming_events: + file_import.merge(event.payload_type.imports(**kwargs)) return file_import def _get_import_type(self, input_path: str) -> ImportType: @@ -143,6 +279,14 @@ def _get_import_type(self, input_path: str) -> ImportType: @classmethod def from_yaml(cls, yaml_data: dict[str, Any], code_model: "CodeModel") -> "Response": + streaming = yaml_data.get("streaming") + if streaming: + return cls( + yaml_data=yaml_data, + code_model=code_model, + headers=[ResponseHeader.from_yaml(header, code_model) for header in yaml_data["headers"]], + type=code_model.lookup_type(id(streaming["itemType"])), + ) type = code_model.lookup_type(id(yaml_data["type"])) if yaml_data.get("type") else None # use ByteIteratorType if we are returning a binary type default_content_type = yaml_data.get("defaultContentType", "application/json") diff --git a/packages/http-client-python/generator/pygen/codegen/serializers/__init__.py b/packages/http-client-python/generator/pygen/codegen/serializers/__init__.py index c5d786a4a6e..e73e3866970 100644 --- a/packages/http-client-python/generator/pygen/codegen/serializers/__init__.py +++ b/packages/http-client-python/generator/pygen/codegen/serializers/__init__.py @@ -525,6 +525,13 @@ def _serialize_and_write_utils_folder(self, env: Environment, namespace: str): general_serializer.serialize_model_base_file(), ) + # write _utils/streaming_base.py (vendored Stream/AsyncStream + JSONL/SSE decoders) + if self.code_model.has_structured_stream: + self.write_file( + utils_folder_path / Path("streaming_base.py"), + general_serializer.serialize_streaming_base_file(), + ) + def _serialize_and_write_top_level_folder(self, env: Environment, namespace: str) -> None: root_dir = self.code_model.get_root_dir() generation_dir = self.code_model.get_generation_dir(namespace) diff --git a/packages/http-client-python/generator/pygen/codegen/serializers/builder_serializer.py b/packages/http-client-python/generator/pygen/codegen/serializers/builder_serializer.py index fe7649a3b89..9cdb21679f9 100644 --- a/packages/http-client-python/generator/pygen/codegen/serializers/builder_serializer.py +++ b/packages/http-client-python/generator/pygen/codegen/serializers/builder_serializer.py @@ -35,12 +35,33 @@ from ..models.utils import NamespaceType, escape_sphinx_field_name from .parameter_serializer import ParameterSerializer, PopKwargType, check_body_optional from ..models.parameter_list import ParameterType +from ..models.response import ( + DEFAULT_SSE_EVENT_TYPE, + StreamingEvent, + get_streaming_event_discriminator, +) from . import utils from ...utils import xml_serializable, json_serializable T = TypeVar("T") OrderedSet = dict[T, None] + +def _sse_event_data_expression(payload_content_type: Optional[str]) -> str: + """Return the generated expression that decodes an SSE event's UTF-8 data field.""" + media_type = payload_content_type.split(";", 1)[0].strip().lower() if payload_content_type else None + is_json = media_type is None or media_type == "application/json" or media_type.endswith("+json") + return "json.loads(_event.data)" if is_json else "_event.data" + + +def _sse_fallback_data_expression(events: list[StreamingEvent]) -> str: + """Decode fallback events only when every candidate payload is JSON.""" + if not events: + return "_event.data" + expressions = {_sse_event_data_expression(event.payload_content_type) for event in events} + return "json.loads(_event.data)" if expressions == {"json.loads(_event.data)"} else "_event.data" + + BuilderType = TypeVar( "BuilderType", bound=Union[ @@ -1021,7 +1042,19 @@ def _call_request_builder_helper( return retval def call_request_builder(self, builder: OperationType, is_paging: bool = False) -> list[str]: - return self._call_request_builder_helper(builder, builder.request_builder, is_paging=is_paging) + retval = self._call_request_builder_helper(builder, builder.request_builder, is_paging=is_paging) + if builder.has_structured_stream_response and any( + response.streaming_kind == "sse" for response in builder.responses + ): + retval.insert(0, '_last_event_id = kwargs.pop("last_event_id", None)') + retval.extend( + [ + "", + "if _last_event_id is not None:", + ' _request.headers["Last-Event-ID"] = _last_event_id', + ] + ) + return retval def response_headers_and_deserialization( self, @@ -1255,15 +1288,140 @@ def handle_error_response( # pylint: disable=too-many-statements, too-many-bran ) return retval - def handle_response(self, builder: OperationType) -> list[str]: - retval: list[str] = ["response = pipeline_response.http_response"] - retval.append("") - retval.extend(self.handle_error_response(builder)) + # pylint: disable=too-many-statements + def handle_structured_stream_response(self, builder: OperationType) -> list[str]: + """Emit the body for an operation returning a structured (JSONL / SSE) stream. + + Produces a per-event deserialization callback and returns a ``Stream`` / + ``AsyncStream`` wrapping the streamed HTTP response. + """ + response = next(r for r in builder.responses if r.is_structured_stream) + item_annotation = response.stream_item_annotation( + is_operation_file=True, serialize_namespace=self.serialize_namespace + ) + stream_class = response.stream_class_name(self.async_mode) # type: ignore[attr-defined] + terminal_event = getattr(response, "terminal_event", None) + terminal_event_names = getattr(response, "terminal_event_names", []) + streaming_events = getattr(response, "streaming_events", []) + unnamed_events = [event for event in streaming_events if event.event_type is None] + unnamed_discriminator = get_streaming_event_discriminator(unnamed_events) + retval: list[str] = [] + retval.append("def _callback(_http_response, _event):") + if response.streaming_kind == "sse": # type: ignore[attr-defined] + named_events = [event for event in streaming_events if event.event_type is not None] + + def emit_deserialized_event(event: StreamingEvent, indent: str) -> None: + event_item_type = event.payload_type + event_annotation = event_item_type.type_annotation( + is_operation_file=True, + serialize_namespace=self.serialize_namespace, + ) + if self.code_model.options["models-mode"] == "msrest": + serialization_type = event_item_type.serialization_type( + serialize_namespace=self.serialize_namespace + ) + retval.append(f"{indent}deserialized = self._deserialize(") + retval.append(f"{indent} '{serialization_type}',") + retval.append(f"{indent} _event_json") + retval.append(f"{indent})") + else: + retval.append(f"{indent}deserialized = _deserialize({event_annotation}, _event_json)") + + def emit_typed_event(event: StreamingEvent, indent: str) -> None: + event_json = _sse_event_data_expression(event.payload_content_type) + retval.append(f"{indent}_event_json = {event_json}") + emit_deserialized_event(event, indent) + + def emit_multiple_unnamed_events(events: list[StreamingEvent], indent: str) -> None: + event_json = _sse_fallback_data_expression(events) + retval.append(f"{indent}_event_json = {event_json}") + if unnamed_discriminator is None or event_json != "json.loads(_event.data)": + retval.append(f"{indent}deserialized = _event_json") + return + discriminator_name, variants = unnamed_discriminator + for index, (discriminator_value, event) in enumerate(variants): + keyword = "if" if index == 0 else "elif" + retval.append( + f"{indent}{keyword} isinstance(_event_json, dict) " + f"and _event_json.get({discriminator_name!r}) == {discriminator_value!r}:" + ) + emit_deserialized_event(event, indent + " ") + retval.append(f"{indent}else:") + retval.append(f"{indent} deserialized = _event_json") + + def emit_message_branch(keyword: str) -> None: + retval.append(f" {keyword} _event.event == {DEFAULT_SSE_EVENT_TYPE!r}:") + if len(unnamed_events) == 1: + emit_typed_event(unnamed_events[0], " ") + else: + emit_multiple_unnamed_events(unnamed_events, " ") + + if named_events: + for index, event in enumerate(named_events): + event_type = event.event_type + keyword = "if" if index == 0 else "elif" + retval.append(f" {keyword} _event.event == {event_type!r}:") + emit_typed_event(event, " ") + if unnamed_events: + emit_message_branch("elif") + elif unnamed_events: + emit_message_branch("if") + + if named_events or unnamed_events: + retval.append(" else:") + retval.append(' raise ValueError(f"Unknown SSE event type: {_event.event!r}")') + else: + retval.append(' raise ValueError(f"Unknown SSE event type: {_event.event!r}")') + else: + retval.append(" _event_json = _event.json()") + if self.code_model.options["models-mode"] == "msrest": + serialization_type = response.stream_item_type.serialization_type( + serialize_namespace=self.serialize_namespace + ) + retval.append(" deserialized = self._deserialize(") + retval.append(f" '{serialization_type}',") + retval.append(" _event_json") + retval.append(" )") + else: + retval.append(f" deserialized = _deserialize({item_annotation}, _event_json)") + retval.append(" return deserialized") retval.append("") - if builder.has_optional_return_type: - retval.append("deserialized = None") - if builder.any_response_has_headers: - retval.append("response_headers = {}") + stream_kwargs = ["response=response", "deserialization_callback=_callback"] + if terminal_event is not None: + stream_kwargs.append(f"terminal_event={terminal_event!r}") + if terminal_event_names: + stream_kwargs.append(f"terminal_event_names={terminal_event_names!r}") + unnamed_terminal_values = ( + [discriminator_value for discriminator_value, event in unnamed_discriminator[1] if event.is_terminal] + if unnamed_discriminator is not None + else [] + ) + if unnamed_terminal_values: + discriminator_name = unnamed_discriminator[0] # type: ignore[index] + retval.append("") + retval.append("def _is_terminal_event(_event):") + retval.append(f" if _event.event != {DEFAULT_SSE_EVENT_TYPE!r}:") + retval.append(" return False") + retval.append(" try:") + retval.append(" _event_json = json.loads(_event.data)") + retval.append(" except (TypeError, ValueError):") + retval.append(" return False") + retval.append( + f" return isinstance(_event_json, dict) " + f"and _event_json.get({discriminator_name!r}) in {unnamed_terminal_values!r}" + ) + stream_kwargs.append("terminal_event_predicate=_is_terminal_event") + retval.append( + f"deserialized: {stream_class}[{item_annotation}] = " + f"{stream_class}({', '.join(stream_kwargs)}) # type: ignore" + ) + retval.append("if cls:") + retval.append(" return cls(pipeline_response, deserialized, {}) # type: ignore") + retval.append("return deserialized") + return retval + + def _handle_response_body(self, builder: OperationType) -> list[str]: + retval: list[str] = [] if builder.has_response_body or builder.any_response_has_headers: # pylint: disable=too-many-nested-blocks if len(builder.responses) > 1: status_codes, res_headers, res_deserialization = [], [], [] @@ -1298,6 +1456,21 @@ def handle_response(self, builder: OperationType) -> list[str]: else: retval.extend(self.response_headers_and_deserialization(builder, builder.responses[0])) retval.append("") + return retval + + def handle_response(self, builder: OperationType) -> list[str]: + retval: list[str] = ["response = pipeline_response.http_response"] + retval.append("") + retval.extend(self.handle_error_response(builder)) + retval.append("") + if builder.has_structured_stream_response: + retval.extend(self.handle_structured_stream_response(builder)) + return retval + if builder.has_optional_return_type: + retval.append("deserialized = None") + if builder.any_response_has_headers: + retval.append("response_headers = {}") + retval.extend(self._handle_response_body(builder)) if ( builder.has_optional_return_type or self.code_model.options["models-mode"] diff --git a/packages/http-client-python/generator/pygen/codegen/serializers/general_serializer.py b/packages/http-client-python/generator/pygen/codegen/serializers/general_serializer.py index d44dbc8bc02..4700b527667 100644 --- a/packages/http-client-python/generator/pygen/codegen/serializers/general_serializer.py +++ b/packages/http-client-python/generator/pygen/codegen/serializers/general_serializer.py @@ -318,6 +318,10 @@ def serialize_model_base_file(self) -> str: template = self.env.get_template("model_base.py.jinja2") return template.render(code_model=self.code_model, file_import=FileImport(self.code_model)) + def serialize_streaming_base_file(self) -> str: + template = self.env.get_template("streaming_base.py.jinja2") + return template.render(code_model=self.code_model, file_import=FileImport(self.code_model)) + def serialize_validation_file(self) -> str: template = self.env.get_template("validation.py.jinja2") return template.render( diff --git a/packages/http-client-python/generator/pygen/codegen/templates/streaming_base.py.jinja2 b/packages/http-client-python/generator/pygen/codegen/templates/streaming_base.py.jinja2 new file mode 100644 index 00000000000..d44e17b408a --- /dev/null +++ b/packages/http-client-python/generator/pygen/codegen/templates/streaming_base.py.jinja2 @@ -0,0 +1,767 @@ +# coding=utf-8 +{% if code_model.license_header %} +{{ code_model.license_header }} +{% endif %} +# pylint: disable=line-too-long,useless-suppression,unnecessary-ellipsis +# -------------------------------------------------------------------------- +# This file is vendored from the core streaming runtime +# ({{ code_model.core_library }}.streaming). It provides the Stream / AsyncStream +# helpers (plus the JSONL / SSE decoders and event types) used by generated +# structured-streaming operations, so the generated package does not take a hard +# dependency on a core runtime that ships the streaming helpers. Do not edit by hand. +# -------------------------------------------------------------------------- +import codecs +import json +from contextlib import aclosing +from types import TracebackType +from typing import ( + Any, + AsyncGenerator, + AsyncIterator, + Callable, + Generator, + Iterator, + List, + Optional, + Protocol, + Sequence, + Type, + TypeVar, + cast, + runtime_checkable, +) + +from typing_extensions import Self + +from {{ code_model.core_library }}.rest import AsyncHttpResponse, HttpResponse + +DecodedType = TypeVar("DecodedType") +ReturnType_co = TypeVar("ReturnType_co", covariant=True) +T_co = TypeVar("T_co", covariant=True) + + +@runtime_checkable +class StreamDecoder(Protocol[T_co]): + """Protocol for stream decoders.""" + + def iter_events(self, iter_bytes: Iterator[bytes]) -> Iterator[T_co]: + """Iterate over events from a byte iterator. + + :param iter_bytes: An iterator of byte chunks. + :type iter_bytes: Iterator[bytes] + :return: An iterator of decoded data. + :rtype: Iterator[DecodedType_co] + """ + ... + + +@runtime_checkable +class AsyncStreamDecoder(Protocol[T_co]): + """Protocol for async stream decoders.""" + + # Why this isn't async def: https://mypy.readthedocs.io/en/stable/more_types.html#asynchronous-iterators + def aiter_events(self, iter_bytes: AsyncIterator[bytes]) -> AsyncIterator[T_co]: + """Asynchronously iterate over events from a byte iterator. + + :param iter_bytes: An asynchronous iterator of byte chunks. + :type iter_bytes: AsyncIterator[bytes] + :return: An asynchronous iterator of decoded data. + :rtype: AsyncIterator[DecodedType_co] + """ + ... + + +class JSONLEvent: + """A single JSON Lines (JSONL) event. + + :ivar data: The raw JSONL record. + :vartype data: str or None + :keyword data: The raw JSONL record. + :paramtype data: str or None + """ + + def __init__( + self, + *, + data: Optional[str] = None, + ) -> None: + self.data = data + + def json(self) -> Any: + """Parse the event data as JSON. + + :return: The parsed JSON value. + :rtype: Any + """ + return json.loads(cast(str, self.data)) + + +class _JSONLLineFramer: + """Incremental JSONL line framer with linear-time behavior. + + JSONL records are separated only by ``"\\n"`` (tolerating ``"\\r\\n"``). Unlike + ``str.splitlines()``, other Unicode boundaries (``\\v``, ``\\f``, ``\\x1c``-``\\x1e``, + ``\\x85``, ``\\u2028``, ``\\u2029``) are preserved because they are valid inside a JSONL + record's string value. + + Rather than re-concatenating and re-splitting the whole pending record on every network chunk + (which is O(n^2) for a single long record fragmented across many chunks), the unfinished record + is held as a list of fragments and joined only when a terminator arrives (or at EOF). Only the + newly decoded text is scanned per chunk, giving O(total) behavior. + """ + + def __init__(self) -> None: + # Fragments of the current, not-yet-terminated record. Never contains a "\n". + self._parts: List[str] = [] + + def push(self, text: str) -> List[str]: + """Feed newly decoded text and return any completed records. + + :param text: Newly decoded text from a single chunk. + :type text: str + :return: Completed records produced by this chunk (may be empty). + :rtype: list[str] + """ + if not text: + return [] + + segments = text.split("\n") + # Fast path: no line terminator, so this is a continuation of the current record. Stash the + # fragment without joining or rescanning the accumulated tail. + if len(segments) == 1: + self._parts.append(text) + return [] + + first = "".join(self._parts) + segments[0] + # All but the final segment are complete records (terminated by "\n"). Strip a trailing "\r" + # to tolerate "\r\n" line endings. + completed = [line[:-1] if line.endswith("\r") else line for line in [first, *segments[1:-1]]] + self._parts = [segments[-1]] + return completed + + def flush(self, extra: str = "") -> List[str]: + """Return the final unterminated record, if any, at end of stream. + + :param extra: Trailing text from finalizing the incremental decoder. + :type extra: str + :return: The final record, if non-empty. + :rtype: list[str] + """ + tail = "".join(self._parts) + extra + self._parts = [] + if not tail: + return [] + return [tail[:-1] if tail.endswith("\r") else tail] + + +def iter_lines(iter_bytes: Iterator[bytes]) -> Iterator[str]: + """Iterate over lines from a byte iterator. + + :param iter_bytes: An iterator of byte chunks. + :type iter_bytes: Iterator[bytes] + :rtype: Iterator[str] + :return: An iterator of lines. + """ + decoder = codecs.getincrementaldecoder("utf-8")() + framer = _JSONLLineFramer() + + for chunk in iter_bytes: + yield from framer.push(decoder.decode(chunk)) + + yield from framer.flush(decoder.decode(b"", final=True)) + + +async def aiter_lines(iter_bytes: AsyncIterator[bytes]) -> AsyncGenerator[str, None]: + """Iterate over lines from a byte iterator. + + :param iter_bytes: An iterator of byte chunks. + :type iter_bytes: Iterator[bytes] + :rtype: Iterator[str] + :return: An iterator of lines. + """ + decoder = codecs.getincrementaldecoder("utf-8")() + framer = _JSONLLineFramer() + + try: + async for chunk in iter_bytes: + for line in framer.push(decoder.decode(chunk)): + yield line + finally: + aclose = getattr(iter_bytes, "aclose", None) + if aclose is not None: + await aclose() + + for line in framer.flush(decoder.decode(b"", final=True)): + yield line + + +class JSONLDecoder: + """Decoder for JSON Lines (JSONL) format. https://jsonlines.org/""" + + def iter_events(self, iter_bytes: Iterator[bytes]) -> Iterator[JSONLEvent]: + """Iterate over JSONL events from a byte iterator. + + :param iter_bytes: An iterator of byte chunks. + :type iter_bytes: Iterator[bytes] + :rtype: Iterator[JSONLEvent] + :return: An iterator of JSONL events. + """ + + yield from (JSONLEvent(data=line) for line in iter_lines(iter_bytes)) + + +class AsyncJSONLDecoder: + """Asynchronous decoder for JSON Lines (JSONL) format. https://jsonlines.org/""" + + # pylint: disable=invalid-overridden-method + async def aiter_events(self, iter_bytes: AsyncIterator[bytes]) -> AsyncIterator[JSONLEvent]: + """Asynchronously iterate over JSONL events from a byte iterator. + + :param iter_bytes: An asynchronous iterator of byte chunks. + :type iter_bytes: AsyncIterator[bytes] + :rtype: AsyncIterator[JSONLEvent] + :return: An asynchronous iterator of JSONL events. + """ + + async with aclosing(aiter_lines(iter_bytes)) as lines: + async for line in lines: + yield JSONLEvent(data=line) + + +class ServerSentEvent: + """A single Server-Sent Event (SSE). + + https://html.spec.whatwg.org/multipage/server-sent-events.html + + :ivar event: The event type. Defaults to ``"message"`` when the stream does not + specify one. + :vartype event: str + :ivar data: The event payload. Multiple ``data`` lines are joined with ``"\\n"``. + Left as a raw string; the caller is responsible for any further parsing. + :vartype data: str + :ivar id: The last event ID. Defaults to an empty string until the stream + provides one. + :vartype id: str + :ivar retry: The reconnection time in milliseconds, if the stream provided one. + :vartype retry: int or None + :keyword event: The event type. Defaults to ``"message"`` when the stream does not + specify one. + :paramtype event: str + :keyword data: The event payload. Multiple ``data`` lines are joined with ``"\\n"``. + Left as a raw string; the caller is responsible for any further parsing. + :paramtype data: str + :keyword id: The last event ID. Defaults to an empty string until the stream + provides one. + :paramtype id: str + :keyword retry: The reconnection time in milliseconds, if the stream provided one. + :paramtype retry: int or None + """ + + def __init__( + self, + *, + event: str = "message", + data: str = "", + id: str = "", # pylint: disable=redefined-builtin + retry: Optional[int] = None, + ) -> None: + self.event = event + self.data = data + self.id = id + self.retry = retry + + def __repr__(self) -> str: + return f"ServerSentEvent(event={self.event!r}, data={self.data!r}, " f"id={self.id!r}, retry={self.retry!r})" + + def __eq__(self, other: object) -> bool: + if not isinstance(other, ServerSentEvent): + return NotImplemented + return (self.event, self.data, self.id, self.retry) == ( + other.event, + other.data, + other.id, + other.retry, + ) + + +class _SSELineFramer: + """Incremental SSE line framer with linear-time behavior. + + Per the SSE spec, lines may be separated by ``"\\r\\n"``, ``"\\r"`` or ``"\\n"``. A lone trailing + ``"\\r"`` at a chunk boundary is ambiguous (it may be the first half of a ``"\\r\\n"``) and its + resolution is deferred until the next chunk (or EOF). + + Rather than re-concatenating and re-scanning the whole pending line on every network chunk + (which is O(n^2) for a single long line fragmented across many chunks), the unfinished line is + held as a list of fragments and joined only when a terminator arrives (or at EOF). Only the + newly decoded text is scanned per chunk, giving O(total) behavior. + """ + + def __init__(self) -> None: + # Fragments of the current, not-yet-terminated line. Never contains a separator. + self._parts: List[str] = [] + # True when the previous chunk ended with a lone "\r" whose "\r\n" status is still unknown. + self._pending_cr = False + + def _emit_current(self) -> str: + line = "".join(self._parts) + self._parts = [] + return line + + def push(self, text: str) -> List[str]: + """Feed newly decoded text and return any completed lines. + + :param text: Newly decoded text from a single chunk. + :type text: str + :return: Completed lines (separators stripped) produced by this chunk. + :rtype: list[str] + """ + if not text: + return [] + + out: List[str] = [] + if self._pending_cr: + # The deferred "\r" terminates the current line now that more data is available. + out.append(self._emit_current()) + self._pending_cr = False + # A leading "\n" is the second half of that "\r\n" pair: consume it. + if text[:1] == "\n": + text = text[1:] + + n = len(text) + start = 0 + i = 0 + while i < n: + char = text[i] + if char == "\n": + self._parts.append(text[start:i]) + out.append(self._emit_current()) + i += 1 + start = i + elif char == "\r": + if i + 1 < n: + self._parts.append(text[start:i]) + out.append(self._emit_current()) + i += 2 if text[i + 1] == "\n" else 1 + start = i + else: + # Trailing lone "\r": defer resolution until the next chunk. + self._parts.append(text[start:i]) + self._pending_cr = True + start = n + break + else: + i += 1 + + if start < n: + self._parts.append(text[start:n]) + return out + + def flush(self, extra: str = "") -> List[str]: + """Return any remaining lines at end of stream. + + A lone trailing ``"\\r"`` is treated as a terminator (its ``"\\r\\n"`` half never arrives), + and a non-empty final unterminated line is emitted; an empty tail is not, so no blank line + (and therefore no spurious event) is invented at EOF. + + :param extra: Trailing text from finalizing the incremental decoder. + :type extra: str + :return: The remaining lines, if any. + :rtype: list[str] + """ + out: List[str] = [] + if self._pending_cr: + out.append(self._emit_current()) + self._pending_cr = False + if extra[:1] == "\n": + extra = extra[1:] + + if extra: + out.extend(self.push(extra)) + + if self._pending_cr: + # 'extra' ended in a lone "\r": at EOF it is a terminator; emit the preceding content. + out.append(self._emit_current()) + self._pending_cr = False + else: + tail = self._emit_current() + if tail: + out.append(tail) + return out + + +def _iter_sse_lines(iter_bytes: Iterator[bytes]) -> Iterator[str]: + """Iterate over SSE lines (line separators stripped) from a byte iterator. + + :param iter_bytes: An iterator of byte chunks. + :type iter_bytes: Iterator[bytes] + :rtype: Iterator[str] + :return: An iterator of decoded lines. + """ + # SSE is always UTF-8 (WHATWG spec). Use utf-8-sig to drop one leading BOM and + # errors="replace" so invalid byte sequences become U+FFFD instead of crashing. + decoder = codecs.getincrementaldecoder("utf-8-sig")(errors="replace") + framer = _SSELineFramer() + + for chunk in iter_bytes: + yield from framer.push(decoder.decode(chunk)) + + yield from framer.flush(decoder.decode(b"", final=True)) + + +async def _aiter_sse_lines(iter_bytes: AsyncIterator[bytes]) -> AsyncGenerator[str, None]: + """Asynchronously iterate over SSE lines (separators stripped) from a byte iterator. + + :param iter_bytes: An asynchronous iterator of byte chunks. + :type iter_bytes: AsyncIterator[bytes] + :rtype: AsyncIterator[str] + :return: An asynchronous iterator of decoded lines. + """ + # SSE is always UTF-8 (WHATWG spec). Use utf-8-sig to drop one leading BOM and + # errors="replace" so invalid byte sequences become U+FFFD instead of crashing. + decoder = codecs.getincrementaldecoder("utf-8-sig")(errors="replace") + framer = _SSELineFramer() + + try: + async for chunk in iter_bytes: + for line in framer.push(decoder.decode(chunk)): + yield line + finally: + aclose = getattr(iter_bytes, "aclose", None) + if aclose is not None: + await aclose() + + for line in framer.flush(decoder.decode(b"", final=True)): + yield line + + +class _SSEEventBuilder: + """Accumulates SSE field lines and builds :class:`ServerSentEvent` instances.""" + + def __init__(self) -> None: + self._data: List[str] = [] + self._event_type = "" + self._last_id = "" + self._retry: Optional[int] = None + + def add_line(self, line: str) -> Optional[ServerSentEvent]: + """Process a single SSE line, dispatching an event on a blank line. + + :param line: A single SSE line with its terminator already stripped. + :type line: str + :return: A :class:`ServerSentEvent` when ``line`` is blank and an event is + pending, otherwise ``None``. + :rtype: ServerSentEvent or None + """ + if line == "": + return self._dispatch() + if line.startswith(":"): + # Comment line, ignored. + return None + + field, _, value = line.partition(":") + if value.startswith(" "): + value = value[1:] + + if field == "event": + self._event_type = value + elif field == "data": + self._data.append(value) + elif field == "id": + if "\x00" not in value: + self._last_id = value + elif field == "retry": + if value.isascii() and value.isdigit(): + try: + self._retry = int(value) + except ValueError: + # All ASCII digits but too long for int() (CPython's int-string + # conversion limit). Ignore rather than crashing the stream. + pass + # Unknown fields are ignored per spec. + return None + + def _dispatch(self) -> Optional[ServerSentEvent]: + if not self._data: + # No data accumulated: reset and dispatch nothing. + self._event_type = "" + return None + event = ServerSentEvent( + event=self._event_type or "message", + data="\n".join(self._data), + id=self._last_id, + retry=self._retry, + ) + self._data = [] + self._event_type = "" + return event + + +class SSEDecoder: + """Decoder for Server-Sent Events (SSE). + + https://html.spec.whatwg.org/multipage/server-sent-events.html + """ + + def iter_events(self, iter_bytes: Iterator[bytes]) -> Iterator[ServerSentEvent]: + """Iterate over SSE events from a byte iterator. + + :param iter_bytes: An iterator of byte chunks. + :type iter_bytes: Iterator[bytes] + :rtype: Iterator[ServerSentEvent] + :return: An iterator of server-sent events. + """ + builder = _SSEEventBuilder() + for line in _iter_sse_lines(iter_bytes): + event = builder.add_line(line) + if event is not None: + yield event + + +class AsyncSSEDecoder: + """Asynchronous decoder for Server-Sent Events (SSE). + + https://html.spec.whatwg.org/multipage/server-sent-events.html + """ + + # pylint: disable=invalid-overridden-method + async def aiter_events(self, iter_bytes: AsyncIterator[bytes]) -> AsyncIterator[ServerSentEvent]: + """Asynchronously iterate over SSE events from a byte iterator. + + :param iter_bytes: An asynchronous iterator of byte chunks. + :type iter_bytes: AsyncIterator[bytes] + :rtype: AsyncIterator[ServerSentEvent] + :return: An asynchronous iterator of server-sent events. + """ + builder = _SSEEventBuilder() + async with aclosing(_aiter_sse_lines(iter_bytes)) as lines: + async for line in lines: + event = builder.add_line(line) + if event is not None: + yield event + + +class Stream(Iterator[ReturnType_co]): + """Stream class for consuming a decoded event stream (e.g. JSONL or SSE). + + :keyword response: The response object. + :paramtype response: ~{{ code_model.core_library }}.rest.HttpResponse + :keyword decoder: A decoder to use for the stream. If omitted, the decoder is + inferred from the response ``Content-Type`` header. + :paramtype decoder: StreamDecoder + :keyword deserialization_callback: A callback that takes the response and the decoded event and + returns a deserialized object. + :paramtype deserialization_callback: Callable[[~{{ code_model.core_library }}.rest.HttpResponse, Any], ReturnType] + :keyword terminal_event: Optional event ``data`` value that terminates the stream (e.g. + ``"[DONE]"``). When an event's ``data`` equals this value, iteration stops and the + event is not passed to ``deserialization_callback``. + :paramtype terminal_event: str or None + :keyword terminal_event_names: Optional event names (the SSE ``event`` field) that terminate + the stream. Unlike ``terminal_event``, such an event carries a payload: it is passed to + ``deserialization_callback`` and yielded, and iteration stops immediately afterwards. + :paramtype terminal_event_names: ~typing.Sequence[str] or None + :keyword terminal_event_predicate: Optional predicate that identifies a payload-bearing + terminal event. The event is yielded before iteration stops. + :paramtype terminal_event_predicate: Callable[[Any], bool] or None + """ + + def __init__( + self, + *, + response: HttpResponse, + deserialization_callback: Callable[[HttpResponse, DecodedType], ReturnType_co], + decoder: Optional[StreamDecoder[DecodedType]] = None, + terminal_event: Optional[str] = None, + terminal_event_names: Optional[Sequence[str]] = None, + terminal_event_predicate: Optional[Callable[[DecodedType], bool]] = None, + ) -> None: + self._response = response + content_type = response.headers.get("Content-Type", "").split(";", 1)[0].strip().lower() + self._decoder: StreamDecoder[Any] = ( + decoder if decoder is not None else (SSEDecoder() if content_type == "text/event-stream" else JSONLDecoder()) + ) + self._deserialization_callback = deserialization_callback + self._terminal_event = terminal_event + self._terminal_event_names = frozenset(terminal_event_names or ()) + self._terminal_event_predicate = terminal_event_predicate + self._last_event_id: Optional[str] = None + self._retry: Optional[int] = None + self._iterator = self._iter_results() + + @property + def last_event_id(self) -> Optional[str]: + """The most recently received SSE event ID, if one was provided.""" + return self._last_event_id + + @property + def retry(self) -> Optional[int]: + """The most recently received valid SSE retry value, if one was provided.""" + return self._retry + + def __next__(self) -> ReturnType_co: + return self._iterator.__next__() + + def __iter__(self) -> Self: + return self + + def _iter_results(self) -> Generator[ReturnType_co, None, None]: + try: + for event in self._decoder.iter_events(self._response.iter_bytes()): + event_id = getattr(event, "id", None) + if event_id is not None: + self._last_event_id = event_id + event_retry = getattr(event, "retry", None) + if event_retry is not None: + self._retry = event_retry + if self._terminal_event is not None and getattr(event, "data", None) == self._terminal_event: + break + result = self._deserialization_callback(self._response, event) + yield result + if ( + self._terminal_event_names + and getattr(event, "event", None) in self._terminal_event_names + ) or ( + self._terminal_event_predicate is not None + and self._terminal_event_predicate(event) + ): + break + finally: + self._response.close() + + def __exit__( + self, + exc_type: Optional[Type[BaseException]] = None, + exc_value: Optional[BaseException] = None, + traceback: Optional[TracebackType] = None, + ) -> None: + self.close() + + def __enter__(self) -> Self: + return self + + def close(self) -> None: + try: + self._iterator.close() + finally: + self._response.close() + + +class AsyncStream(AsyncIterator[ReturnType_co]): + """AsyncStream class for asynchronously consuming a decoded event stream (e.g. JSONL or SSE). + + :keyword response: The response object. + :paramtype response: ~{{ code_model.core_library }}.rest.AsyncHttpResponse + :keyword decoder: A decoder to use for the stream. If omitted, the decoder is + inferred from the response ``Content-Type`` header. + :paramtype decoder: AsyncStreamDecoder + :keyword deserialization_callback: A callback that takes the response and the decoded event and + returns a deserialized object. + :paramtype deserialization_callback: Callable[[~{{ code_model.core_library }}.rest.AsyncHttpResponse, Any], ReturnType] + :keyword terminal_event: Optional event ``data`` value that terminates the stream (e.g. + ``"[DONE]"``). When an event's ``data`` equals this value, iteration stops and the + event is not passed to ``deserialization_callback``. + :paramtype terminal_event: str or None + :keyword terminal_event_names: Optional event names (the SSE ``event`` field) that terminate + the stream. Unlike ``terminal_event``, such an event carries a payload: it is passed to + ``deserialization_callback`` and yielded, and iteration stops immediately afterwards. + :paramtype terminal_event_names: ~typing.Sequence[str] or None + :keyword terminal_event_predicate: Optional predicate that identifies a payload-bearing + terminal event. The event is yielded before iteration stops. + :paramtype terminal_event_predicate: Callable[[Any], bool] or None + """ + + def __init__( + self, + *, + response: AsyncHttpResponse, + deserialization_callback: Callable[[AsyncHttpResponse, DecodedType], ReturnType_co], + decoder: Optional[AsyncStreamDecoder[DecodedType]] = None, + terminal_event: Optional[str] = None, + terminal_event_names: Optional[Sequence[str]] = None, + terminal_event_predicate: Optional[Callable[[DecodedType], bool]] = None, + ) -> None: + self._response = response + content_type = response.headers.get("Content-Type", "").split(";", 1)[0].strip().lower() + self._decoder: AsyncStreamDecoder[Any] = ( + decoder + if decoder is not None + else (AsyncSSEDecoder() if content_type == "text/event-stream" else AsyncJSONLDecoder()) + ) + self._deserialization_callback = deserialization_callback + self._terminal_event = terminal_event + self._terminal_event_names = frozenset(terminal_event_names or ()) + self._terminal_event_predicate = terminal_event_predicate + self._last_event_id: Optional[str] = None + self._retry: Optional[int] = None + self._iterator = self._iter_results() + + @property + def last_event_id(self) -> Optional[str]: + """The most recently received SSE event ID, if one was provided.""" + return self._last_event_id + + @property + def retry(self) -> Optional[int]: + """The most recently received valid SSE retry value, if one was provided.""" + return self._retry + + async def __anext__(self) -> ReturnType_co: + return await self._iterator.__anext__() + + def __aiter__(self) -> Self: + return self + + async def _iter_results(self) -> AsyncGenerator[ReturnType_co, None]: + events = self._decoder.aiter_events(self._response.iter_bytes()) + try: + async for event in events: + event_id = getattr(event, "id", None) + if event_id is not None: + self._last_event_id = event_id + event_retry = getattr(event, "retry", None) + if event_retry is not None: + self._retry = event_retry + if self._terminal_event is not None and getattr(event, "data", None) == self._terminal_event: + break + result = self._deserialization_callback(self._response, event) + yield result + if ( + self._terminal_event_names + and getattr(event, "event", None) in self._terminal_event_names + ) or ( + self._terminal_event_predicate is not None + and self._terminal_event_predicate(event) + ): + break + finally: + try: + aclose = getattr(events, "aclose", None) + if aclose is not None: + await aclose() + finally: + await self._response.close() + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]] = None, + exc_value: Optional[BaseException] = None, + traceback: Optional[TracebackType] = None, + ) -> None: + await self.close() + + async def __aenter__(self) -> Self: + return self + + async def close(self) -> None: + try: + await self._iterator.aclose() + finally: + await self._response.close() + + +__all__ = [ + "Stream", + "AsyncStream", + "JSONLEvent", + "ServerSentEvent", +] diff --git a/packages/http-client-python/generator/pygen/preprocess/__init__.py b/packages/http-client-python/generator/pygen/preprocess/__init__.py index ff3d6f094e3..d79fcb42235 100644 --- a/packages/http-client-python/generator/pygen/preprocess/__init__.py +++ b/packages/http-client-python/generator/pygen/preprocess/__init__.py @@ -35,6 +35,8 @@ def update_overload_section( for overload_s, original_s in zip(overload[section], yaml_data[section]): if overload_s.get("type"): overload_s["type"] = original_s["type"] + if overload_s.get("streaming"): + overload_s["streaming"] = original_s["streaming"] if overload_s.get("headers"): for overload_h, original_h in zip(overload_s["headers"], original_s["headers"]): if overload_h.get("type"): diff --git a/packages/http-client-python/package-lock.json b/packages/http-client-python/package-lock.json index 64e2b328755..2f0ef1b4c5e 100644 --- a/packages/http-client-python/package-lock.json +++ b/packages/http-client-python/package-lock.json @@ -29,11 +29,11 @@ "@typespec/compiler": "^1.15.0", "@typespec/events": "~0.85.0", "@typespec/http": "^1.15.0", - "@typespec/http-specs": "0.1.0-alpha.41", + "@typespec/http-specs": "0.1.0-alpha.42-dev.2", "@typespec/openapi": "^1.15.0", "@typespec/rest": "~0.85.0", "@typespec/spec-api": "0.1.0-alpha.16", - "@typespec/spector": "0.1.0-alpha.28", + "@typespec/spector": "0.1.0-alpha.29-dev.0", "@typespec/sse": "~0.85.0", "@typespec/streams": "~0.85.0", "@typespec/versioning": "~0.85.0", @@ -2472,26 +2472,26 @@ } }, "node_modules/@typespec/http-specs": { - "version": "0.1.0-alpha.41", - "resolved": "https://pkgs.dev.azure.com/azure-sdk/public/_packaging/azure-sdk-for-js/npm/registry/@typespec/http-specs/-/http-specs-0.1.0-alpha.41.tgz", - "integrity": "sha1-f6jsbsXAIlEFa1+1A/FQqcvxZkk=", + "version": "0.1.0-alpha.42-dev.2", + "resolved": "https://pkgs.dev.azure.com/azure-sdk/public/_packaging/azure-sdk-for-js/npm/registry/@typespec/http-specs/-/http-specs-0.1.0-alpha.42-dev.2.tgz", + "integrity": "sha1-EhDehpyxXzA+dUEwWHHzLWJBQbM=", "dev": true, "license": "MIT", "dependencies": { - "@typespec/spec-api": "^0.1.0-alpha.16", - "@typespec/spector": "^0.1.0-alpha.28" + "@typespec/spec-api": "^0.1.0-alpha.16 || >= 0.1.0-dev.0", + "@typespec/spector": "^0.1.0-alpha.28 || >= 0.1.0-alpha.29-dev.0" }, "engines": { "node": ">=22.0.0" }, "peerDependencies": { - "@typespec/compiler": "^1.15.0", - "@typespec/events": "^0.85.0", - "@typespec/http": "^1.15.0", - "@typespec/rest": "^0.85.0", - "@typespec/sse": "^0.85.0", - "@typespec/versioning": "^0.85.0", - "@typespec/xml": "^0.85.0" + "@typespec/compiler": "^1.15.0 || >= 1.16.0-dev.0", + "@typespec/events": "^0.85.0 || >= 0.86.0-dev.0", + "@typespec/http": "^1.15.0 || >= 1.16.0-dev.0", + "@typespec/rest": "^0.85.0 || >= 0.86.0-dev.0", + "@typespec/sse": "^0.85.0 || >= 0.86.0-dev.0", + "@typespec/versioning": "^0.85.0 || >= 0.86.0-dev.0", + "@typespec/xml": "^0.85.0 || >= 0.86.0-dev.0" } }, "node_modules/@typespec/openapi": { @@ -2584,19 +2584,19 @@ "license": "MIT" }, "node_modules/@typespec/spector": { - "version": "0.1.0-alpha.28", - "resolved": "https://pkgs.dev.azure.com/azure-sdk/public/_packaging/azure-sdk-for-js/npm/registry/@typespec/spector/-/spector-0.1.0-alpha.28.tgz", - "integrity": "sha1-/4kblwKNH+UZXSzz1SmZds8nt88=", + "version": "0.1.0-alpha.29-dev.0", + "resolved": "https://pkgs.dev.azure.com/azure-sdk/public/_packaging/azure-sdk-for-js/npm/registry/@typespec/spector/-/spector-0.1.0-alpha.29-dev.0.tgz", + "integrity": "sha1-Y76ZtxW9/gmRAbGeEFfohnKnUhk=", "dev": true, "license": "MIT", "dependencies": { "@azure/identity": "^4.13.1", - "@typespec/compiler": "^1.15.0", - "@typespec/http": "^1.15.0", - "@typespec/rest": "^0.85.0", - "@typespec/spec-api": "^0.1.0-alpha.16", - "@typespec/spec-coverage-sdk": "^0.1.0-alpha.16", - "@typespec/versioning": "^0.85.0", + "@typespec/compiler": "^1.15.0 || >= 1.16.0-dev.0", + "@typespec/http": "^1.15.0 || >= 1.16.0-dev.0", + "@typespec/rest": "^0.85.0 || >= 0.86.0-dev.0", + "@typespec/spec-api": "^0.1.0-alpha.16 || >= 0.1.0-dev.0", + "@typespec/spec-coverage-sdk": "^0.1.0-alpha.16 || >= 0.1.0-dev.0", + "@typespec/versioning": "^0.85.0 || >= 0.86.0-dev.0", "ajv": "^8.18.0", "express": "^5.2.1", "micromatch": "^4.0.8", diff --git a/packages/http-client-python/package.json b/packages/http-client-python/package.json index b530962f347..17f49dd347b 100644 --- a/packages/http-client-python/package.json +++ b/packages/http-client-python/package.json @@ -116,12 +116,12 @@ "@typespec/rest": "~0.85.0", "@typespec/versioning": "~0.85.0", "@typespec/events": "~0.85.0", - "@typespec/spector": "0.1.0-alpha.28", + "@typespec/spector": "0.1.0-alpha.29-dev.0", "@typespec/spec-api": "0.1.0-alpha.16", "@typespec/sse": "~0.85.0", "@typespec/streams": "~0.85.0", "@typespec/xml": "~0.85.0", - "@typespec/http-specs": "0.1.0-alpha.41", + "@typespec/http-specs": "0.1.0-alpha.42-dev.2", "@types/js-yaml": "~4.0.5", "@types/node": "~25.0.2", "@types/semver": "7.5.8", diff --git a/packages/http-client-python/tests/mock_api/shared/asynctests/test_streaming_jsonl_async.py b/packages/http-client-python/tests/mock_api/shared/asynctests/test_streaming_jsonl_async.py index 74e05cebd14..d46dd1c41b9 100644 --- a/packages/http-client-python/tests/mock_api/shared/asynctests/test_streaming_jsonl_async.py +++ b/packages/http-client-python/tests/mock_api/shared/asynctests/test_streaming_jsonl_async.py @@ -7,6 +7,8 @@ import pytest_asyncio from streaming.jsonl.aio import JsonlClient +from streaming.jsonl._utils.streaming_base import AsyncStream +from streaming.jsonl.basic.models import Info @pytest_asyncio.fixture @@ -25,4 +27,8 @@ async def test_basic_send(client: JsonlClient): @pytest.mark.asyncio async def test_basic_recv(client: JsonlClient): - assert b"".join([d async for d in (await client.basic.receive())]) == JSONL + stream = await client.basic.receive() + assert isinstance(stream, AsyncStream) + items = [item async for item in stream] + assert all(isinstance(item, Info) for item in items) + assert [item.desc for item in items] == ["one", "two", "three"] diff --git a/packages/http-client-python/tests/mock_api/shared/asynctests/test_streaming_sse_async.py b/packages/http-client-python/tests/mock_api/shared/asynctests/test_streaming_sse_async.py new file mode 100644 index 00000000000..66c5f1abb0a --- /dev/null +++ b/packages/http-client-python/tests/mock_api/shared/asynctests/test_streaming_sse_async.py @@ -0,0 +1,289 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- +import json + +import pytest +import pytest_asyncio + +from streaming.sse.aio import SseClient +from streaming.sse._utils.streaming_base import AsyncStream +from streaming.sse.named.models import ResponseCreated, ResponseDelta +from streaming.sse.protocol.data.models import WithEnvelope1 +from streaming.sse.protocol.models import Info as ProtocolInfo +from streaming.sse.retrieve.models import FinalResult, PartialResult, RetrievalRequest +from streaming.sse.unnamed.models import Info + + +@pytest_asyncio.fixture +async def client(): + async with SseClient(endpoint="http://localhost:3000") as client: + yield client + + +@pytest.mark.asyncio +async def test_unnamed_receive(client: SseClient): + async with await client.unnamed.receive() as stream: + assert isinstance(stream, AsyncStream) + items = [await stream.__anext__() for _ in range(3)] + assert all(isinstance(item, Info) for item in items) + assert [item.desc for item in items] == ["one", "two", "three"] + + +@pytest.mark.asyncio +async def test_named_receive(client: SseClient): + stream = await client.named.receive() + assert isinstance(stream, AsyncStream) + items = [item async for item in stream] + # The terminal "[DONE]" event stops iteration and is not yielded. + assert len(items) == 3 + assert isinstance(items[0], ResponseCreated) and items[0].id == "resp_1" + assert isinstance(items[1], ResponseDelta) and items[1].delta == "Hello" + assert isinstance(items[2], ResponseDelta) and items[2].delta == " world" + + +@pytest.mark.asyncio +async def test_retrieve_stream(client: SseClient): + stream = await client.retrieve.stream(RetrievalRequest(query="what is typespec?")) + assert isinstance(stream, AsyncStream) + items = [item async for item in stream] + # The terminal "[DONE]" event stops iteration and is not yielded. + assert len(items) == 3 + assert isinstance(items[0], PartialResult) and items[0].text == "partial one" + assert isinstance(items[1], PartialResult) and items[1].text == "partial two" + assert isinstance(items[2], FinalResult) and items[2].references == ["doc1", "doc2"] + + +@pytest.mark.asyncio +async def test_protocol_data_with_envelope(client: SseClient): + async with await client.protocol.data.with_envelope() as stream: + assert await stream.__anext__() == "hello" + + +@pytest.mark.asyncio +async def test_protocol_data_without_envelope(client: SseClient): + async with await client.protocol.data.without_envelope() as stream: + item = await stream.__anext__() + assert isinstance(item, WithEnvelope1) + assert item.metadata == {"source": "test"} + assert item.contents == "world" + + +@pytest.mark.asyncio +async def test_protocol_event_id(client: SseClient): + async with await client.protocol.id() as stream: + item = await stream.__anext__() + assert isinstance(item, ProtocolInfo) + assert item.message == "hello" + assert stream.last_event_id == "event-1" + + +@pytest.mark.asyncio +async def test_protocol_invalid_event_id(client: SseClient): + async with await client.protocol.invalid_id() as stream: + item = await stream.__anext__() + assert isinstance(item, ProtocolInfo) + assert item.message == "hello" + assert stream.last_event_id == "" + + +@pytest.mark.asyncio +async def test_protocol_retry(client: SseClient): + async with await client.protocol.retry() as stream: + item = await stream.__anext__() + assert isinstance(item, ProtocolInfo) + assert item.message == "hello" + assert stream.retry == 1000 + + +@pytest.mark.asyncio +async def test_protocol_invalid_retry(client: SseClient): + async with await client.protocol.invalid_retry() as stream: + item = await stream.__anext__() + assert isinstance(item, ProtocolInfo) + assert item.message == "hello" + assert stream.retry is None + + +@pytest.mark.asyncio +async def test_protocol_reconnect(client: SseClient): + async with await client.protocol.reconnect() as stream: + first = await stream.__anext__() + assert isinstance(first, ProtocolInfo) + assert first.message == "hello" + assert stream.last_event_id == "event-1" + + async with await client.protocol.reconnect(last_event_id="event-1") as stream: + second = await stream.__anext__() + assert isinstance(second, ProtocolInfo) + assert second.message == "world" + assert stream.last_event_id == "event-2" + + +# --------------------------------------------------------------------------- +# Named / model terminal events (yield-then-stop) -- see the sync test module +# for the rationale. Driven through the generated ``AsyncStream`` with a fake +# response because no published Spector spec produces named-model terminals. +# --------------------------------------------------------------------------- + + +class _FakeAsyncResponse: + """A minimal AsyncHttpResponse-shaped stand-in that replays SSE bytes.""" + + def __init__(self, body: bytes): + self.headers = {"Content-Type": "text/event-stream"} + self._body = body + self.closed = False + + def iter_bytes(self): + async def gen(): + for index in range(0, len(self._body), 8): + yield self._body[index : index + 8] + + return gen() + + async def close(self): + self.closed = True + + +class _FailingCloseAsyncResponse(_FakeAsyncResponse): + def iter_bytes(self): + class _Chunks: + def __init__(self, body): + self._body = body + self._done = False + + def __aiter__(self): + return self + + async def __anext__(self): + if self._done: + raise StopAsyncIteration + self._done = True + return self._body + + async def aclose(self): + raise RuntimeError("iterator cleanup failed") + + return _Chunks(self._body) + + +def _event_kind(_response, event): + return (event.event, json.loads(event.data)) + + +_NAMED_TERMINAL_SSE = ( + b'event: response.partial\ndata: {"text": "one"}\n\n' + b'event: response.delta\ndata: {"delta": "hi"}\n\n' + b'event: response.completed\ndata: {"references": []}\n\n' + b'event: response.delta\ndata: {"delta": "AFTER-TERMINAL"}\n\n' +) + +_TERMINAL_EVENT_NAMES = ["response.completed", "error"] + + +@pytest.mark.asyncio +async def test_named_terminal_event_yields_then_stops(): + response = _FakeAsyncResponse(_NAMED_TERMINAL_SSE) + stream = AsyncStream( + response=response, + deserialization_callback=_event_kind, + terminal_event_names=_TERMINAL_EVENT_NAMES, + ) + items = [item async for item in stream] + # The named terminal `response.completed` IS yielded, then iteration stops. + assert items == [ + ("response.partial", {"text": "one"}), + ("response.delta", {"delta": "hi"}), + ("response.completed", {"references": []}), + ] + assert response.closed + + +@pytest.mark.asyncio +async def test_unnamed_terminal_event_yields_then_stops(): + body = b'data: {"status": "done"}\n\n' b'data: {"status": "AFTER-TERMINAL"}\n\n' + response = _FakeAsyncResponse(body) + stream = AsyncStream( + response=response, + deserialization_callback=_event_kind, + terminal_event_names=["message"], + ) + + assert [item async for item in stream] == [("message", {"status": "done"})] + assert response.closed + + +@pytest.mark.asyncio +async def test_unnamed_terminal_predicate_preserves_service_event_type(): + body = b'data: {"kind": "progress"}\n\n' b'data: {"kind": "complete"}\n\n' b'data: {"kind": "after"}\n\n' + response = _FakeAsyncResponse(body) + stream = AsyncStream( + response=response, + deserialization_callback=_event_kind, + terminal_event_predicate=lambda event: json.loads(event.data).get("kind") == "complete", + ) + + assert [item async for item in stream] == [ + ("message", {"kind": "progress"}), + ("message", {"kind": "complete"}), + ] + assert response.closed + + +@pytest.mark.asyncio +async def test_response_closes_when_async_iterator_cleanup_fails(): + response = _FailingCloseAsyncResponse(b"data: first\n\n") + stream = AsyncStream( + response=response, + deserialization_callback=lambda _response, event: event.data, + ) + + assert await stream.__anext__() == "first" + with pytest.raises(RuntimeError, match="iterator cleanup failed"): + await stream.__anext__() + assert response.closed + + +@pytest.mark.asyncio +async def test_sentinel_and_named_terminal_coexist(): + body = ( + b'event: response.delta\ndata: {"delta": "a"}\n\n' + b"data: [DONE]\n\n" + b'event: response.delta\ndata: {"delta": "AFTER-DONE"}\n\n' + ) + response = _FakeAsyncResponse(body) + stream = AsyncStream( + response=response, + deserialization_callback=_event_kind, + terminal_event="[DONE]", + terminal_event_names=_TERMINAL_EVENT_NAMES, + ) + items = [item async for item in stream] + # The bare `[DONE]` sentinel stops iteration WITHOUT being yielded. + assert items == [("response.delta", {"delta": "a"})] + assert response.closed + + +@pytest.mark.asyncio +async def test_sse_protocol_metadata_is_available(): + body = b'id: event-1\nretry: 1000\nevent: message\ndata: {"message": "hello"}\n\n' + response = _FakeAsyncResponse(body) + stream = AsyncStream(response=response, deserialization_callback=lambda _response, event: event.data) + + assert [item async for item in stream] == ['{"message": "hello"}'] + assert stream.last_event_id == "event-1" + assert stream.retry == 1000 + + +@pytest.mark.asyncio +async def test_sse_protocol_invalid_metadata_is_ignored(): + body = b"id: invalid\x00id\nretry: not-a-number\nevent: message\ndata: hello\n\n" + response = _FakeAsyncResponse(body) + stream = AsyncStream(response=response, deserialization_callback=lambda _response, event: event.data) + + assert [item async for item in stream] == ["hello"] + assert stream.last_event_id == "" + assert stream.retry is None diff --git a/packages/http-client-python/tests/mock_api/shared/test_streaming_jsonl.py b/packages/http-client-python/tests/mock_api/shared/test_streaming_jsonl.py index 494c17a3493..bc2febf91d2 100644 --- a/packages/http-client-python/tests/mock_api/shared/test_streaming_jsonl.py +++ b/packages/http-client-python/tests/mock_api/shared/test_streaming_jsonl.py @@ -6,6 +6,8 @@ import pytest from streaming.jsonl import JsonlClient +from streaming.jsonl._utils.streaming_base import Stream +from streaming.jsonl.basic.models import Info @pytest.fixture @@ -22,4 +24,8 @@ def test_basic_send(client: JsonlClient): def test_basic_recv(client: JsonlClient): - assert b"".join(client.basic.receive()) == JSONL + stream = client.basic.receive() + assert isinstance(stream, Stream) + items = list(stream) + assert all(isinstance(item, Info) for item in items) + assert [item.desc for item in items] == ["one", "two", "three"] diff --git a/packages/http-client-python/tests/mock_api/shared/test_streaming_sse.py b/packages/http-client-python/tests/mock_api/shared/test_streaming_sse.py new file mode 100644 index 00000000000..405b44983e7 --- /dev/null +++ b/packages/http-client-python/tests/mock_api/shared/test_streaming_sse.py @@ -0,0 +1,253 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- +import json + +import pytest + +from streaming.sse import SseClient +from streaming.sse._utils.streaming_base import Stream +from streaming.sse.named.models import ResponseCreated, ResponseDelta +from streaming.sse.protocol.data.models import WithEnvelope1 +from streaming.sse.protocol.models import Info as ProtocolInfo +from streaming.sse.retrieve.models import FinalResult, PartialResult, RetrievalRequest +from streaming.sse.unnamed.models import Info + + +@pytest.fixture +def client(): + with SseClient(endpoint="http://localhost:3000") as client: + yield client + + +def test_unnamed_receive(client: SseClient): + with client.unnamed.receive() as stream: + assert isinstance(stream, Stream) + items = [next(stream) for _ in range(3)] + assert all(isinstance(item, Info) for item in items) + assert [item.desc for item in items] == ["one", "two", "three"] + + +def test_named_receive(client: SseClient): + stream = client.named.receive() + assert isinstance(stream, Stream) + items = list(stream) + # The terminal "[DONE]" event stops iteration and is not yielded. + assert len(items) == 3 + assert isinstance(items[0], ResponseCreated) and items[0].id == "resp_1" + assert isinstance(items[1], ResponseDelta) and items[1].delta == "Hello" + assert isinstance(items[2], ResponseDelta) and items[2].delta == " world" + + +def test_retrieve_stream(client: SseClient): + stream = client.retrieve.stream(RetrievalRequest(query="what is typespec?")) + assert isinstance(stream, Stream) + items = list(stream) + # The terminal "[DONE]" event stops iteration and is not yielded. + assert len(items) == 3 + assert isinstance(items[0], PartialResult) and items[0].text == "partial one" + assert isinstance(items[1], PartialResult) and items[1].text == "partial two" + assert isinstance(items[2], FinalResult) and items[2].references == ["doc1", "doc2"] + + +def test_protocol_data_with_envelope(client: SseClient): + with client.protocol.data.with_envelope() as stream: + assert next(stream) == "hello" + + +def test_protocol_data_without_envelope(client: SseClient): + with client.protocol.data.without_envelope() as stream: + item = next(stream) + assert isinstance(item, WithEnvelope1) + assert item.metadata == {"source": "test"} + assert item.contents == "world" + + +def test_protocol_event_id(client: SseClient): + with client.protocol.id() as stream: + item = next(stream) + assert isinstance(item, ProtocolInfo) + assert item.message == "hello" + assert stream.last_event_id == "event-1" + + +def test_protocol_invalid_event_id(client: SseClient): + with client.protocol.invalid_id() as stream: + item = next(stream) + assert isinstance(item, ProtocolInfo) + assert item.message == "hello" + assert stream.last_event_id == "" + + +def test_protocol_retry(client: SseClient): + with client.protocol.retry() as stream: + item = next(stream) + assert isinstance(item, ProtocolInfo) + assert item.message == "hello" + assert stream.retry == 1000 + + +def test_protocol_invalid_retry(client: SseClient): + with client.protocol.invalid_retry() as stream: + item = next(stream) + assert isinstance(item, ProtocolInfo) + assert item.message == "hello" + assert stream.retry is None + + +def test_protocol_reconnect(client: SseClient): + with client.protocol.reconnect() as stream: + first = next(stream) + assert isinstance(first, ProtocolInfo) + assert first.message == "hello" + assert stream.last_event_id == "event-1" + + with client.protocol.reconnect(last_event_id="event-1") as stream: + second = next(stream) + assert isinstance(second, ProtocolInfo) + assert second.message == "world" + assert stream.last_event_id == "event-2" + + +# --------------------------------------------------------------------------- +# Named / model terminal events (yield-then-stop). +# +# The published Spector SSE specs only cover the bare string-constant `[DONE]` +# sentinel (drop-and-stop, exercised above). A ``@terminalEvent`` can also be a +# *named* event carrying a model payload (e.g. `response.completed`, `error`): +# such an event IS deserialized and yielded, and iteration stops immediately +# afterwards. No published spec produces that shape, so we drive the generated +# ``Stream`` runtime directly with a fake response instead of the mock server. +# --------------------------------------------------------------------------- + + +class _FakeResponse: + """A minimal HttpResponse-shaped stand-in that replays SSE bytes.""" + + def __init__(self, body: bytes): + self.headers = {"Content-Type": "text/event-stream"} + self._body = body + self.closed = False + + def iter_bytes(self): + # Emit in small chunks so incremental SSE framing is exercised. + for index in range(0, len(self._body), 8): + yield self._body[index : index + 8] + + def close(self): + self.closed = True + + +def _event_kind(_response, event): + return (event.event, json.loads(event.data)) + + +_NAMED_TERMINAL_SSE = ( + b'event: response.partial\ndata: {"text": "one"}\n\n' + b'event: response.delta\ndata: {"delta": "hi"}\n\n' + b'event: response.completed\ndata: {"references": []}\n\n' + b'event: response.delta\ndata: {"delta": "AFTER-TERMINAL"}\n\n' +) + +_TERMINAL_EVENT_NAMES = ["response.completed", "error"] + + +def test_named_terminal_event_yields_then_stops(): + response = _FakeResponse(_NAMED_TERMINAL_SSE) + stream = Stream( + response=response, + deserialization_callback=_event_kind, + terminal_event_names=_TERMINAL_EVENT_NAMES, + ) + items = list(stream) + # The named terminal `response.completed` IS yielded (it carries a payload), + # then iteration stops -- the trailing `response.delta` must not appear. + assert items == [ + ("response.partial", {"text": "one"}), + ("response.delta", {"delta": "hi"}), + ("response.completed", {"references": []}), + ] + assert response.closed + + +def test_named_terminal_event_first_stops_immediately(): + body = b'event: error\ndata: {"code": "boom"}\n\n' b'event: response.delta\ndata: {"delta": "AFTER-ERROR"}\n\n' + response = _FakeResponse(body) + stream = Stream( + response=response, + deserialization_callback=_event_kind, + terminal_event_names=_TERMINAL_EVENT_NAMES, + ) + items = list(stream) + assert items == [("error", {"code": "boom"})] + assert response.closed + + +def test_unnamed_terminal_event_yields_then_stops(): + body = b'data: {"status": "done"}\n\n' b'data: {"status": "AFTER-TERMINAL"}\n\n' + response = _FakeResponse(body) + stream = Stream( + response=response, + deserialization_callback=_event_kind, + terminal_event_names=["message"], + ) + + assert list(stream) == [("message", {"status": "done"})] + assert response.closed + + +def test_unnamed_terminal_predicate_preserves_service_event_type(): + body = b'data: {"kind": "progress"}\n\n' b'data: {"kind": "complete"}\n\n' b'data: {"kind": "after"}\n\n' + response = _FakeResponse(body) + stream = Stream( + response=response, + deserialization_callback=_event_kind, + terminal_event_predicate=lambda event: json.loads(event.data).get("kind") == "complete", + ) + + assert list(stream) == [ + ("message", {"kind": "progress"}), + ("message", {"kind": "complete"}), + ] + assert response.closed + + +def test_sentinel_and_named_terminal_coexist(): + body = ( + b'event: response.delta\ndata: {"delta": "a"}\n\n' + b"data: [DONE]\n\n" + b'event: response.delta\ndata: {"delta": "AFTER-DONE"}\n\n' + ) + response = _FakeResponse(body) + stream = Stream( + response=response, + deserialization_callback=_event_kind, + terminal_event="[DONE]", + terminal_event_names=_TERMINAL_EVENT_NAMES, + ) + items = list(stream) + # The bare `[DONE]` sentinel stops iteration WITHOUT being yielded. + assert items == [("response.delta", {"delta": "a"})] + assert response.closed + + +def test_sse_protocol_metadata_is_available(): + body = b'id: event-1\nretry: 1000\nevent: message\ndata: {"message": "hello"}\n\n' + response = _FakeResponse(body) + stream = Stream(response=response, deserialization_callback=lambda _response, event: event.data) + + assert list(stream) == ['{"message": "hello"}'] + assert stream.last_event_id == "event-1" + assert stream.retry == 1000 + + +def test_sse_protocol_invalid_metadata_is_ignored(): + body = b"id: invalid\x00id\nretry: not-a-number\nevent: message\ndata: hello\n\n" + response = _FakeResponse(body) + stream = Stream(response=response, deserialization_callback=lambda _response, event: event.data) + + assert list(stream) == ["hello"] + assert stream.last_event_id == "" + assert stream.retry is None diff --git a/packages/http-client-python/tests/unit/test_sse_streaming.py b/packages/http-client-python/tests/unit/test_sse_streaming.py new file mode 100644 index 00000000000..04795806926 --- /dev/null +++ b/packages/http-client-python/tests/unit/test_sse_streaming.py @@ -0,0 +1,143 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- +from types import SimpleNamespace + +import pytest + +from pygen.codegen.models.response import ( + StreamingEvent, + _get_terminal_event_names, + get_streaming_event_discriminator, +) +from pygen.codegen.serializers.builder_serializer import ( + OperationSerializer, + _sse_fallback_data_expression, +) + + +@pytest.fixture(scope="session", autouse=True) +def testserver(): + yield + + +def _streaming_event(event_type=None, *, is_terminal=False, payload_type=None): + return StreamingEvent( + event_type=event_type, + payload_type=payload_type, # type: ignore[arg-type] + payload_content_type="application/json", + is_terminal=is_terminal, + ) + + +def _discriminated_payload(value, annotation="Payload"): + return SimpleNamespace( + discriminator_property=SimpleNamespace(wire_name="kind"), + discriminator_value=value, + type_annotation=lambda **kwargs: annotation, + ) + + +def test_single_unnamed_terminal_event_uses_default_sse_event_name(): + events = [ + _streaming_event(is_terminal=True), + ] + + assert _get_terminal_event_names(events) == ["message"] + + +def test_multiple_unnamed_events_cannot_identify_terminal_by_event_name(): + events = [ + _streaming_event(), + _streaming_event(is_terminal=True), + ] + + assert _get_terminal_event_names(events) == [] + + +def test_multiple_unnamed_events_expose_common_discriminator(): + connected = _streaming_event(payload_type=_discriminated_payload("connected")) + disconnected = _streaming_event( + payload_type=_discriminated_payload("disconnected"), + is_terminal=True, + ) + + assert get_streaming_event_discriminator([connected, disconnected]) == ( + "kind", + [("connected", connected), ("disconnected", disconnected)], + ) + + +def test_inconsistent_unnamed_discriminators_are_ambiguous(): + events = [ + _streaming_event(payload_type=_discriminated_payload("connected")), + _streaming_event(payload_type=SimpleNamespace(discriminator_property=None)), + ] + + assert get_streaming_event_discriminator(events) is None + + +@pytest.mark.parametrize( + ("async_mode", "stream_class"), + [(False, "Stream"), (True, "AsyncStream")], +) +def test_generated_unnamed_discriminator_dispatch_and_terminal_predicate(async_mode, stream_class): + connected = _streaming_event(payload_type=_discriminated_payload("connected", "_models.Connected")) + disconnected = _streaming_event( + payload_type=_discriminated_payload("disconnected", "_models.Disconnected"), + is_terminal=True, + ) + response = SimpleNamespace( + is_structured_stream=True, + streaming_kind="sse", + streaming_events=[connected, disconnected], + terminal_event=None, + terminal_event_names=[], + stream_item_annotation=lambda **kwargs: "Union[_models.Connected, _models.Disconnected]", + stream_class_name=lambda is_async: "AsyncStream" if is_async else "Stream", + ) + code_model = SimpleNamespace( + options={"models-mode": "dpg"}, + get_serialize_namespace=lambda *args, **kwargs: "test", + ) + serializer = OperationSerializer(code_model, async_mode=async_mode, client_namespace="test") + + generated = "\n".join(serializer.handle_structured_stream_response(SimpleNamespace(responses=[response]))) + + assert ("if isinstance(_event_json, dict) and " "_event_json.get('kind') == 'connected':") in generated + assert "_deserialize(_models.Connected, _event_json)" in generated + assert ("elif isinstance(_event_json, dict) and " "_event_json.get('kind') == 'disconnected':") in generated + assert "_deserialize(_models.Disconnected, _event_json)" in generated + assert "def _is_terminal_event(_event):" in generated + assert "_event_json.get('kind') in ['disconnected']" in generated + assert f"deserialized: {stream_class}[" in generated + assert "terminal_event_predicate=_is_terminal_event" in generated + assert generated.count("return cls(pipeline_response, deserialized, {})") == 1 + assert 'raise ValueError(f"Unknown SSE event type: {_event.event!r}")' in generated + assert not any(line.strip().startswith("_event.event =") for line in generated.splitlines()) + + +def test_named_terminal_event_uses_explicit_event_name(): + events = [ + _streaming_event("progress"), + _streaming_event("complete", is_terminal=True), + ] + + assert _get_terminal_event_names(events) == ["complete"] + + +@pytest.mark.parametrize( + ("content_types", "expected"), + [ + ([], "_event.data"), + (["application/json", "application/vnd.example+json"], "json.loads(_event.data)"), + (["text/plain", "text/csv"], "_event.data"), + (["application/json", "text/plain"], "_event.data"), + ], +) +def test_sse_fallback_data_expression(content_types, expected): + events = [SimpleNamespace(payload_content_type=content_type) for content_type in content_types] + + assert _sse_fallback_data_expression(events) == expected