agentfield 0.1.22rc2__py3-none-any.whl

agentfield/http_connection_manager.py

@@ -0,0 +1,429 @@
+"""
+HTTP Connection Manager for async execution.
+
+This module provides aiohttp session pooling with configurable connection limits,
+connection reuse, proper cleanup, timeout handling, and connection health monitoring.
+Supports both single requests and batch operations for the AgentField SDK async execution.
+"""
+
+import asyncio
+import time
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Union
+
+import aiohttp
+
+from .async_config import AsyncConfig
+from .logger import get_logger
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class ConnectionMetrics:
+    """Metrics for connection pool monitoring."""
+
+    total_requests: int = 0
+    successful_requests: int = 0
+    failed_requests: int = 0
+    timeout_requests: int = 0
+    active_connections: int = 0
+    pool_size: int = 0
+    created_at: float = field(default_factory=time.time)
+
+    @property
+    def success_rate(self) -> float:
+        """Calculate success rate as a percentage."""
+        if self.total_requests == 0:
+            return 0.0
+        return (self.successful_requests / self.total_requests) * 100
+
+    @property
+    def uptime(self) -> float:
+        """Get uptime in seconds."""
+        return time.time() - self.created_at
+
+    def record_request(self, success: bool, timeout: bool = False) -> None:
+        """Record a request attempt."""
+        self.total_requests += 1
+        if success:
+            self.successful_requests += 1
+        else:
+            self.failed_requests += 1
+            if timeout:
+                self.timeout_requests += 1
+
+
+@dataclass
+class ConnectionHealth:
+    """Health status of connection pool."""
+
+    is_healthy: bool = True
+    last_check: float = field(default_factory=time.time)
+    consecutive_failures: int = 0
+    last_error: Optional[str] = None
+
+    def mark_healthy(self) -> None:
+        """Mark connection as healthy."""
+        self.is_healthy = True
+        self.consecutive_failures = 0
+        self.last_error = None
+        self.last_check = time.time()
+
+    def mark_unhealthy(self, error: str) -> None:
+        """Mark connection as unhealthy."""
+        self.is_healthy = False
+        self.consecutive_failures += 1
+        self.last_error = error
+        self.last_check = time.time()
+
+
+class ConnectionManager:
+    """
+    HTTP Connection Manager with aiohttp session pooling.
+
+    Provides efficient HTTP connection management for async execution with:
+    - Configurable connection limits and timeouts
+    - Connection reuse and proper cleanup
+    - Health monitoring and metrics
+    - Support for single requests and batch operations
+    - Thread-safe operations for concurrent access
+    """
+
+    def __init__(self, config: Optional[AsyncConfig] = None):
+        """
+        Initialize the connection manager.
+
+        Args:
+            config: AsyncConfig instance for configuration parameters
+        """
+        self.config = config or AsyncConfig()
+        self._session: Optional[aiohttp.ClientSession] = None
+        self._connector: Optional[aiohttp.TCPConnector] = None
+        self._lock = asyncio.Lock()
+        self._closed = False
+
+        # Metrics and health monitoring
+        self.metrics = ConnectionMetrics()
+        self.health = ConnectionHealth()
+
+        # Background tasks
+        self._health_check_task: Optional[asyncio.Task] = None
+        self._cleanup_task: Optional[asyncio.Task] = None
+
+        logger.debug(f"ConnectionManager initialized with config: {self.config}")
+
+    async def __aenter__(self):
+        """Async context manager entry."""
+        await self.start()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+        await self.close()
+
+    async def start(self) -> None:
+        """
+        Start the connection manager and initialize session.
+
+        Raises:
+            RuntimeError: If manager is already started or closed
+        """
+        async with self._lock:
+            if self._session is not None:
+                raise RuntimeError("ConnectionManager is already started")
+
+            if self._closed:
+                raise RuntimeError(
+                    "ConnectionManager is closed and cannot be restarted"
+                )
+
+            # Create TCP connector with configuration
+            self._connector = aiohttp.TCPConnector(
+                limit=self.config.connection_pool_size,
+                limit_per_host=self.config.connection_pool_per_host,
+                ttl_dns_cache=300,  # 5 minutes DNS cache
+                use_dns_cache=True,
+                keepalive_timeout=30,
+                enable_cleanup_closed=True,
+                force_close=False,
+            )
+
+            # Create session with timeout configuration
+            timeout = aiohttp.ClientTimeout(
+                total=self.config.max_execution_timeout,
+                connect=self.config.polling_timeout,
+                sock_read=self.config.polling_timeout,
+            )
+
+            self._session = aiohttp.ClientSession(
+                connector=self._connector,
+                timeout=timeout,
+                headers={
+                    "User-Agent": "AgentField-SDK-AsyncClient/1.0",
+                    "Accept": "application/json",
+                    "Content-Type": "application/json",
+                },
+            )
+
+            # Update metrics
+            self.metrics.pool_size = self.config.connection_pool_size
+
+            # Start background tasks if enabled
+            if self.config.enable_performance_logging:
+                self._health_check_task = asyncio.create_task(self._health_check_loop())
+                self._cleanup_task = asyncio.create_task(self._cleanup_loop())
+
+            logger.info(
+                f"ConnectionManager started with pool size {self.config.connection_pool_size}"
+            )
+
+    async def close(self) -> None:
+        """
+        Close the connection manager and cleanup resources.
+        """
+        async with self._lock:
+            if self._closed:
+                return
+
+            self._closed = True
+
+            # Cancel background tasks
+            if self._health_check_task:
+                self._health_check_task.cancel()
+                try:
+                    await self._health_check_task
+                except asyncio.CancelledError:
+                    pass
+
+            if self._cleanup_task:
+                self._cleanup_task.cancel()
+                try:
+                    await self._cleanup_task
+                except asyncio.CancelledError:
+                    pass
+
+            # Close session and connector
+            if self._session:
+                await self._session.close()
+                self._session = None
+
+            if self._connector:
+                await self._connector.close()
+                self._connector = None
+
+            logger.info("ConnectionManager closed")
+
+    @asynccontextmanager
+    async def get_session(self):
+        """
+        Get an aiohttp session for making requests.
+
+        Yields:
+            aiohttp.ClientSession: Active session for making requests
+
+        Raises:
+            RuntimeError: If manager is not started or is closed
+        """
+        if self._session is None:
+            raise RuntimeError("ConnectionManager is not started. Call start() first.")
+
+        if self._closed:
+            raise RuntimeError("ConnectionManager is closed")
+
+        try:
+            yield self._session
+        except Exception as e:
+            self.health.mark_unhealthy(str(e))
+            raise
+
+    async def request(self, method: str, url: str, **kwargs) -> aiohttp.ClientResponse:
+        """
+        Make a single HTTP request.
+
+        Args:
+            method: HTTP method (GET, POST, etc.)
+            url: Request URL
+            **kwargs: Additional arguments for aiohttp request
+
+        Returns:
+            aiohttp.ClientResponse: Response object
+
+        Raises:
+            aiohttp.ClientError: For HTTP-related errors
+            asyncio.TimeoutError: For timeout errors
+        """
+        start_time = time.time()
+        success = False
+        timeout_occurred = False
+
+        try:
+            async with self.get_session() as session:
+                response = await session.request(method, url, **kwargs)
+                success = True
+                self.health.mark_healthy()
+                return response
+
+        except asyncio.TimeoutError:
+            timeout_occurred = True
+            logger.warn(f"Request timeout for {method} {url}")
+            raise
+        except Exception as e:
+            self.health.mark_unhealthy(str(e))
+            logger.error(f"Request failed for {method} {url}: {e}")
+            raise
+        finally:
+            # Record metrics
+            self.metrics.record_request(success, timeout_occurred)
+
+            # Log slow requests
+            duration = time.time() - start_time
+            if (
+                self.config.log_slow_executions
+                and duration > self.config.slow_execution_threshold
+            ):
+                logger.warn(f"Slow request: {method} {url} took {duration:.2f}s")
+
+    async def batch_request(
+        self, requests: List[Dict[str, Any]]
+    ) -> List[Union[aiohttp.ClientResponse, Exception]]:
+        """
+        Make multiple HTTP requests concurrently.
+
+        Args:
+            requests: List of request dictionaries with 'method', 'url', and optional kwargs
+
+        Returns:
+            List of responses or exceptions for each request
+        """
+        if not requests:
+            return []
+
+        # Limit concurrent requests
+        semaphore = asyncio.Semaphore(self.config.max_active_polls)
+
+        async def make_request(
+            req_data: Dict[str, Any],
+        ) -> Union[aiohttp.ClientResponse, Exception]:
+            async with semaphore:
+                try:
+                    method = req_data.pop("method")
+                    url = req_data.pop("url")
+                    return await self.request(method, url, **req_data)
+                except Exception as e:
+                    return e
+
+        # Execute all requests concurrently
+        tasks = [make_request(req.copy()) for req in requests]
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+
+        logger.debug(f"Batch request completed: {len(requests)} requests")
+        return results
+
+    async def health_check(self) -> bool:
+        """
+        Perform a health check on the connection pool.
+
+        Returns:
+            bool: True if healthy, False otherwise
+        """
+        try:
+            # Simple health check - try to create a request
+            if self._session is None or self._session.closed:
+                self.health.mark_unhealthy("Session is closed")
+                return False
+
+            # Check connector health
+            if self._connector is None or self._connector.closed:
+                self.health.mark_unhealthy("Connector is closed")
+                return False
+
+            self.health.mark_healthy()
+            return True
+
+        except Exception as e:
+            self.health.mark_unhealthy(str(e))
+            return False
+
+    async def _health_check_loop(self) -> None:
+        """Background task for periodic health checks."""
+        while not self._closed:
+            try:
+                await asyncio.sleep(60)  # Check every minute
+                await self.health_check()
+
+                # Log health status if unhealthy
+                if not self.health.is_healthy:
+                    logger.warn(f"Connection pool unhealthy: {self.health.last_error}")
+
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"Health check loop error: {e}")
+
+    async def _cleanup_loop(self) -> None:
+        """Background task for periodic cleanup."""
+        while not self._closed:
+            try:
+                await asyncio.sleep(self.config.cleanup_interval)
+
+                # Update active connections metric
+                if self._connector:
+                    self.metrics.active_connections = len(self._connector._conns)
+
+                # Log metrics if performance logging is enabled
+                if self.config.enable_performance_logging:
+                    logger.debug(
+                        f"Connection metrics: {self.metrics.total_requests} total, "
+                        f"{self.metrics.success_rate:.1f}% success, "
+                        f"{self.metrics.active_connections} active"
+                    )
+
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"Cleanup loop error: {e}")
+
+    def get_metrics(self) -> ConnectionMetrics:
+        """
+        Get current connection metrics.
+
+        Returns:
+            ConnectionMetrics: Current metrics snapshot
+        """
+        # Update active connections if connector is available
+        if self._connector:
+            self.metrics.active_connections = len(self._connector._conns)
+
+        return self.metrics
+
+    def get_health(self) -> ConnectionHealth:
+        """
+        Get current health status.
+
+        Returns:
+            ConnectionHealth: Current health status
+        """
+        return self.health
+
+    @property
+    def is_healthy(self) -> bool:
+        """Check if connection manager is healthy."""
+        return self.health.is_healthy and not self._closed
+
+    @property
+    def is_closed(self) -> bool:
+        """Check if connection manager is closed."""
+        return self._closed
+
+    def __repr__(self) -> str:
+        """String representation of the connection manager."""
+        return (
+            f"ConnectionManager("
+            f"pool_size={self.config.connection_pool_size}, "
+            f"healthy={self.is_healthy}, "
+            f"closed={self.is_closed}, "
+            f"requests={self.metrics.total_requests}"
+            f")"
+        )

