ai-lib-python 0.7.0__py3-none-any.whl → 0.8.2__py3-none-any.whl

ai_lib_python/__init__.py

@@ -27,7 +27,7 @@ from ai_lib_python.types.message import (
 )
 from ai_lib_python.types.tool import ToolCall, ToolDefinition
 
-__version__ = "0.7.0"
+__version__ = "0.7.5"
 
 __all__ = [
     # Client

ai_lib_python/batch/collector.py

@@ -191,7 +191,10 @@ class BatchCollector(Generic[T, R]):
 
         try:
             # Execute batch
-            results = await self._executor(data_list)
+            executor = self._executor
+            if executor is None:
+                raise RuntimeError("No executor set")
+            results = await executor(data_list)
 
             # Resolve futures
             for request, result in zip(requests, results, strict=False):

ai_lib_python/cache/manager.py

@@ -130,6 +130,7 @@ class CacheManager:
         self._key_generator = CacheKeyGenerator()
         self._stats = CacheStats()
 
+        self._backend: CacheBackend
         if not self._config.enabled:
             self._backend = NullCache()
         else:

ai_lib_python/client/__init__.py

@@ -10,8 +10,8 @@ This module provides:
 
 from ai_lib_python.client.builder import AiClientBuilder, ChatRequestBuilder
 from ai_lib_python.client.cancel import (
-    CancellableStream,
     CancelHandle,
+    CancellableStream,
     CancelReason,
     CancelState,
     CancelToken,
@@ -25,11 +25,11 @@ __all__ = [
     "AiClient",
     "AiClientBuilder",
     "CallStats",
-    "CancellableStream",
     "CancelHandle",
     "CancelReason",
     "CancelState",
     "CancelToken",
+    "CancellableStream",
     "ChatRequestBuilder",
     "ChatResponse",
     "create_cancel_pair",

ai_lib_python/client/builder.py

@@ -52,6 +52,7 @@ class AiClientBuilder:
         self._rate_limit_config: RateLimiterConfig | None = None
         self._circuit_breaker_config: CircuitBreakerConfig | None = None
         self._resilient_config: ResilientConfig | None = None
+        self._api_keys: dict[str, str] = {}
 
     def model(self, model_id: str) -> AiClientBuilder:
         """Set the model to use.
@@ -149,6 +150,37 @@ class AiClientBuilder:
         self._max_inflight = n
         return self
 
+    def retry(self, max_attempts: int = 3, backoff: float = 1.0) -> AiClientBuilder:
+        """Configure retry policy with simple parameters.
+
+        Args:
+            max_attempts: Maximum number of attempts (including initial)
+            backoff: Backoff base in seconds
+
+        Returns:
+            Self for chaining
+        """
+        from ai_lib_python.resilience import RetryConfig
+
+        self._retry_config = RetryConfig(
+            max_retries=max_attempts - 1,
+            min_delay_ms=int(backoff * 1000),
+        )
+        return self
+
+    def api_key_for(self, model_id: str, key: str) -> AiClientBuilder:
+        """Set API key for a specific fallback model.
+
+        Args:
+            model_id: Model identifier
+            key: API key to use for this model
+
+        Returns:
+            Self for chaining
+        """
+        self._api_keys[model_id] = key
+        return self
+
     def with_retry(self, config: RetryConfig) -> AiClientBuilder:
         """Configure retry policy.
 
@@ -252,6 +284,7 @@ class AiClientBuilder:
             timeout=self._timeout,
             hot_reload=self._hot_reload,
             resilient_config=resilient_config,
+            api_keys=self._api_keys,
         )
 
 

ai_lib_python/client/cancel.py

@@ -77,8 +77,12 @@ class CancelToken:
 
     def _start_timeout(self) -> None:
         """Start the timeout task."""
+        timeout = self._timeout
+        if timeout is None:
+            return
+
         async def timeout_handler() -> None:
-            await asyncio.sleep(self._timeout)  # type: ignore
+            await asyncio.sleep(timeout)
             if not self._state.cancelled:
                 self.cancel(CancelReason.TIMEOUT)
 
@@ -125,7 +129,7 @@ class CancelToken:
             try:
                 result = callback(reason)
                 if asyncio.iscoroutine(result):
-                    _ = asyncio.create_task(result)  # type: ignore  # noqa: RUF006
+                    _ = asyncio.create_task(result)  # noqa: RUF006
             except Exception:
                 pass
 
@@ -185,7 +189,7 @@ class CancelToken:
             try:
                 result = callback(self._state.reason)
                 if asyncio.iscoroutine(result):
-                    _ = asyncio.create_task(result)  # type: ignore  # noqa: RUF006
+                    _ = asyncio.create_task(result)  # noqa: RUF006
             except Exception:
                 pass
         return self

ai_lib_python/client/core.py

@@ -56,6 +56,10 @@ class AiClient:
         model_id: str,
         fallbacks: list[str] | None = None,
         executor: ResilientExecutor | None = None,
+        loader: ProtocolLoader | None = None,
+        api_keys: dict[str, str] | None = None,
+        base_url_override: str | None = None,
+        timeout: float | None = None,
     ) -> None:
         """Initialize the client (internal use).
 
@@ -67,6 +71,10 @@ class AiClient:
         self._model_id = model_id
         self._fallbacks = fallbacks or []
         self._executor = executor
+        self._loader = loader
+        self._api_keys = api_keys or {}
+        self._base_url_override = base_url_override
+        self._timeout = timeout
 
     @classmethod
     async def create(
@@ -131,6 +139,7 @@ class AiClient:
         timeout: float | None = None,
         hot_reload: bool = False,
         resilient_config: ResilientConfig | None = None,
+        api_keys: dict[str, str] | None = None,
     ) -> AiClient:
         """Internal creation method.
 
@@ -181,9 +190,13 @@ class AiClient:
             manifest=manifest,
             transport=transport,
             pipeline=pipeline,
-            model_id=model_id,
+            model_id=model,  # Keep the full model name including provider
             fallbacks=fallbacks,
             executor=executor,
+            loader=loader,
+            api_keys=api_keys,
+            base_url_override=base_url_override,
+            timeout=timeout,
         )
 
     def chat(self) -> ChatRequestBuilder:
