Azure / Azure/azure-sdk-for-python

add support for SSE

Abierto
#44,906 0 comentarios 0 reacciones 1 asignado Reclamado por @iscai-msft Ver en GitHub
Azure.Core
Lenguaje dominante
Python
Estrellas
5.6k
Forks
3.4k
Merge medio
2 d 2 h
PR fusionados (30 d)
202

Descripción

copy paste of @kristapratico GitHub issue :*

# Streaming

This gist provides a draft of what the streaming implementation could look like for JSONL and Server Sent Events (SSE) implemented in Python.

Stream response examples: [Examples](https://gist.github.com/kristapratico/0c7aad76a0fe19959ac7d1063b02a050)

TypeSpec Reference: [TypeSpec Streams Sync](https://gist.github.com/chrisradek/b2656d7ee5db1b29768d50db0628f033)

## Stream interface

Provide a generic, stream-agnostic `Stream` iterator in azure-core. It iterates over http_response.iter_bytes(), parsing and decoding chunks using the appropriate decoder class. When an event is ready to be dispatched, it calls deserialization_callback to deserialize the JSON to the model returned to the user.

In azure.core:

> Note: below code is just a prototype I was playing with to try design concept with SSE streaming. It does not cover all scenarios.

```python
from typing import Iterator, TypeVar, Type

ReturnType = TypeVar("ReturnType")

class Stream(Iterator[ReturnType]):
"""Stream class.

:keyword response: The response object.
:paramtype response: ~azure.core.pipeline.PipelineResponse
:keyword deserialization_callback: A callback that takes HttpResponse and returns a deserialized object
:paramtype deserialization_callback: Callable
:keyword terminal_event: A terminal event that indicates the end of the SSE stream.
:paramtype terminal_event: Optional[str]
"""
def __init__(
self,
*,
response: PipelineResponse,
deserialization_callback: Callable[[Any, Any], ReturnType],
terminal_event: Optional[str] = None,
) -> None:
self._response = response.http_response
self._decoder = SSEDecoder() if self._response.headers.get("Content-Type") == "text/event-stream" else JSONLDecoder()
self._deserialization_callback = deserialization_callback
self._terminal_event = terminal_event
self._iterator = self._iter_events()

def __next__(self) -> ReturnType:
return self._iterator.__next__()

def __iter__(self) -> Iterator[ReturnType]:
yield from self._iterator

def _iter_events(self) -> Iterator[ReturnType]:
for line in self._parse_chunk(self._response.iter_bytes()):
for data in line.splitlines():
if data:
self._decoder.decode(data)
else:
event = self._decoder.event()
if self._terminal_event:
if event.data == self._terminal_event:
break

event_model = self._deserialization_callback(self._response, event.json())

yield event_model

def _parse_chunk(self, iter_bytes: Iterator[bytes]) -> Iterator[str]:
data = b''
for chunk in iter_bytes:
for line in chunk.splitlines(keepends=True):
data += line
if data.endswith((b'\r\r', b'\n\n', b'\r\n\r\n')):
yield data.decode("utf-8")
data = b''

def __exit__(
self,
exc_type: type[BaseException] | None,
exc: BaseException | None,
exc_tb: TracebackType | None,
) -> None:
self.close()

def __enter__(self) -> Self:
return self

def close(self) -> None:
self._response.close()
```

## Decoders

Provide separate decoder classes for JSONL and SSE. These aim to collect and decode the data per stream chunk. Each decoder has a `decode` method which takes str and decodes per the spec, and an `event` method which creates and returns a SSE or JSONL event.

**SSEDecoder**

```python
class SSEDecoder:
def __init__(self) -> None:
self.data: list[str] = []
self.last_event_id = None
self.event = None
self.retry = None

def decode(self, line: str) -> None:
if line.startswith(":"):
# comment, ignore the line
return None

if ":" in line:
field, _, value = line.partition(":")
if value.startswith(" "):
# data:test and data: test are equivalent
value = value[1:]
else:
field = line
value = ""

if field == "data":
self.data.append(value)
elif field == "event":
self.event = value
elif field == "id":
if "\0" in value:
pass
else:
self.last_event_id = value
elif field == "retry":
try:
self.retry = int(value)
except (TypeError, ValueError):
pass

# else: ignore the field

def event(self) -> ServerSentEvent:
sse = ServerSentEvent(
event=self.event,
data="\n".join(self.data),
id=self.last_event_id,
retry=self.retry,
)

self.data = []
self.event = None
self.retry = None
return sse
```

**JSONLDecoder**

```python
class JSONLDecoder:
def __init__(self) -> None:
self.data: list[str] = []

def decode(self, line: str) -> None:
self.data.append(line)

def event(self) -> JSONLEvent:
jsonl = JSONLEvent(data="\n".join(self.data))
self.data = []
return jsonl
```

## Event types

We have 2 types which represent a `ServerSentEvent` and a `JSONLEvent`.

**ServerSentEvent**

```python
class ServerSentEvent:
def __init__(
self,
*,
data: Optional[str] = None,
event: Optional[str] = None,
id: Optional[str] = None,
retry: Optional[int] = None,
) -> None:
self.data = data
self.event = event
self.id = id
self.retry = retry

def json(self) -> Any:
return json.loads(self.data)
```

**JSONLEvent**

```python
class JSONLEvent:
def __init__(
self,
*,
data: Optional[str] = None,
) -> None:
self.data = data

def json(self) -> Any:
return json.loads(self.data)
```

> Note: the only type exposed to users is the `Stream` class. The rest of the above classes are internal implementation details.

## Library-specific code

### Generated for homogeneous events

This example shows JSONL streaming, but similar could work for SSE.
Homogeneous events return the same type for every event.

Models:

```python
class Notification:
id: str
content: str
timestamp: datetime.datetime
```

Operations:

```python

class OperationMixin:
def subscribe(self, **kwargs: Any) -> Stream[Notification]:
pipeline_response: PipelineResponse = self._client._pipeline.run(
request, stream=True, **kwargs
)

def callback(pipeline_response, model_json):
deserialized = _deserialize(Notification, model_json)
if cls:
return cls(pipeline_response, deserialized, {})
return deserialized

return Stream[Notification](
response=pipeline_response,
deserialization_callback=callback
)
```

User code:

```python

response = client.subscribe(..., stream=True)
for item in response:
print(item) # Notification
```

### Generated for heterogeneous events

Heterogeneous events can have more than one event type.

Models:

```python
class ChannelEvent(_model_base.Model):
kind: str

class UserConnect(ChannelEvent, discriminator="userconnect"):
kind: Literal["userconnect"]
username: str
time: str

class UserMessage(ChannelEvent, discriminator="usermessage"):
kind: Literal["usermessage"]
username: str
time: str
text: str

class UserDisconnect(ChannelEvent, discriminator="userdisconnect"):
kind: Literal["userdisconnect"]
username: str
time: str

ChannelEvents: TypeAlias = Union[UserConnect, UserMessage, UserDisconnect]
```

Operations:

```python

class OperationMixin:
def subscribe(self, **kwargs: Any) -> Stream[ChannelEvents]:
pipeline_response: PipelineResponse = self._client._pipeline.run(
request, stream=True, **kwargs
)

def callback(pipeline_response, model_json):
deserialized = _deserialize(ChannelEvent, model_json)
if cls:
return cls(pipeline_response, deserialized, {})
return deserialized

return Stream[ChannelEvents](
response=pipeline_response,
deserialization_callback=callback,
terminal_event="[DONE]",
)
```

User code:

```python
events: ChannelEvents = client.subscribe(..., stream=True)
for event in events:
if event.kind == "userconnect":
print("User is connecting...")
elif event.kind == "usermessage":
print("hello")
elif event.kind == "userdisconnect":
print("goodbye")
```

## Handling Events

> **NOTE**: For the initial implementation we will not be adding any convenience for handling events like outlined below. We will wait for feedback and assess later.

**1) Allow the user to register callbacks for event types**

Provide a method on the `Stream` class which lets you register a callback for a specific event type. When there is an event of that type, the provided callback is called with the ChannelEvents variant.

```python
class Stream:
# ...insert implementation from above section...

def add_callback(self, event_type: str, on_event: Callable[[ReturnType], None]) -> None:
self._callbacks[event_type] = callback

def iter_events(self) -> ReturnType:
"""Convenience method"""
for event in self._iterator:
yield event
```

User code:

```python
def say_goodbye(user_disconnect):
print("Goodbye!")

response = client.subscribe(..., stream=True)
response.add_callback(event_type="userdisconnect", say_goodbye)

# Provide convenience method that iterates over events so users don't have to iterate themselves
response.iter_events()
```

**2) Provide an EventHandler interface which a user can subclass and access specific events when they are dispatched**

