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

prompture/drivers/airllm_driver.py

@@ -1,6 +1,7 @@
 import logging
+from typing import Any, Optional
+
 from ..driver import Driver
-from typing import Any, Dict, Optional
 
 logger = logging.getLogger(__name__)
 
@@ -13,12 +14,9 @@ class AirLLMDriver(Driver):
     ``generate()`` call so the rest of Prompture works without it installed.
     """
 
-    MODEL_PRICING = {
-        "default": {"prompt": 0.0, "completion": 0.0}
-    }
+    MODEL_PRICING = {"default": {"prompt": 0.0, "completion": 0.0}}
 
-    def __init__(self, model: str = "meta-llama/Llama-2-7b-hf",
-                 compression: Optional[str] = None):
+    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"``).
@@ -26,7 +24,7 @@ class AirLLMDriver(Driver):
         """
         self.model = model
         self.compression = compression
-        self.options: Dict[str, Any] = {}
+        self.options: dict[str, Any] = {}
         self._llm = None
         self._tokenizer = None
 
@@ -42,9 +40,8 @@ class AirLLMDriver(Driver):
             from airllm import AutoModel
         except ImportError:
             raise ImportError(
-                "The 'airllm' package is required for the AirLLM driver. "
-                "Install it with: pip install prompture[airllm]"
-            )
+                "The 'airllm' package is required for the AirLLM driver. Install it with: pip install prompture[airllm]"
+            ) from None
 
         try:
             from transformers import AutoTokenizer
@@ -52,12 +49,11 @@ class AirLLMDriver(Driver):
             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} "
-                     f"(compression={self.compression})")
+        logger.info(f"Loading AirLLM model: {self.model} (compression={self.compression})")
 
-        load_kwargs: Dict[str, Any] = {}
+        load_kwargs: dict[str, Any] = {}
         if self.compression:
             load_kwargs["compression"] = self.compression
 
@@ -68,7 +64,7 @@ class AirLLMDriver(Driver):
     # ------------------------------------------------------------------
     # Driver interface
     # ------------------------------------------------------------------
-    def generate(self, prompt: str, options: Dict[str, Any] = None) -> Dict[str, Any]:
+    def generate(self, prompt: str, options: Optional[dict[str, Any]] = None) -> dict[str, Any]:
         self._ensure_loaded()
 
         merged_options = self.options.copy()
@@ -78,14 +74,11 @@ class AirLLMDriver(Driver):
         max_new_tokens = merged_options.get("max_new_tokens", 256)
 
         # Tokenize
