kubiya-control-plane-api 0.1.0__py3-none-any.whl → 0.3.4__py3-none-any.whl

control_plane_api/worker/utils/chunk_batcher.py

@@ -0,0 +1,305 @@
+"""
+Smart chunk batching for streaming to reduce HTTP requests.
+
+Instead of sending one POST per chunk (50-70 requests), batch chunks
+with configurable time/size windows (5-10 requests).
+
+Batching Strategy:
+- Time window: Flush after X ms (default: 100ms)
+- Size window: Flush when batch reaches Y bytes (default: 100 bytes)
+- Immediate flush: On tool events, errors, or completion
+
+This provides:
+- 90%+ reduction in HTTP requests
+- Still feels real-time (100ms is imperceptible)
+- Lower latency (fewer round trips)
+- Better serverless performance (fewer cold starts)
+- Lower costs (fewer invocations)
+"""
+
+import asyncio
+import time
+from typing import Dict, Any, Optional, Callable
+from dataclasses import dataclass, field
+import structlog
+
+logger = structlog.get_logger()
+
+
+@dataclass
+class BatchConfig:
+    """Configuration for chunk batching behavior."""
+
+    # Time-based batching: flush after this many milliseconds
+    time_window_ms: int = 100
+
+    # Size-based batching: flush when accumulated content reaches this size
+    size_window_bytes: int = 100
+
+    # Maximum batch size before forced flush (safety limit)
+    max_batch_size_bytes: int = 1000
+
+    # Enable/disable batching (for testing/debugging)
+    enabled: bool = True
+
+    @classmethod
+    def from_env(cls) -> "BatchConfig":
+        """
+        Create configuration from environment variables.
+
+        Environment variables:
+            CHUNK_BATCHING_ENABLED: Enable/disable batching (default: true)
+            CHUNK_BATCHING_TIME_WINDOW_MS: Time window in ms (default: 100)
+            CHUNK_BATCHING_SIZE_WINDOW_BYTES: Size window in bytes (default: 100)
+            CHUNK_BATCHING_MAX_SIZE_BYTES: Max batch size in bytes (default: 1000)
+
+        Returns:
+            BatchConfig instance with values from environment
+        """
+        import os
+
+        return cls(
+            enabled=os.getenv("CHUNK_BATCHING_ENABLED", "true").lower() == "true",
+            time_window_ms=int(os.getenv("CHUNK_BATCHING_TIME_WINDOW_MS", "100")),
+            size_window_bytes=int(os.getenv("CHUNK_BATCHING_SIZE_WINDOW_BYTES", "100")),
+            max_batch_size_bytes=int(os.getenv("CHUNK_BATCHING_MAX_SIZE_BYTES", "1000")),
+        )
+
+
+@dataclass
+class ContentBatch:
+    """Accumulated content chunks waiting to be flushed."""
+
+    chunks: list[str] = field(default_factory=list)
+    total_size: int = 0
+    first_chunk_time: Optional[float] = None
+
+    def add(self, content: str) -> None:
+        """Add content to the batch."""
+        self.chunks.append(content)
+        self.total_size += len(content.encode('utf-8'))
+
+        if self.first_chunk_time is None:
+            self.first_chunk_time = time.time()
+
+    def get_combined_content(self) -> str:
+        """Get all chunks combined into single string."""
+        return ''.join(self.chunks)
+
+    def clear(self) -> None:
+        """Clear the batch after flushing."""
+        self.chunks.clear()
+        self.total_size = 0
+        self.first_chunk_time = None
+
+    def is_empty(self) -> bool:
+        """Check if batch is empty."""
+        return len(self.chunks) == 0
+
+    def age_ms(self) -> float:
+        """Get age of batch in milliseconds."""
+        if self.first_chunk_time is None:
+            return 0
+        return (time.time() - self.first_chunk_time) * 1000
+
+
+class ChunkBatcher:
+    """
+    Smart batching for streaming chunks to reduce HTTP requests.
+
+    Usage:
+        batcher = ChunkBatcher(
+            publish_func=control_plane.publish_event,
+            execution_id=execution_id,
+            message_id=message_id,
+            config=BatchConfig(time_window_ms=100, size_window_bytes=100)
+        )
+
+        # Add chunks as they arrive
+        await batcher.add_chunk("Hello")
+        await batcher.add_chunk(" world")
+
+        # Flush remaining chunks when done
+        await batcher.flush()
+    """
+
+    def __init__(
+        self,
+        publish_func: Callable,
+        execution_id: str,
+        message_id: str,
+        config: Optional[BatchConfig] = None
+    ):
+        self.publish_func = publish_func
+        self.execution_id = execution_id
+        self.message_id = message_id
+        self.config = config or BatchConfig()
+
+        self.batch = ContentBatch()
+        self._flush_task: Optional[asyncio.Task] = None
+        self._stats = {
+            "chunks_received": 0,
+            "batches_sent": 0,
+            "bytes_sent": 0,
+            "flushes_by_time": 0,
+            "flushes_by_size": 0,
+            "flushes_manual": 0,
+        }
+
+    async def add_chunk(self, content: str) -> None:
+        """
+        Add a chunk to the batch.
+
+        Automatically flushes if:
+        - Batch size exceeds size_window_bytes
+        - Batch age exceeds time_window_ms
+        - Max batch size is reached (safety)
+        """
+        if not self.config.enabled:
+            # Batching disabled - send immediately
+            await self._publish_batch([content])
+            return
+
+        self._stats["chunks_received"] += 1
+        self.batch.add(content)
+
+        # Check if we should flush immediately due to size
+        should_flush_size = self.batch.total_size >= self.config.size_window_bytes
+        should_flush_max = self.batch.total_size >= self.config.max_batch_size_bytes
+
+        if should_flush_max:
+            # Safety: flush immediately if max size reached
+            logger.debug(
+                "Flushing batch (max size reached)",
+                execution_id=self.execution_id[:8],
+                batch_size=self.batch.total_size,
+                chunk_count=len(self.batch.chunks),
+            )
+            await self.flush(reason="max_size")
+        elif should_flush_size:
+            # Size threshold reached - flush now
+            await self.flush(reason="size")
+        else:
+            # Start/reset timer for time-based flush
+            await self._schedule_time_flush()
+
+    async def _schedule_time_flush(self) -> None:
+        """Schedule a time-based flush if not already scheduled."""
+        if self._flush_task is not None and not self._flush_task.done():
+            # Timer already running
+            return
+
+        self._flush_task = asyncio.create_task(self._time_based_flush())
+
+    async def _time_based_flush(self) -> None:
+        """Wait for time window, then flush."""
+        await asyncio.sleep(self.config.time_window_ms / 1000.0)
+
+        if not self.batch.is_empty():
+            await self.flush(reason="time")
+
+    async def flush(self, reason: str = "manual") -> None:
+        """
+        Flush current batch immediately.
+
+        Args:
+            reason: Why flush was triggered (for stats/debugging)
+        """
+        if self.batch.is_empty():
+            return
+
+        # Cancel pending timer if any
+        if self._flush_task is not None and not self._flush_task.done():
+            self._flush_task.cancel()
+            try:
+                await self._flush_task
+            except asyncio.CancelledError:
+                pass
+
+        # Publish the batch
+        chunks = self.batch.chunks.copy()
+        await self._publish_batch(chunks)
+
+        # Update stats
+        if reason == "time":
+            self._stats["flushes_by_time"] += 1
+        elif reason == "size" or reason == "max_size":
+            self._stats["flushes_by_size"] += 1
+        else:
+            self._stats["flushes_manual"] += 1
+
+        # Clear batch
+        self.batch.clear()
+
+    async def _publish_batch(self, chunks: list[str]) -> None:
+        """Publish a batch of chunks as single event."""
+        combined_content = ''.join(chunks)
+
+        try:
+            # Call publish_func (could be sync or async)
+            result = self.publish_func(
+                execution_id=self.execution_id,
+                event_type="message_chunk",
+                data={
+                    "role": "assistant",
+                    "content": combined_content,
+                    "is_chunk": True,
+                    "message_id": self.message_id,
+                    # Metadata for debugging
+                    "batch_info": {
+                        "chunk_count": len(chunks),
+                        "batch_size": len(combined_content.encode('utf-8')),
+                    } if len(chunks) > 1 else None,
+                }
+            )
+
+            # Await if it's a coroutine
+            if asyncio.iscoroutine(result):
+                await result
+
+            self._stats["batches_sent"] += 1
+            self._stats["bytes_sent"] += len(combined_content.encode('utf-8'))
+
+        except Exception as e:
+            logger.warning(
+                "Failed to publish batch",
+                execution_id=self.execution_id[:8],
+                error=str(e),
+                chunk_count=len(chunks),
+            )
+
+    def get_stats(self) -> Dict[str, Any]:
+        """
+        Get batching statistics.
+
+        Returns:
+            Dict with stats about batching performance
+        """
+        chunks_received = self._stats["chunks_received"]
+        batches_sent = self._stats["batches_sent"]
+
+        return {
+            **self._stats,
+            "reduction_percent": round(
+                (1 - batches_sent / max(chunks_received, 1)) * 100, 1
+            ) if chunks_received > 0 else 0,
+            "avg_batch_size": round(
+                chunks_received / max(batches_sent, 1), 1
+            ) if batches_sent > 0 else 0,
+        }
+
+    async def close(self) -> None:
+        """
+        Close the batcher and flush remaining chunks.
+
+        Call this when streaming is complete.
+        """
+        await self.flush(reason="close")
+
+        # Log stats
+        stats = self.get_stats()
+        logger.info(
+            "Chunk batching stats",
+            execution_id=self.execution_id[:8],
+            **stats
+        )

