feat: add streaming support utilities

Author: LifeJiggy

src/gradient/_utils/__init__.py

@@ -37,6 +37,8 @@
     BatchProcessor as BatchProcessor,
     DataExporter as DataExporter,
     PaginationHelper as PaginationHelper,
+    StreamProcessor as StreamProcessor,
+    StreamCollector as StreamCollector,
 )
 from ._compat import (
     get_args as get_args,

src/gradient/_utils/_utils.py

@@ -838,6 +838,142 @@ async def paginate_async(self, fetch_func: Callable[[dict[str, Any]], Any], **kw
         return all_items
 
 
+# Streaming Support Classes
+class StreamProcessor:
+    """Utility for processing streaming API responses with custom handlers."""
+
+    def __init__(self) -> None:
+        """Initialize stream processor."""
+        self._handlers: dict[str, Callable[[Any], Any]] = {}
+
+    def add_handler(self, event_type: str, handler: Callable[[Any], Any]) -> None:
+        """Add event handler for specific event type.
+
+        Args:
+            event_type: Type of event to handle
+            handler: Function to process the event
+        """
+        self._handlers[event_type] = handler
+
+    def remove_handler(self, event_type: str) -> None:
+        """Remove handler for specific event type."""
+        self._handlers.pop(event_type, None)
+
+    def process_event(self, event: Any) -> Any | None:
+        """Process a single streaming event.
+
+        Args:
+            event: The event data to process
+
+        Returns:
+            Result of handler if one exists, None otherwise
+        """
+        event_type = self._get_event_type(event)
+        handler = self._handlers.get(event_type)
+        if handler:
+            return handler(event)
+        return None
+
+    def process_stream(self, stream: Any) -> list[Any]:
+        """Process entire streaming response.
+
+        Args:
+            stream: The stream to process
+
+        Returns:
+            List of all processed event results
+        """
+        results = []
+        for event in stream:
+            result = self.process_event(event)
+            if result is not None:
+                results.append(result)
+        return results
+
+    async def process_stream_async(self, stream: Any) -> list[Any]:
+        """Async version of process_stream."""
+        results = []
+        async for event in stream:
+            result = self.process_event(event)
+            if result is not None:
+                results.append(result)
+        return results
+
+    def _get_event_type(self, event: Any) -> str:
+        """Extract event type from event data."""
+        # Handle different event formats
+        if hasattr(event, 'event') and event.event:
+            return event.event
+        elif hasattr(event, 'type') and event.type:
+            return event.type
+        elif isinstance(event, dict):
+            return event.get('event') or event.get('type') or 'unknown'
+        else:
+            return 'unknown'
+
+
+class StreamCollector:
+    """Utility for collecting and aggregating streaming events."""
+
+    def __init__(self) -> None:
+        """Initialize stream collector."""
+        self._events: list[Any] = []
+        self._aggregated: dict[str, Any] = {}
+
+    def collect(self, event: Any) -> None:
+        """Collect a streaming event."""
+        self._events.append(event)
+        self._aggregate_event(event)
+
+    def get_events(self, event_type: str | None = None) -> list[Any]:
+        """Get collected events, optionally filtered by type."""
+        if event_type is None:
+            return self._events.copy()
+
+        return [e for e in self._events if self._get_event_type(e) == event_type]
+
+    def get_aggregated(self) -> dict[str, Any]:
+        """Get aggregated event data."""
+        return self._aggregated.copy()
+
+    def clear(self) -> None:
+        """Clear all collected events and aggregated data."""
+        self._events.clear()
+        self._aggregated.clear()
+
+    def count_events(self, event_type: str | None = None) -> int:
+        """Count events, optionally filtered by type."""
+        if event_type is None:
+            return len(self._events)
+        return len(self.get_events(event_type))
+
+    def _aggregate_event(self, event: Any) -> None:
+        """Aggregate event data for summary statistics."""
+        event_type = self._get_event_type(event)
+
+        if event_type not in self._aggregated:
+            self._aggregated[event_type] = {
+                'count': 0,
+                'events': [],
+                'last_event': None
+            }
+
+        self._aggregated[event_type]['count'] += 1
+        self._aggregated[event_type]['events'].append(event)
+        self._aggregated[event_type]['last_event'] = event
+
+    def _get_event_type(self, event: Any) -> str:
+        """Extract event type from event data."""
+        if hasattr(event, 'event') and event.event:
+            return event.event
+        elif hasattr(event, 'type') and event.type:
+            return event.type
+        elif isinstance(event, dict):
+            return event.get('event') or event.get('type') or 'unknown'
+        else:
+            return 'unknown'
+
+
 # API Key Validation Functions
 def validate_api_key(api_key: str | None) -> bool:
     """Validate an API key format.

tests/test_streaming_support.py

@@ -0,0 +1,162 @@
+"""Tests for streaming support functionality."""
+
+import pytest
+from gradient._utils import StreamProcessor, StreamCollector
+
+
+class TestStreamProcessor:
+    """Test stream processor functionality."""
+
+    def test_stream_processor_basic(self):
+        """Test basic stream processor functionality."""
+        processor = StreamProcessor()
+
+        # Add handler
+        results = []
+        def text_handler(event):
+            results.append(f"processed: {event.get('text', '')}")
+            return f"processed: {event.get('text', '')}"
+
+        processor.add_handler("text", text_handler)
+
+        # Process events
+        event1 = {"type": "text", "text": "Hello"}
+        event2 = {"type": "other", "data": "ignored"}
+        event3 = {"type": "text", "text": "World"}
+
+        result1 = processor.process_event(event1)
+        result2 = processor.process_event(event2)
+        result3 = processor.process_event(event3)
+
+        assert result1 == "processed: Hello"
+        assert result2 is None  # No handler for "other"
+        assert result3 == "processed: World"
+        assert results == ["processed: Hello", "processed: World"]
+
+    def test_stream_processor_remove_handler(self):
+        """Test removing event handlers."""
+        processor = StreamProcessor()
+
+        def handler(event):
+            return "handled"
+
+        processor.add_handler("test", handler)
+        assert processor.process_event({"type": "test"}) == "handled"
+
+        processor.remove_handler("test")
+        assert processor.process_event({"type": "test"}) is None
+
+    def test_stream_processor_process_stream(self):
+        """Test processing entire stream."""
+        processor = StreamProcessor()
+
+        def text_handler(event):
+            return event.get("text", "").upper()
+
+        processor.add_handler("text", text_handler)
+
+        stream = [
+            {"type": "text", "text": "hello"},
+            {"type": "other", "data": "ignored"},
+            {"type": "text", "text": "world"}
+        ]
+
+        results = processor.process_stream(stream)
+        assert results == ["HELLO", "WORLD"]
+
+    def test_stream_processor_event_type_extraction(self):
+        """Test event type extraction from different formats."""
+        processor = StreamProcessor()
+
+        # Test different event formats
+        event1 = {"type": "custom"}
+        event2 = type('MockEvent', (), {'event': 'mock'})()
+        event3 = {"event": "dict_event"}
+        event4 = "unknown_format"
+
+        assert processor._get_event_type(event1) == "custom"
+        assert processor._get_event_type(event2) == "mock"
+        assert processor._get_event_type(event3) == "dict_event"
+        assert processor._get_event_type(event4) == "unknown"
+
+
+class TestStreamCollector:
+    """Test stream collector functionality."""
+
+    def test_stream_collector_basic(self):
+        """Test basic stream collector functionality."""
+        collector = StreamCollector()
+
+        # Collect events
+        event1 = {"type": "text", "text": "Hello"}
+        event2 = {"type": "text", "text": "World"}
+        event3 = {"type": "error", "message": "Something went wrong"}
+
+        collector.collect(event1)
+        collector.collect(event2)
+        collector.collect(event3)
+
+        # Check all events
+        all_events = collector.get_events()
+        assert len(all_events) == 3
+
+        # Check filtered events
+        text_events = collector.get_events("text")
+        assert len(text_events) == 2
+        assert all(e["type"] == "text" for e in text_events)
+
+        error_events = collector.get_events("error")
+        assert len(error_events) == 1
+        assert error_events[0]["type"] == "error"
+
+    def test_stream_collector_aggregation(self):
+        """Test event aggregation."""
+        collector = StreamCollector()
+
+        # Collect events
+        collector.collect({"type": "text", "text": "Hello"})
+        collector.collect({"type": "text", "text": "World"})
+        collector.collect({"type": "error", "message": "Error 1"})
+        collector.collect({"type": "error", "message": "Error 2"})
+        collector.collect({"type": "text", "text": "Again"})
+
+        aggregated = collector.get_aggregated()
+
+        # Check text events aggregation
+        assert aggregated["text"]["count"] == 3
+        assert len(aggregated["text"]["events"]) == 3
+        assert aggregated["text"]["last_event"]["text"] == "Again"
+
+        # Check error events aggregation
+        assert aggregated["error"]["count"] == 2
+        assert len(aggregated["error"]["events"]) == 2
+        assert aggregated["error"]["last_event"]["message"] == "Error 2"
+
+    def test_stream_collector_count_events(self):
+        """Test event counting."""
+        collector = StreamCollector()
+
+        collector.collect({"type": "text"})
+        collector.collect({"type": "text"})
+        collector.collect({"type": "error"})
+        collector.collect({"type": "text"})
+
+        assert collector.count_events() == 4
+        assert collector.count_events("text") == 3
+        assert collector.count_events("error") == 1
+        assert collector.count_events("unknown") == 0
+
+    def test_stream_collector_clear(self):
+        """Test clearing collected events."""
+        collector = StreamCollector()
+
+        collector.collect({"type": "text"})
+        collector.collect({"type": "error"})
+
+        assert collector.count_events() == 2
+        assert len(collector.get_aggregated()) == 2
+
+        collector.clear()
+
+        assert collector.count_events() == 0
+        assert len(collector.get_aggregated()) == 0