Coverage for agentos/core/streaming_optimizer.py: 0%

1""" 

2AgentOS Streaming Optimizer — SSE Stream Processing & Backpressure 

3━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 

4 

5Production-grade streaming optimization for LLM responses: 

6 

7 - Chunk aggregation (avoid 1-token-at-a-time jitter) 

8 - Backpressure handling (pause/resume flow control) 

9 - Stream transformation pipeline (chains of transformers) 

10 - Token counting and rate estimation 

11 - Adaptive chunk sizing based on network conditions 

12 - Stream cancellation and cleanup 

13 

14Usage: 

15 optimizer = StreamingOptimizer() 

16 async for chunk in optimizer.optimize(llm_stream): 

17 yield chunk 

18""" 

19 

20from __future__ import annotations 

21 

22import asyncio 

23import time 

24from collections.abc import AsyncIterable, AsyncIterator, Callable 

25from dataclasses import dataclass, field 

26from enum import StrEnum 

27from typing import Any 

28 

29# --------------------------------------------------------------------------- 

30# Stream Chunk 

31# --------------------------------------------------------------------------- 

32 

33 

34@dataclass 

35class StreamChunk: 

36 """A single chunk from an LLM stream.""" 

37 

38 content: str 

39 index: int 

40 timestamp: float = field(default_factory=time.time) 

41 is_final: bool = False 

42 metadata: dict[str, Any] = field(default_factory=dict) 

43 

44 

45# --------------------------------------------------------------------------- 

46# Configuration 

47# --------------------------------------------------------------------------- 

48 

49 

50class AggregationStrategy(StrEnum): 

51 """How to aggregate small chunks.""" 

52 

53 NONE = "none" # Pass through as-is 

54 FIXED_SIZE = "fixed" # Wait for N tokens before emitting 

55 TIME_WINDOW = "time" # Emit every T milliseconds 

56 ADAPTIVE = "adaptive" # Adjust based on latency 

57 

58 

59@dataclass 

60class StreamConfig: 

61 """Configuration for streaming optimization.""" 

62 

63 strategy: AggregationStrategy = AggregationStrategy.ADAPTIVE 

64 min_chunk_size: int = 3 # Minimum tokens before emitting 

65 max_chunk_size: int = 50 # Maximum tokens in a single chunk 

66 time_window_ms: int = 50 # Max wait time between emissions 

67 buffer_max_chunks: int = 100 # Max unacknowledged chunks (backpressure) 

68 adaptive_latency_target_ms: int = 100 # Target round-trip latency 

69 adaptive_min_chunk_size: int = 1 

70 adaptive_max_chunk_size: int = 100 

71 enable_compression: bool = False 

72 track_performance: bool = True 

73 

74 

75# --------------------------------------------------------------------------- 

76# Performance Tracker 

77# --------------------------------------------------------------------------- 

78 

79 

80@dataclass 

81class StreamMetrics: 

82 """Stream performance metrics.""" 

83 

84 total_chunks_received: int = 0 

85 total_chunks_emitted: int = 0 

86 total_tokens_received: int = 0 

87 total_tokens_emitted: int = 0 

88 avg_chunk_latency_ms: float = 0.0 

89 total_wall_time_ms: float = 0.0 

90 backpressure_events: int = 0 

91 peak_buffer_size: int = 0 

92 tokens_per_second: float = 0.0 

93 

94 @property 

95 def aggregation_ratio(self) -> float: 

96 if self.total_chunks_received == 0: 

97 return 1.0 

98 return self.total_chunks_received / max(1, self.total_chunks_emitted) 

99 

100 

101class MetricsCollector: 

102 """Collect streaming performance metrics.""" 

103 

104 def __init__(self): 

105 self._metrics = StreamMetrics() 

106 self._latencies: list[float] = [] 

107 

108 def record_received(self, tokens: int = 1) -> None: 

109 self._metrics.total_chunks_received += 1 

110 self._metrics.total_tokens_received += tokens 

111 

112 def record_emitted(self, tokens: int = 1, latency_ms: float = 0.0) -> None: 

113 self._metrics.total_chunks_emitted += 1 

114 self._metrics.total_tokens_emitted += tokens 

115 self._latencies.append(latency_ms) 

116 

117 def record_backpressure(self) -> None: 

