prompture 0.0.33.dev1__py3-none-any.whl → 0.0.34__py3-none-any.whl

prompture/cost_mixin.py

@@ -0,0 +1,51 @@
+"""Shared cost-calculation mixin for LLM drivers."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+class CostMixin:
+    """Mixin that provides ``_calculate_cost`` to sync and async drivers.
+
+    Drivers that charge per-token should inherit from this mixin alongside
+    their base class (``Driver`` or ``AsyncDriver``).  Free/local drivers
+    (Ollama, LM Studio, LocalHTTP, HuggingFace, AirLLM) can skip it.
+    """
+
+    # Subclasses must define MODEL_PRICING as a class attribute.
+    MODEL_PRICING: dict[str, dict[str, Any]] = {}
+
+    # Divisor for hardcoded MODEL_PRICING rates.
+    # Most drivers use per-1K-token pricing (1_000).
+    # Grok uses per-1M-token pricing (1_000_000).
+    # Google uses per-1M-character pricing (1_000_000).
+    _PRICING_UNIT: int = 1_000
+
+    def _calculate_cost(
+        self,
+        provider: str,
+        model: str,
+        prompt_tokens: int | float,
+        completion_tokens: int | float,
+    ) -> float:
+        """Calculate USD cost for a generation call.
+
+        Resolution order:
+        1. Live rates from ``model_rates.get_model_rates()`` (per 1M tokens).
+        2. Hardcoded ``MODEL_PRICING`` on the driver class (unit set by ``_PRICING_UNIT``).
+        3. Zero if neither source has data.
+        """
+        from .model_rates import get_model_rates
+
+        live_rates = get_model_rates(provider, model)
+        if live_rates:
+            prompt_cost = (prompt_tokens / 1_000_000) * live_rates["input"]
+            completion_cost = (completion_tokens / 1_000_000) * live_rates["output"]
+        else:
+            unit = self._PRICING_UNIT
+            model_pricing = self.MODEL_PRICING.get(model, {"prompt": 0, "completion": 0})
+            prompt_cost = (prompt_tokens / unit) * model_pricing["prompt"]
+            completion_cost = (completion_tokens / unit) * model_pricing["completion"]
+
+        return round(prompt_cost + completion_cost, 6)

prompture/discovery.py

@@ -1,39 +1,41 @@
 """Discovery module for auto-detecting available models."""
+
+import logging
 import os
+
 import requests
-import logging
-from typing import List, Dict, Any, Set
 
 from .drivers import (
-    DRIVER_REGISTRY,
-    OpenAIDriver,
     AzureDriver,
     ClaudeDriver,
     GoogleDriver,
-    GroqDriver,
-    OpenRouterDriver,
     GrokDriver,
-    OllamaDriver,
+    GroqDriver,
     LMStudioDriver,
     LocalHTTPDriver,
+    OllamaDriver,
+    OpenAIDriver,
+    OpenRouterDriver,
 )
 from .settings import settings
 
 logger = logging.getLogger(__name__)
 
-def get_available_models() -> List[str]:
+
+def get_available_models() -> list[str]:
     """
     Auto-detects all available models based on configured drivers and environment variables.
-    
+
     Iterates through supported providers and checks if they are configured (e.g. API key present).
     For static drivers, returns models from their MODEL_PRICING keys.
     For dynamic drivers (like Ollama), attempts to fetch available models from the endpoint.
