-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathstreaming.py
More file actions
212 lines (172 loc) · 7.48 KB
/
Copy pathstreaming.py
File metadata and controls
212 lines (172 loc) · 7.48 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
"""Streaming support for Hawk chat responses via SSE."""
from __future__ import annotations
from typing import TYPE_CHECKING
from .types import StreamEvent, StreamEventType, ToolCall
if TYPE_CHECKING:
from collections.abc import AsyncIterator, Iterator
import httpx
# Cap accumulated bytes for a single SSE event to guard against a
# misbehaving or malicious daemon/proxy sending an unterminated stream of
# `data:` lines and exhausting client memory.
MAX_EVENT_SIZE = 10 * 1024 * 1024
class StreamEventTooLargeError(ValueError):
"""Raised when a single SSE event's accumulated data exceeds MAX_EVENT_SIZE."""
class StreamReader:
"""Reads SSE events from a streaming chat response (sync).
Wraps an httpx response that returns server-sent events.
"""
def __init__(self, response: httpx.Response) -> None:
self._response = response
self._lines: Iterator[str] = response.iter_lines()
self._closed = False
def __iter__(self) -> Iterator[StreamEvent]:
"""Iterate over raw stream events."""
return self.events()
def __enter__(self) -> StreamReader:
return self
def __exit__(self, *args: object) -> None:
self.close()
def events(self) -> Iterator[StreamEvent]:
"""Yield typed StreamEvent objects from the SSE stream.
The underlying httpx response is closed when iteration completes,
when the caller breaks early, or when iteration raises — so a
``break`` inside a ``for event in reader.events()`` loop does not
leak the connection.
"""
current_event: str | None = None
current_data: str | None = None
try:
for line in self._lines:
if line == "":
# Empty line signals end of an event
if current_data is not None:
yield StreamEvent(event=current_event, data=current_data)
current_event = None
current_data = None
continue
if line.startswith("event: "):
current_event = line[7:]
elif line.startswith("data: "):
if current_data is not None:
current_data += "\n" + line[6:]
else:
current_data = line[6:]
elif line == "data:":
if current_data is not None:
current_data += "\n"
else:
current_data = ""
if current_data is not None and len(current_data) > MAX_EVENT_SIZE:
raise StreamEventTooLargeError(
f"SSE event data exceeded {MAX_EVENT_SIZE} bytes without a terminating blank line"
)
finally:
self.close()
def collect_text(self) -> str:
"""Consume the entire stream and return concatenated text content."""
parts: list[str] = []
for event in self.events():
if event.event is None or event.event == StreamEventType.CONTENT:
parts.append(event.data)
return "".join(parts)
def collect_tool_calls(self) -> list[ToolCall]:
"""Consume the stream and assemble tool call deltas into complete calls."""
import json
tool_calls: list[ToolCall] = []
for event in self.events():
if event.event == StreamEventType.TOOL_CALL:
try:
data = json.loads(event.data)
tool_calls.append(
ToolCall(
id=data.get("id", ""),
name=data.get("name", ""),
arguments=data.get("arguments", {}),
)
)
except (json.JSONDecodeError, KeyError, AttributeError, TypeError):
continue
return tool_calls
def close(self) -> None:
"""Close the underlying response."""
if not self._closed:
self._response.close()
self._closed = True
class AsyncStreamReader:
"""Reads SSE events from an async streaming chat response.
Wraps an httpx async response that returns server-sent events.
"""
def __init__(self, response: httpx.Response) -> None:
self._response = response
self._lines: AsyncIterator[str] = response.aiter_lines()
self._closed = False
def __aiter__(self) -> AsyncIterator[StreamEvent]:
"""Iterate over raw stream events."""
return self.events()
async def __aenter__(self) -> AsyncStreamReader:
return self
async def __aexit__(self, *args: object) -> None:
await self.close()
async def events(self) -> AsyncIterator[StreamEvent]:
"""Yield typed StreamEvent objects from the SSE stream.
The underlying httpx response is closed when iteration completes,
when the caller breaks early, or when iteration raises.
"""
current_event: str | None = None
current_data: str | None = None
try:
async for line in self._lines:
if line == "":
if current_data is not None:
yield StreamEvent(event=current_event, data=current_data)
current_event = None
current_data = None
continue
if line.startswith("event: "):
current_event = line[7:]
elif line.startswith("data: "):
if current_data is not None:
current_data += "\n" + line[6:]
else:
current_data = line[6:]
elif line == "data:":
if current_data is not None:
current_data += "\n"
else:
current_data = ""
if current_data is not None and len(current_data) > MAX_EVENT_SIZE:
raise StreamEventTooLargeError(
f"SSE event data exceeded {MAX_EVENT_SIZE} bytes without a terminating blank line"
)
finally:
await self.close()
async def collect_text(self) -> str:
"""Consume the entire stream and return concatenated text content."""
parts: list[str] = []
async for event in self.events():
if event.event is None or event.event == StreamEventType.CONTENT:
parts.append(event.data)
return "".join(parts)
async def collect_tool_calls(self) -> list[ToolCall]:
"""Consume the stream and assemble tool call deltas."""
import json
tool_calls: list[ToolCall] = []
async for event in self.events():
if event.event == StreamEventType.TOOL_CALL:
try:
data = json.loads(event.data)
tool_calls.append(
ToolCall(
id=data.get("id", ""),
name=data.get("name", ""),
arguments=data.get("arguments", {}),
)
)
except (json.JSONDecodeError, KeyError, AttributeError, TypeError):
continue
return tool_calls
async def close(self) -> None:
"""Close the underlying response."""
if not self._closed:
await self._response.aclose()
self._closed = True