For services that expose multiple events, we can provide additional convenience by generating a tailored EventHandler class (the method names come from the event types, may need to do some language-specific conversion to snake_case, camelCase, etc).

Models:

```python
class ChannelEventHandler:

def on_userconnect(self, userconnect: UserConnect) -> None:
pass

def on_usermessage(self, usermessage: UserMessage) -> None:
pass

def on_userdisconnect(self, userdisconnect: UserDisconnect) -> None:
pass
```

Operations:

```python
class OperationsMixin:

def subscribe(self, **kwargs: Any) -> Stream[ChannelEvents]:
pipeline_response: PipelineResponse = self._client._pipeline.run(
request, stream=True, **kwargs
)

def callback(pipeline_response, model_json):
deserialized = _deserialize(ChannelEvent, model_json)
if cls:
return cls(pipeline_response, deserialized, {})
return deserialized

return Stream[ChannelEvents](
response=pipeline_response,
deserialization_callback=callback,
event_hander=event_handler or ChannelEventHandler(),
)

```

User code:

A user will define their own Event handler subclass and implement the methods for the events they want act on.

```python

class CustomEventHandler(ChannelEventHandler):

def on_usermessage(self, usermessage: UserMessage) -> None:
# my custom logic for the event
...

stream = client.subscribe(..., event_handler=CustomEventHandler())
# events handled in CustomEventHandler
stream.iter_events()
```

Guía de contribución

Abrir la guía de contribución

Evaluación

Este issue todavía no se ha evaluado.

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.