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

prompture/drivers/claude_driver.py

@@ -1,81 +1,128 @@
 """Driver for Anthropic's Claude models. Requires the `anthropic` library.
 Use with API key in CLAUDE_API_KEY env var or provide directly.
 """
+
+import json
 import os
-from typing import Any, Dict
+from typing import Any
+
 try:
     import anthropic
 except Exception:
     anthropic = None
 
+from ..cost_mixin import CostMixin
 from ..driver import Driver
 
-class ClaudeDriver(Driver):
+
+class ClaudeDriver(CostMixin, Driver):
+    supports_json_mode = True
+    supports_json_schema = True
+
     # Claude pricing per 1000 tokens (prices should be kept current with Anthropic's pricing)
     MODEL_PRICING = {
         # Claude Opus 4.1
         "claude-opus-4-1-20250805": {
-            "prompt": 0.015,      # $15 per 1M prompt tokens
-            "completion": 0.075,   # $75 per 1M completion tokens
+            "prompt": 0.015,  # $15 per 1M prompt tokens
+            "completion": 0.075,  # $75 per 1M completion tokens
         },
         # Claude Opus 4.0
         "claude-opus-4-20250514": {
-            "prompt": 0.015,      # $15 per 1M prompt tokens
-            "completion": 0.075,   # $75 per 1M completion tokens
+            "prompt": 0.015,  # $15 per 1M prompt tokens
+            "completion": 0.075,  # $75 per 1M completion tokens
         },
         # Claude Sonnet 4.0
         "claude-sonnet-4-20250514": {
-            "prompt": 0.003,      # $3 per 1M prompt tokens
-            "completion": 0.015,   # $15 per 1M completion tokens
+            "prompt": 0.003,  # $3 per 1M prompt tokens
+            "completion": 0.015,  # $15 per 1M completion tokens
         },
         # Claude Sonnet 3.7
         "claude-3-7-sonnet-20250219": {
-            "prompt": 0.003,      # $3 per 1M prompt tokens
-            "completion": 0.015,   # $15 per 1M completion tokens
+            "prompt": 0.003,  # $3 per 1M prompt tokens
+            "completion": 0.015,  # $15 per 1M completion tokens
         },
         # Claude Haiku 3.5
         "claude-3-5-haiku-20241022": {
-            "prompt": 0.0008,     # $0.80 per 1M prompt tokens
-            "completion": 0.004,   # $4 per 1M completion tokens
-        }
+            "prompt": 0.0008,  # $0.80 per 1M prompt tokens
+            "completion": 0.004,  # $4 per 1M completion tokens
+        },
     }
 
     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")
 
-    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, 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 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.Anthropic(api_key=self.api_key)
-        resp = client.messages.create(
-            model=model,
-            messages=[{"role": "user", "content": prompt}],
-            temperature=opts["temperature"],
-            max_tokens=opts["max_tokens"]
-        )
-        
+
+        # 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 = 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 = client.messages.create(**common_kwargs)
+                text = resp.content[0].text
+        else:
+            resp = client.messages.create(**common_kwargs)
+            text = resp.content[0].text
+
         # Extract token usage from Claude response
         prompt_tokens = resp.usage.input_tokens
         completion_tokens = resp.usage.output_tokens
         total_tokens = prompt_tokens + completion_tokens