118 self._metrics.backpressure_events += 1 

119 

120 def record_buffer_size(self, size: int) -> None: 

121 if size > self._metrics.peak_buffer_size: 

122 self._metrics.peak_buffer_size = size 

123 

124 def finalize(self, wall_time_ms: float) -> StreamMetrics: 

125 self._metrics.total_wall_time_ms = wall_time_ms 

126 if self._latencies: 

127 self._metrics.avg_chunk_latency_ms = sum(self._latencies) / len(self._latencies) 

128 if wall_time_ms > 0: 

129 self._metrics.tokens_per_second = self._metrics.total_tokens_emitted / ( 

130 wall_time_ms / 1000 

131 ) 

132 return self._metrics 

133 

134 

135# --------------------------------------------------------------------------- 

136# Stream Transformer 

137# --------------------------------------------------------------------------- 

138 

139 

140StreamTransformer = Callable[[StreamChunk], StreamChunk | None] 

141 

142 

143class TransformerPipeline: 

144 """Chain of stream transformers applied to each chunk.""" 

145 

146 def __init__(self): 

147 self._transformers: list[StreamTransformer] = [] 

148 

149 def add(self, transformer: StreamTransformer) -> TransformerPipeline: 

150 self._transformers.append(transformer) 

151 return self 

152 

153 def apply(self, chunk: StreamChunk) -> StreamChunk | None: 

154 current = chunk 

155 for t in self._transformers: 

156 if current is None: 

157 return None 

158 current = t(current) 

159 return current 

160 

161 @property 

162 def size(self) -> int: 

163 return len(self._transformers) 

164 

165 

166# Built-in transformers 

167 

168 

169def strip_leading_whitespace() -> StreamTransformer: 

170 """Remove leading whitespace from first chunk.""" 

171 first = True 

172 

173 def transformer(chunk: StreamChunk) -> StreamChunk: 

174 nonlocal first 

175 if first: 

176 chunk.content = chunk.content.lstrip() 

177 first = False 

178 return chunk 

179 

180 return transformer 

181 

182 

183def normalize_newlines() -> StreamTransformer: 

184 """Normalize all line endings to '\n'.""" 

185 

186 def transformer(chunk: StreamChunk) -> StreamChunk: 

187 chunk.content = chunk.content.replace("\r\n", "\n").replace("\r", "\n") 

188 return chunk 

189 

190 return transformer 

191 

192 

193def filter_empty_chunks() -> StreamTransformer: 

194 """Drop chunks with empty content.""" 

195 

196 def transformer(chunk: StreamChunk) -> StreamChunk | None: 

197 return chunk if chunk.content else None 

198 

199 return transformer 

200 

201 

202def add_token_count() -> StreamTransformer: 

203 """Add approximate token count to metadata.""" 

204 

205 def transformer(chunk: StreamChunk) -> StreamChunk: 

206 # Rough estimate: ~1.3 chars per token 

207 chunk.metadata["approx_tokens"] = max(1, int(len(chunk.content) / 1.3)) 

208 return chunk 

209 

210 return transformer 

211 

212 

213# --------------------------------------------------------------------------- 

214# Streaming Optimizer 

215# --------------------------------------------------------------------------- 

216 

217 

218class StreamingOptimizer: 

219 """ 

220 Optimize LLM streaming with aggregation, backpressure, and metrics. 

221 

222 Usage: 

223 opt = StreamingOptimizer() 

224 async for chunk in opt.optimize(llm_stream): 

225 yield chunk.content 

226 

227 With transformers: 

228 opt = StreamingOptimizer() 

229 opt.pipeline.add(strip_leading_whitespace()) 

230 async for chunk in opt.optimize(llm_stream): 

231 ... 

232 """ 

233 

234 def __init__(self, config: StreamConfig | None = None): 

235 self._config = config or StreamConfig() 

236 self._pipeline = TransformerPipeline() 

237 self._buffer: list[StreamChunk] = [] 

238 self._token_accumulator: list[str] = [] 

239 self._metrics = MetricsCollector() 

240 self._start_time: float | None = None 

241 self._paused = False 

242 

243 @property 

244 def pipeline(self) -> TransformerPipeline: 

245 return self._pipeline 

246 

247 @property 

248 def config(self) -> StreamConfig: 

249 return self._config 

