askai-python 0.5.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

ask_ai/__init__.py

@@ -2,7 +2,7 @@
 ask-ai: A minimal Python SDK to switch between LLM providers in one line.
 """
 
-__version__ = "0.5.0"
+__version__ = "0.6.0"
 
 from .base import Response
 from .capabilities import Capabilities
@@ -17,6 +17,7 @@ from .exceptions import (
     MediaSecurityError,
     MediaTypeNotSupportedError,
     ProviderAuthError,
+    ProviderConnectionError,
     ProviderError,
     ProviderInvalidRequestError,
     ProviderRateLimitError,
@@ -54,6 +55,7 @@ __all__ = [
     "AskAIParsingError",
     "ProviderError",
     "ProviderAuthError",
+    "ProviderConnectionError",
     "ProviderRateLimitError",
     "ProviderTimeoutError",
     "ProviderInvalidRequestError",

ask_ai/base.py

@@ -7,13 +7,14 @@ from collections.abc import AsyncIterator, Iterator
 from typing import Any, Optional, Union
 
 from .capabilities import Capabilities
-from .config import AdvancedConfig
+from .config import CONFIG_OPTION_NAMES, AdvancedConfig, validate_option_names
 from .events import Completed, Error, TextDelta, ToolCall, UsageUpdate
 from .exceptions import StructuredValidationError
 from .media import AudioObject, ImageObject
 from .observability import Observability, ResponseMetadata, Usage
 
 logger = logging.getLogger("ask_ai")
+DEFAULT_RETRY_COUNT = 2
 
 
 class Response:
@@ -26,12 +27,14 @@ class Response:
         media: Union[ImageObject, AudioObject, None] = None,
         response_model: Any = None,
         metadata: Any = None,
+        usage: Optional[Usage] = None,
         **kwargs,
     ):
         self._raw_text = text
         self.media = media
         self.response_model = response_model
         self.metadata = metadata
+        self.usage = usage
         self._kwargs = kwargs
 
     @property
@@ -110,7 +113,18 @@ class BaseProvider:
     Abstract base class for all AI providers.
     Implements the core 'ask' logic and configuration management.
     """
-    def __init__(self, api_key: str = None, model: str = None, persona: str = None, **kwargs):
+    option_names = frozenset()
+
+    def __init__(
+        self,
+        api_key: str = None,
+        model: str = None,
+        persona: str = None,
+        strict_options: bool = False,
+        **kwargs,
+    ):
+        self.strict_options = strict_options
+        self._validate_option_names(kwargs, CONFIG_OPTION_NAMES | self.option_names, strict_options)
         # Zero-Config: Try env var if key not provided
         self.api_key = api_key or os.environ.get(self._get_api_key_env_var())
         self.model = model or self._get_default_model()
@@ -118,14 +132,33 @@ class BaseProvider:
         self.config = AdvancedConfig(**kwargs)
         self._middleware: list[Observability] = []
 
+    @staticmethod
+    def _validate_option_names(options, allowed, strict=False):
+        validate_option_names(options, allowed, strict)
+
     def use(self, middleware: Observability) -> "BaseProvider":
         """Attach optional observability middleware and return this provider."""
         self._middleware.append(middleware)
         return self
 
-    def session(self, max_messages: Optional[int] = None):
+    def session(
+        self,
+        max_messages: Optional[int] = None,
+        strict_options: Optional[bool] = None,
+        **kwargs,
+    ):
+        strict = (
+            getattr(self, "strict_options", False)
+            if strict_options is None
+            else strict_options
+        )
+        validate_option_names(kwargs, {"max_messages", "strict_options"}, strict)
         from .session import ConversationSession
-        return ConversationSession(self, max_messages=max_messages)
+        return ConversationSession(
+            self,
+            max_messages=max_messages,
+            strict_options=strict,
+        )
 
     def capabilities(self) -> Capabilities:
         """Return conservative capabilities for custom providers."""
@@ -197,6 +230,13 @@ class BaseProvider:
                 response = self._execute_with_retry(
                     self._send_request, request_messages, config, output_type, retry
                 )
+                response._kwargs.update(
+                    {
+                        option: config.extra[option]
+                        for option in ("clean", "strip", "code")
+                        if option in config.extra
+                    }
+                )
                 retry_count += getattr(self, "_retry_count", 0)
                 if response_model is None:
                     return self._attach_metadata(response, started_at, retry_count)
@@ -234,6 +274,13 @@ class BaseProvider:
                 response = await self._execute_with_retry_async(
                     self._send_request_async, request_messages, config, output_type, retry
                 )
