prompture 0.0.29.dev8__py3-none-any.whl → 0.0.35__py3-none-any.whl

prompture/drivers/openai_driver.py

@@ -1,17 +1,27 @@
 """Minimal OpenAI driver (migrated to openai>=1.0.0).
 Requires the `openai` package. Uses OPENAI_API_KEY env var.
 """
+
+import json
 import os
-from typing import Any, Dict
+from collections.abc import Iterator
+from typing import Any
+
 try:
     from openai import OpenAI
 except Exception:
     OpenAI = None
 
+from ..cost_mixin import CostMixin
 from ..driver import Driver
 
 
-class OpenAIDriver(Driver):
+class OpenAIDriver(CostMixin, Driver):
+    supports_json_mode = True
+    supports_json_schema = True
+    supports_tool_use = True
+    supports_streaming = True
+
     # Approximate pricing per 1K tokens (keep updated with OpenAI's official pricing)
     # Each model entry also defines which token parameter it supports and
     # whether it accepts temperature.
@@ -62,7 +72,16 @@ class OpenAIDriver(Driver):
         else:
             self.client = None
 
-    def generate(self, prompt: str, options: Dict[str, Any]) -> Dict[str, Any]:
+    supports_messages = True
+
+    def generate(self, prompt: str, options: dict[str, Any]) -> dict[str, Any]:
+        messages = [{"role": "user", "content": prompt}]
+        return self._do_generate(messages, options)
+
+    def generate_messages(self, messages: list[dict[str, Any]], options: dict[str, Any]) -> dict[str, Any]:
+        return self._do_generate(messages, options)
+
+    def _do_generate(self, messages: list[dict[str, Any]], options: dict[str, Any]) -> dict[str, Any]:
         if self.client is None:
             raise RuntimeError("openai package (>=1.0.0) is not installed")
 
@@ -79,7 +98,7 @@ class OpenAIDriver(Driver):
         # Base kwargs
         kwargs = {
             "model": model,
-            "messages": [{"role": "user", "content": prompt}],
+            "messages": messages,
         }
 
         # Assign token limit with the correct parameter name
@@ -89,6 +108,21 @@ class OpenAIDriver(Driver):
         if supports_temperature and "temperature" in opts:
             kwargs["temperature"] = opts["temperature"]
 
+        # Native JSON mode support
+        if options.get("json_mode"):
+            json_schema = options.get("json_schema")
+            if json_schema:
+                kwargs["response_format"] = {
+                    "type": "json_schema",
+                    "json_schema": {
+                        "name": "extraction",
+                        "strict": True,
+                        "schema": json_schema,
+                    },
+                }
+            else:
+                kwargs["response_format"] = {"type": "json_object"}
+
         resp = self.client.chat.completions.create(**kwargs)
 
         # Extract usage info
@@ -97,11 +131,8 @@ class OpenAIDriver(Driver):
         completion_tokens = getattr(usage, "completion_tokens", 0)
         total_tokens = getattr(usage, "total_tokens", 0)
 
