prompture 0.0.38.dev2__py3-none-any.whl → 0.0.42__py3-none-any.whl

prompture/drivers/openai_driver.py

@@ -12,7 +12,7 @@ try:
 except Exception:
     OpenAI = None
 
-from ..cost_mixin import CostMixin
+from ..cost_mixin import CostMixin, prepare_strict_schema
 from ..driver import Driver
 
 
@@ -93,10 +93,17 @@ class OpenAIDriver(CostMixin, Driver):
 
         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")
-        supports_temperature = model_info.get("supports_temperature", True)
+        # Lookup model-specific config (live models.dev data + hardcoded fallback)
+        model_config = self._get_model_config("openai", model)
+        tokens_param = model_config["tokens_param"]
+        supports_temperature = model_config["supports_temperature"]
+
+        # Validate capabilities against models.dev metadata
+        self._validate_model_capabilities(
+            "openai",
+            model,
+            using_json_schema=bool(options.get("json_schema")),
+        )
 
         # Defaults
         opts = {"temperature": 1.0, "max_tokens": 512, **options}
@@ -118,12 +125,13 @@ class OpenAIDriver(CostMixin, Driver):
         if options.get("json_mode"):
             json_schema = options.get("json_schema")
             if json_schema:
+                schema_copy = prepare_strict_schema(json_schema)
                 kwargs["response_format"] = {
                     "type": "json_schema",
                     "json_schema": {
                         "name": "extraction",
                         "strict": True,
-                        "schema": json_schema,
+                        "schema": schema_copy,
                     },
                 }
             else:
@@ -168,9 +176,11 @@ class OpenAIDriver(CostMixin, Driver):
             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)
+        model_config = self._get_model_config("openai", model)
+        tokens_param = model_config["tokens_param"]
+        supports_temperature = model_config["supports_temperature"]
+
+        self._validate_model_capabilities("openai", model, using_tool_use=True)
 
         opts = {"temperature": 1.0, "max_tokens": 512, **options}
 
@@ -239,9 +249,9 @@ class OpenAIDriver(CostMixin, Driver):
             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)
+        model_config = self._get_model_config("openai", model)
+        tokens_param = model_config["tokens_param"]
+        supports_temperature = model_config["supports_temperature"]
 
         opts = {"temperature": 1.0, "max_tokens": 512, **options}
 

prompture/drivers/openrouter_driver.py

@@ -2,54 +2,66 @@
 Requires the `requests` package. Uses OPENROUTER_API_KEY env var.
 """
 
+import contextlib
+import json
 import os
+from collections.abc import Iterator
 from typing import Any
 
 import requests
 
-from ..cost_mixin import CostMixin
+from ..cost_mixin import CostMixin, prepare_strict_schema
 from ..driver import Driver
 
 
 class OpenRouterDriver(CostMixin, Driver):
     supports_json_mode = True
+    supports_json_schema = True
+    supports_tool_use = True
+    supports_streaming = True
     supports_vision = True
 
     # Approximate pricing per 1K tokens based on OpenRouter's pricing
     # https://openrouter.ai/docs#pricing
     MODEL_PRICING = {
-        "openai/gpt-3.5-turbo": {
-            "prompt": 0.0015,
-            "completion": 0.002,
+        "openai/gpt-4o": {
+            "prompt": 0.005,
+            "completion": 0.015,
             "tokens_param": "max_tokens",
             "supports_temperature": True,
         },
-        "anthropic/claude-2": {
-            "prompt": 0.008,
-            "completion": 0.024,
+        "openai/gpt-4o-mini": {
+            "prompt": 0.00015,
+            "completion": 0.0006,
             "tokens_param": "max_tokens",
             "supports_temperature": True,
         },
-        "google/palm-2-chat-bison": {
-            "prompt": 0.0005,
-            "completion": 0.0005,
+        "anthropic/claude-sonnet-4-20250514": {
+            "prompt": 0.003,
+            "completion": 0.015,
             "tokens_param": "max_tokens",
             "supports_temperature": True,
         },
-        "meta-llama/llama-2-70b-chat": {
-            "prompt": 0.0007,
-            "completion": 0.0007,
+        "google/gemini-2.0-flash-001": {
+            "prompt": 0.0001,
+            "completion": 0.0004,
+            "tokens_param": "max_tokens",
+            "supports_temperature": True,
+        },
+        "meta-llama/llama-3.1-70b-instruct": {
+            "prompt": 0.0004,
+            "completion": 0.0004,
             "tokens_param": "max_tokens",
             "supports_temperature": True,
         },
     }
 
-    def __init__(self, api_key: str | None = None, model: str = "openai/gpt-3.5-turbo"):
+    def __init__(self, api_key: str | None = None, model: str = "openai/gpt-4o-mini"):
         """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