agentfield/litellm_adapters.py

@@ -0,0 +1,140 @@
+"""
+LiteLLM Provider Adapters
+
+This module centralizes provider-specific parameter transformations and patches
+required to ensure compatibility across different LLM providers.
+
+Each patch should be:
+1. Well-documented with the reason for its existence
+2. Tied to specific providers/models that require it
+3. Transparent about what transformation is being applied
+
+This abstraction allows the core SDK to remain clean while handling necessary
+provider-specific quirks in one maintainable location.
+"""
+
+from typing import Dict, Any
+
+
+def get_provider_from_model(model: str) -> str:
+    """
+    Extract provider name from model string.
+
+    LiteLLM uses the format "provider/model-name" (e.g., "openai/gpt-4o").
+    This function extracts the provider prefix.
+
+    Args:
+        model: Model string in LiteLLM format
+
+    Returns:
+        Provider name (e.g., "openai", "anthropic", "cohere")
+        Returns "unknown" if format doesn't match
+
+    Examples:
+        >>> get_provider_from_model("openai/gpt-4o")
+        'openai'
+        >>> get_provider_from_model("anthropic/claude-3-opus")
+        'anthropic'
+        >>> get_provider_from_model("gpt-4o")
+        'unknown'
+    """
+    if "/" in model:
+        return model.split("/")[0]
+    return "unknown"
+
+
+def apply_openai_patches(params: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Apply OpenAI-specific parameter patches.
+
+    **Patch 1: max_tokens → max_completion_tokens**
+
+    Reason: OpenAI's newer models (gpt-4o, gpt-4o-mini, etc.) use
+    `max_completion_tokens` instead of `max_tokens` to disambiguate between
+    input tokens and output tokens. LiteLLM may not always handle this
+    transformation automatically for all OpenAI models.
+
+    This patch ensures compatibility by renaming the parameter when targeting
+    OpenAI models.
+
+    Reference: https://platform.openai.com/docs/api-reference/chat/create
+
+    Args:
+        params: Parameter dictionary to transform
+
+    Returns:
+        Transformed parameter dictionary
+    """
+    # Create a copy to avoid mutating the original
+    patched = params.copy()
+
+    # Patch: max_tokens → max_completion_tokens for OpenAI
+    if "max_tokens" in patched:
+        patched["max_completion_tokens"] = patched.pop("max_tokens")
+
+    return patched
+
+
+def apply_provider_patches(params: Dict[str, Any], model: str) -> Dict[str, Any]:
+    """
+    Apply provider-specific parameter transformations.
+
+    This is the main entry point for all provider-specific patches. It detects
+    the provider from the model string and applies appropriate transformations.
+
+    **When to add a new patch:**
+    1. A specific provider requires a different parameter name
+    2. A provider has parameter constraints that differ from LiteLLM defaults
+    3. There's a known incompatibility that needs a workaround
+
+    **How to add a new patch:**
+    1. Create a new function: `apply_{provider}_patches(params)`
+    2. Document the patch reason and affected models
+    3. Add a new elif branch in this function
+
+    Args:
+        params: Parameter dictionary from AIConfig.get_litellm_params()
+        model: Model string (e.g., "openai/gpt-4o")
+
+    Returns:
+        Transformed parameter dictionary with provider-specific patches applied
+
+    Examples:
+        >>> params = {"model": "openai/gpt-4o", "max_tokens": 1000}
+        >>> apply_provider_patches(params, "openai/gpt-4o")
+        {'model': 'openai/gpt-4o', 'max_completion_tokens': 1000}
+    """
+    provider = get_provider_from_model(model)
+
+    # Apply provider-specific patches
+    if provider == "openai":
+        return apply_openai_patches(params)
+
+    # Add more providers here as needed:
+    # elif provider == "anthropic":
+    #     return apply_anthropic_patches(params)
+    # elif provider == "cohere":
+    #     return apply_cohere_patches(params)
+
+    # No patches needed for this provider
+    return params
+
+
+def filter_none_values(params: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Remove None values from parameter dictionary.
+
+    This ensures we only pass explicitly set parameters to LiteLLM,
+    allowing it to use its own defaults for unset values.
+
+    Args:
+        params: Parameter dictionary potentially containing None values
+
+    Returns:
+        Dictionary with None values removed
+
+    Examples:
+        >>> filter_none_values({"a": 1, "b": None, "c": "test"})
+        {'a': 1, 'c': 'test'}
+    """
+    return {k: v for k, v in params.items() if v is not None}