prompture 0.0.29.dev8__py3-none-any.whl → 0.0.38.dev2__py3-none-any.whl

prompture/drivers/__init__.py

@@ -1,71 +1,149 @@
-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
-    ),
-    "claude": lambda model=None: ClaudeDriver(
-        api_key=settings.claude_api_key,
-        model=model or settings.claude_model
-    ),
-    "lmstudio": lambda model=None: LMStudioDriver(
+# 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
-    ),
-    "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
+        model=model or settings.lmstudio_model,
+        api_key=settings.lmstudio_api_key,
     ),
-    "openrouter": lambda model=None: OpenRouterDriver(
-        api_key=settings.openrouter_api_key,
-        model=model or settings.openrouter_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
     ),
-    "grok": lambda model=None: GrokDriver(
-        api_key=settings.grok_api_key,
-        model=model or settings.grok_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,
     ),
-}
+    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 +151,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 +174,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",
 ]

prompture/drivers/airllm_driver.py

@@ -0,0 +1,109 @@
+import logging
+from typing import Any, Optional
+
+from ..driver import Driver
+
+logger = logging.getLogger(__name__)
+
+
+class AirLLMDriver(Driver):
+    """Driver for AirLLM — run large models (70B+) on consumer GPUs via
+    layer-by-layer memory management.
+
+    The ``airllm`` package is a lazy dependency: it is imported on first
+    ``generate()`` call so the rest of Prompture works without it installed.
+    """
+
+    MODEL_PRICING = {"default": {"prompt": 0.0, "completion": 0.0}}
+
+    def __init__(self, model: str = "meta-llama/Llama-2-7b-hf", compression: Optional[str] = None):
+        """
+        Args:
+            model: HuggingFace repo ID (e.g. ``"meta-llama/Llama-2-70b-hf"``).
+            compression: Optional quantization mode — ``"4bit"`` or ``"8bit"``.
+        """
+        self.model = model
+        self.compression = compression
+        self.options: dict[str, Any] = {}
+        self._llm = None
+        self._tokenizer = None
+
+    # ------------------------------------------------------------------
+    # Lazy model loading
+    # ------------------------------------------------------------------
+    def _ensure_loaded(self):
+        """Load the AirLLM model and tokenizer on first use."""
+        if self._llm is not None:
+            return
+
+        try:
+            from airllm import AutoModel
+        except ImportError:
+            raise ImportError(
+                "The 'airllm' package is required for the AirLLM driver. Install it with: pip install prompture[airllm]"
+            ) from None
+
+        try:
+            from transformers import AutoTokenizer
+        except ImportError:
+            raise ImportError(
+                "The 'transformers' package is required for the AirLLM driver. "
+                "Install it with: pip install transformers"
+            ) from None
+
+        logger.info(f"Loading AirLLM model: {self.model} (compression={self.compression})")
+
+        load_kwargs: dict[str, Any] = {}
+        if self.compression:
+            load_kwargs["compression"] = self.compression
+
+        self._llm = AutoModel.from_pretrained(self.model, **load_kwargs)
+        self._tokenizer = AutoTokenizer.from_pretrained(self.model)
+        logger.info("AirLLM model loaded successfully")
+
+    # ------------------------------------------------------------------
+    # Driver interface
+    # ------------------------------------------------------------------
+    def generate(self, prompt: str, options: Optional[dict[str, Any]] = None) -> dict[str, Any]:
+        self._ensure_loaded()
+
+        merged_options = self.options.copy()
+        if options:
+            merged_options.update(options)
+
+        max_new_tokens = merged_options.get("max_new_tokens", 256)
+
+        # Tokenize
+        input_ids = self._tokenizer(prompt, return_tensors="pt").input_ids
+
+        prompt_tokens = input_ids.shape[1]
+
+        logger.debug(f"AirLLM generating with max_new_tokens={max_new_tokens}, prompt_tokens={prompt_tokens}")
+
+        # Generate
+        output_ids = self._llm.generate(
+            input_ids,
+            max_new_tokens=max_new_tokens,
+        )
+
+        # Decode only the newly generated tokens (strip the prompt prefix)
+        new_tokens = output_ids[0, prompt_tokens:]
+        completion_tokens = len(new_tokens)
+        text = self._tokenizer.decode(new_tokens, skip_special_tokens=True)
+
+        total_tokens = prompt_tokens + completion_tokens
+
+        meta = {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": total_tokens,
+            "cost": 0.0,
+            "raw_response": {
+                "model": self.model,
+                "compression": self.compression,
+                "max_new_tokens": max_new_tokens,
+            },
+            "model_name": self.model,
+        }
+
+        return {"text": text, "meta": meta}