-        input_ids = self._tokenizer(
-            prompt, return_tensors="pt"
-        ).input_ids
+        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}, "
-                      f"prompt_tokens={prompt_tokens}")
+        logger.debug(f"AirLLM generating with max_new_tokens={max_new_tokens}, prompt_tokens={prompt_tokens}")
 
         # Generate
         output_ids = self._llm.generate(

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,117 @@
+"""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
+
+    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
+
+    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(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,107 @@
+"""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
+
+    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
+
+    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(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}

prompture/drivers/async_google_driver.py

@@ -0,0 +1,132 @@
+"""Async Google Generative AI (Gemini) driver."""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any
+
+import google.generativeai as genai
+
+from ..async_driver import AsyncDriver
+from ..cost_mixin import CostMixin
+from .google_driver import GoogleDriver
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncGoogleDriver(CostMixin, AsyncDriver):
+    """Async driver for Google's Generative AI API (Gemini)."""
+
+    supports_json_mode = True
+    supports_json_schema = True
+
+    MODEL_PRICING = GoogleDriver.MODEL_PRICING
+    _PRICING_UNIT = 1_000_000
+
+    def __init__(self, api_key: str | None = None, model: str = "gemini-1.5-pro"):
+        self.api_key = api_key or os.getenv("GOOGLE_API_KEY")
+        if not self.api_key:
+            raise ValueError("Google API key not found. Set GOOGLE_API_KEY env var or pass api_key to constructor")
+        self.model = model
+        genai.configure(api_key=self.api_key)
+        self.options: dict[str, Any] = {}
+
+    def _calculate_cost_chars(self, prompt_chars: int, completion_chars: int) -> float:
+        """Calculate cost from character counts (same logic as sync GoogleDriver)."""
+        from ..model_rates import get_model_rates
+
+        live_rates = get_model_rates("google", self.model)
+        if live_rates:
+            est_prompt_tokens = prompt_chars / 4
+            est_completion_tokens = completion_chars / 4
+            prompt_cost = (est_prompt_tokens / 1_000_000) * live_rates["input"]
+            completion_cost = (est_completion_tokens / 1_000_000) * live_rates["output"]
+        else:
+            model_pricing = self.MODEL_PRICING.get(self.model, {"prompt": 0, "completion": 0})
+            prompt_cost = (prompt_chars / 1_000_000) * model_pricing["prompt"]
+            completion_cost = (completion_chars / 1_000_000) * model_pricing["completion"]
+        return round(prompt_cost + completion_cost, 6)
+
+    supports_messages = True
+
+    async def generate(self, prompt: str, options: dict[str, Any] | None = None) -> 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(messages, options)
+
+    async def _do_generate(
+        self, messages: list[dict[str, str]], options: dict[str, Any] | None = None
+    ) -> dict[str, Any]:
+        merged_options = self.options.copy()
+        if options:
+            merged_options.update(options)
+
+        generation_config = merged_options.get("generation_config", {})
+        safety_settings = merged_options.get("safety_settings", {})
+
+        if "temperature" in merged_options and "temperature" not in generation_config:
+            generation_config["temperature"] = merged_options["temperature"]
+        if "max_tokens" in merged_options and "max_output_tokens" not in generation_config:
+            generation_config["max_output_tokens"] = merged_options["max_tokens"]
+        if "top_p" in merged_options and "top_p" not in generation_config:
+            generation_config["top_p"] = merged_options["top_p"]
+        if "top_k" in merged_options and "top_k" not in generation_config:
+            generation_config["top_k"] = merged_options["top_k"]
+
+        # Native JSON mode support
+        if merged_options.get("json_mode"):
+            generation_config["response_mime_type"] = "application/json"
+            json_schema = merged_options.get("json_schema")
+            if json_schema:
+                generation_config["response_schema"] = json_schema
+
+        # Convert messages to Gemini format
+        system_instruction = None
+        contents: list[dict[str, Any]] = []
+        for msg in messages:
+            role = msg.get("role", "user")
+            content = msg.get("content", "")
+            if role == "system":
+                system_instruction = content
+            else:
+                gemini_role = "model" if role == "assistant" else "user"
+                contents.append({"role": gemini_role, "parts": [content]})
+
+        try:
+            model_kwargs: dict[str, Any] = {}
+            if system_instruction:
+                model_kwargs["system_instruction"] = system_instruction
+            model = genai.GenerativeModel(self.model, **model_kwargs)
+
+            gen_input: Any = contents if len(contents) != 1 else contents[0]["parts"][0]
+            response = await model.generate_content_async(
+                gen_input,
+                generation_config=generation_config if generation_config else None,
+                safety_settings=safety_settings if safety_settings else None,
+            )
+
+            if not response.text:
+                raise ValueError("Empty response from model")
+
+            total_prompt_chars = sum(len(msg.get("content", "")) for msg in messages)
+            completion_chars = len(response.text)
+
+            total_cost = self._calculate_cost_chars(total_prompt_chars, completion_chars)
+
+            meta = {
+                "prompt_chars": total_prompt_chars,
+                "completion_chars": completion_chars,
+                "total_chars": total_prompt_chars + completion_chars,
+                "cost": total_cost,
+                "raw_response": response.prompt_feedback if hasattr(response, "prompt_feedback") else None,
+                "model_name": self.model,
+            }
+
+            return {"text": response.text, "meta": meta}
+
+        except Exception as e:
+            logger.error(f"Google API request failed: {e}")
+            raise RuntimeError(f"Google API request failed: {e}") from e

prompture/drivers/async_grok_driver.py

@@ -0,0 +1,91 @@
+"""Async xAI Grok driver using httpx."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import httpx
+
+from ..async_driver import AsyncDriver
+from ..cost_mixin import CostMixin
+from .grok_driver import GrokDriver
+
+
+class AsyncGrokDriver(CostMixin, AsyncDriver):
+    supports_json_mode = True
+
+    MODEL_PRICING = GrokDriver.MODEL_PRICING
+    _PRICING_UNIT = 1_000_000
+
+    def __init__(self, api_key: str | None = None, model: str = "grok-4-fast-reasoning"):
+        self.api_key = api_key or os.getenv("GROK_API_KEY")
+        self.model = model
+        self.api_base = "https://api.x.ai/v1"
+
+    supports_messages = True
+
+    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(messages, options)
+
+    async def _do_generate(self, messages: list[dict[str, str]], options: dict[str, Any]) -> dict[str, Any]:
+        if not self.api_key:
+            raise RuntimeError("GROK_API_KEY environment variable is required")
+
+        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}
+
+        payload = {
+            "model": model,
+            "messages": messages,
+        }
+        payload[tokens_param] = opts.get("max_tokens", 512)
+
+        if supports_temperature and "temperature" in opts:
+            payload["temperature"] = opts["temperature"]
+
+        # Native JSON mode support
+        if options.get("json_mode"):
+            payload["response_format"] = {"type": "json_object"}
+
+        headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
+
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.post(
+                    f"{self.api_base}/chat/completions", headers=headers, json=payload, timeout=120
+                )
+                response.raise_for_status()
+                resp = response.json()
+            except httpx.HTTPStatusError as e:
+                raise RuntimeError(f"Grok API request failed: {e!s}") from e
+            except Exception as e:
+                raise RuntimeError(f"Grok API request failed: {e!s}") from e
+
+        usage = resp.get("usage", {})
+        prompt_tokens = usage.get("prompt_tokens", 0)
+        completion_tokens = usage.get("completion_tokens", 0)
+        total_tokens = usage.get("total_tokens", 0)
+
+        total_cost = self._calculate_cost("grok", 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_name": model,
+        }
+
+        text = resp["choices"][0]["message"]["content"]
+        return {"text": text, "meta": meta}

prompture/drivers/async_groq_driver.py

@@ -0,0 +1,84 @@
+"""Async Groq driver. Requires the ``groq`` package."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+try:
+    import groq
+except Exception:
+    groq = None
+
+from ..async_driver import AsyncDriver
+from ..cost_mixin import CostMixin
+from .groq_driver import GroqDriver
+
+
+class AsyncGroqDriver(CostMixin, AsyncDriver):
+    supports_json_mode = True
+
+    MODEL_PRICING = GroqDriver.MODEL_PRICING
+
+    def __init__(self, api_key: str | None = None, model: str = "llama2-70b-4096"):
+        self.api_key = api_key or os.getenv("GROQ_API_KEY")
+        self.model = model
+        if groq:
+            self.client = groq.AsyncClient(api_key=self.api_key)
+        else:
+            self.client = None
+
+    supports_messages = True
+
+    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(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("groq package 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": 0.7, "max_tokens": 512, **options}
+
+        kwargs = {
+            "model": model,
+            "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"):
+            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("groq", 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,
+        }
+
+        text = resp.choices[0].message.content
+        return {"text": text, "meta": meta}