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 import deque 

25from dataclasses import dataclass, field 

26from enum import Enum 

27from typing import Any, AsyncIterable, AsyncIterator, Callable, Dict, List, Optional, Set 

28 

29 

30# --------------------------------------------------------------------------- 

31# Stream Chunk 

32# --------------------------------------------------------------------------- 

33 

34 

35@dataclass 

36class StreamChunk: 

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

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(str, Enum): 

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

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

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

54 TIME_WINDOW = "time" # Emit every T milliseconds 

55 ADAPTIVE = "adaptive" # Adjust based on latency 

56 

57 

58@dataclass 

59class StreamConfig: 

60 """Configuration for streaming optimization.""" 

61 strategy: AggregationStrategy = AggregationStrategy.ADAPTIVE 

62 min_chunk_size: int = 3 # Minimum tokens before emitting 

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

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

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

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

67 adaptive_min_chunk_size: int = 1 

68 adaptive_max_chunk_size: int = 100 

69 enable_compression: bool = False 

70 track_performance: bool = True 

71 

72 

73# --------------------------------------------------------------------------- 

74# Performance Tracker 

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

76 

77 

78@dataclass 

79class StreamMetrics: 

80 """Stream performance metrics.""" 

81 total_chunks_received: int = 0 

82 total_chunks_emitted: int = 0 

83 total_tokens_received: int = 0 

84 total_tokens_emitted: int = 0 

85 avg_chunk_latency_ms: float = 0.0 

86 total_wall_time_ms: float = 0.0 

87 backpressure_events: int = 0 

88 peak_buffer_size: int = 0 

89 tokens_per_second: float = 0.0 

90 

91 @property 

92 def aggregation_ratio(self) -> float: 

93 if self.total_chunks_received == 0: 

94 return 1.0 

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

96 

97 

98class MetricsCollector: 

99 """Collect streaming performance metrics.""" 

100 

101 def __init__(self): 

102 self._metrics = StreamMetrics() 

103 self._latencies: List[float] = [] 

104 

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

106 self._metrics.total_chunks_received += 1 

107 self._metrics.total_tokens_received += tokens 

108 

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

110 self._metrics.total_chunks_emitted += 1 

111 self._metrics.total_tokens_emitted += tokens 

112 self._latencies.append(latency_ms) 

113 

114 def record_backpressure(self) -> None: 

115 self._metrics.backpressure_events += 1 

116 

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

118 if size > self._metrics.peak_buffer_size: 

119 self._metrics.peak_buffer_size = size 

120 

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

122 self._metrics.total_wall_time_ms = wall_time_ms 

123 if self._latencies: 

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

125 if wall_time_ms > 0: 

126 self._metrics.tokens_per_second = ( 

127 self._metrics.total_tokens_emitted / (wall_time_ms / 1000) 

128 ) 

129 return self._metrics 

130 

131 

132# --------------------------------------------------------------------------- 

133# Stream Transformer 

134# --------------------------------------------------------------------------- 

135 

136 

137StreamTransformer = Callable[[StreamChunk], Optional[StreamChunk]] 

138 

139 

140class TransformerPipeline: 

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

142 

143 def __init__(self): 

144 self._transformers: List[StreamTransformer] = [] 

145 

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

147 self._transformers.append(transformer) 

148 return self 

149 

150 def apply(self, chunk: StreamChunk) -> Optional[StreamChunk]: 

151 current = chunk 

152 for t in self._transformers: 

153 if current is None: 

154 return None 

155 current = t(current) 

156 return current 

157 

158 @property 

159 def size(self) -> int: 

160 return len(self._transformers) 

161 

162 

163# Built-in transformers 

164 

165 

166def strip_leading_whitespace() -> StreamTransformer: 

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

168 first = True 

169 

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

171 nonlocal first 

172 if first: 

173 chunk.content = chunk.content.lstrip() 

174 first = False 

175 return chunk 

176 

177 return transformer 

178 

179 

180def normalize_newlines() -> StreamTransformer: 

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

182 

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

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

185 return chunk 

186 

187 return transformer 

188 

189 

190def filter_empty_chunks() -> StreamTransformer: 

191 """Drop chunks with empty content.""" 

192 def transformer(chunk: StreamChunk) -> Optional[StreamChunk]: 