prompture/drivers/async_airllm_driver.py

@@ -0,0 +1,26 @@
+"""Async AirLLM driver — wraps the sync GPU-bound driver with asyncio.to_thread."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+from ..async_driver import AsyncDriver
+from .airllm_driver import AirLLMDriver
+
+
+class AsyncAirLLMDriver(AsyncDriver):
+    """Async wrapper around :class:`AirLLMDriver`.
+
+    AirLLM is GPU-bound with no native async API, so we delegate to
+    ``asyncio.to_thread()`` to avoid blocking the event loop.
+    """
+
+    MODEL_PRICING = AirLLMDriver.MODEL_PRICING
+
+    def __init__(self, model: str = "meta-llama/Llama-2-7b-hf", compression: str | None = None):
+        self.model = model
+        self._sync_driver = AirLLMDriver(model=model, compression=compression)
+
+    async def generate(self, prompt: str, options: dict[str, Any] | None = None) -> dict[str, Any]:
+        return await asyncio.to_thread(self._sync_driver.generate, prompt, options)

prompture/drivers/async_azure_driver.py

@@ -0,0 +1,123 @@
+"""Async Azure OpenAI driver. Requires the ``openai`` package (>=1.0.0)."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+try:
+    from openai import AsyncAzureOpenAI
+except Exception:
+    AsyncAzureOpenAI = None
+
+from ..async_driver import AsyncDriver
+from ..cost_mixin import CostMixin
+from .azure_driver import AzureDriver
+
+
+class AsyncAzureDriver(CostMixin, AsyncDriver):
+    supports_json_mode = True
+    supports_json_schema = True
+    supports_vision = True
+
+    MODEL_PRICING = AzureDriver.MODEL_PRICING
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        endpoint: str | None = None,
+        deployment_id: str | None = None,
+        model: str = "gpt-4o-mini",
+    ):
+        self.api_key = api_key or os.getenv("AZURE_API_KEY")
+        self.endpoint = endpoint or os.getenv("AZURE_API_ENDPOINT")
+        self.deployment_id = deployment_id or os.getenv("AZURE_DEPLOYMENT_ID")
+        self.api_version = os.getenv("AZURE_API_VERSION", "2023-07-01-preview")
+        self.model = model
+
+        if not self.api_key:
+            raise ValueError("Missing Azure API key (AZURE_API_KEY).")
+        if not self.endpoint:
+            raise ValueError("Missing Azure API endpoint (AZURE_API_ENDPOINT).")
+        if not self.deployment_id:
+            raise ValueError("Missing Azure deployment ID (AZURE_DEPLOYMENT_ID).")
+
+        if AsyncAzureOpenAI:
+            self.client = AsyncAzureOpenAI(
+                api_key=self.api_key,
+                api_version=self.api_version,
+                azure_endpoint=self.endpoint,
+            )
+        else:
+            self.client = None
+
+    supports_messages = True
+
+    def _prepare_messages(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        from .vision_helpers import _prepare_openai_vision_messages
+
+        return _prepare_openai_vision_messages(messages)
+
+    async def generate(self, prompt: str, options: dict[str, Any]) -> dict[str, Any]:
+        messages = [{"role": "user", "content": prompt}]
+        return await self._do_generate(messages, options)
+
+    async def generate_messages(self, messages: list[dict[str, str]], options: dict[str, Any]) -> dict[str, Any]:
+        return await self._do_generate(self._prepare_messages(messages), options)
+
+    async def _do_generate(self, messages: list[dict[str, str]], options: dict[str, Any]) -> dict[str, Any]:
+        if self.client is None:
+            raise RuntimeError("openai package (>=1.0.0) with AsyncAzureOpenAI 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 = {
+            "model": self.deployment_id,
+            "messages": messages,
+        }
+        kwargs[tokens_param] = opts.get("max_tokens", 512)
+
+        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 = await 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("azure", model, prompt_tokens, completion_tokens)
+
+        meta = {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": total_tokens,
+            "cost": total_cost,
+            "raw_response": resp.model_dump(),
+            "model_name": model,
+            "deployment_id": self.deployment_id,
+        }
+
+        text = resp.choices[0].message.content
+        return {"text": text, "meta": meta}

prompture/drivers/async_claude_driver.py

@@ -0,0 +1,113 @@
+"""Async Anthropic Claude driver. Requires the ``anthropic`` package."""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import Any
+
+try:
+    import anthropic
+except Exception:
+    anthropic = None
+
+from ..async_driver import AsyncDriver
+from ..cost_mixin import CostMixin
+from .claude_driver import ClaudeDriver
+
+
+class AsyncClaudeDriver(CostMixin, AsyncDriver):
+    supports_json_mode = True
+    supports_json_schema = True
+    supports_vision = True
+
+    MODEL_PRICING = ClaudeDriver.MODEL_PRICING
+
+    def __init__(self, api_key: str | None = None, model: str = "claude-3-5-haiku-20241022"):
+        self.api_key = api_key or os.getenv("CLAUDE_API_KEY")
+        self.model = model or os.getenv("CLAUDE_MODEL_NAME", "claude-3-5-haiku-20241022")
+
+    supports_messages = True
+
+    def _prepare_messages(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        from .vision_helpers import _prepare_claude_vision_messages
+
+        return _prepare_claude_vision_messages(messages)
+
+    async def generate(self, prompt: str, options: dict[str, Any]) -> dict[str, Any]:
+        messages = [{"role": "user", "content": prompt}]
+        return await self._do_generate(messages, options)
+
+    async def generate_messages(self, messages: list[dict[str, str]], options: dict[str, Any]) -> dict[str, Any]:
+        return await self._do_generate(self._prepare_messages(messages), options)
+
+    async def _do_generate(self, messages: list[dict[str, str]], options: dict[str, Any]) -> dict[str, Any]:
+        if anthropic is None:
+            raise RuntimeError("anthropic package not installed")
+
+        opts = {**{"temperature": 0.0, "max_tokens": 512}, **options}
+        model = options.get("model", self.model)
+
+        client = anthropic.AsyncAnthropic(api_key=self.api_key)
+
+        # Anthropic requires system messages as a top-level parameter
+        system_content = None
+        api_messages = []
+        for msg in messages:
+            if msg.get("role") == "system":
+                system_content = msg.get("content", "")
+            else:
+                api_messages.append(msg)
+
+        # Build common kwargs
+        common_kwargs: dict[str, Any] = {
+            "model": model,
+            "messages": api_messages,
+            "temperature": opts["temperature"],
+            "max_tokens": opts["max_tokens"],
+        }
+        if system_content:
+            common_kwargs["system"] = system_content
+
+        # Native JSON mode: use tool-use for schema enforcement
+        if options.get("json_mode"):
+            json_schema = options.get("json_schema")
+            if json_schema:
+                tool_def = {
+                    "name": "extract_json",
+                    "description": "Extract structured data matching the schema",
+                    "input_schema": json_schema,
+                }
+                resp = await client.messages.create(
+                    **common_kwargs,
+                    tools=[tool_def],
+                    tool_choice={"type": "tool", "name": "extract_json"},
+                )
+                text = ""
+                for block in resp.content:
+                    if block.type == "tool_use":
+                        text = json.dumps(block.input)
+                        break
+            else:
+                resp = await client.messages.create(**common_kwargs)
+                text = resp.content[0].text
+        else:
+            resp = await client.messages.create(**common_kwargs)
+            text = resp.content[0].text
+
+        prompt_tokens = resp.usage.input_tokens
+        completion_tokens = resp.usage.output_tokens
+        total_tokens = prompt_tokens + completion_tokens
+
+        total_cost = self._calculate_cost("claude", model, prompt_tokens, completion_tokens)
+
+        meta = {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": total_tokens,
+            "cost": total_cost,
+            "raw_response": dict(resp),
+            "model_name": model,
+        }
+
+        return {"text": text, "meta": meta}