-    
+
     Returns:
         A list of unique model strings in the format "provider/model_id".
     """
-    available_models: Set[str] = set()
-    
+    available_models: set[str] = set()
+    configured_providers: set[str] = set()
+
     # Map of provider name to driver class
     # We need to map the registry keys to the actual classes to check MODEL_PRICING
     # and instantiate for dynamic checks if needed.
@@ -56,16 +58,18 @@ def get_available_models() -> List[str]:
             # We can check this by looking at the settings or env vars that the driver uses.
             # A simple way is to try to instantiate it with defaults, but that might fail if keys are missing.
             # Instead, let's check the specific requirements for each known provider.
-            
+
             is_configured = False
-            
+
             if provider == "openai":
                 if settings.openai_api_key or os.getenv("OPENAI_API_KEY"):
                     is_configured = True
             elif provider == "azure":
-                if (settings.azure_api_key or os.getenv("AZURE_API_KEY")) and \
-                   (settings.azure_api_endpoint or os.getenv("AZURE_API_ENDPOINT")) and \
-                   (settings.azure_deployment_id or os.getenv("AZURE_DEPLOYMENT_ID")):
+                if (
+                    (settings.azure_api_key or os.getenv("AZURE_API_KEY"))
+                    and (settings.azure_api_endpoint or os.getenv("AZURE_API_ENDPOINT"))
+                    and (settings.azure_deployment_id or os.getenv("AZURE_DEPLOYMENT_ID"))
+                ):
                     is_configured = True
             elif provider == "claude":
                 if settings.claude_api_key or os.getenv("CLAUDE_API_KEY"):
@@ -87,46 +91,49 @@ def get_available_models() -> List[str]:
                 # We will check connectivity later
                 is_configured = True
             elif provider == "lmstudio":
-                 # LM Studio is similar to Ollama, defaults to localhost
-                 is_configured = True
-            elif provider == "local_http":
-                 if settings.local_http_endpoint or os.getenv("LOCAL_HTTP_ENDPOINT"):
-                     is_configured = True
+                # LM Studio is similar to Ollama, defaults to localhost
+                is_configured = True
+            elif provider == "local_http" and (settings.local_http_endpoint or os.getenv("LOCAL_HTTP_ENDPOINT")):
+                is_configured = True
 
             if not is_configured:
                 continue
 
+            configured_providers.add(provider)
+
             # 2. Static Detection: Get models from MODEL_PRICING
             if hasattr(driver_cls, "MODEL_PRICING"):
                 pricing = driver_cls.MODEL_PRICING
-                for model_id in pricing.keys():
+                for model_id in pricing:
                     # Skip "default" or generic keys if they exist
                     if model_id == "default":
                         continue
-                    
-                    # For Azure, the model_id in pricing is usually the base model name, 
-                    # but the user needs to use the deployment ID. 
-                    # However, our Azure driver implementation uses the deployment_id from init 
-                    # as the "model" for the request, but expects the user to pass a model name 
-                    # that maps to pricing? 
+
+                    # For Azure, the model_id in pricing is usually the base model name,
+                    # but the user needs to use the deployment ID.
+                    # However, our Azure driver implementation uses the deployment_id from init
+                    # as the "model" for the request, but expects the user to pass a model name
+                    # that maps to pricing?
                     # Looking at AzureDriver:
                     # kwargs = {"model": self.deployment_id, ...}
                     # model = options.get("model", self.model) -> used for pricing lookup
-                    # So we should list the keys in MODEL_PRICING as available "models" 
+                    # So we should list the keys in MODEL_PRICING as available "models"
                     # even though for Azure specifically it's a bit weird because of deployment IDs.
                     # But for general discovery, listing supported models is correct.
-                    
+
                     available_models.add(f"{provider}/{model_id}")
 
             # 3. Dynamic Detection: Specific logic for Ollama
             if provider == "ollama":
                 try:
-                    endpoint = settings.ollama_endpoint or os.getenv("OLLAMA_ENDPOINT", "http://localhost:11434/api/generate")
+                    endpoint = settings.ollama_endpoint or os.getenv(
+                        "OLLAMA_ENDPOINT", "http://localhost:11434/api/generate"
+                    )
                     # We need the base URL for tags, usually http://localhost:11434/api/tags
                     # The configured endpoint might be .../api/generate or .../api/chat
                     base_url = endpoint.split("/api/")[0]
                     tags_url = f"{base_url}/api/tags"
-                    
+
                     resp = requests.get(tags_url, timeout=2)
                     if resp.status_code == 200:
                         data = resp.json()
@@ -139,11 +146,19 @@ def get_available_models() -> List[str]:
                                 available_models.add(f"ollama/{name}")
                 except Exception as e:
                     logger.debug(f"Failed to fetch Ollama models: {e}")
-            
+
             # Future: Add dynamic detection for LM Studio if they have an endpoint for listing models
-            
+
         except Exception as e:
             logger.warning(f"Error detecting models for provider {provider}: {e}")
             continue
 
+    # Enrich with live model list from models.dev cache
+    from .model_rates import PROVIDER_MAP, get_all_provider_models
+
+    for prompture_name, api_name in PROVIDER_MAP.items():
+        if prompture_name in configured_providers:
+            for model_id in get_all_provider_models(api_name):
+                available_models.add(f"{prompture_name}/{model_id}")
+
     return sorted(list(available_models))

prompture/driver.py

@@ -1,7 +1,15 @@
-"""Driver base class for LLM adapters.
-"""
+"""Driver base class for LLM adapters."""
+
 from __future__ import annotations
-from typing import Any, Dict
+
+import logging
+import time
+from typing import Any
+
+from .callbacks import DriverCallbacks
+
+logger = logging.getLogger("prompture.driver")
+
 
 class Driver:
     """Adapter base. Implementar generate(prompt, options) -> {"text": ... , "meta": {...}}
