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

ask_ai/__init__.py

@@ -2,29 +2,36 @@
 ask-ai: A minimal Python SDK to switch between LLM providers in one line.
 """
 
-__version__ = "0.4.1"
+__version__ = "0.5.0"
 
-from .providers import (
-    OpenAI,
-    Anthropic,
-    Google,
-    Groq,
-    Azure,
-    OpenRouter
-)
-from .config import AdvancedConfig
 from .base import Response
-from .media import ImageObject, AudioObject
+from .capabilities import Capabilities
+from .config import AdvancedConfig
+from .events import Completed, Error, TextDelta, ToolCall, UsageUpdate
 from .exceptions import (
-    AskAIError,
     AskAIConfigError,
+    AskAIError,
     AskAINetworkError,
-    AskAIRateLimitError,
     AskAIParsingError,
-    ProviderError,
+    AskAIRateLimitError,
+    MediaSecurityError,
     MediaTypeNotSupportedError,
+    ProviderAuthError,
+    ProviderError,
+    ProviderInvalidRequestError,
+    ProviderRateLimitError,
+    ProviderTimeoutError,
+    ProviderUnavailableError,
     StreamNotSupportedError,
+    StructuredValidationError,
 )
+from .media import AudioObject, ImageObject
+from .media_security import MediaPolicy, download_media
+from .observability import Observability, ResponseMetadata, Usage
+from .providers import Anthropic, Azure, Google, Groq, OpenAI, OpenRouter
+from .router import Router
+from .session import ConversationSession
+from .structured import parse_structured
 
 __all__ = [
     # Providers
@@ -46,6 +53,27 @@ __all__ = [
     "AskAIRateLimitError",
     "AskAIParsingError",
     "ProviderError",
+    "ProviderAuthError",
+    "ProviderRateLimitError",
+    "ProviderTimeoutError",
+    "ProviderInvalidRequestError",
+    "ProviderUnavailableError",
     "MediaTypeNotSupportedError",
+    "MediaSecurityError",
     "StreamNotSupportedError",
+    "StructuredValidationError",
+    "parse_structured",
+    "TextDelta",
+    "ToolCall",
+    "UsageUpdate",
+    "Completed",
+    "Error",
+    "Observability",
+    "ResponseMetadata",
+    "Usage",
+    "Capabilities",
+    "Router",
+    "ConversationSession",
+    "MediaPolicy",
+    "download_media",
 ]

ask_ai/base.py

@@ -1,12 +1,17 @@
-import os
-import time
-import random
 import asyncio
 import logging
-from typing import Union, List, Dict, Any, Iterator, AsyncIterator
+import os
+import random
+import time
+from collections.abc import AsyncIterator, Iterator
+from typing import Any, Optional, Union
 
+from .capabilities import Capabilities
 from .config import AdvancedConfig
-from .media import ImageObject, AudioObject
+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")
 
@@ -15,10 +20,18 @@ class Response:
     """
     Unified response object for all ask-ai requests.
     """
-    def __init__(self, text: str = "", media: Union[ImageObject, AudioObject, None] = None, response_model: Any = None, **kwargs):
+    def __init__(
+        self,
+        text: str = "",
+        media: Union[ImageObject, AudioObject, None] = None,
+        response_model: Any = None,
+        metadata: Any = None,
+        **kwargs,
+    ):
         self._raw_text = text
         self.media = media
         self.response_model = response_model
+        self.metadata = metadata
         self._kwargs = kwargs
 
     @property
@@ -36,7 +49,20 @@ class Response:
         return text
 
     @property
-    def json(self) -> Union[Dict[str, Any], list]:
+    def raw(self) -> str:
+        """Return the unmodified provider text."""
+        return self._raw_text
+
+    @property
+    def parsed(self) -> Any:
+        """Return the validated structured response, if one was requested."""
+        if not self.response_model:
+            raise ValueError("No response_model was specified for this request.")
+        from .structured import parse_structured
+        return parse_structured(self._raw_text, self.response_model)
+
+    @property
+    def json(self) -> Union[dict[str, Any], list]:
         """Returns parsed JSON. If 'json=True' was passed, this is highly reliable."""
         from .utils import parse_json
         return parse_json(self._raw_text)
@@ -46,10 +72,13 @@ class Response:
         """Returns parsed Pydantic model instance if response_model was provided."""
         if not self.response_model:
             raise ValueError("No response_model was specified for this request.")
-        return self.to_model(self.response_model)
+        return self.parsed
 
     def to_model(self, model_class: Any) -> Any:
         """Parse the response text into a specified Pydantic model."""
+        from .structured import parse_structured
+        if isinstance(model_class, dict) or hasattr(model_class, "model_validate") or hasattr(model_class, "parse_obj"):
+            return parse_structured(self._raw_text, model_class)
         text_content = self.text
         if hasattr(model_class, "model_validate_json"):
             return model_class.model_validate_json(text_content)
@@ -87,6 +116,141 @@ class BaseProvider:
         self.model = model or self._get_default_model()
         self.persona = persona
         self.config = AdvancedConfig(**kwargs)
+        self._middleware: list[Observability] = []
+
+    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):
+        from .session import ConversationSession
+        return ConversationSession(self, max_messages=max_messages)
+
+    def capabilities(self) -> Capabilities:
+        """Return conservative capabilities for custom providers."""
+        return Capabilities()
+
+    def _attach_metadata(self, response: Response, started_at: float, retry_count: int) -> Response:
+        usage = getattr(response, "usage", None)
+        if isinstance(usage, dict):
+            usage = Usage(**usage)
+        metadata = ResponseMetadata(
+            provider=self.__class__.__name__,
+            model=self.model,
+            latency_ms=(time.perf_counter() - started_at) * 1000,
+            retry_count=retry_count,
+            usage=usage,
+            estimated_cost=None,
+        )
+        response.metadata = metadata
+        for middleware in getattr(self, "_middleware", []):
+            metadata = ResponseMetadata(
+                **{**metadata.__dict__, "estimated_cost": middleware.cost(self.model, usage)}
+            )
+            response.metadata = metadata
+            middleware.emit({"type": "completed", "metadata": metadata, "response": response})
+        return response
+
+    def _attach_failure_metadata(self, error: Exception, started_at: float) -> None:
+        metadata = ResponseMetadata(
+            provider=self.__class__.__name__,
+            model=self.model,
+            latency_ms=(time.perf_counter() - started_at) * 1000,
+            retry_count=getattr(self, "_retry_count", 0),
+            failure_reason=str(error),
+        )
+        for middleware in getattr(self, "_middleware", []):
+            middleware.emit({"type": "failed", "metadata": metadata, "error": error})
+
+    def _validation_retry_messages(
+        self, messages: list[dict[str, str]], response: Response, error: StructuredValidationError
+    ) -> list[dict[str, str]]:
+        return [
+            *messages,
+            {"role": "assistant", "content": response.raw},
+            {
+                "role": "user",
+                "content": (
+                    "Your previous response failed schema validation: "
+                    f"{error}. Return a corrected JSON response only."
+                ),
+            },
+        ]
+
+    def _execute_response(
+        self,
+        messages: list[dict[str, str]],
+        config: AdvancedConfig,
+        output_type: Optional[str],
+        retry: int,
+        response_model: Any,
+        validation_retries: int,
+        started_at: float,
+    ) -> Response:
+        if validation_retries < 0:
+            raise ValueError("validation_retries must be non-negative.")
+        request_messages = messages
+        retry_count = 0
+        try:
+            for attempt in range(validation_retries + 1):
+                response = self._execute_with_retry(
+                    self._send_request, request_messages, config, output_type, retry
+                )
+                retry_count += getattr(self, "_retry_count", 0)
+                if response_model is None:
+                    return self._attach_metadata(response, started_at, retry_count)
+                response.response_model = response_model
+                try:
+                    response.parsed
+                except StructuredValidationError as error:
+                    if attempt == validation_retries:
+                        raise
+                    request_messages = self._validation_retry_messages(messages, response, error)
+                else:
+                    return self._attach_metadata(response, started_at, retry_count)
+            raise RuntimeError("Structured validation retry loop exited unexpectedly.")
+        except Exception as error:
+            self._retry_count = retry_count + getattr(self, "_retry_count", 0)
+            self._attach_failure_metadata(error, started_at)
+            raise
+
+    async def _execute_response_async(
+        self,
+        messages: list[dict[str, str]],
+        config: AdvancedConfig,
+        output_type: Optional[str],
+        retry: int,
+        response_model: Any,
+        validation_retries: int,
+        started_at: float,
+    ) -> Response:
+        if validation_retries < 0:
+            raise ValueError("validation_retries must be non-negative.")
+        request_messages = messages
+        retry_count = 0
+        try:
+            for attempt in range(validation_retries + 1):
+                response = await self._execute_with_retry_async(
+                    self._send_request_async, request_messages, config, output_type, retry
+                )
+                retry_count += getattr(self, "_retry_count", 0)
+                if response_model is None:
+                    return self._attach_metadata(response, started_at, retry_count)
+                response.response_model = response_model
+                try:
+                    response.parsed
+                except StructuredValidationError as error:
+                    if attempt == validation_retries:
+                        raise
+                    request_messages = self._validation_retry_messages(messages, response, error)
+                else:
+                    return self._attach_metadata(response, started_at, retry_count)
+            raise RuntimeError("Structured validation retry loop exited unexpectedly.")
+        except Exception as error:
+            self._retry_count = retry_count + getattr(self, "_retry_count", 0)
+            self._attach_failure_metadata(error, started_at)
+            raise
 
     def _get_api_key_env_var(self) -> str:
         """Subclasses should return the env var name, e.g. 'OPENAI_API_KEY'"""
@@ -103,7 +267,9 @@ class BaseProvider:
         """
         new_conf = AdvancedConfig(**kwargs)
         for k, v in new_conf.__dict__.items():