+                response._kwargs.update(
+                    {
+                        option: config.extra[option]
+                        for option in ("clean", "strip", "code")
+                        if option in config.extra
+                    }
+                )
                 retry_count += getattr(self, "_retry_count", 0)
                 if response_model is None:
                     return self._attach_metadata(response, started_at, retry_count)
@@ -265,6 +312,8 @@ class BaseProvider:
         Update global advanced settings for this instance.
         Merges new settings with existing ones.
         """
+        strict = kwargs.pop("strict_options", self.strict_options)
+        self._validate_option_names(kwargs, CONFIG_OPTION_NAMES | self.option_names, strict)
         new_conf = AdvancedConfig(**kwargs)
         for k, v in new_conf.__dict__.items():
             if k in ('extra', '_safe_mode_provided'):
@@ -308,7 +357,9 @@ class BaseProvider:
         final_config = self._build_final_config(kwargs)
         messages = self._prepare_messages(query, final_config)
         output_type = kwargs.get('output_type')
-        retry = int(kwargs.get('retry', final_config.extra.get('retry', 0)) or 0)
+        retry = int(
+            kwargs.get("retry", final_config.extra.get("retry", DEFAULT_RETRY_COUNT)) or 0
+        )
         response_model = kwargs.get('response_model')
         validation_retries = int(kwargs.get("validation_retries", 0) or 0)
 
@@ -352,7 +403,14 @@ class BaseProvider:
 
         final_config = self._build_final_config(kwargs)
         messages = self._prepare_messages(query, final_config)
-        yield from self._send_request_stream(messages, final_config)
+        retry = int(
+            kwargs.get("retry", final_config.extra.get("retry", DEFAULT_RETRY_COUNT))
+            or 0
+        )
+        yield from self._iter_with_retry(
+            lambda: self._send_request_stream(messages, final_config),
+            retry,
+        )
 
     def ask_stream_events(self, query: str, **kwargs) -> Iterator[Any]:
         """Yield typed events while preserving the legacy text stream API."""
@@ -360,18 +418,44 @@ class BaseProvider:
         if fallback_providers:
             if not isinstance(fallback_providers, (list, tuple)):
                 fallback_providers = [fallback_providers]
+            last_error = None
             for provider in fallback_providers:
-                provider_instance = provider() if isinstance(provider, type) else provider
-                yield from provider_instance.ask_stream_events(query, **kwargs)
-                return
+                emitted = False
+                try:
+                    provider_instance = provider() if isinstance(provider, type) else provider
+                    for event in provider_instance.ask_stream_events(query, **kwargs):
+                        if isinstance(event, Error):
+                            if not emitted and event.emitted_chunks == 0:
+                                last_error = event
+                                break
+                            yield event
+                            return
+                        emitted = True
+                        yield event
+                    else:
+                        return
+                except Exception as error:
+                    if emitted:
+                        raise
+                    last_error = Error(error=error, emitted_chunks=0)
+            if last_error is not None:
+                yield last_error
+            return
 
         emitted_chunks = 0
         text_parts = []
         usage = None
         final_config = self._build_final_config(kwargs)
         messages = self._prepare_messages(query, final_config)
+        retry = int(
+            kwargs.get("retry", final_config.extra.get("retry", DEFAULT_RETRY_COUNT))
+            or 0
+        )
         try:
-            for event in self._send_request_stream_events(messages, final_config):
+            for event in self._iter_with_retry(
+                lambda: self._send_request_stream_events(messages, final_config),
+                retry,
+            ):
                 if isinstance(event, TextDelta):
                     emitted_chunks += 1
                     text_parts.append(event.text)
@@ -419,7 +503,9 @@ class BaseProvider:
         final_config = self._build_final_config(kwargs)
         messages = self._prepare_messages(query, final_config)
         output_type = kwargs.get('output_type')
-        retry = int(kwargs.get('retry', final_config.extra.get('retry', 0)) or 0)
+        retry = int(
+            kwargs.get("retry", final_config.extra.get("retry", DEFAULT_RETRY_COUNT)) or 0
+        )
         response_model = kwargs.get('response_model')
         validation_retries = int(kwargs.get("validation_retries", 0) or 0)
 
@@ -463,7 +549,14 @@ class BaseProvider:
 
         final_config = self._build_final_config(kwargs)
         messages = self._prepare_messages(query, final_config)
-        async for chunk in self._send_request_stream_async(messages, final_config):
+        retry = int(
+            kwargs.get("retry", final_config.extra.get("retry", DEFAULT_RETRY_COUNT))
+            or 0
+        )
+        async for chunk in self._aiter_with_retry(
+            lambda: self._send_request_stream_async(messages, final_config),
+            retry,
+        ):
             yield chunk
 
     async def ask_stream_events_async(self, query: str, **kwargs) -> AsyncIterator[Any]:
@@ -472,19 +565,44 @@ class BaseProvider:
         if fallback_providers:
             if not isinstance(fallback_providers, (list, tuple)):
                 fallback_providers = [fallback_providers]
+            last_error = None
             for provider in fallback_providers:
-                provider_instance = provider() if isinstance(provider, type) else provider
-                async for event in provider_instance.ask_stream_events_async(query, **kwargs):
-                    yield event
-                return
+                emitted = False
+                try:
+                    provider_instance = provider() if isinstance(provider, type) else provider
+                    async for event in provider_instance.ask_stream_events_async(query, **kwargs):
+                        if isinstance(event, Error):
+                            if not emitted and event.emitted_chunks == 0:
+                                last_error = event
+                                break
+                            yield event
+                            return
+                        emitted = True
+                        yield event
+                    else:
+                        return
+                except Exception as error:
+                    if emitted:
+                        raise
+                    last_error = Error(error=error, emitted_chunks=0)
+            if last_error is not None:
+                yield last_error
+            return
 
         emitted_chunks = 0
         text_parts = []
         usage = None
         final_config = self._build_final_config(kwargs)
         messages = self._prepare_messages(query, final_config)
+        retry = int(
+            kwargs.get("retry", final_config.extra.get("retry", DEFAULT_RETRY_COUNT))
+            or 0
+        )
         try:
-            async for event in self._send_request_stream_events_async(messages, final_config):
+            async for event in self._aiter_with_retry(
+                lambda: self._send_request_stream_events_async(messages, final_config),
+                retry,
+            ):
                 if isinstance(event, TextDelta):
                     emitted_chunks += 1
                     text_parts.append(event.text)
@@ -504,11 +622,20 @@ class BaseProvider:
 
     def _build_final_config(self, kwargs: dict[str, Any]) -> AdvancedConfig:
         """Build final config by merging global config with per-request kwargs."""
-        request_config = AdvancedConfig(**kwargs)
+        options = dict(kwargs)
+        strict = options.pop("strict_options", getattr(self, "strict_options", False))
+        validate_option_names(
+            options,
+            CONFIG_OPTION_NAMES
+            | self.option_names
+            | {"strict_options", "providers", "fallback"},
+            strict,
+        )
+        request_config = AdvancedConfig(**options)
         final_config = self.config.merge(request_config)
 
         # Ensure timeout has a default
-        timeout = kwargs.get('timeout', final_config.extra.get('timeout', 30))
+        timeout = options.get('timeout', final_config.extra.get('timeout', 30))
         final_config.extra['timeout'] = float(timeout) if timeout is not None else 30.0
 
         return final_config
@@ -521,14 +648,10 @@ class BaseProvider:
 
         # Smart capability: if json=True or response_model is specified
         response_model = config.extra.get('response_model')
-        if response_model:
-            if hasattr(response_model, "model_json_schema"):
-                schema = response_model.model_json_schema()
-            elif hasattr(response_model, "schema"):
-                schema = response_model.schema()
-            else:
-                schema = None
+        if response_model is not None:
+            from .structured import json_schema_for_model
 
+            schema = json_schema_for_model(response_model)
             if schema:
                 import json
                 schema_str = json.dumps(schema)
@@ -555,6 +678,58 @@ class BaseProvider:
 
     # ─── Retry Logic ────────────────────────────────────────
 
+    def _iter_with_retry(self, iterator_factory, retry: int):
+        attempts = 0
+        self._retry_count = 0
+        while True:
+            emitted = False
+            try:
+                for item in iterator_factory():
+                    emitted = True
+                    yield item
+                self._retry_count = attempts
+                return
+            except Exception as error:
+                if emitted or attempts >= retry or not self._is_retryable_error(error):
+                    raise
+                attempts += 1
+                self._retry_count = attempts
+                delay = min(2 ** attempts + random.uniform(0, 1), 30)
+                logger.warning(
+                    "Transient streaming error (%s). Retrying in %.1fs... (%d/%d)",
+                    error.__class__.__name__,
+                    delay,
+                    attempts,
+                    retry,
+                )
+                time.sleep(delay)
+
+    async def _aiter_with_retry(self, iterator_factory, retry: int):
+        attempts = 0
+        self._retry_count = 0
+        while True:
+            emitted = False
+            try:
+                async for item in iterator_factory():
+                    emitted = True
+                    yield item
+                self._retry_count = attempts
+                return
+            except Exception as error:
+                if emitted or attempts >= retry or not self._is_retryable_error(error):
+                    raise
+                attempts += 1
+                self._retry_count = attempts
+                delay = min(2 ** attempts + random.uniform(0, 1), 30)
+                logger.warning(
+                    "Transient async streaming error (%s). Retrying in %.1fs... (%d/%d)",
+                    error.__class__.__name__,
+                    delay,
+                    attempts,
+                    retry,
+                )
+                await asyncio.sleep(delay)
+
     def _execute_with_retry(self, func, messages, config, output_type, retry):
         """Execute a sync function with exponential backoff + jitter."""
         attempts = 0

ask_ai/config.py

@@ -1,7 +1,49 @@
-from typing import Any, Optional
+import difflib
+import warnings
+from typing import Any, Iterable, Optional
 
 _NOT_PROVIDED = object()
 
+CONFIG_OPTION_NAMES = frozenset(
+    {
+        "temperature",
+        "max_tokens",
+        "top_p",
+        "frequency_penalty",
+        "presence_penalty",
+        "stop",
+        "system_message",
+        "prompt",
+        "system",
+        "safe_mode",
+        "timeout",
+        "retry",
+        "response_model",
+        "validation_retries",
+        "output_type",
+        "clean",
+        "strip",
+        "code",
+        "json",
+    }
+)
+
+
+def validate_option_names(
+    options: Iterable[str],
+    allowed: Iterable[str],
+    strict: bool = False,
+) -> None:
+    allowed_set = set(allowed)
+    allowed_names = sorted(allowed_set)
+    for name in sorted(set(options) - allowed_set):
+        suggestions = difflib.get_close_matches(name, allowed_names, n=1, cutoff=0.6)
+        suggestion = f" Did you mean '{suggestions[0]}'?" if suggestions else ""
+        message = f"Unknown option '{name}'.{suggestion}"
+        if strict:
+            raise ValueError(message)
+        warnings.warn(message, UserWarning, stacklevel=3)
+
 
 class AdvancedConfig:
     """
