prompture 0.0.29.dev8__py3-none-any.whl → 0.0.38.dev2__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,187 @@
+"""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}")
+
+            # Dynamic Detection: LM Studio loaded models
+            if provider == "lmstudio":
+                try:
+                    endpoint = settings.lmstudio_endpoint or os.getenv(
+                        "LMSTUDIO_ENDPOINT", "http://127.0.0.1:1234/v1/chat/completions"
+                    )
+                    base_url = endpoint.split("/v1/")[0]
+                    models_url = f"{base_url}/v1/models"
+
+                    headers: dict[str, str] = {}
+                    api_key = settings.lmstudio_api_key or os.getenv("LMSTUDIO_API_KEY")
+                    if api_key:
+                        headers["Authorization"] = f"Bearer {api_key}"
+
+                    resp = requests.get(models_url, headers=headers, timeout=2)
+                    if resp.status_code == 200:
+                        data = resp.json()
+                        models = data.get("data", [])
+                        for model in models:
+                            model_id = model.get("id")
+                            if model_id:
+                                available_models.add(f"lmstudio/{model_id}")
+                except Exception as e:
+                    logger.debug(f"Failed to fetch LM Studio models: {e}")
+
+        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,197 @@ 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
+    supports_vision: 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``.
+        """
+        self._check_vision_support(messages)
+        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)
+
+    def _check_vision_support(self, messages: list[dict[str, Any]]) -> None:
+        """Raise if messages contain image blocks and the driver lacks vision support."""
+        if self.supports_vision:
+            return
+        for msg in messages:
+            content = msg.get("content")
+            if isinstance(content, list):
+                for block in content:
+                    if isinstance(block, dict) and block.get("type") == "image":
+                        raise NotImplementedError(
+                            f"{self.__class__.__name__} does not support vision/image inputs. "
+                            "Use a vision-capable model."
+                        )
+
+    def _prepare_messages(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        """Transform universal message format into provider-specific wire format.
+
+        Vision-capable drivers override this to convert the universal image
+        blocks into their provider-specific format.  The base implementation
+        validates vision support and returns messages unchanged.
+        """
+        self._check_vision_support(messages)
+        return messages
+
+    @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", "")
+            # Handle content that is a list of blocks (vision messages)
+            if isinstance(content, list):
+                text_parts = []
+                for block in content:
+                    if isinstance(block, dict):
+                        if block.get("type") == "text":
+                            text_parts.append(block.get("text", ""))
+                        elif block.get("type") == "image":
+                            text_parts.append("[image]")
+                    elif isinstance(block, str):
+                        text_parts.append(block)
+                content = " ".join(text_parts)
+            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)