@@ -203,7 +216,7 @@ class AiClient:
         return ChatRequestBuilder(self)
 
     async def _execute_chat(self, builder: ChatRequestBuilder) -> ChatResponse:
-        """Execute a non-streaming chat request.
+        """Execute a non-streaming chat request with fallback support.
 
         Args:
             builder: Configured request builder
@@ -211,19 +224,89 @@ class AiClient:
         Returns:
             ChatResponse with the completion
         """
-        async def do_request() -> ChatResponse:
-            payload = builder.build_payload()
-            endpoint = self._manifest.get_chat_endpoint()
+        models_to_try = [self._model_id, *self._fallbacks]
+        last_error = None
 
-            response = await self._transport.post(endpoint, json=payload)
-            data = response.json()
-
-            return self._parse_response(data)
-
-        # Use executor if available for resilience
-        if self._executor:
-            return await self._executor.execute(do_request)
-        return await do_request()
+        for model in models_to_try:
+            try:
+                # 1. Resolve manifest and transport for this model
+                if model == self._model_id:
+                    manifest = self._manifest
+                    transport = self._transport
+                    pipeline = self._pipeline
+                else:
+                    # Dynamic load for fallback
+                    if not self._loader:
+                        raise ValueError("ProtocolLoader missing for fallback")
+                    manifest = await self._loader.load_model(model)
+
+                    parts = model.split("/")
+                    m_id = parts[1] if len(parts) >= 2 else model
+
+                    # Resolve key for this model
+                    m_key = self._api_keys.get(model)
+
+                    from ai_lib_python.transport import HttpTransport
+                    transport = HttpTransport(
+                        manifest=manifest,
+                        model_id=m_id,
+                        api_key=m_key,
+                        base_url_override=self._base_url_override,
+                        timeout=self._timeout,
+                    )
+                    pipeline = Pipeline.from_manifest(manifest)
+
+                async def do_request(
+                    m: ProtocolManifest = manifest,
+                    t: HttpTransport = transport,
+                    p: Pipeline = pipeline,
+                    mid: str = model,
+                ) -> ChatResponse:
+                    # Debug print for model being used
+                    print(f"DEBUG: Executing request for model: {mid}, manifest ID: {m.id}")
+
+                    # Update builder's client temporary context?
+                    # Actually builder.build_payload() uses self._client._model_id
+                    # This is tricky as builder is bound to the primary client.
+                    # We need to temporarily override the client context in the builder.
+
+                    # Create a temporary builder/payload
+                    # For simplicity, we'll just manually build the payload here or
+                    # temporarily swap self._model_id (hacky but it works for this pattern)
+                    original_model_id = self._model_id
+                    original_manifest = self._manifest
+                    try:
+                        self._model_id = mid
+                        self._manifest = m
+                        payload = builder.build_payload()
+                        print(f"DEBUG: Payload model: {payload.get('model')}")
+                    finally:
+                        self._model_id = original_model_id
+                        self._manifest = original_manifest
+
+                    endpoint = m.get_chat_endpoint()
+                    print(f"DEBUG: Endpoint: {endpoint}")
+                    response = await t.post(endpoint, json=payload)
+                    data = response.json()
+
+                    # Parse using the correct pipeline
+                    return self._parse_response(data)
+
+                # Use executor if available for resilience
+                if self._executor:
+                    return await self._executor.execute(do_request)
+                return await do_request()
+
+            except Exception as e:
+                from ai_lib_python.errors import is_fallbackable
+                # Check if we should fallback
+                error_class = getattr(e, "error_class", None)
+                if model != models_to_try[-1] and (error_class is None or is_fallbackable(error_class)):
+                    last_error = e
+                    continue
+                raise e
+
+        raise last_error or RuntimeError("Fallback exhausted")
 
     async def _execute_chat_with_stats(
         self, builder: ChatRequestBuilder

ai_lib_python/computer_use/__init__.py

@@ -14,7 +14,6 @@ from enum import Enum
 from typing import Any
 from urllib.parse import urlparse
 
-
 # ─── Normalized Action Types ────────────────────────────────────────────────
 
 

ai_lib_python/drivers/anthropic.py

@@ -162,6 +162,16 @@ class AnthropicDriver(ProviderDriver):
                     system_parts.append(m.content)
                 continue
 
+            if role == "tool":
+                # Anthropic: tool results as user message with tool_result block
+                tool_id = getattr(m, "tool_call_id", None)
+                if tool_id and isinstance(m.content, str):
+                    msgs.append({
+                        "role": "user",
+                        "content": [{"type": "tool_result", "tool_use_id": tool_id, "content": m.content}],
+                    })
+                continue
+
             if isinstance(m.content, str):
                 content: Any = [{"type": "text", "text": m.content}]
             else:

ai_lib_python/drivers/gemini.py

@@ -55,7 +55,7 @@ class GeminiDriver(ProviderDriver):
     def build_request(
         self,
         messages: list[Message],
-        model: str,
+        _model: str,
         *,
         temperature: float | None = None,
         max_tokens: int | None = None,
@@ -161,6 +161,16 @@ class GeminiDriver(ProviderDriver):
                     system_parts.append(m.content)
                 continue
 
+            if role == "tool":
+                # Gemini: function_response with name (tool_call_id) and response
+                tool_id = getattr(m, "tool_call_id", None)
+                if tool_id and isinstance(m.content, str):
+                    contents.append({
+                        "role": "user",
+                        "parts": [{"functionResponse": {"name": tool_id, "response": {"result": m.content}}}],
+                    })
+                continue
+
             gemini_role = "model" if role == "assistant" else "user"
 
             if isinstance(m.content, str):

ai_lib_python/drivers/openai.py

@@ -18,7 +18,7 @@ from ai_lib_python.drivers import (
 from ai_lib_python.protocol.v2.capabilities import Capability
 from ai_lib_python.protocol.v2.manifest import ApiStyle
 from ai_lib_python.types.events import StreamingEvent
-from ai_lib_python.types.message import ContentBlock, Message
+from ai_lib_python.types.message import Message
 
 
 class OpenAiDriver(ProviderDriver):
@@ -120,14 +120,18 @@ class OpenAiDriver(ProviderDriver):
         # role is stored as str because model uses use_enum_values=True
         role = m.role if isinstance(m.role, str) else m.role.value
         if isinstance(m.content, str):
-            return {"role": role, "content": m.content}
-        # list[ContentBlock] → OpenAI content array
-        blocks = []
-        for b in m.content:
-            if b.type == "text":
-                blocks.append({"type": "text", "text": b.text})
-            elif b.type == "image":
-                blocks.append(b.model_dump(by_alias=True))
-            else:
-                blocks.append(b.model_dump(by_alias=True))
-        return {"role": role, "content": blocks}
+            out: dict[str, Any] = {"role": role, "content": m.content}
+        else:
+            # list[ContentBlock] → OpenAI content array
+            blocks = []
+            for b in m.content:
+                if b.type == "text":
+                    blocks.append({"type": "text", "text": b.text})
+                elif b.type == "image":
+                    blocks.append(b.model_dump(by_alias=True))
+                else:
+                    blocks.append(b.model_dump(by_alias=True))
+            out = {"role": role, "content": blocks}
+        if role == "tool" and getattr(m, "tool_call_id", None):
+            out["tool_call_id"] = m.tool_call_id
+        return out

ai_lib_python/embeddings/client.py

@@ -193,9 +193,12 @@ class EmbeddingClient:
         Returns:
             Endpoint path
         """