@@ -37,6 +79,12 @@ class AdvancedConfig:
             kwargs.pop("system", None)
 
         self.system_message = system_message
+        if safe_mode is not _NOT_PROVIDED:
+            warnings.warn(
+                "safe_mode is deprecated and has no defined behavior; it will be removed in a future release.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
         self.safe_mode = False if safe_mode is _NOT_PROVIDED else safe_mode
         self._safe_mode_provided = safe_mode is not _NOT_PROVIDED
         self.extra: dict[str, Any] = kwargs

ask_ai/exceptions.py

@@ -38,6 +38,11 @@ class ProviderTimeoutError(AskAINetworkError, ProviderError):
     pass
 
 
+class ProviderConnectionError(AskAINetworkError, ProviderError):
+    """Raised when a provider request fails due to a network connection error."""
+    pass
+
+
 class ProviderInvalidRequestError(ProviderError):
     """Raised when a provider rejects request parameters."""
     pass
@@ -83,6 +88,8 @@ def map_provider_exception(error: Exception, provider: str, operation: str = "re
     if status_code is None:
         response = getattr(error, "response", None)
         status_code = getattr(response, "status_code", None)
+    if status_code is None:
+        status_code = getattr(error, "code", None)
 
     error_text = str(error).lower()
     error_name = error.__class__.__name__.lower()
@@ -98,11 +105,19 @@ def map_provider_exception(error: Exception, provider: str, operation: str = "re
         for marker in ("ratelimit", "rate_limit", "rate limit", "too many requests")
     ):
         return ProviderRateLimitError(f"{prefix}: {error}")
+    if any(marker in error_name or marker in error_text for marker in ("timeout", "timed out")):
+        return ProviderTimeoutError(f"{prefix}: {error}")
     if any(
         marker in error_name or marker in error_text
-        for marker in ("timeout", "timed out", "connecterror", "connectionerror", "connection reset")
+        for marker in (
+            "connecterror",
+            "connectionerror",
+            "connection reset",
+            "connection refused",
+            "connection error",
+        )
     ):
-        return ProviderTimeoutError(f"{prefix}: {error}")
+        return ProviderConnectionError(f"{prefix}: {error}")
     if isinstance(status_code, int) and status_code >= 500:
         return ProviderUnavailableError(f"{prefix}: {error}")
     if status_code in (400, 404, 422) or any(

ask_ai/media_security.py

@@ -16,7 +16,7 @@ class MediaPolicy:
     timeout: float = 30.0
     allow_private: bool = False
 
-    def validate_url(self, url: str) -> None:
+    def _validated_addresses(self, url: str) -> tuple[str, list[str]]:
         parsed = urlparse(url)
         if parsed.scheme != "https" or not parsed.hostname:
             raise MediaSecurityError("Media URLs must use HTTPS and include a hostname.")
@@ -53,6 +53,10 @@ class MediaPolicy:
                     raise MediaSecurityError(
                         f"Media host resolves to a non-public IP address: {address}"
                     )
+        return parsed.hostname, [str(address) for address in addresses]
+
+    def validate_url(self, url: str) -> None:
+        self._validated_addresses(url)
 
     def validate_mime(self, mime_type: str) -> None:
         if self.allowed_mime_types and mime_type not in self.allowed_mime_types:
@@ -62,14 +66,19 @@ class MediaPolicy:
 def download_media(url: str, policy: Optional[MediaPolicy] = None) -> bytes:
     """Stream media bytes with URL, MIME, size, and timeout validation."""
     policy = policy or MediaPolicy()
-    policy.validate_url(url)
+    hostname, addresses = policy._validated_addresses(url)
     try:
         import httpx
     except ImportError as error:
         raise MediaSecurityError("httpx is required for secure media downloads.") from error
 
     try:
-        with httpx.Client(timeout=policy.timeout, follow_redirects=False) as client:
+        transport = _create_pinned_transport(httpx, hostname, addresses)
+        with httpx.Client(
+            timeout=policy.timeout,
+            follow_redirects=False,
+            transport=transport,
+        ) as client:
             with client.stream("GET", url) as response:
                 if 300 <= response.status_code < 400:
                     raise MediaSecurityError(
@@ -103,3 +112,41 @@ def download_media(url: str, policy: Optional[MediaPolicy] = None) -> bytes:
         raise
     except Exception as error:
         raise MediaSecurityError(f"Secure media download failed: {error}") from error
+
+
+def _create_pinned_transport(httpx, hostname: str, addresses: list[str]):
+    import httpcore
+
+    normalized_hostname = hostname.encode("idna").decode("ascii").lower().rstrip(".")
+
+    class PinnedNetworkBackend(httpcore.SyncBackend):
+        def connect_tcp(
+            self,
+            host,
+            port,
+            timeout=None,
+            local_address=None,
+            socket_options=None,
+        ):
+            if host.lower().rstrip(".") != normalized_hostname:
+                raise httpcore.ConnectError(
+                    "Media transport attempted to connect to an unvalidated host."
+                )
+
+            last_error = None
+            for address in addresses:
+                try:
+                    return super().connect_tcp(
+                        address,
+                        port,
+                        timeout=timeout,
+                        local_address=local_address,
+                        socket_options=socket_options,
+                    )
+                except (httpcore.ConnectError, httpcore.ConnectTimeout) as error:
+                    last_error = error
+            raise last_error
+
+    transport = httpx.HTTPTransport()
+    transport._pool._network_backend = PinnedNetworkBackend()
+    return transport

ask_ai/observability.py

@@ -9,6 +9,37 @@ class Usage:
     total_tokens: Optional[int] = None
 
 
+def usage_from_fields(
+    source: Any,
+    input_field: str,
+    output_field: str,
+    total_field: Optional[str] = None,
+) -> Optional[Usage]:
+    """Normalize provider-specific token usage fields without inventing missing counts."""
+    if source is None:
+        return None
+
+    def field(name: Optional[str]) -> Optional[int]:
+        if name is None:
+            return None
+        if isinstance(source, dict):
+            return source.get(name)
+        return getattr(source, name, None)
+
+    input_tokens = field(input_field)
+    output_tokens = field(output_field)
+    total_tokens = field(total_field)
+    if input_tokens is None and output_tokens is None and total_tokens is None:
+        return None
+    if total_tokens is None and input_tokens is not None and output_tokens is not None:
+        total_tokens = input_tokens + output_tokens
+    return Usage(
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        total_tokens=total_tokens,
+    )
+
+
 @dataclass(frozen=True)
 class ResponseMetadata:
     provider: str