-        # Calculate cost
-        model_pricing = self.MODEL_PRICING.get(model, {"prompt": 0, "completion": 0})
-        prompt_cost = (prompt_tokens / 1000) * model_pricing["prompt"]
-        completion_cost = (completion_tokens / 1000) * model_pricing["completion"]
-        total_cost = prompt_cost + completion_cost
+        # Calculate cost via shared mixin
+        total_cost = self._calculate_cost("openai", model, prompt_tokens, completion_tokens)
 
         # Standardized meta object
         meta = {
@@ -115,3 +146,141 @@ class OpenAIDriver(Driver):
 
         text = resp.choices[0].message.content
         return {"text": text, "meta": meta}
+
+    # ------------------------------------------------------------------
+    # 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."""
+        if self.client is None:
+            raise RuntimeError("openai package (>=1.0.0) is not installed")
+
+        model = options.get("model", self.model)
+        model_info = self.MODEL_PRICING.get(model, {})
+        tokens_param = model_info.get("tokens_param", "max_tokens")
+        supports_temperature = model_info.get("supports_temperature", True)
+
+        opts = {"temperature": 1.0, "max_tokens": 512, **options}
+
+        kwargs: dict[str, Any] = {
+            "model": model,
+            "messages": messages,
+            "tools": tools,
+        }
+        kwargs[tokens_param] = opts.get("max_tokens", 512)
+
+        if supports_temperature and "temperature" in opts:
+            kwargs["temperature"] = opts["temperature"]
+
+        resp = self.client.chat.completions.create(**kwargs)
+
+        usage = getattr(resp, "usage", None)
+        prompt_tokens = getattr(usage, "prompt_tokens", 0)
+        completion_tokens = getattr(usage, "completion_tokens", 0)
+        total_tokens = getattr(usage, "total_tokens", 0)
+        total_cost = self._calculate_cost("openai", model, prompt_tokens, completion_tokens)
+
+        meta = {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": total_tokens,
+            "cost": round(total_cost, 6),
+            "raw_response": resp.model_dump(),
+            "model_name": model,
+        }
+
+        choice = resp.choices[0]
+        text = choice.message.content or ""
+        stop_reason = choice.finish_reason
+
+        tool_calls_out: list[dict[str, Any]] = []
+        if choice.message.tool_calls:
+            for tc in choice.message.tool_calls:
+                try:
+                    args = json.loads(tc.function.arguments)
+                except (json.JSONDecodeError, TypeError):
+                    args = {}
+                tool_calls_out.append({
+                    "id": tc.id,
+                    "name": tc.function.name,
+                    "arguments": args,
+                })
+
+        return {
+            "text": text,
+            "meta": meta,
+            "tool_calls": tool_calls_out,
+            "stop_reason": stop_reason,
+        }
+
+    # ------------------------------------------------------------------
+    # Streaming
+    # ------------------------------------------------------------------
+
+    def generate_messages_stream(
+        self,
+        messages: list[dict[str, Any]],
+        options: dict[str, Any],
+    ) -> Iterator[dict[str, Any]]:
+        """Yield response chunks via OpenAI streaming API."""
+        if self.client is None:
+            raise RuntimeError("openai package (>=1.0.0) is not installed")
+
+        model = options.get("model", self.model)
+        model_info = self.MODEL_PRICING.get(model, {})
+        tokens_param = model_info.get("tokens_param", "max_tokens")
+        supports_temperature = model_info.get("supports_temperature", True)
+
+        opts = {"temperature": 1.0, "max_tokens": 512, **options}
+
+        kwargs: dict[str, Any] = {
+            "model": model,
+            "messages": messages,
+            "stream": True,
+            "stream_options": {"include_usage": True},
+        }
+        kwargs[tokens_param] = opts.get("max_tokens", 512)
+
+        if supports_temperature and "temperature" in opts:
+            kwargs["temperature"] = opts["temperature"]
+
+        stream = self.client.chat.completions.create(**kwargs)
+
+        full_text = ""
+        prompt_tokens = 0
+        completion_tokens = 0
+
+        for chunk in stream:
+            # Usage comes in the final chunk
+            if getattr(chunk, "usage", None):
+                prompt_tokens = chunk.usage.prompt_tokens or 0
+                completion_tokens = chunk.usage.completion_tokens or 0
+
+            if chunk.choices:
+                delta = chunk.choices[0].delta
+                content = getattr(delta, "content", None) or ""
+                if content:
+                    full_text += content
+                    yield {"type": "delta", "text": content}
+
+        total_tokens = prompt_tokens + completion_tokens
+        total_cost = self._calculate_cost("openai", model, prompt_tokens, completion_tokens)
+
+        yield {
+            "type": "done",
+            "text": full_text,
+            "meta": {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": total_tokens,
+                "cost": round(total_cost, 6),
+                "raw_response": {},
+                "model_name": model,
+            },
+        }

prompture/drivers/openrouter_driver.py

@@ -1,14 +1,19 @@
 """OpenRouter driver implementation.
 Requires the `requests` package. Uses OPENROUTER_API_KEY env var.
 """
+
 import os
-from typing import Any, Dict
+from typing import Any
+
 import requests
 
+from ..cost_mixin import CostMixin
 from ..driver import Driver
 
 
-class OpenRouterDriver(Driver):
+class OpenRouterDriver(CostMixin, Driver):
+    supports_json_mode = True
+
     # Approximate pricing per 1K tokens based on OpenRouter's pricing
     # https://openrouter.ai/docs#pricing
     MODEL_PRICING = {
@@ -40,7 +45,7 @@ class OpenRouterDriver(Driver):
 
     def __init__(self, api_key: str | None = None, model: str = "openai/gpt-3.5-turbo"):
         """Initialize OpenRouter driver.
-        
+
         Args:
             api_key: OpenRouter API key. If not provided, will look for OPENROUTER_API_KEY env var
             model: Model to use. Defaults to openai/gpt-3.5-turbo
@@ -48,10 +53,10 @@ class OpenRouterDriver(Driver):
         self.api_key = api_key or os.getenv("OPENROUTER_API_KEY")
         if not self.api_key:
             raise ValueError("OpenRouter API key not found. Set OPENROUTER_API_KEY env var.")
-        
+
         self.model = model
         self.base_url = "https://openrouter.ai/api/v1"
-        
+
         # Required headers for OpenRouter
         self.headers = {
             "Authorization": f"Bearer {self.api_key}",
@@ -59,21 +64,21 @@ class OpenRouterDriver(Driver):
             "Content-Type": "application/json",
         }
 
-    def generate(self, prompt: str, options: Dict[str, Any]) -> Dict[str, Any]:
-        """Generate completion using OpenRouter API.
-        
-        Args:
-            prompt: The prompt text
-            options: Generation options
-            
-        Returns:
-            Dict containing generated text and metadata
-        """
+    supports_messages = True
+
+    def generate(self, prompt: str, options: dict[str, Any]) -> dict[str, Any]:
+        messages = [{"role": "user", "content": prompt}]
+        return self._do_generate(messages, options)
+
+    def generate_messages(self, messages: list[dict[str, str]], options: dict[str, Any]) -> dict[str, Any]:
+        return self._do_generate(messages, options)
+
+    def _do_generate(self, messages: list[dict[str, str]], options: dict[str, Any]) -> dict[str, Any]:
         if not self.api_key:
             raise RuntimeError("OpenRouter API key not found")
 
         model = options.get("model", self.model)
-        
+
         # Lookup model-specific config
         model_info = self.MODEL_PRICING.get(model, {})
         tokens_param = model_info.get("tokens_param", "max_tokens")
@@ -85,7 +90,7 @@ class OpenRouterDriver(Driver):
         # Base request data
         data = {
             "model": model,
-            "messages": [{"role": "user", "content": prompt}],
+            "messages": messages,
         }
 
         # Add token limit with correct parameter name
@@ -95,6 +100,10 @@ class OpenRouterDriver(Driver):
         if supports_temperature and "temperature" in opts:
             data["temperature"] = opts["temperature"]
 
+        # Native JSON mode support
+        if options.get("json_mode"):
+            data["response_format"] = {"type": "json_object"}
+
         try:
             response = requests.post(
                 f"{self.base_url}/chat/completions",
@@ -110,11 +119,8 @@ class OpenRouterDriver(Driver):
             completion_tokens = usage.get("completion_tokens", 0)
             total_tokens = usage.get("total_tokens", 0)
 
-            # Calculate cost
-            model_pricing = self.MODEL_PRICING.get(model, {"prompt": 0, "completion": 0})
-            prompt_cost = (prompt_tokens / 1000) * model_pricing["prompt"]
-            completion_cost = (completion_tokens / 1000) * model_pricing["completion"]
-            total_cost = prompt_cost + completion_cost
+            # Calculate cost via shared mixin
+            total_cost = self._calculate_cost("openrouter", model, prompt_tokens, completion_tokens)
 
             # Standardized meta object
             meta = {
@@ -130,11 +136,11 @@ class OpenRouterDriver(Driver):
             return {"text": text, "meta": meta}
 
         except requests.exceptions.RequestException as e:
-            error_msg = f"OpenRouter API request failed: {str(e)}"
-            if hasattr(e.response, 'json'):
+            error_msg = f"OpenRouter API request failed: {e!s}"
+            if hasattr(e.response, "json"):
                 try:
                     error_details = e.response.json()
                     error_msg = f"{error_msg} - {error_details.get('error', {}).get('message', '')}"
                 except Exception:
                     pass
-            raise RuntimeError(error_msg) from e
+            raise RuntimeError(error_msg) from e

prompture/drivers/registry.py

@@ -0,0 +1,306 @@
+"""Driver registry with plugin support.
+
+This module provides a public API for registering custom drivers and
+supports auto-discovery of drivers via Python entry points.
+
+Example usage:
+    # Register a custom driver
+    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")
+
+For entry point discovery, add to your package's pyproject.toml:
+    [project.entry-points."prompture.drivers"]
+    my_provider = "my_package.drivers:my_driver_factory"
+
+    [project.entry-points."prompture.async_drivers"]
+    my_provider = "my_package.drivers:my_async_driver_factory"
+"""
+
+from __future__ import annotations
+
+import logging
+import sys
+from typing import Callable
+
+logger = logging.getLogger("prompture.drivers.registry")
+
+# Type alias for driver factory functions
+# A factory takes an optional model name and returns a driver instance
+DriverFactory = Callable[[str | None], object]
+
+# Internal registries - populated by built-in drivers and plugins
+_SYNC_REGISTRY: dict[str, DriverFactory] = {}
+_ASYNC_REGISTRY: dict[str, DriverFactory] = {}
+
+# Track whether entry points have been loaded
+_entry_points_loaded = False
+
+
+def register_driver(name: str, factory: DriverFactory, *, overwrite: bool = False) -> None:
+    """Register a custom driver factory for a provider name.
+
+    Args:
+        name: Provider name (e.g., "my_provider"). Will be lowercased.
+        factory: A callable that takes an optional model name and returns
+                 a driver instance. The driver must implement the
+                 ``Driver`` interface (specifically ``generate()``).
+        overwrite: If True, allow overwriting an existing registration.
+                   Defaults to False.
+
+    Raises:
+        ValueError: If a driver with this name is already registered
+                    and overwrite=False.
+
+    Example:
+        >>> def my_factory(model=None):
+        ...     return MyDriver(model=model or "default-model")
+        >>> register_driver("my_provider", my_factory)
+        >>> driver = get_driver_for_model("my_provider/custom-model")
+    """
+    name = name.lower()
+    if name in _SYNC_REGISTRY and not overwrite:
+        raise ValueError(f"Driver '{name}' is already registered. Use overwrite=True to replace it.")
+    _SYNC_REGISTRY[name] = factory
+    logger.debug("Registered sync driver: %s", name)
+
+
+def register_async_driver(name: str, factory: DriverFactory, *, overwrite: bool = False) -> None:
+    """Register a custom async driver factory for a provider name.
+
+    Args:
+        name: Provider name (e.g., "my_provider"). Will be lowercased.
+        factory: A callable that takes an optional model name and returns
+                 an async driver instance. The driver must implement the
+                 ``AsyncDriver`` interface (specifically ``async generate()``).
+        overwrite: If True, allow overwriting an existing registration.
+                   Defaults to False.
+
+    Raises:
+        ValueError: If an async driver with this name is already registered
+                    and overwrite=False.
+
+    Example:
+        >>> def my_async_factory(model=None):
+        ...     return MyAsyncDriver(model=model or "default-model")
+        >>> register_async_driver("my_provider", my_async_factory)
+        >>> driver = get_async_driver_for_model("my_provider/custom-model")
+    """
+    name = name.lower()
+    if name in _ASYNC_REGISTRY and not overwrite:
+        raise ValueError(f"Async driver '{name}' is already registered. Use overwrite=True to replace it.")
+    _ASYNC_REGISTRY[name] = factory
+    logger.debug("Registered async driver: %s", name)
+
+
+def unregister_driver(name: str) -> bool:
+    """Unregister a sync driver by name.
+
+    Args:
+        name: Provider name to unregister.
+
+    Returns:
+        True if the driver was unregistered, False if it wasn't registered.
+    """
+    name = name.lower()
+    if name in _SYNC_REGISTRY:
+        del _SYNC_REGISTRY[name]
+        logger.debug("Unregistered sync driver: %s", name)
+        return True
+    return False
+
+
+def unregister_async_driver(name: str) -> bool:
+    """Unregister an async driver by name.
+
+    Args:
+        name: Provider name to unregister.
+
+    Returns:
+        True if the driver was unregistered, False if it wasn't registered.
+    """
+    name = name.lower()
+    if name in _ASYNC_REGISTRY:
+        del _ASYNC_REGISTRY[name]
+        logger.debug("Unregistered async driver: %s", name)
+        return True
+    return False
+
+
+def list_registered_drivers() -> list[str]:
+    """Return a sorted list of registered sync driver names."""
+    _ensure_entry_points_loaded()
+    return sorted(_SYNC_REGISTRY.keys())
+
+
+def list_registered_async_drivers() -> list[str]:
+    """Return a sorted list of registered async driver names."""
+    _ensure_entry_points_loaded()
+    return sorted(_ASYNC_REGISTRY.keys())
+
+
+def is_driver_registered(name: str) -> bool:
+    """Check if a sync driver is registered.
+
+    Args:
+        name: Provider name to check.
+
+    Returns:
+        True if the driver is registered.
+    """
+    _ensure_entry_points_loaded()
+    return name.lower() in _SYNC_REGISTRY
+
+
+def is_async_driver_registered(name: str) -> bool:
+    """Check if an async driver is registered.
+
+    Args:
+        name: Provider name to check.
+
+    Returns:
+        True if the async driver is registered.
+    """
+    _ensure_entry_points_loaded()
+    return name.lower() in _ASYNC_REGISTRY
+
+
+def get_driver_factory(name: str) -> DriverFactory:
+    """Get a registered sync driver factory by name.
+
+    Args:
+        name: Provider name.
+
+    Returns:
+        The factory function.
+
+    Raises:
+        ValueError: If the driver is not registered.
+    """
+    _ensure_entry_points_loaded()
+    name = name.lower()
+    if name not in _SYNC_REGISTRY:
+        raise ValueError(f"Unsupported provider '{name}'")
+    return _SYNC_REGISTRY[name]
+
+
+def get_async_driver_factory(name: str) -> DriverFactory:
+    """Get a registered async driver factory by name.
+
+    Args:
+        name: Provider name.
+
+    Returns:
+        The factory function.
+
+    Raises:
+        ValueError: If the async driver is not registered.
+    """
+    _ensure_entry_points_loaded()
+    name = name.lower()
+    if name not in _ASYNC_REGISTRY:
+        raise ValueError(f"Unsupported provider '{name}'")
+    return _ASYNC_REGISTRY[name]
+
+
+def load_entry_point_drivers() -> tuple[int, int]:
+    """Load drivers from installed packages via entry points.
+
+    This function scans for packages that define entry points in the
+    ``prompture.drivers`` and ``prompture.async_drivers`` groups.
+
+    Returns:
+        A tuple of (sync_count, async_count) indicating how many drivers
+        were loaded from entry points.
+
+    Example pyproject.toml for a plugin package:
+        [project.entry-points."prompture.drivers"]
+        my_provider = "my_package.drivers:create_my_driver"
+
+        [project.entry-points."prompture.async_drivers"]
+        my_provider = "my_package.drivers:create_my_async_driver"
+    """
+    global _entry_points_loaded
+
+    sync_count = 0
+    async_count = 0
+
+    # Python 3.9+ has importlib.metadata in stdlib
+    # Python 3.8 needs importlib_metadata backport
+    if sys.version_info >= (3, 10):
+        from importlib.metadata import entry_points
+
+        sync_eps = entry_points(group="prompture.drivers")
+        async_eps = entry_points(group="prompture.async_drivers")
+    else:
+        from importlib.metadata import entry_points
+
+        all_eps = entry_points()
+        sync_eps = all_eps.get("prompture.drivers", [])
+        async_eps = all_eps.get("prompture.async_drivers", [])
+
+    # Load sync drivers
+    for ep in sync_eps:
+        try:
+            # Skip if already registered (built-in drivers take precedence)
+            if ep.name.lower() in _SYNC_REGISTRY:
+                logger.debug("Skipping entry point driver '%s' (already registered)", ep.name)
+                continue
+
+            factory = ep.load()
+            _SYNC_REGISTRY[ep.name.lower()] = factory
+            sync_count += 1
+            logger.info("Loaded sync driver from entry point: %s", ep.name)
+        except Exception:
+            logger.exception("Failed to load sync driver entry point: %s", ep.name)
+
+    # Load async drivers
+    for ep in async_eps:
+        try:
+            # Skip if already registered (built-in drivers take precedence)
+            if ep.name.lower() in _ASYNC_REGISTRY:
+                logger.debug("Skipping entry point async driver '%s' (already registered)", ep.name)
+                continue
+
+            factory = ep.load()
+            _ASYNC_REGISTRY[ep.name.lower()] = factory
+            async_count += 1
+            logger.info("Loaded async driver from entry point: %s", ep.name)
+        except Exception:
+            logger.exception("Failed to load async driver entry point: %s", ep.name)
+
+    _entry_points_loaded = True
+    return (sync_count, async_count)
+
+
+def _ensure_entry_points_loaded() -> None:
+    """Ensure entry points have been loaded (lazy initialization)."""
+    global _entry_points_loaded
+    if not _entry_points_loaded:
+        load_entry_point_drivers()
+
+
+def _get_sync_registry() -> dict[str, DriverFactory]:
+    """Get the internal sync registry dict (for internal use by drivers/__init__.py)."""
+    _ensure_entry_points_loaded()
+    return _SYNC_REGISTRY
+
+
+def _get_async_registry() -> dict[str, DriverFactory]:
+    """Get the internal async registry dict (for internal use by drivers/async_registry.py)."""
+    _ensure_entry_points_loaded()
+    return _ASYNC_REGISTRY
+
+
+def _reset_registries() -> None:
+    """Reset registries to empty state (for testing only)."""
+    global _entry_points_loaded
+    _SYNC_REGISTRY.clear()
+    _ASYNC_REGISTRY.clear()
+    _entry_points_loaded = False