-        # Try to get from manifest, default to OpenAI-style
-        if hasattr(self._manifest, "embedding_endpoint"):
-            return self._manifest.embedding_endpoint
+        # Prefer manifest endpoint mapping when present.
+        embedding_cfg = self._manifest.endpoints.get("embeddings")
+        if isinstance(embedding_cfg, dict):
+            path = embedding_cfg.get("path")
+            if isinstance(path, str):
+                return path
         return "/v1/embeddings"
 
     @property
@@ -325,8 +328,9 @@ class EmbeddingClientBuilder:
         manifest = await loader.load_provider(provider_id)
 
         # Create transport
-        transport = HttpTransport.from_manifest(
-            manifest,
+        transport = HttpTransport(
+            manifest=manifest,
+            model_id=model_id,
             api_key=self._api_key,
             base_url_override=self._base_url,
             timeout=self._timeout,

ai_lib_python/guardrails/__init__.py

@@ -7,14 +7,14 @@ both user inputs and AI model outputs to ensure safety and compliance.
 Core principle: All logic is operators, all configuration is protocol.
 """
 
-from ai_lib_python.guardrails.base import Guardrail, GuardrailViolation, GuardrailResult
+from ai_lib_python.guardrails.base import Guardrail, GuardrailResult, GuardrailViolation
 from ai_lib_python.guardrails.filters import (
+    EmailFilter,
     KeywordFilter,
-    RegexFilter,
     LengthFilter,
     ProfanityFilter,
+    RegexFilter,
     UrlFilter,
-    EmailFilter,
 )
 from ai_lib_python.guardrails.validators import ContentValidator
 

ai_lib_python/guardrails/base.py

@@ -6,7 +6,7 @@ from __future__ import annotations
 
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -29,14 +29,14 @@ class GuardrailViolation:
     message: str
     severity: GuardrailSeverity
     matched_text: str | None = None
-    metadata: dict = field(default_factory=dict)
+    metadata: dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self) -> None:
         """Validate severity."""
         if not isinstance(self.severity, GuardrailSeverity):
             self.severity = GuardrailSeverity(self.severity)
 
-    def to_dict(self) -> dict:
+    def to_dict(self) -> dict[str, Any]:
         """Convert violation to dictionary."""
         return {
             "rule_id": self.rule_id,
@@ -54,10 +54,10 @@ class GuardrailResult:
     is_safe: bool
     violations: list[GuardrailViolation] = field(default_factory=list)
     filtered_content: str | None = None
-    metadata: dict = field(default_factory=dict)
+    metadata: dict[str, Any] = field(default_factory=dict)
 
     @classmethod
-    def safe(cls, content: str | None = None) -> "GuardrailResult":
+    def safe(cls, content: str | None = None) -> GuardrailResult:
         """Create a safe result."""
         return cls(is_safe=True, violations=[], filtered_content=content)
 
@@ -66,7 +66,7 @@ class GuardrailResult:
         cls,
         violations: list[GuardrailViolation],
         filtered_content: str | None = None,
-    ) -> "GuardrailResult":
+    ) -> GuardrailResult:
         """Create a violated result."""
         return cls(
             is_safe=False,
@@ -74,7 +74,7 @@ class GuardrailResult:
             filtered_content=filtered_content,
         )
 
-    def to_dict(self) -> dict:
+    def to_dict(self) -> dict[str, Any]:
         """Convert result to dictionary."""
         return {
             "is_safe": self.is_safe,
@@ -303,7 +303,7 @@ class ConditionalGuardrail(Guardrail):
         self,
         rule_id: str,
         guardrail: Guardrail,
-        condition: Callable[[dict], bool],
+        condition: Callable[[dict[str, Any]], bool],
         severity: GuardrailSeverity = GuardrailSeverity.WARNING,
     ) -> None:
         """Initialize conditional guardrail.
@@ -317,9 +317,9 @@ class ConditionalGuardrail(Guardrail):
         super().__init__(rule_id, severity)
         self._guardrail = guardrail
         self._condition = condition
-        self._context: dict = {}
+        self._context: dict[str, Any] = {}
 
-    def set_context(self, context: dict) -> None:
+    def set_context(self, context: dict[str, Any]) -> None:
         """Set the context for condition evaluation."""
         self._context = context
 

ai_lib_python/guardrails/filters.py

@@ -5,14 +5,13 @@ Concrete filter implementations for common guardrail use cases.
 from __future__ import annotations
 
 import re
-import string
 from typing import TYPE_CHECKING
 
 from ai_lib_python.guardrails.base import (
-    CompositeGuardrail,
     Guardrail,
     GuardrailResult,
     GuardrailSeverity,
+    GuardrailViolation,
 )
 
 if TYPE_CHECKING:
@@ -121,19 +120,13 @@ class KeywordFilter(Guardrail):
         self,
         message: str,
         matched_text: str,
-    ) -> "GuardrailResult":
-        """Create a violation result."""
-        from ai_lib_python.guardrails.base import GuardrailViolation
-
-        return GuardrailResult.violated(
-            [
-                GuardrailViolation(
-                    rule_id=self._rule_id,
-                    message=message,
-                    severity=self._severity,
-                    matched_text=matched_text,
-                )
-            ]
+    ) -> GuardrailViolation:
+        """Create a violation."""
+        return GuardrailViolation(
+            rule_id=self._rule_id,
+            message=message,
+            severity=self._severity,
+            matched_text=matched_text,
         )
 
 
@@ -176,7 +169,7 @@ class RegexFilter(Guardrail):
         self._message = message
 
     @property
-    def pattern(self) -> re.Pattern:
+    def pattern(self) -> re.Pattern[str]:
         """Get the compiled pattern."""
         return self._pattern
 
@@ -253,6 +246,7 @@ class LengthFilter(Guardrail):
         self._min_length = min_length
         self._max_length = max_length
         self._count_mode = count_mode
+        self._counter: Callable[[str], int]
 
         if count_mode == "chars":
             self._counter = len
@@ -273,8 +267,6 @@ class LengthFilter(Guardrail):
         violations = []
 
         if self._min_length is not None and length < self._min_length:
-            from ai_lib_python.guardrails.base import GuardrailViolation
-
             violations.append(
                 GuardrailViolation(
                     rule_id=self._rule_id,
@@ -284,8 +276,6 @@ class LengthFilter(Guardrail):
             )
 
         if self._max_length is not None and length > self._max_length:
-            from ai_lib_python.guardrails.base import GuardrailViolation
-
             violations.append(
                 GuardrailViolation(
                     rule_id=self._rule_id,
@@ -352,12 +342,12 @@ class ProfanityFilter(Guardrail):
                 # Find the actual matched text (preserving case)
                 pattern = re.escape(keyword)
                 if not self._case_sensitive:
-                    pattern = re.compile(pattern, re.IGNORECASE)
+                    pattern_obj = re.compile(pattern, re.IGNORECASE)
+                else:
+                    pattern_obj = re.compile(pattern)
 
-                match = pattern.search(content)
+                match = pattern_obj.search(content)
                 if match:
-                    from ai_lib_python.guardrails.base import GuardrailViolation
-
                     violations.append(
                         GuardrailViolation(
                             rule_id=self._rule_id,
@@ -463,19 +453,13 @@ class UrlFilter(Guardrail):
 
         return GuardrailResult.safe(content=content)
 
-    def _create_url_violation(self, message: str, url: str) -> "GuardrailResult":
-        """Create a URL violation result."""
-        from ai_lib_python.guardrails.base import GuardrailViolation
-
-        return GuardrailResult.violated(
-            [
-                GuardrailViolation(
-                    rule_id=self._rule_id,
-                    message=message,
-                    severity=self._severity,
-                    matched_text=url,
-                )
-            ]
+    def _create_url_violation(self, message: str, url: str) -> GuardrailViolation:
+        """Create a URL violation."""
+        return GuardrailViolation(
+            rule_id=self._rule_id,
+            message=message,
+            severity=self._severity,
+            matched_text=url,
         )
 
 
@@ -567,17 +551,11 @@ class EmailFilter(Guardrail):
         """Replace email addresses."""
         return self._EMAIL_PATTERN.sub(self._replacement, content)
 
-    def _create_email_violation(self, message: str, email: str) -> "GuardrailResult":
-        """Create an email violation result."""
-        from ai_lib_python.guardrails.base import GuardrailViolation
-
-        return GuardrailResult.violated(
-            [
-                GuardrailViolation(
-                    rule_id=self._rule_id,
-                    message=message,
-                    severity=self._severity,
-                    matched_text=email,
-                )
-            ]
+    def _create_email_violation(self, message: str, email: str) -> GuardrailViolation:
+        """Create an email violation."""
+        return GuardrailViolation(
+            rule_id=self._rule_id,
+            message=message,
+            severity=self._severity,
+            matched_text=email,
         )