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..5fa79ea8015 --- /dev/null +++ b/.chronus/changes/structured-streaming-2026-0-0.md @@ -0,0 +1,15 @@ +--- +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: + ... +``` diff --git a/cspell.yaml b/cspell.yaml index fce7a455653..ff5a09eecf9 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 @@ -120,6 +126,7 @@ words: - intrinsics - ints - IOHTTP + - isascii - isdigit - isinstance - issecret diff --git a/packages/http-client-python/README.md b/packages/http-client-python/README.md index 0c3335df5fa..c896aa6a990 100644 --- a/packages/http-client-python/README.md +++ b/packages/http-client-python/README.md @@ -153,3 +153,24 @@ Whether to clear the output folder before generating the code. Defaults to `fals **Type:** `boolean` Emit YAML code model only, without running Python generator. For batch processing. + +## Structured streaming (JSONL / SSE) + +Operations whose HTTP response is a JSONL (`application/jsonl`) or SSE (`text/event-stream`) stream generate client methods that return `Stream[T]` (sync) / `AsyncStream[T]` (async), yielding deserialized model instances instead of raw bytes. This is driven by the TCGC response stream metadata (the response stream type) — there is no opt-in emitter option. This applies to both the Azure and unbranded flavors. + +For an operation returning `JsonlStream`, the generated method returns `Stream[Thing]` (sync) / `AsyncStream[Thing]` (async), yielding deserialized `Thing` instances as each JSONL line arrives. Similarly, `SSEStream` produces a `Stream` / `AsyncStream` over the SSE event payloads. + +Generated packages re-export `Stream` and `AsyncStream` from their base namespace: + +```python +from your_sdk import Stream +from your_sdk.models import Thing + +stream: Stream[Thing] = client.receive() +for thing in stream: + ... +``` + +The `Stream` / `AsyncStream` runtime (plus the JSONL and SSE decoders) is **vendored** into the generated package at `_utils/streaming_base.py` (alongside `_utils/model_base.py`). It depends only on the released core runtime for the flavor — `azure.core.rest` for the Azure flavor and `corehttp.rest` for the unbranded flavor — so no unreleased `azure.core.streaming` dependency is required at runtime. + +SSE `@events` unions use TCGC event metadata to deserialize each named event into its corresponding generated model. Events marked with `@terminalEvent` stop iteration without being yielded. diff --git a/packages/http-client-python/emitter/src/http.ts b/packages/http-client-python/emitter/src/http.ts index 31d970d4b65..5526dfbb47a 100644 --- a/packages/http-client-python/emitter/src/http.ts +++ b/packages/http-client-python/emitter/src/http.ts @@ -42,6 +42,157 @@ export enum ReferredByOperationTypes { NonPagingOnly = 2, } +type StructuredStreamKind = "jsonl" | "sse"; +type EmittedType = ReturnType; + +interface StructuredStreamEvent { + eventType: string | undefined; + /** + * Payload type for this one SSE event. 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}. + */ + itemType: 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; +} + +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.itemType}: 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; +} + +/** The subset of an {@link SdkSseMetadata} event that terminal-event partitioning depends on. */ +interface SseEventLike { + eventType?: string | undefined; + isTerminalEvent: boolean; + type: SdkType; + payloadType: SdkType; +} + +/** + * 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`. + * + * `toItemType` maps an event payload to its emitted type; 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 SseEventLike[], + toItemType: (payloadType: 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, + itemType: toItemType(event.payloadType), + isTerminal: true, + }); + continue; + } + dispatch.push({ eventType: event.eventType, itemType: toItemType(event.payloadType) }); + } + 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), + }; + + const sseMetadata = response.sseMetadata; + if (!sseMetadata) return streaming; + + const { events, terminalEvent } = partitionSseEvents(sseMetadata.events, (payloadType) => + getType(context, payloadType), + ); + 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 +833,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..f0c03f4ac52 --- /dev/null +++ b/packages/http-client-python/emitter/test/streaming.test.ts @@ -0,0 +1,135 @@ +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].itemType, completed); + strictEqual(events[2].eventType, "error"); + strictEqual(events[2].isTerminal, true); + strictEqual(events[2].itemType, 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); + }); + }); +}); 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..5db0e8aed85 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) 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..724db0db50a 100644 --- a/packages/http-client-python/generator/pygen/codegen/models/response.py +++ b/packages/http-client-python/generator/pygen/codegen/models/response.py @@ -58,6 +58,24 @@ 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[tuple[Optional[str], BaseType]] = [] + self._streaming_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 (``_streaming_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 = [ + (event.get("eventType"), self.code_model.lookup_type(id(event["itemType"]))) + for event in streaming.get("events", []) + ] + self.terminal_event_names = [ + event["eventType"] + for event in streaming.get("events", []) + if event.get("isTerminal") and event.get("eventType") is not None + ] @property def result_property(self) -> str: @@ -92,12 +110,65 @@ 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 + + @property + def terminal_event(self) -> Optional[str]: + """Terminal event marker for a heterogeneous SSE stream, if any. + + TCGC ``sseMetadata`` supplies this marker directly. For compatibility with + older metadata, a string-literal member of the item union is used as a fallback. + """ + if self.streaming_kind != "sse": + return None + if self._streaming_terminal_event is not None: + return self._streaming_terminal_event + if not isinstance(self.type, CombinedType): + return None + from .constant_type import ConstantType + + for member in self.type.types: + if isinstance(member, ConstantType) and isinstance(member.value, str): + return member.value + return 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]: + if len(self.streaming_events) == 1: + return self.streaming_events[0][1] + 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 +180,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_type, event_item_type in self.streaming_events: + file_import.merge(event_item_type.imports(**kwargs)) return file_import def _get_import_type(self, input_path: str) -> ImportType: @@ -143,6 +247,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..a09bc2e96a9 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 @@ -1255,15 +1255,119 @@ 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", []) + retval: list[str] = [] + retval.append("def _callback(_http_response, _event):") + if response.streaming_kind == "sse": # type: ignore[attr-defined] + retval.append(" _event_json = json.loads(_event.data)") + named_events = [ + (event_type, event_item_type) + for event_type, event_item_type in streaming_events + if event_type is not None + ] + unnamed_events = [event_item_type for event_type, event_item_type in streaming_events if event_type is None] + if named_events: + for index, (event_type, event_item_type) in enumerate(named_events): + event_annotation = event_item_type.type_annotation( + is_operation_file=True, + serialize_namespace=self.serialize_namespace, + ) + keyword = "if" if index == 0 else "elif" + retval.append(f" {keyword} _event.event == {event_type!r}:") + if self.code_model.options["models-mode"] == "msrest": + serialization_type = event_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({event_annotation}, _event_json)") + retval.append(" else:") + if len(unnamed_events) == 1: + event_annotation = unnamed_events[0].type_annotation( + is_operation_file=True, + serialize_namespace=self.serialize_namespace, + ) + if self.code_model.options["models-mode"] == "msrest": + serialization_type = unnamed_events[0].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({event_annotation}, _event_json)") + else: + retval.append(" deserialized = _event_json") + elif len(unnamed_events) == 1: + event_annotation = unnamed_events[0].type_annotation( + is_operation_file=True, + serialize_namespace=self.serialize_namespace, + ) + if self.code_model.options["models-mode"] == "msrest": + serialization_type = unnamed_events[0].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({event_annotation}, _event_json)") + else: + 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)") + 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(" if cls:") + retval.append(" return cls(pipeline_response, deserialized, {}) # type: ignore") + 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}") + retval.append(f"return {stream_class}({', '.join(stream_kwargs)}) # type: ignore") + 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 +1402,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..5246167751e --- /dev/null +++ b/packages/http-client-python/generator/pygen/codegen/templates/streaming_base.py.jinja2 @@ -0,0 +1,707 @@ +# 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 + """ + + 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, + ) -> 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._iterator = self._iter_results() + + 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()): + 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: + 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 + """ + + 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, + ) -> 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._iterator = self._iter_results() + + 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: + 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: + break + finally: + aclose = getattr(events, "aclose", None) + if aclose is not None: + await aclose() + 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/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..e8eff299908 --- /dev/null +++ b/packages/http-client-python/tests/mock_api/shared/asynctests/test_streaming_sse_async.py @@ -0,0 +1,132 @@ +# ------------------------------------------------------------------------- +# 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.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): + stream = await client.unnamed.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"] + + +@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"] + + +# --------------------------------------------------------------------------- +# 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 + + +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_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 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..5be85663b1f --- /dev/null +++ b/packages/http-client-python/tests/mock_api/shared/test_streaming_sse.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. +# -------------------------------------------------------------------------- +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.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): + stream = client.unnamed.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"] + + +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"] + + +# --------------------------------------------------------------------------- +# 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_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