250 

251 async def optimize(self, stream: AsyncIterable[StreamChunk]) -> AsyncIterator[StreamChunk]: 

252 """ 

253 Optimize a stream of chunks. 

254 

255 Applies aggregation and transformer pipeline. 

256 """ 

257 self._start_time = time.time() 

258 

259 async for chunk in stream: 

260 self._metrics.record_received() 

261 

262 # Backpressure: wait if buffer is full 

263 while len(self._buffer) >= self._config.buffer_max_chunks: 

264 self._metrics.record_backpressure() 

265 await asyncio.sleep(0.01) 

266 

267 self._buffer.append(chunk) 

268 self._metrics.record_buffer_size(len(self._buffer)) 

269 

270 # Check if we should emit 

271 result = await self._maybe_emit() 

272 if result is not None: 

273 yield result 

274 

275 # Flush remaining buffer 

276 async for chunk in self._flush(): 

277 yield chunk 

278 

279 self._metrics.finalize((time.time() - self._start_time) * 1000) 

280 

281 async def optimize_simple(self, text_stream: AsyncIterable[str]) -> AsyncIterator[str]: 

282 """ 

283 Optimize a simple text stream (strings instead of StreamChunk objects). 

284 """ 

285 async for chunk in self.optimize( 

286 StreamChunk(content=text, index=i) for i, text in enumerate(text_stream) 

287 ): 

288 yield chunk.content 

289 

290 async def _maybe_emit(self) -> StreamChunk | None: 

291 """Check if we should emit an aggregated chunk.""" 

292 if not self._buffer: 

293 return None 

294 

295 tokens = sum(chunk.metadata.get("approx_tokens", 1) for chunk in self._buffer) 

296 

297 should_emit = False 

298 

299 if self._config.strategy == AggregationStrategy.NONE: 

300 should_emit = True 

301 elif self._config.strategy == AggregationStrategy.FIXED_SIZE: 

302 should_emit = tokens >= self._config.min_chunk_size 

303 elif self._config.strategy == AggregationStrategy.TIME_WINDOW: 

304 if self._buffer: 

305 elapsed = (time.time() - self._buffer[0].timestamp) * 1000 

306 should_emit = elapsed >= self._config.time_window_ms 

307 elif self._config.strategy == AggregationStrategy.ADAPTIVE: 

308 should_emit = tokens >= self._config.adaptive_min_chunk_size and ( 

309 tokens >= self._config.max_chunk_size or self._buffer[-1].is_final 

310 ) 

311 

312 if not should_emit: 

313 return None 

314 

315 return self._emit_aggregated() 

316 

317 def _emit_aggregated(self) -> StreamChunk | None: 

318 """Aggregate buffered chunks into a single emission.""" 

319 if not self._buffer: 

320 return None 

321 

322 # Aggregate content 

323 content = "".join(c.content for c in self._buffer) 

324 is_final = self._buffer[-1].is_final 

325 index = self._buffer[-1].index 

326 timestamp = time.time() 

327 

328 # Clear buffer 

329 count = len(self._buffer) 

330 self._buffer.clear() 

331 

332 # Build aggregated chunk 

333 chunk = StreamChunk( 

334 content=content, 

335 index=index, 

336 timestamp=timestamp, 

337 is_final=is_final, 

338 metadata={"aggregated_from": count}, 

339 ) 

340 

341 # Apply transformer pipeline 

342 chunk = self._pipeline.apply(chunk) 

343 if chunk is None: 

344 return None 

345 

346 # Record metrics 

347 latency_ms = (timestamp - self._start_time) * 1000 if self._start_time else 0 

348 approx_tokens = chunk.metadata.get("approx_tokens", 1) 

349 self._metrics.record_emitted(tokens=approx_tokens, latency_ms=latency_ms) 

350 

351 return chunk 

352 

353 async def _flush(self) -> AsyncIterator[StreamChunk]: 

354 """Flush remaining buffered chunks.""" 

355 while self._buffer: 

356 chunk = self._emit_aggregated() 

357 if chunk is not None: 

358 yield chunk 

359 

360 def get_metrics(self) -> StreamMetrics: 

361 return self._metrics._metrics 

362 

363 def reset_metrics(self) -> None: 

364 self._metrics = MetricsCollector() 

365 self._start_time = None