@@ -20,5 +28,117 @@ class Driver:
     additional provider-specific metadata while the core fields provide
     standardized access to token usage and cost information.
     """
-    def generate(self, prompt: str, options: Dict[str,Any]) -> Dict[str,Any]:
-        raise NotImplementedError
+
+    supports_json_mode: bool = False
+    supports_json_schema: bool = False
+    supports_messages: bool = False
+
+    callbacks: DriverCallbacks | None = None
+
+    def generate(self, prompt: str, options: dict[str, Any]) -> dict[str, Any]:
+        raise NotImplementedError
+
+    def generate_messages(self, messages: list[dict[str, str]], options: dict[str, Any]) -> dict[str, Any]:
+        """Generate a response from a list of conversation messages.
+
+        Each message is a dict with ``"role"`` (``"system"``, ``"user"``, or
+        ``"assistant"``) and ``"content"`` keys.
+
+        The default implementation flattens the messages into a single prompt
+        string and delegates to :meth:`generate`.  Drivers that natively
+        support message arrays should override this method and set
+        ``supports_messages = True``.
+        """
+        prompt = self._flatten_messages(messages)
+        return self.generate(prompt, options)
+
+    # ------------------------------------------------------------------
+    # Hook-aware wrappers
+    # ------------------------------------------------------------------
+
+    def generate_with_hooks(self, prompt: str, options: dict[str, Any]) -> dict[str, Any]:
+        """Wrap :meth:`generate` with on_request / on_response / on_error callbacks."""
+        driver_name = getattr(self, "model", self.__class__.__name__)
+        self._fire_callback(
+            "on_request",
+            {"prompt": prompt, "messages": None, "options": options, "driver": driver_name},
+        )
+        t0 = time.perf_counter()
+        try:
+            resp = self.generate(prompt, options)
+        except Exception as exc:
+            self._fire_callback(
+                "on_error",
+                {"error": exc, "prompt": prompt, "messages": None, "options": options, "driver": driver_name},
+            )
+            raise
+        elapsed_ms = (time.perf_counter() - t0) * 1000
+        self._fire_callback(
+            "on_response",
+            {
+                "text": resp.get("text", ""),
+                "meta": resp.get("meta", {}),
+                "driver": driver_name,
+                "elapsed_ms": elapsed_ms,
+            },
+        )
+        return resp
+
+    def generate_messages_with_hooks(self, messages: list[dict[str, str]], options: dict[str, Any]) -> dict[str, Any]:
+        """Wrap :meth:`generate_messages` with callbacks."""
+        driver_name = getattr(self, "model", self.__class__.__name__)
+        self._fire_callback(
+            "on_request",
+            {"prompt": None, "messages": messages, "options": options, "driver": driver_name},
+        )
+        t0 = time.perf_counter()
+        try:
+            resp = self.generate_messages(messages, options)
+        except Exception as exc:
+            self._fire_callback(
+                "on_error",
+                {"error": exc, "prompt": None, "messages": messages, "options": options, "driver": driver_name},
+            )
+            raise
+        elapsed_ms = (time.perf_counter() - t0) * 1000
+        self._fire_callback(
+            "on_response",
+            {
+                "text": resp.get("text", ""),
+                "meta": resp.get("meta", {}),
+                "driver": driver_name,
+                "elapsed_ms": elapsed_ms,
+            },
+        )
+        return resp
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _fire_callback(self, event: str, payload: dict[str, Any]) -> None:
+        """Invoke a single callback, swallowing and logging any exception."""
+        if self.callbacks is None:
+            return
+        cb = getattr(self.callbacks, event, None)
+        if cb is None:
+            return
+        try:
+            cb(payload)
+        except Exception:
+            logger.exception("Callback %s raised an exception", event)
+
+    @staticmethod
+    def _flatten_messages(messages: list[dict[str, str]]) -> str:
+        """Join messages into a single prompt string with role prefixes."""
+        parts: list[str] = []
+        for msg in messages:
+            role = msg.get("role", "user")
+            content = msg.get("content", "")
+            if role == "system":
+                parts.append(f"[System]: {content}")
+            elif role == "assistant":
+                parts.append(f"[Assistant]: {content}")
+            else:
+                parts.append(f"[User]: {content}")
+        return "\n\n".join(parts)

prompture/drivers/__init__.py

@@ -1,76 +1,145 @@
-from .openai_driver import OpenAIDriver
-from .local_http_driver import LocalHTTPDriver
-from .ollama_driver import OllamaDriver
-from .claude_driver import ClaudeDriver
+"""Driver registry and factory functions.
+
+This module provides:
+- Built-in drivers for popular LLM providers
+- A pluggable registry system for custom drivers
+- Factory functions to instantiate drivers by provider/model name
+
+Custom Driver Registration:
+    from prompture import register_driver
+
+    def my_driver_factory(model=None):
+        return MyCustomDriver(model=model)
+
+    register_driver("my_provider", my_driver_factory)
+
+    # Now you can use it
+    driver = get_driver_for_model("my_provider/my-model")
+
+Entry Point Discovery:
+    Third-party packages can register drivers via entry points.
+    Add to your pyproject.toml:
+
+    [project.entry-points."prompture.drivers"]
+    my_provider = "my_package.drivers:my_driver_factory"
+"""
+
+from typing import Optional
+
+from ..settings import settings
+from .airllm_driver import AirLLMDriver
+from .async_airllm_driver import AsyncAirLLMDriver
+from .async_azure_driver import AsyncAzureDriver
+from .async_claude_driver import AsyncClaudeDriver
+from .async_google_driver import AsyncGoogleDriver
+from .async_grok_driver import AsyncGrokDriver
+from .async_groq_driver import AsyncGroqDriver
+from .async_hugging_driver import AsyncHuggingFaceDriver
+from .async_lmstudio_driver import AsyncLMStudioDriver
+from .async_local_http_driver import AsyncLocalHTTPDriver
+from .async_ollama_driver import AsyncOllamaDriver
+from .async_openai_driver import AsyncOpenAIDriver
+from .async_openrouter_driver import AsyncOpenRouterDriver
+from .async_registry import ASYNC_DRIVER_REGISTRY, get_async_driver, get_async_driver_for_model
 from .azure_driver import AzureDriver
-from .lmstudio_driver import LMStudioDriver
+from .claude_driver import ClaudeDriver
 from .google_driver import GoogleDriver
+from .grok_driver import GrokDriver
 from .groq_driver import GroqDriver
+from .lmstudio_driver import LMStudioDriver
+from .local_http_driver import LocalHTTPDriver
+from .ollama_driver import OllamaDriver
+from .openai_driver import OpenAIDriver
 from .openrouter_driver import OpenRouterDriver
-from .grok_driver import GrokDriver
-from .airllm_driver import AirLLMDriver
-from ..settings import settings
+from .registry import (
+    _get_sync_registry,
+    get_async_driver_factory,
+    get_driver_factory,
+    is_async_driver_registered,
+    is_driver_registered,
+    list_registered_async_drivers,
+    list_registered_drivers,
+    load_entry_point_drivers,
+    register_async_driver,
+    register_driver,
+    unregister_async_driver,
+    unregister_driver,
+)
 
-
-# Central registry: maps provider → factory function
-DRIVER_REGISTRY = {
-    "openai": lambda model=None: OpenAIDriver(
-        api_key=settings.openai_api_key,
-        model=model or settings.openai_model
-    ),
-    "ollama": lambda model=None: OllamaDriver(
-        endpoint=settings.ollama_endpoint,
-        model=model or settings.ollama_model
-    ),
-    "claude": lambda model=None: ClaudeDriver(
-        api_key=settings.claude_api_key,
-        model=model or settings.claude_model
+# Register built-in sync drivers
+register_driver(
+    "openai",
+    lambda model=None: OpenAIDriver(api_key=settings.openai_api_key, model=model or settings.openai_model),
+    overwrite=True,
+)
+register_driver(
+    "ollama",
+    lambda model=None: OllamaDriver(endpoint=settings.ollama_endpoint, model=model or settings.ollama_model),
+    overwrite=True,
+)
+register_driver(
+    "claude",
+    lambda model=None: ClaudeDriver(api_key=settings.claude_api_key, model=model or settings.claude_model),
+    overwrite=True,
+)
+register_driver(
+    "lmstudio",
+    lambda model=None: LMStudioDriver(endpoint=settings.lmstudio_endpoint, model=model or settings.lmstudio_model),
+    overwrite=True,
+)
+register_driver(
+    "azure",
+    lambda model=None: AzureDriver(
+        api_key=settings.azure_api_key, endpoint=settings.azure_api_endpoint, deployment_id=settings.azure_deployment_id
     ),
-    "lmstudio": lambda model=None: LMStudioDriver(
-        endpoint=settings.lmstudio_endpoint,
-        model=model or settings.lmstudio_model
-    ),
-    "azure": lambda model=None: AzureDriver(
-        api_key=settings.azure_api_key,
-        endpoint=settings.azure_api_endpoint,
-        deployment_id=settings.azure_deployment_id
-    ),
-    "local_http": lambda model=None: LocalHTTPDriver(
-        endpoint=settings.local_http_endpoint,
-        model=model
-    ),
-    "google": lambda model=None: GoogleDriver(
-        api_key=settings.google_api_key,
-        model=model or settings.google_model
-    ),
-    "groq": lambda model=None: GroqDriver(
-        api_key=settings.groq_api_key,
-        model=model or settings.groq_model
-    ),
-    "openrouter": lambda model=None: OpenRouterDriver(
-        api_key=settings.openrouter_api_key,
-        model=model or settings.openrouter_model
-    ),
-    "grok": lambda model=None: GrokDriver(
-        api_key=settings.grok_api_key,
-        model=model or settings.grok_model
-    ),
-    "airllm": lambda model=None: AirLLMDriver(
+    overwrite=True,
+)
+register_driver(
+    "local_http",
+    lambda model=None: LocalHTTPDriver(endpoint=settings.local_http_endpoint, model=model),
+    overwrite=True,
+)
+register_driver(
+    "google",
+    lambda model=None: GoogleDriver(api_key=settings.google_api_key, model=model or settings.google_model),
+    overwrite=True,
+)
+register_driver(
+    "groq",
+    lambda model=None: GroqDriver(api_key=settings.groq_api_key, model=model or settings.groq_model),
+    overwrite=True,
+)
+register_driver(
+    "openrouter",
+    lambda model=None: OpenRouterDriver(api_key=settings.openrouter_api_key, model=model or settings.openrouter_model),
+    overwrite=True,
+)
+register_driver(
+    "grok",
+    lambda model=None: GrokDriver(api_key=settings.grok_api_key, model=model or settings.grok_model),
+    overwrite=True,
+)
+register_driver(
+    "airllm",
+    lambda model=None: AirLLMDriver(
         model=model or settings.airllm_model,
         compression=settings.airllm_compression,
     ),
-}
+    overwrite=True,
+)
 
+# Backwards compatibility: expose registry dict (read-only view recommended)
+DRIVER_REGISTRY = _get_sync_registry()
 
-def get_driver(provider_name: str = None):
+
+def get_driver(provider_name: Optional[str] = None):
     """
     Factory to get a driver instance based on the provider name (legacy style).
     Uses default model from settings if not overridden.
     """
     provider = (provider_name or settings.ai_provider or "ollama").strip().lower()
-    if provider not in DRIVER_REGISTRY:
-        raise ValueError(f"Unknown provider: {provider_name}")
-    return DRIVER_REGISTRY[provider]()  # use default model from settings
+    factory = get_driver_factory(provider)
+    return factory()  # use default model from settings
 
 
 def get_driver_for_model(model_str: str):
@@ -78,21 +147,21 @@ def get_driver_for_model(model_str: str):
     Factory to get a driver instance based on a full model string.
     Format: provider/model_id
     Example: "openai/gpt-4-turbo-preview"
-    
+
     Args:
         model_str: Model identifier string. Can be either:
                    - Full format: "provider/model" (e.g. "openai/gpt-4")
                    - Provider only: "provider" (e.g. "openai")
-    
+
     Returns:
         A configured driver instance for the specified provider/model.
-        
+
     Raises:
         ValueError: If provider is invalid or format is incorrect.
     """
     if not isinstance(model_str, str):
         raise ValueError("Model string must be a string, got {type(model_str)}")
