Track statistics about a streaming response.
| 364 | # ================= STREAMING TESTS ================= |
| 365 | |
| 366 | class StreamStats: |
| 367 | """Track statistics about a streaming response.""" |
| 368 | |
| 369 | def __init__(self): |
| 370 | self.event_types = set() |
| 371 | self.event_counts = {} |
| 372 | self.first_event_time = None |
| 373 | self.last_event_time = None |
| 374 | self.total_chunks = 0 |
| 375 | self.events = [] |
| 376 | self.text_content = "" |
| 377 | self.content_blocks = {} |
| 378 | self.has_tool_use = False |
| 379 | self.has_error = False |
| 380 | self.error_message = "" |
| 381 | self.text_content_by_block = {} |
| 382 | |
| 383 | def add_event(self, event_data): |
| 384 | """Track information about each received event.""" |
| 385 | now = datetime.now() |
| 386 | if self.first_event_time is None: |
| 387 | self.first_event_time = now |
| 388 | self.last_event_time = now |
| 389 | |
| 390 | self.total_chunks += 1 |
| 391 | |
| 392 | # Record event type and increment count |
| 393 | if "type" in event_data: |
| 394 | event_type = event_data["type"] |
| 395 | self.event_types.add(event_type) |
| 396 | self.event_counts[event_type] = self.event_counts.get(event_type, 0) + 1 |
| 397 | |
| 398 | # Track specific event data |
| 399 | if event_type == "content_block_start": |
| 400 | block_idx = event_data.get("index") |
| 401 | content_block = event_data.get("content_block", {}) |
| 402 | if content_block.get("type") == "tool_use": |
| 403 | self.has_tool_use = True |
| 404 | self.content_blocks[block_idx] = content_block |
| 405 | self.text_content_by_block[block_idx] = "" |
| 406 | |
| 407 | elif event_type == "content_block_delta": |
| 408 | block_idx = event_data.get("index") |
| 409 | delta = event_data.get("delta", {}) |
| 410 | if delta.get("type") == "text_delta": |
| 411 | text = delta.get("text", "") |
| 412 | self.text_content += text |
| 413 | # Also track text by block ID |
| 414 | if block_idx in self.text_content_by_block: |
| 415 | self.text_content_by_block[block_idx] += text |
| 416 | |
| 417 | # Keep track of all events for debugging |
| 418 | self.events.append(event_data) |
| 419 | |
| 420 | def get_duration(self): |
| 421 | """Calculate the total duration of the stream in seconds.""" |
| 422 | if self.first_event_time is None or self.last_event_time is None: |
| 423 | return 0 |
no outgoing calls