ai-lib-python 0.5.0__py3-none-any.whl

ai_lib_python/__init__.py

@@ -0,0 +1,43 @@
+"""
+ai-lib-python: Official Python Runtime for AI-Protocol
+
+The canonical Pythonic implementation for unified AI model interaction.
+Core principle: All logic is operators, all configuration is protocol.
+"""
+
+from ai_lib_python.client import AiClient, AiClientBuilder, CallStats, ChatResponse
+from ai_lib_python.errors import AiLibError, ProtocolError, TransportError
+from ai_lib_python.types.events import StreamingEvent
+from ai_lib_python.types.message import (
+    ContentBlock,
+    Message,
+    MessageContent,
+    MessageRole,
+)
+from ai_lib_python.types.tool import ToolCall, ToolDefinition
+
+__version__ = "0.5.0"
+
+__all__ = [
+    # Client
+    "AiClient",
+    "AiClientBuilder",
+    # Errors
+    "AiLibError",
+    "CallStats",
+    "ChatResponse",
+    "ContentBlock",
+    # Types - Message
+    "Message",
+    "MessageContent",
+    "MessageRole",
+    "ProtocolError",
+    # Types - Events
+    "StreamingEvent",
+    "ToolCall",
+    # Types - Tool
+    "ToolDefinition",
+    "TransportError",
+    # Version
+    "__version__",
+]

ai_lib_python/batch/__init__.py

@@ -0,0 +1,15 @@
+"""
+Batch processing module for ai-lib-python.
+
+Provides request batching and batch execution utilities.
+"""
+
+from ai_lib_python.batch.collector import BatchCollector, BatchConfig
+from ai_lib_python.batch.executor import BatchExecutor, BatchResult
+
+__all__ = [
+    "BatchCollector",
+    "BatchConfig",
+    "BatchExecutor",
+    "BatchResult",
+]

ai_lib_python/batch/collector.py