+            model: Model to use. Defaults to openai/gpt-4o-mini
         """
         self.api_key = api_key or os.getenv("OPENROUTER_API_KEY")
         if not self.api_key:
@@ -85,10 +97,17 @@ class OpenRouterDriver(CostMixin, Driver):
 
         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")
-        supports_temperature = model_info.get("supports_temperature", True)
+        # Lookup model-specific config (live models.dev data + hardcoded fallback)
+        model_config = self._get_model_config("openrouter", model)
+        tokens_param = model_config["tokens_param"]
+        supports_temperature = model_config["supports_temperature"]
+
+        # Validate capabilities against models.dev metadata
+        self._validate_model_capabilities(
+            "openrouter",
+            model,
+            using_json_schema=bool(options.get("json_schema")),
+        )
 
         # Defaults
         opts = {"temperature": 1.0, "max_tokens": 512, **options}
@@ -108,45 +127,230 @@ class OpenRouterDriver(CostMixin, Driver):
 
         # Native JSON mode support
         if options.get("json_mode"):
-            data["response_format"] = {"type": "json_object"}
+            json_schema = options.get("json_schema")
+            if json_schema:
+                schema_copy = prepare_strict_schema(json_schema)
+                data["response_format"] = {
+                    "type": "json_schema",
+                    "json_schema": {
+                        "name": "extraction",
+                        "strict": True,
+                        "schema": schema_copy,
+                    },
+                }
+            else:
+                data["response_format"] = {"type": "json_object"}
+
+        try:
+            response = requests.post(
+                f"{self.base_url}/chat/completions",
+                headers=self.headers,
+                json=data,
+                timeout=120,
+            )
+            response.raise_for_status()
+            resp = response.json()
+        except requests.exceptions.HTTPError as e:
+            body = ""
+            if e.response is not None:
+                with contextlib.suppress(Exception):
+                    body = e.response.text
+            error_msg = f"OpenRouter API request failed: {e!s}"
+            if body:
+                error_msg += f"\nResponse: {body}"
+            raise RuntimeError(error_msg) from e
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"OpenRouter 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)
+        total_tokens = usage.get("total_tokens", 0)
+
+        # Calculate cost via shared mixin
+        total_cost = self._calculate_cost("openrouter", model, prompt_tokens, completion_tokens)
+
+        # Standardized meta object
+        meta = {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": total_tokens,
+            "cost": round(total_cost, 6),
+            "raw_response": resp,
+            "model_name": model,
+        }
+
+        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 not self.api_key:
+            raise RuntimeError("OpenRouter API key not found")
+
+        model = options.get("model", self.model)
+        model_config = self._get_model_config("openrouter", model)
+        tokens_param = model_config["tokens_param"]
+        supports_temperature = model_config["supports_temperature"]
+
+        self._validate_model_capabilities("openrouter", model, using_tool_use=True)
+
+        opts = {"temperature": 1.0, "max_tokens": 512, **options}
+
+        data: dict[str, Any] = {
+            "model": model,
+            "messages": messages,
+            "tools": tools,
+        }
+        data[tokens_param] = opts.get("max_tokens", 512)
+
+        if supports_temperature and "temperature" in opts:
+            data["temperature"] = opts["temperature"]
 
         try:
             response = requests.post(
                 f"{self.base_url}/chat/completions",
                 headers=self.headers,
                 json=data,
+                timeout=120,
             )
             response.raise_for_status()
             resp = response.json()
+        except requests.exceptions.HTTPError as e:
+            error_msg = f"OpenRouter API request failed: {e!s}"
+            raise RuntimeError(error_msg) from e
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"OpenRouter 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("openrouter", model, prompt_tokens, completion_tokens)
 
-            # Extract usage info
-            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)
+        meta = {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": total_tokens,
+            "cost": round(total_cost, 6),
+            "raw_response": resp,
+            "model_name": model,
+        }
+
+        choice = resp["choices"][0]
+        text = choice["message"].get("content") or ""
+        stop_reason = choice.get("finish_reason")
+
+        tool_calls_out: list[dict[str, Any]] = []
+        for tc in choice["message"].get("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,
+        }
 
-            # Calculate cost via shared mixin
-            total_cost = self._calculate_cost("openrouter", model, prompt_tokens, completion_tokens)
+    # ------------------------------------------------------------------
+    # Streaming
+    # ------------------------------------------------------------------
 
-            # Standardized meta object
-            meta = {
+    def generate_messages_stream(
+        self,
+        messages: list[dict[str, Any]],
+        options: dict[str, Any],
+    ) -> Iterator[dict[str, Any]]:
+        """Yield response chunks via OpenRouter streaming API."""
+        if not self.api_key:
+            raise RuntimeError("OpenRouter API key not found")
+
+        model = options.get("model", self.model)
+        model_config = self._get_model_config("openrouter", model)
+        tokens_param = model_config["tokens_param"]
+        supports_temperature = model_config["supports_temperature"]
+
+        opts = {"temperature": 1.0, "max_tokens": 512, **options}
+
+        data: dict[str, Any] = {
+            "model": model,
+            "messages": messages,
+            "stream": True,
+            "stream_options": {"include_usage": True},
+        }
+        data[tokens_param] = opts.get("max_tokens", 512)
+
+        if supports_temperature and "temperature" in opts:
+            data["temperature"] = opts["temperature"]
+
+        response = requests.post(
+            f"{self.base_url}/chat/completions",
+            headers=self.headers,
+            json=data,
+            stream=True,
+            timeout=120,
+        )
+        response.raise_for_status()
+
+        full_text = ""
+        prompt_tokens = 0
+        completion_tokens = 0
+
+        for line in response.iter_lines(decode_unicode=True):
+            if not line or not line.startswith("data: "):
+                continue
+            payload = line[len("data: "):]
+            if payload.strip() == "[DONE]":
+                break
+            try:
+                chunk = json.loads(payload)
+            except json.JSONDecodeError:
+                continue
+
+            # Usage comes in the final chunk
+            usage = chunk.get("usage")
+            if usage:
+                prompt_tokens = usage.get("prompt_tokens", 0)
+                completion_tokens = usage.get("completion_tokens", 0)
+
+            choices = chunk.get("choices", [])
+            if choices:
+                delta = choices[0].get("delta", {})
+                content = delta.get("content", "")
+                if content:
+                    full_text += content
+                    yield {"type": "delta", "text": content}
+
+        total_tokens = prompt_tokens + completion_tokens
+        total_cost = self._calculate_cost("openrouter", 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": resp,
+                "raw_response": {},
                 "model_name": model,
-            }
-
-            text = resp["choices"][0]["message"]["content"]
-            return {"text": text, "meta": meta}
-
-        except requests.exceptions.RequestException as e:
-            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
+            },
+        }