control_plane_api/worker/utils/retry_utils.py

@@ -0,0 +1,60 @@
+"""Retry utilities with exponential backoff"""
+
+import time
+import httpx
+from typing import Callable, TypeVar, Optional
+from functools import wraps
+import structlog
+
+logger = structlog.get_logger()
+
+T = TypeVar('T')
+
+
+def retry_with_backoff(
+    max_retries: int = 3,
+    initial_delay: float = 1.0,
+    backoff_factor: float = 2.0,
+    exceptions: tuple = (httpx.TimeoutException, httpx.ConnectError)
+):
+    """
+    Decorator for retrying functions with exponential backoff
+
+    Args:
+        max_retries: Maximum number of retry attempts
+        initial_delay: Initial delay in seconds
+        backoff_factor: Multiplier for each retry (exponential)
+        exceptions: Tuple of exceptions to catch and retry
+    """
+    def decorator(func: Callable[..., T]) -> Callable[..., T]:
+        @wraps(func)
+        def wrapper(*args, **kwargs) -> T:
+            last_exception = None
+
+            for attempt in range(max_retries):
+                try:
+                    return func(*args, **kwargs)
+                except exceptions as e:
+                    last_exception = e
+                    if attempt < max_retries - 1:
+                        delay = initial_delay * (backoff_factor ** attempt)
+                        logger.warning(
+                            f"{func.__name__} failed, retrying...",
+                            extra={
+                                "attempt": attempt + 1,
+                                "max_retries": max_retries,
+                                "delay": delay,
+                                "error": str(e)
+                            }
+                        )
+                        time.sleep(delay)
+                    else:
+                        logger.error(
+                            f"{func.__name__} failed after {max_retries} attempts",
+                            extra={"error": str(e)}
+                        )
+
+            raise last_exception
+
+        return wrapper
+    return decorator