-            if k == 'extra':
+            if k in ('extra', '_safe_mode_provided'):
+                continue
+            if k == 'safe_mode' and not new_conf._safe_mode_provided:
                 continue
             if v is not None:
                 setattr(self.config, k, v)
@@ -144,13 +310,13 @@ class BaseProvider:
         output_type = kwargs.get('output_type')
         retry = int(kwargs.get('retry', final_config.extra.get('retry', 0)) or 0)
         response_model = kwargs.get('response_model')
+        validation_retries = int(kwargs.get("validation_retries", 0) or 0)
 
-        response = self._execute_with_retry(
-            self._send_request, messages, final_config, output_type, retry
+        started_at = time.perf_counter()
+        return self._execute_response(
+            messages, final_config, output_type, retry,
+            response_model, validation_retries, started_at,
         )
-        if response_model:
-            response.response_model = response_model
-        return response
 
     def ask_stream(self, query: str, **kwargs) -> Iterator[str]:
         """
@@ -171,9 +337,14 @@ class BaseProvider:
                 else:
                     p_inst = p
                 try:
-                    yield from p_inst.ask_stream(query, **kwargs)
+                    emitted = False
+                    for chunk in p_inst.ask_stream(query, **kwargs):
+                        emitted = True
+                        yield chunk
                     return
                 except Exception as e:
+                    if emitted:
+                        raise
                     errors.append(f"{p_inst.__class__.__name__}: {e}")
                     logger.warning("Provider fallback failed for %s: %s", p_inst.__class__.__name__, e)
             from .exceptions import ProviderError
@@ -183,6 +354,39 @@ class BaseProvider:
         messages = self._prepare_messages(query, final_config)
         yield from self._send_request_stream(messages, final_config)
 
+    def ask_stream_events(self, query: str, **kwargs) -> Iterator[Any]:
+        """Yield typed events while preserving the legacy text stream API."""
+        fallback_providers = kwargs.pop("providers", None) or kwargs.pop("fallback", None)
+        if fallback_providers:
+            if not isinstance(fallback_providers, (list, tuple)):
+                fallback_providers = [fallback_providers]
+            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_chunks = 0
+        text_parts = []
+        usage = None
+        final_config = self._build_final_config(kwargs)
+        messages = self._prepare_messages(query, final_config)
+        try:
+            for event in self._send_request_stream_events(messages, final_config):
+                if isinstance(event, TextDelta):
+                    emitted_chunks += 1
+                    text_parts.append(event.text)
+                    yield event
+                elif isinstance(event, UsageUpdate):
+                    usage = event
+                    yield event
+                elif isinstance(event, ToolCall):
+                    yield event
+                else:
+                    yield event
+            yield Completed(response=Response(text="".join(text_parts)), usage=usage)
+        except Exception as error:
+            yield Error(error=error, emitted_chunks=emitted_chunks)
+
     # ─── Async API ──────────────────────────────────────────
 
     async def ask_async(self, query: str, **kwargs) -> Response:
@@ -217,13 +421,13 @@ class BaseProvider:
         output_type = kwargs.get('output_type')
         retry = int(kwargs.get('retry', final_config.extra.get('retry', 0)) or 0)
         response_model = kwargs.get('response_model')
+        validation_retries = int(kwargs.get("validation_retries", 0) or 0)
 
-        response = await self._execute_with_retry_async(
-            self._send_request_async, messages, final_config, output_type, retry
+        started_at = time.perf_counter()
+        return await self._execute_response_async(
+            messages, final_config, output_type, retry,
+            response_model, validation_retries, started_at,
         )
-        if response_model:
-            response.response_model = response_model
-        return response
 
     async def ask_stream_async(self, query: str, **kwargs) -> AsyncIterator[str]:
         """
@@ -244,10 +448,14 @@ class BaseProvider:
                 else:
                     p_inst = p
                 try:
-                    async for chunk in await p_inst.ask_stream_async(query, **kwargs):
+                    emitted = False
+                    async for chunk in p_inst.ask_stream_async(query, **kwargs):
+                        emitted = True
                         yield chunk
                     return
                 except Exception as e:
+                    if emitted:
+                        raise
                     errors.append(f"{p_inst.__class__.__name__}: {e}")
                     logger.warning("Provider fallback failed for %s: %s", p_inst.__class__.__name__, e)
             from .exceptions import ProviderError
@@ -258,9 +466,43 @@ class BaseProvider:
         async for chunk in self._send_request_stream_async(messages, final_config):
             yield chunk
 
+    async def ask_stream_events_async(self, query: str, **kwargs) -> AsyncIterator[Any]:
+        """Yield typed events while preserving the legacy async text stream API."""
+        fallback_providers = kwargs.pop("providers", None) or kwargs.pop("fallback", None)
+        if fallback_providers:
+            if not isinstance(fallback_providers, (list, tuple)):
+                fallback_providers = [fallback_providers]
+            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_chunks = 0
+        text_parts = []
+        usage = None
+        final_config = self._build_final_config(kwargs)
+        messages = self._prepare_messages(query, final_config)
+        try:
+            async for event in self._send_request_stream_events_async(messages, final_config):
+                if isinstance(event, TextDelta):
+                    emitted_chunks += 1
+                    text_parts.append(event.text)
+                    yield event
+                elif isinstance(event, UsageUpdate):
+                    usage = event
+                    yield event
+                elif isinstance(event, ToolCall):
+                    yield event
+                else:
+                    yield event
+            yield Completed(response=Response(text="".join(text_parts)), usage=usage)
+        except Exception as error:
+            yield Error(error=error, emitted_chunks=emitted_chunks)
+
     # ─── Config Helpers ─────────────────────────────────────
 
-    def _build_final_config(self, kwargs: Dict[str, Any]) -> AdvancedConfig:
+    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)
         final_config = self.config.merge(request_config)