-        
-        # Calculate cost — try live rates first (per 1M tokens), fall back to hardcoded (per 1K tokens)
-        from ..model_rates import get_model_rates
-        live_rates = get_model_rates("claude", 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:
-            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("claude", model, prompt_tokens, completion_tokens)
+
         # Create standardized meta object
         meta = {
             "prompt_tokens": prompt_tokens,
@@ -83,8 +130,7 @@ class ClaudeDriver(Driver):
             "total_tokens": total_tokens,
             "cost": round(total_cost, 6),  # Round to 6 decimal places
             "raw_response": dict(resp),
-            "model_name": model
+            "model_name": model,
         }
-        
-        text = resp.content[0].text
-        return {"text": text, "meta": meta}
+
+        return {"text": text, "meta": meta}

prompture/drivers/google_driver.py

@@ -1,60 +1,55 @@
-import os
 import logging
-import google.generativeai as genai
-from typing import Any, Dict
-from ..driver import Driver
-
 import os
-import logging
+from typing import Any, Optional
+
 import google.generativeai as genai
-from typing import Any, Dict
+
+from ..cost_mixin import CostMixin
 from ..driver import Driver
 
 logger = logging.getLogger(__name__)
 
 
-class GoogleDriver(Driver):
+class GoogleDriver(CostMixin, Driver):
     """Driver for Google's Generative AI API (Gemini)."""
 
+    supports_json_mode = True
+    supports_json_schema = True
+
     # Based on current Gemini pricing (as of 2025)
     # Source: https://cloud.google.com/vertex-ai/pricing#gemini_models
+    _PRICING_UNIT = 1_000_000
     MODEL_PRICING = {
         "gemini-1.5-pro": {
             "prompt": 0.00025,  # $0.25/1M chars input
-            "completion": 0.0005  # $0.50/1M chars output
+            "completion": 0.0005,  # $0.50/1M chars output
         },
         "gemini-1.5-pro-vision": {
             "prompt": 0.00025,  # $0.25/1M chars input
-            "completion": 0.0005  # $0.50/1M chars output
+            "completion": 0.0005,  # $0.50/1M chars output
         },
         "gemini-2.5-pro": {
             "prompt": 0.0004,  # $0.40/1M chars input
-            "completion": 0.0008  # $0.80/1M chars output
+            "completion": 0.0008,  # $0.80/1M chars output
         },
         "gemini-2.5-flash": {
             "prompt": 0.0004,  # $0.40/1M chars input
-            "completion": 0.0008  # $0.80/1M chars output
+            "completion": 0.0008,  # $0.80/1M chars output
         },
         "gemini-2.5-flash-lite": {
             "prompt": 0.0002,  # $0.20/1M chars input
-            "completion": 0.0004  # $0.40/1M chars output
+            "completion": 0.0004,  # $0.40/1M chars output
         },
-         "gemini-2.0-flash": {
+        "gemini-2.0-flash": {
             "prompt": 0.0004,  # $0.40/1M chars input
-            "completion": 0.0008  # $0.80/1M chars output
+            "completion": 0.0008,  # $0.80/1M chars output
         },
         "gemini-2.0-flash-lite": {
             "prompt": 0.0002,  # $0.20/1M chars input
-            "completion": 0.0004  # $0.40/1M chars output
-        },
-        "gemini-1.5-flash": {
-            "prompt": 0.00001875,
-            "completion": 0.000075
+            "completion": 0.0004,  # $0.40/1M chars output
         },
-        "gemini-1.5-flash-8b": {
-            "prompt": 0.00001,
-            "completion": 0.00004
-        }
+        "gemini-1.5-flash": {"prompt": 0.00001875, "completion": 0.000075},
+        "gemini-1.5-flash-8b": {"prompt": 0.00001, "completion": 0.00004},
     }
 
     def __init__(self, api_key: str | None = None, model: str = "gemini-1.5-pro"):
@@ -75,8 +70,8 @@ class GoogleDriver(Driver):
 
         # Configure google.generativeai
         genai.configure(api_key=self.api_key)
-        self.options: Dict[str, Any] = {}
-        
+        self.options: dict[str, Any] = {}
+
         # Validate connection and model availability
         self._validate_connection()
 
@@ -90,16 +85,36 @@ class GoogleDriver(Driver):
             logger.warning(f"Could not validate connection to Google API: {e}")
             raise
 
-    def generate(self, prompt: str, options: Dict[str, Any] = None) -> Dict[str, Any]:
-        """Generate text using Google's Generative AI.
-
-        Args:
-            prompt: The input prompt
-            options: Additional options to pass to the model
+    def _calculate_cost_chars(self, prompt_chars: int, completion_chars: int) -> float:
+        """Calculate cost from character counts.
 
-        Returns:
-            Dict containing generated text and metadata
+        Live rates use token-based pricing (estimate ~4 chars/token).
+        Hardcoded MODEL_PRICING uses per-1M-character rates.
         """
+        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
+
+    def generate(self, prompt: str, options: Optional[dict[str, Any]] = None) -> 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: Optional[dict[str, Any]] = None) -> dict[str, Any]:
         merged_options = self.options.copy()
         if options:
             merged_options.update(options)
@@ -107,7 +122,7 @@ class GoogleDriver(Driver):
         # Extract specific options for Google's API
         generation_config = merged_options.get("generation_config", {})
         safety_settings = merged_options.get("safety_settings", {})
-        
+
         # Map common options to generation_config if not present
         if "temperature" in merged_options and "temperature" not in generation_config:
             generation_config["temperature"] = merged_options["temperature"]
@@ -118,44 +133,57 @@ class GoogleDriver(Driver):
         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 uses "model" for assistant role
+                gemini_role = "model" if role == "assistant" else "user"
+                contents.append({"role": gemini_role, "parts": [content]})
+
         try:
             logger.debug(f"Initializing {self.model} for generation")
-            model = genai.GenerativeModel(self.model)
+            model_kwargs: dict[str, Any] = {}
+            if system_instruction:
+                model_kwargs["system_instruction"] = system_instruction
+            model = genai.GenerativeModel(self.model, **model_kwargs)
 
             # Generate response
-            logger.debug(f"Generating with prompt: {prompt}")
+            logger.debug(f"Generating with {len(contents)} content parts")
+            # If single user message, pass content directly for backward compatibility
+            gen_input: Any = contents if len(contents) != 1 else contents[0]["parts"][0]
             response = model.generate_content(
-                prompt,
+                gen_input,
                 generation_config=generation_config if generation_config else None,
-                safety_settings=safety_settings if safety_settings else None
+                safety_settings=safety_settings if safety_settings else None,
             )
-            
+
             if not response.text:
                 raise ValueError("Empty response from model")
 
             # Calculate token usage and cost
-            prompt_chars = len(prompt)
+            total_prompt_chars = sum(len(msg.get("content", "")) for msg in messages)
             completion_chars = len(response.text)
 
-            # Try live rates first (per 1M tokens), fall back to hardcoded character-based pricing
-            from ..model_rates import get_model_rates
-            live_rates = get_model_rates("google", self.model)
-            if live_rates:
-                # models.dev reports token-based pricing; estimate tokens from chars (~4 chars/token)
-                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"]
-            total_cost = prompt_cost + completion_cost
+            # Google uses character-based cost estimation
+            total_cost = self._calculate_cost_chars(total_prompt_chars, completion_chars)
 
             meta = {
-                "prompt_chars": prompt_chars,
+                "prompt_chars": total_prompt_chars,
                 "completion_chars": completion_chars,
-                "total_chars": prompt_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,
@@ -165,4 +193,4 @@ class GoogleDriver(Driver):
 
         except Exception as e:
             logger.error(f"Google API request failed: {e}")
-            raise RuntimeError(f"Google API request failed: {e}")
+            raise RuntimeError(f"Google API request failed: {e}") from e

prompture/drivers/grok_driver.py

@@ -1,15 +1,21 @@
 """xAI Grok driver.
 Requires the `requests` package. Uses GROK_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 GrokDriver(Driver):
+class GrokDriver(CostMixin, Driver):
+    supports_json_mode = True
+
     # Pricing per 1M tokens based on xAI's documentation
+    _PRICING_UNIT = 1_000_000
     MODEL_PRICING = {
         "grok-code-fast-1": {
             "prompt": 0.20,
@@ -72,19 +78,16 @@ class GrokDriver(Driver):
         self.model = model
         self.api_base = "https://api.x.ai/v1"
 
-    def generate(self, prompt: str, options: Dict[str, Any]) -> Dict[str, Any]:
-        """Generate completion using Grok API.
+    supports_messages = True
 
-        Args:
-            prompt: Input prompt
-            options: Generation options
-
-        Returns:
-            Dict containing generated text and metadata
-        
-        Raises:
-            RuntimeError: If API key is missing or request fails
-        """
+    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("GROK_API_KEY environment variable is required")
 
@@ -101,7 +104,7 @@ class GrokDriver(Driver):
         # Base request payload
         payload = {
             "model": model,
-            "messages": [{"role": "user", "content": prompt}],
+            "messages": messages,
         }
 
         # Add token limit with correct parameter name
@@ -111,39 +114,27 @@ class GrokDriver(Driver):
         if supports_temperature and "temperature" in opts:
             payload["temperature"] = opts["temperature"]
 
-        headers = {
-            "Authorization": f"Bearer {self.api_key}",
-            "Content-Type": "application/json"
-        }
+        # 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"}
 
         try:
-            response = requests.post(
-                f"{self.api_base}/chat/completions",
-                headers=headers,
-                json=payload
-            )
+            response = requests.post(f"{self.api_base}/chat/completions", headers=headers, json=payload)
             response.raise_for_status()
             resp = response.json()
         except requests.exceptions.RequestException as e:
-            raise RuntimeError(f"Grok API request failed: {str(e)}")
+            raise RuntimeError(f"Grok API request failed: {e!s}") from e
 
         # Extract usage info
         usage = resp.get("usage", {})
         prompt_tokens = usage.get("prompt_tokens", 0)
-        completion_tokens = usage.get("completion_tokens", 0) 
+        completion_tokens = usage.get("completion_tokens", 0)
         total_tokens = usage.get("total_tokens", 0)
 
-        # Calculate cost — try live rates first (per 1M tokens), fall back to hardcoded (per 1M tokens)
-        from ..model_rates import get_model_rates
-        live_rates = get_model_rates("grok", 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:
-            model_pricing = self.MODEL_PRICING.get(model, {"prompt": 0, "completion": 0})
-            prompt_cost = (prompt_tokens / 1_000_000) * model_pricing["prompt"]
-            completion_cost = (completion_tokens / 1_000_000) * model_pricing["completion"]
-        total_cost = prompt_cost + completion_cost
+        # Calculate cost via shared mixin
+        total_cost = self._calculate_cost("grok", model, prompt_tokens, completion_tokens)
 
         # Standardized meta object
         meta = {
@@ -156,4 +147,4 @@ class GrokDriver(Driver):
         }
 
         text = resp["choices"][0]["message"]["content"]
-        return {"text": text, "meta": meta}
+        return {"text": text, "meta": meta}

prompture/drivers/groq_driver.py

@@ -1,18 +1,22 @@
 """Groq driver for prompture.
 Requires the `groq` package. Uses GROQ_API_KEY env var.
 """
+
 import os
-from typing import Any, Dict
+from typing import Any
 
 try:
     import groq
 except Exception:
     groq = None
 
+from ..cost_mixin import CostMixin
 from ..driver import Driver
 
 
-class GroqDriver(Driver):
+class GroqDriver(CostMixin, Driver):
+    supports_json_mode = True
+
     # Approximate pricing per 1K tokens (to be updated with official pricing)
     # Each model entry defines token parameters and temperature support
     MODEL_PRICING = {
@@ -32,7 +36,7 @@ class GroqDriver(Driver):
 
     def __init__(self, api_key: str | None = None, model: str = "llama2-70b-4096"):
         """Initialize Groq driver.
-        
+
         Args:
             api_key: Groq API key (defaults to GROQ_API_KEY env var)
             model: Model to use (defaults to llama2-70b-4096)
@@ -44,20 +48,16 @@ class GroqDriver(Driver):
         else:
             self.client = None
 
-    def generate(self, prompt: str, options: Dict[str, Any]) -> Dict[str, Any]:
-        """Generate completion using Groq API.
-        
-        Args:
-            prompt: Input prompt
-            options: Generation options
-            
-        Returns:
-            Dict containing generated text and metadata
-            
-        Raises:
-            RuntimeError: If groq package is not installed
-            groq.error.*: Various Groq API errors
-        """
+    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 self.client is None:
             raise RuntimeError("groq package is not installed")
 
@@ -74,7 +74,7 @@ class GroqDriver(Driver):
         # Base kwargs for API call
         kwargs = {
             "model": model,
-            "messages": [{"role": "user", "content": prompt}],
+            "messages": messages,
         }
 
         # Set token limit with correct parameter name
@@ -84,29 +84,24 @@ class GroqDriver(Driver):
         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"}
+
         try:
             resp = self.client.chat.completions.create(**kwargs)
-        except Exception as e:
+        except Exception:
             # Re-raise any Groq API errors
             raise
 
         # Extract usage statistics
         usage = getattr(resp, "usage", None)
         prompt_tokens = getattr(usage, "prompt_tokens", 0)
-        completion_tokens = getattr(usage, "completion_tokens", 0) 
+        completion_tokens = getattr(usage, "completion_tokens", 0)
         total_tokens = getattr(usage, "total_tokens", 0)
 
-        # Calculate costs — try live rates first (per 1M tokens), fall back to hardcoded (per 1K tokens)
-        from ..model_rates import get_model_rates
-        live_rates = get_model_rates("groq", 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:
-            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("groq", model, prompt_tokens, completion_tokens)
 
         # Standard metadata object
         meta = {
@@ -120,4 +115,4 @@ class GroqDriver(Driver):
 
         # Extract generated text
         text = resp.choices[0].message.content
-        return {"text": text, "meta": meta}
+        return {"text": text, "meta": meta}