193 return chunk if chunk.content else None 

194 return transformer 

195 

196 

197def add_token_count() -> StreamTransformer: 

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

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

200 # Rough estimate: ~1.3 chars per token 

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

202 return chunk 

203 return transformer 

204 

205 

206# --------------------------------------------------------------------------- 

207# Streaming Optimizer 

208# --------------------------------------------------------------------------- 

209 

210 

211class StreamingOptimizer: 

212 """ 

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

214 

215 Usage: 

216 opt = StreamingOptimizer() 

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

218 yield chunk.content 

219 

220 With transformers: 

221 opt = StreamingOptimizer() 

222 opt.pipeline.add(strip_leading_whitespace()) 

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

224 ... 

225 """ 

226 

227 def __init__(self, config: Optional[StreamConfig] = None): 

228 self._config = config or StreamConfig() 

229 self._pipeline = TransformerPipeline() 

230 self._buffer: List[StreamChunk] = [] 

231 self._token_accumulator: List[str] = [] 

232 self._metrics = MetricsCollector() 

233 self._start_time: Optional[float] = None 

234 self._paused = False 

235 

236 @property 

237 def pipeline(self) -> TransformerPipeline: 

238 return self._pipeline 

239 

240 @property 

241 def config(self) -> StreamConfig: 

242 return self._config 

243 

244 async def optimize( 

245 self, stream: AsyncIterable[StreamChunk] 

246 ) -> AsyncIterator[StreamChunk]: 

247 """ 

248 Optimize a stream of chunks. 

249 

250 Applies aggregation and transformer pipeline. 

251 """ 

252 self._start_time = time.time() 

253 

254 async for chunk in stream: 

255 self._metrics.record_received() 

256 

257 # Backpressure: wait if buffer is full 

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

259 self._metrics.record_backpressure() 

260 await asyncio.sleep(0.01) 

261 

262 self._buffer.append(chunk) 

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

264 

265 # Check if we should emit 

266 result = await self._maybe_emit() 

267 if result is not None: 

268 yield result 

269 

270 # Flush remaining buffer 

271 async for chunk in self._flush(): 

272 yield chunk 

273 

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

275 

276 async def optimize_simple( 

277 self, text_stream: AsyncIterable[str] 

278 ) -> AsyncIterator[str]: 

279 """ 

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

281 """ 

282 async for chunk in self.optimize( 

283 StreamChunk(content=text, index=i) 

284 for i, text in enumerate(text_stream) 

285 ): 

286 yield chunk.content 

287 

288 async def _maybe_emit(self) -> Optional[StreamChunk]: 

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

290 if not self._buffer: 

291 return None 

292 

293 tokens = sum( 

294 chunk.metadata.get("approx_tokens", 1) for chunk in self._buffer 

295 ) 

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 = ( 

309 tokens >= self._config.adaptive_min_chunk_size 

310 and ( 

311 tokens >= self._config.max_chunk_size 

312 or self._buffer[-1].is_final 

313 ) 

314 ) 

315 

316 if not should_emit: 

317 return None 

318 

319 return self._emit_aggregated() 

320 

321 def _emit_aggregated(self) -> Optional[StreamChunk]: 

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

323 if not self._buffer: 

324 return None 

325 

326 # Aggregate content 

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

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

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

330 timestamp = time.time() 

331 

332 # Clear buffer 

333 count = len(self._buffer) 

334 self._buffer.clear() 

335 

336 # Build aggregated chunk 

337 chunk = StreamChunk( 

338 content=content, 

339 index=index, 

340 timestamp=timestamp, 

341 is_final=is_final, 

342 metadata={"aggregated_from": count}, 

343 ) 

344 

345 # Apply transformer pipeline 

346 chunk = self._pipeline.apply(chunk) 

347 if chunk is None: 

348 return None 

349 

350 # Record metrics 

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

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

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

354 

355 return chunk 

356 

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

358 """Flush remaining buffered chunks.""" 

359 while self._buffer: 

360 chunk = self._emit_aggregated() 

361 if chunk is not None: 

362 yield chunk 

363 

364 def get_metrics(self) -> StreamMetrics: 

365 return self._metrics._metrics 

366 

367 def reset_metrics(self) -> None: 

368 self._metrics = MetricsCollector() 

369 self._start_time = None