-        
+
     if not model_str:
         raise ValueError("Model string cannot be empty")
 
@@ -101,26 +170,55 @@ def get_driver_for_model(model_str: str):
     provider = parts[0].lower()
     model_id = parts[1] if len(parts) > 1 else None
 
-    # Validate provider
-    if provider not in DRIVER_REGISTRY:
-        raise ValueError(f"Unsupported provider '{provider}'")
-        
+    # Get factory (validates provider exists)
+    factory = get_driver_factory(provider)
+
     # Create driver with model ID if provided, otherwise use default
-    return DRIVER_REGISTRY[provider](model_id)
+    return factory(model_id)
 
 
 __all__ = [
-    "OpenAIDriver",
-    "LocalHTTPDriver",
-    "OllamaDriver",
-    "ClaudeDriver",
-    "LMStudioDriver",
+    "ASYNC_DRIVER_REGISTRY",
+    # Legacy registry dicts (for backwards compatibility)
+    "DRIVER_REGISTRY",
+    # Sync drivers
+    "AirLLMDriver",
+    # Async drivers
+    "AsyncAirLLMDriver",
+    "AsyncAzureDriver",
+    "AsyncClaudeDriver",
+    "AsyncGoogleDriver",
+    "AsyncGrokDriver",
+    "AsyncGroqDriver",
+    "AsyncHuggingFaceDriver",
+    "AsyncLMStudioDriver",
+    "AsyncLocalHTTPDriver",
+    "AsyncOllamaDriver",
+    "AsyncOpenAIDriver",
+    "AsyncOpenRouterDriver",
     "AzureDriver",
+    "ClaudeDriver",
     "GoogleDriver",
+    "GrokDriver",
     "GroqDriver",
+    "LMStudioDriver",
+    "LocalHTTPDriver",
+    "OllamaDriver",
+    "OpenAIDriver",
     "OpenRouterDriver",
-    "GrokDriver",
-    "AirLLMDriver",
+    "get_async_driver",
+    "get_async_driver_for_model",
+    # Factory functions
     "get_driver",
     "get_driver_for_model",
+    "is_async_driver_registered",
+    "is_driver_registered",
+    "list_registered_async_drivers",
+    "list_registered_drivers",
+    "load_entry_point_drivers",
+    "register_async_driver",
+    # Registry functions (public API)
+    "register_driver",
+    "unregister_async_driver",
+    "unregister_driver",
 ]