@@ -0,0 +1,244 @@
+"""
+Batch collector for grouping requests.
+
+Collects requests and groups them for batch execution.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+@dataclass
+class BatchConfig:
+    """Configuration for batch collection.
+
+    Attributes:
+        max_batch_size: Maximum requests per batch
+        max_wait_ms: Maximum wait time before flushing batch
+        group_by: Function to determine grouping key
+    """
+
+    max_batch_size: int = 10
+    max_wait_ms: float = 100.0
+    group_by: Callable[[Any], str] | None = None
+
+    @classmethod
+    def default(cls) -> BatchConfig:
+        """Create default configuration."""
+        return cls()
+
+    @classmethod
+    def for_embeddings(cls) -> BatchConfig:
+        """Create configuration optimized for embeddings."""
+        return cls(
+            max_batch_size=100,
+            max_wait_ms=50.0,
+        )
+
+    @classmethod
+    def for_chat(cls) -> BatchConfig:
+        """Create configuration for chat completions."""
+        return cls(
+            max_batch_size=5,
+            max_wait_ms=10.0,
+        )
+
+
+@dataclass
+class PendingRequest(Generic[T]):
+    """A pending request waiting for batch execution.
+
+    Attributes:
+        data: Request data
+        future: Future to resolve with result
+        added_at: Timestamp when added
+        group_key: Grouping key
+    """
+
+    data: T
+    future: asyncio.Future[Any]
+    added_at: float = field(default_factory=time.time)
+    group_key: str = "_default_"
+
+
+class BatchCollector(Generic[T, R]):
+    """Collects requests for batch processing.
+
+    Accumulates requests until batch size or time limit is reached,
+    then triggers batch execution.
+
+    Example:
+        >>> async def process_batch(items):
+        ...     return [f"result_{i}" for i in range(len(items))]
+        ...
+        >>> collector = BatchCollector(
+        ...     config=BatchConfig(max_batch_size=5),
+        ...     executor=process_batch,
+        ... )
+        >>> result = await collector.add("request1")
+    """
+
+    def __init__(
+        self,
+        config: BatchConfig | None = None,
+        executor: Callable[[list[T]], Awaitable[list[R]]] | None = None,
+    ) -> None:
+        """Initialize batch collector.
+
+        Args:
+            config: Batch configuration
+            executor: Function to execute batches
+        """
+        self._config = config or BatchConfig.default()
+        self._executor = executor
+        self._pending: dict[str, list[PendingRequest[T]]] = {}
+        self._lock = asyncio.Lock()
+        self._timers: dict[str, asyncio.Task[None]] = {}
+        self._running = True
+
+    def set_executor(
+        self, executor: Callable[[list[T]], Awaitable[list[R]]]
+    ) -> None:
+        """Set the batch executor function.
+
+        Args:
+            executor: Function to execute batches
+        """
+        self._executor = executor
+
+    async def add(self, data: T) -> R:
+        """Add a request to the batch.
+
+        Args:
+            data: Request data
+
+        Returns:
+            Result from batch execution
+        """
+        if not self._running:
+            raise RuntimeError("Batch collector is stopped")
+
+        if self._executor is None:
+            raise RuntimeError("No executor set")
+
+        # Determine group key
+        group_key = (
+            self._config.group_by(data) if self._config.group_by else "_default_"
+        )
+
+        # Create future for result
+        loop = asyncio.get_event_loop()
+        future: asyncio.Future[R] = loop.create_future()
+
+        async with self._lock:
+            # Add to pending
+            if group_key not in self._pending:
+                self._pending[group_key] = []
+
+            self._pending[group_key].append(
+                PendingRequest(data=data, future=future, group_key=group_key)
+            )
+
+            # Check if batch is full
+            if len(self._pending[group_key]) >= self._config.max_batch_size:
+                await self._flush_group(group_key)
+            else:
+                # Start timer if not already running
+                if group_key not in self._timers or self._timers[group_key].done():
+                    self._timers[group_key] = asyncio.create_task(
+                        self._timer_flush(group_key)
+                    )
+
+        return await future
+
+    async def _timer_flush(self, group_key: str) -> None:
+        """Flush group after timeout."""
+        await asyncio.sleep(self._config.max_wait_ms / 1000.0)
+
+        async with self._lock:
+            if self._pending.get(group_key):
+                await self._flush_group(group_key)
+
+    async def _flush_group(self, group_key: str) -> None:
+        """Flush a specific group.
+
+        Args:
+            group_key: Group to flush
+        """
+        if group_key not in self._pending or not self._pending[group_key]:
+            return
+
+        # Get pending requests
+        requests = self._pending.pop(group_key)
+
+        # Cancel timer
+        if group_key in self._timers:
+            self._timers[group_key].cancel()
+            del self._timers[group_key]
+
+        # Extract data
+        data_list = [r.data for r in requests]
+
+        try:
+            # Execute batch
+            results = await self._executor(data_list)
+
+            # Resolve futures
+            for request, result in zip(requests, results, strict=False):
+                if not request.future.done():
+                    request.future.set_result(result)
+
+        except Exception as e:
+            # Reject all futures
+            for request in requests:
+                if not request.future.done():
+                    request.future.set_exception(e)
+
+    async def flush(self) -> None:
+        """Flush all pending batches."""
+        async with self._lock:
+            for group_key in list(self._pending.keys()):
+                await self._flush_group(group_key)
+
+    async def stop(self) -> None:
+        """Stop the collector and flush pending requests."""
+        self._running = False
+        await self.flush()
+
+        # Cancel all timers
+        for timer in self._timers.values():
+            timer.cancel()
+        self._timers.clear()
+
+    def get_pending_count(self, group_key: str | None = None) -> int:
+        """Get count of pending requests.
+
+        Args:
+            group_key: Group to count (all if None)
+
+        Returns:
+            Number of pending requests
+        """
+        if group_key:
+            return len(self._pending.get(group_key, []))
+        return sum(len(p) for p in self._pending.values())
+
+    @property
+    def config(self) -> BatchConfig:
+        """Get batch configuration."""
+        return self._config
+
+    @property
+    def is_running(self) -> bool:
+        """Check if collector is running."""
+        return self._running

ai_lib_python/batch/executor.py

@@ -0,0 +1,224 @@
+"""
+Batch executor for parallel request execution.
+
+Executes multiple requests concurrently with rate limiting.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Generic, TypeVar
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable, Sequence
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+@dataclass
+class BatchResult(Generic[R]):
+    """Result of a batch execution.
+
+    Attributes:
+        results: List of results (None for failed requests)
+        errors: List of errors (None for successful requests)
+        total_time_ms: Total execution time in milliseconds
+        successful_count: Number of successful requests
+        failed_count: Number of failed requests
+    """
+
+    results: list[R | None] = field(default_factory=list)
+    errors: list[Exception | None] = field(default_factory=list)
+    total_time_ms: float = 0.0
+
+    @property
+    def successful_count(self) -> int:
+        """Get count of successful results."""
+        return sum(1 for e in self.errors if e is None)
+
+    @property
+    def failed_count(self) -> int:
+        """Get count of failed results."""
+        return sum(1 for e in self.errors if e is not None)
+
+    @property
+    def all_successful(self) -> bool:
+        """Check if all requests succeeded."""
+        return all(e is None for e in self.errors)
+
+    def get_successful_results(self) -> list[R]:
+        """Get only successful results."""
+        return [r for r, e in zip(self.results, self.errors, strict=False) if e is None and r is not None]
+
+    def get_errors(self) -> list[tuple[int, Exception]]:
+        """Get errors with their indices."""
+        return [(i, e) for i, e in enumerate(self.errors) if e is not None]
+
+
+class BatchExecutor(Generic[T, R]):
+    """Executes batches of requests concurrently.
+
+    Provides controlled parallel execution with configurable
+    concurrency limits.
+
+    Example:
+        >>> async def call_api(item):
+        ...     return await client.chat().user(item).execute()
+        ...
+        >>> executor = BatchExecutor(call_api, max_concurrent=5)
+        >>> result = await executor.execute(["Q1", "Q2", "Q3"])
+        >>> print(result.successful_count)
+    """
+
+    def __init__(
+        self,
+        operation: Callable[[T], Awaitable[R]],
+        max_concurrent: int = 10,
+        fail_fast: bool = False,
+    ) -> None:
+        """Initialize batch executor.
+
+        Args:
+            operation: Async function to execute for each item
+            max_concurrent: Maximum concurrent operations
+            fail_fast: Stop on first error
+        """
+        self._operation = operation
+        self._max_concurrent = max_concurrent
+        self._fail_fast = fail_fast
+        self._semaphore = asyncio.Semaphore(max_concurrent)
+
+    async def execute(self, items: Sequence[T]) -> BatchResult[R]:
+        """Execute operation for all items.
+
+        Args:
+            items: Items to process
+
+        Returns:
+            BatchResult with results and errors
+        """
+        start_time = time.time()
+        result = BatchResult[R]()
+
+        # Initialize result lists
+        result.results = [None] * len(items)
+        result.errors = [None] * len(items)
+
+        # Create tasks
+        tasks = [
+            asyncio.create_task(self._execute_one(item, idx, result))
+            for idx, item in enumerate(items)
+        ]
+
+        # Execute all tasks
+        if self._fail_fast:
+            # Stop on first error
+            try:
+                await asyncio.gather(*tasks)
+            except Exception:
+                # Cancel remaining tasks
+                for task in tasks:
+                    if not task.done():
+                        task.cancel()
+                # Wait for cancellation to complete
+                await asyncio.gather(*tasks, return_exceptions=True)
+        else:
+            # Continue on errors
+            await asyncio.gather(*tasks, return_exceptions=True)
+
+        result.total_time_ms = (time.time() - start_time) * 1000
+        return result
+
+    async def _execute_one(
+        self,
+        item: T,
+        index: int,
+        result: BatchResult[R],
+    ) -> None:
+        """Execute operation for a single item.
+
+        Args:
+            item: Item to process
+            index: Item index
+            result: Result object to update
+        """
+        async with self._semaphore:
+            try:
+                value = await self._operation(item)
+                result.results[index] = value
+            except Exception as e:
+                result.errors[index] = e
+                if self._fail_fast:
+                    raise
+
+    async def execute_with_progress(
+        self,
+        items: Sequence[T],
+        on_progress: Callable[[int, int], None] | None = None,
+    ) -> BatchResult[R]:
+        """Execute with progress callback.
+
+        Args:
+            items: Items to process
+            on_progress: Callback(completed, total)
+
+        Returns:
+            BatchResult with results and errors
+        """
+        start_time = time.time()
+        result = BatchResult[R]()
+        completed = 0
+        total = len(items)
+
+        # Initialize result lists
+        result.results = [None] * total
+        result.errors = [None] * total
+
+        async def execute_with_callback(item: T, index: int) -> None:
+            nonlocal completed
+            async with self._semaphore:
+                try:
+                    value = await self._operation(item)
+                    result.results[index] = value
+                except Exception as e:
+                    result.errors[index] = e
+
+                completed += 1
+                if on_progress:
+                    on_progress(completed, total)
+
+        # Execute all tasks
+        tasks = [execute_with_callback(item, idx) for idx, item in enumerate(items)]
+        await asyncio.gather(*tasks, return_exceptions=True)
+
+        result.total_time_ms = (time.time() - start_time) * 1000
+        return result
+
+    @property
+    def max_concurrent(self) -> int:
+        """Get maximum concurrent operations."""
+        return self._max_concurrent
+
+
+async def batch_execute(
+    items: Sequence[T],
+    operation: Callable[[T], Awaitable[R]],
+    max_concurrent: int = 10,
+    fail_fast: bool = False,
+) -> BatchResult[R]:
+    """Convenience function for batch execution.
+
+    Args:
+        items: Items to process
+        operation: Async function for each item
+        max_concurrent: Maximum concurrent operations
+        fail_fast: Stop on first error
+
+    Returns:
+        BatchResult with results and errors
+    """
+    executor = BatchExecutor(operation, max_concurrent, fail_fast)
+    return await executor.execute(items)

ai_lib_python/cache/__init__.py

@@ -0,0 +1,26 @@
+"""
+Response caching module for ai-lib-python.
+
+Provides caching for AI responses with TTL and multiple backends.
+"""
+
+from ai_lib_python.cache.backends import (
+    CacheBackend,
+    DiskCache,
+    MemoryCache,
+    NullCache,
+)
+from ai_lib_python.cache.key import CacheKey, CacheKeyGenerator
+from ai_lib_python.cache.manager import CacheConfig, CacheManager, CacheStats
+
+__all__ = [
+    "CacheBackend",
+    "CacheConfig",
+    "CacheKey",
+    "CacheKeyGenerator",
+    "CacheManager",
+    "CacheStats",
+    "DiskCache",
+    "MemoryCache",
+    "NullCache",
+]