@@ -271,7 +513,7 @@ class BaseProvider:
 
         return final_config
 
-    def _prepare_messages(self, query: str, config: AdvancedConfig) -> List[Dict[str, str]]:
+    def _prepare_messages(self, query: str, config: AdvancedConfig) -> list[dict[str, str]]:
         """Build the messages list from query, persona, and config."""
         messages = []
 
@@ -317,15 +559,18 @@ class BaseProvider:
         """Execute a sync function with exponential backoff + jitter."""
         attempts = 0
         last_error = None
+        self._retry_count = 0
 
         while attempts <= retry:
             try:
                 response = func(messages, config, output_type)
+                self._retry_count = attempts
                 return response
             except Exception as e:
                 last_error = e
                 if self._is_retryable_error(e) and attempts < retry:
                     attempts += 1
+                    self._retry_count = attempts
                     delay = min(2 ** attempts + random.uniform(0, 1), 30)
                     logger.warning(
                         "Transient error (%s). Retrying in %.1fs... (%d/%d)",
@@ -342,15 +587,18 @@ class BaseProvider:
         """Execute an async function with exponential backoff + jitter."""
         attempts = 0
         last_error = None
+        self._retry_count = 0
 
         while attempts <= retry:
             try:
                 response = await func(messages, config, output_type)
+                self._retry_count = attempts
                 return response
             except Exception as e:
                 last_error = e
                 if self._is_retryable_error(e) and attempts < retry:
                     attempts += 1
+                    self._retry_count = attempts
                     delay = min(2 ** attempts + random.uniform(0, 1), 30)
                     logger.warning(
                         "Transient error (%s). Retrying in %.1fs... (%d/%d)",
@@ -365,43 +613,57 @@ class BaseProvider:
 
     def _is_retryable_error(self, e: Exception) -> bool:
         """Dynamically detect rate limits and network timeouts across different provider SDKs."""
-        from .exceptions import AskAINetworkError, AskAIRateLimitError
-        if isinstance(e, (AskAINetworkError, AskAIRateLimitError)):
+        from .exceptions import (
+            AskAINetworkError,
+            AskAIRateLimitError,
+            ProviderRateLimitError,
+            ProviderTimeoutError,
+            ProviderUnavailableError,
+        )
+        if isinstance(e, (AskAINetworkError, AskAIRateLimitError, ProviderRateLimitError,
+                          ProviderTimeoutError, ProviderUnavailableError)):
             return True
-
-        err_str = str(e).lower()
-        err_cls = e.__class__.__name__.lower()
-
-        retryable_indicators = ("ratelimit", "rate_limit", "429", "timeout", "connection", "connecterror")
-        return any(indicator in err_cls or indicator in err_str for indicator in retryable_indicators)
+        return False
 
     # ─── Abstract Methods (Subclasses Must Implement) ──────
 
-    def _send_request(self, messages: List[Dict[str, str]], config: AdvancedConfig, output_type: str = None) -> Response:
+    def _send_request(self, messages: list[dict[str, str]], config: AdvancedConfig, output_type: str = None) -> Response:
         """Sync request — subclasses must implement."""
         raise NotImplementedError
 
-    async def _send_request_async(self, messages: List[Dict[str, str]], config: AdvancedConfig, output_type: str = None) -> Response:
+    async def _send_request_async(self, messages: list[dict[str, str]], config: AdvancedConfig, output_type: str = None) -> Response:
         """Async request — subclasses should override for native async."""
         # Default fallback: run sync in thread pool
         loop = asyncio.get_running_loop()
         return await loop.run_in_executor(None, self._send_request, messages, config, output_type)
 
-    def _send_request_stream(self, messages: List[Dict[str, str]], config: AdvancedConfig) -> Iterator[str]:
+    def _send_request_stream(self, messages: list[dict[str, str]], config: AdvancedConfig) -> Iterator[str]:
         """Sync streaming — subclasses should override."""
         from .exceptions import StreamNotSupportedError
         raise StreamNotSupportedError(f"{self.__class__.__name__} does not support streaming.")
 
-    async def _send_request_stream_async(self, messages: List[Dict[str, str]], config: AdvancedConfig) -> AsyncIterator[str]:
+    def _send_request_stream_events(self, messages: list[dict[str, str]], config: AdvancedConfig) -> Iterator[Any]:
+        for index, chunk in enumerate(self._send_request_stream(messages, config)):
+            yield TextDelta(text=chunk, index=index)
+
+    async def _send_request_stream_async(self, messages: list[dict[str, str]], config: AdvancedConfig) -> AsyncIterator[str]:
         """Async streaming — subclasses should override for native async."""
         # Default fallback: wrap sync stream
         loop = asyncio.get_running_loop()
         for chunk in await loop.run_in_executor(None, lambda: list(self._send_request_stream(messages, config))):
             yield chunk
 
+    async def _send_request_stream_events_async(
+        self, messages: list[dict[str, str]], config: AdvancedConfig
+    ) -> AsyncIterator[Any]:
+        index = 0
+        async for chunk in self._send_request_stream_async(messages, config):
+            yield TextDelta(text=chunk, index=index)
+            index += 1
+
     # ─── Tool Definitions for Smart Intent ─────────────────
 
-    def _get_media_tools(self) -> List[Dict[str, Any]]:
+    def _get_media_tools(self) -> list[dict[str, Any]]:
         return [
             {
                 "type": "function",

ask_ai/capabilities.py

@@ -0,0 +1,17 @@
+from dataclasses import asdict, dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class Capabilities:
+    text: bool = True
+    streaming: bool = False
+    async_streaming: bool = False
+    structured_output: bool = False
+    tool_calls: bool = False
+    image_generation: bool = False
+    audio_generation: bool = False
+    vision: bool = False
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)

ask_ai/cli.py

@@ -1,16 +1,7 @@
-import sys
 import argparse
-from . import __version__
-from . import (
-    OpenAI,
-    Groq,
-    Google,
-    OpenRouter,
-    Azure,
-    Anthropic,
-    AskAIError
-)
+import sys
 
+from . import Anthropic, AskAIError, Azure, Google, Groq, OpenAI, OpenRouter, __version__
 
 PROVIDER_MAP = {
     "openai": OpenAI,
@@ -70,14 +61,9 @@ def main():
 
         # Stream or regular
         if args.stream:
-            try:
-                for chunk in ai.ask_stream(args.query, **ask_kwargs):
-                    print(chunk, end="", flush=True)
-                print()  # newline at end
-            except Exception:
-                # Fallback to non-streaming if provider doesn't support it
-                response = ai.ask(args.query, **ask_kwargs)
-                print(response.text)
+            for chunk in ai.ask_stream(args.query, **ask_kwargs):
+                print(chunk, end="", flush=True)
+            print()  # newline at end
         else:
             response = ai.ask(args.query, **ask_kwargs)