prompture 0.0.29.dev8__py3-none-any.whl → 0.0.35__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

@@ -0,0 +1,164 @@
+"""Discovery module for auto-detecting available models."""
+
+import logging
+import os
+
+import requests
+
+from .drivers import (
+    AzureDriver,
+    ClaudeDriver,
+    GoogleDriver,
+    GrokDriver,
+    GroqDriver,
+    LMStudioDriver,
+    LocalHTTPDriver,
+    OllamaDriver,
+    OpenAIDriver,
+    OpenRouterDriver,
+)
+from .settings import settings
+
+logger = logging.getLogger(__name__)
+
+
+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()
+    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.
+    provider_classes = {
+        "openai": OpenAIDriver,
+        "azure": AzureDriver,
+        "claude": ClaudeDriver,
+        "google": GoogleDriver,
+        "groq": GroqDriver,
+        "openrouter": OpenRouterDriver,
+        "grok": GrokDriver,
+        "ollama": OllamaDriver,
+        "lmstudio": LMStudioDriver,
+        "local_http": LocalHTTPDriver,
+    }
+
+    for provider, driver_cls in provider_classes.items():
+        try:
+            # 1. Check if the provider is configured (has API key or endpoint)
+            # 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"))
+                ):
+                    is_configured = True
+            elif provider == "claude":
+                if settings.claude_api_key or os.getenv("CLAUDE_API_KEY"):
+                    is_configured = True
+            elif provider == "google":
+                if settings.google_api_key or os.getenv("GOOGLE_API_KEY"):
+                    is_configured = True
+            elif provider == "groq":
+                if settings.groq_api_key or os.getenv("GROQ_API_KEY"):
+                    is_configured = True
+            elif provider == "openrouter":
+                if settings.openrouter_api_key or os.getenv("OPENROUTER_API_KEY"):
+                    is_configured = True
+            elif provider == "grok":
+                if settings.grok_api_key or os.getenv("GROK_API_KEY"):
+                    is_configured = True
+            elif provider == "ollama":
+                # Ollama is always considered "configured" as it defaults to localhost
+                # 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" 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:
+                    # 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?
+                    # 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"
+                    # 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"
+                    )
+                    # 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()
+                        models = data.get("models", [])
+                        for model in models:
+                            name = model.get("name")
+                            if name:
+                                # Ollama model names often include tags like "llama3:latest"
+                                # We can keep them as is.
+                                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,16 @@
-"""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 collections.abc import Iterator
+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 +29,159 @@ 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
+    supports_tool_use: bool = False
+    supports_streaming: 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, Any]], 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)
+
+    # ------------------------------------------------------------------
+    # Tool use
+    # ------------------------------------------------------------------
+
+    def generate_messages_with_tools(
+        self,
+        messages: list[dict[str, Any]],
+        tools: list[dict[str, Any]],
+        options: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Generate a response that may include tool calls.
+
+        Returns a dict with keys: ``text``, ``meta``, ``tool_calls``, ``stop_reason``.
+        ``tool_calls`` is a list of ``{"id": str, "name": str, "arguments": dict}``.
+
+        Drivers that support tool use should override this method and set
+        ``supports_tool_use = True``.
+        """
+        raise NotImplementedError(f"{self.__class__.__name__} does not support tool use")
+
+    # ------------------------------------------------------------------
+    # Streaming
+    # ------------------------------------------------------------------
+
+    def generate_messages_stream(
+        self,
+        messages: list[dict[str, Any]],
+        options: dict[str, Any],
+    ) -> Iterator[dict[str, Any]]:
+        """Yield response chunks incrementally.
+
+        Each chunk is a dict:
+        - ``{"type": "delta", "text": str}`` for content fragments
+        - ``{"type": "done", "text": str, "meta": dict}`` for the final summary
+
+        Drivers that support streaming should override this method and set
+        ``supports_streaming = True``.
+        """
+        raise NotImplementedError(f"{self.__class__.__name__} does not support streaming")
+
+    # ------------------------------------------------------------------
+    # 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, Any]], 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, Any]]) -> 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,71 +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 ..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
+# 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
     ),
-    "claude": lambda model=None: ClaudeDriver(
-        api_key=settings.claude_api_key,
-        model=model or settings.claude_model
+    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,
     ),
-    "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
-    ),
-}
+    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):
@@ -73,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")
 
@@ -96,25 +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",
+    "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",
 ]