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

prompture/drivers/azure_driver.py

@@ -1,17 +1,24 @@
 """Driver for Azure OpenAI Service (migrated to openai>=1.0.0).
 Requires the `openai` package.
 """
+
 import os
-from typing import Any, Dict
+from typing import Any
+
 try:
     from openai import AzureOpenAI
 except Exception:
     AzureOpenAI = None
 
+from ..cost_mixin import CostMixin
 from ..driver import Driver
 
 
-class AzureDriver(Driver):
+class AzureDriver(CostMixin, Driver):
+    supports_json_mode = True
+    supports_json_schema = True
+    supports_vision = True
+
     # Pricing per 1K tokens (adjust if your Azure pricing differs from OpenAI defaults)
     MODEL_PRICING = {
         "gpt-5-mini": {
@@ -82,7 +89,21 @@ class AzureDriver(Driver):
         else:
             self.client = None
 
-    def generate(self, prompt: str, options: Dict[str, Any]) -> Dict[str, Any]:
+    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)
+
+    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(self._prepare_messages(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("openai package (>=1.0.0) with AzureOpenAI not installed")
 
@@ -96,13 +117,28 @@ class AzureDriver(Driver):
         # Build request kwargs
         kwargs = {
             "model": self.deployment_id,  # for Azure, use deployment name
-            "messages": [{"role": "user", "content": prompt}],
+            "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 = self.client.chat.completions.create(**kwargs)
 
         # Extract usage
@@ -111,11 +147,8 @@ class AzureDriver(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("azure", model, prompt_tokens, completion_tokens)
 
         # Standardized meta object
         meta = {

prompture/drivers/claude_driver.py

@@ -1,75 +1,137 @@
 """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 collections.abc import Iterator
+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
+    supports_tool_use = True
+    supports_streaming = True
+    supports_vision = 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 _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)
+
+    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(self._prepare_messages(messages), options)
+
+    def _do_generate(self, messages: list[dict[str, Any]], 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 based on model pricing
-        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,
@@ -77,8 +139,169 @@ 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,
+        }
+
+        return {"text": text, "meta": meta}
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _extract_system_and_messages(
+        self, messages: list[dict[str, Any]]
+    ) -> tuple[str | None, list[dict[str, Any]]]:
+        """Separate system message from conversation messages for Anthropic API."""
+        system_content = None
+        api_messages: list[dict[str, Any]] = []
+        for msg in messages:
+            if msg.get("role") == "system":
+                system_content = msg.get("content", "")
+            else:
+                api_messages.append(msg)
+        return system_content, api_messages
+
+    # ------------------------------------------------------------------
+    # 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 (Anthropic)."""
+        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)
+
+        system_content, api_messages = self._extract_system_and_messages(messages)
+
+        # Convert tools from OpenAI format to Anthropic format if needed
+        anthropic_tools = []
+        for t in tools:
+            if "type" in t and t["type"] == "function":
+                # OpenAI format -> Anthropic format
+                fn = t["function"]
+                anthropic_tools.append({
+                    "name": fn["name"],
+                    "description": fn.get("description", ""),
+                    "input_schema": fn.get("parameters", {"type": "object", "properties": {}}),
+                })
+            elif "input_schema" in t:
+                # Already Anthropic format
+                anthropic_tools.append(t)
+            else:
+                anthropic_tools.append(t)
+
+        kwargs: dict[str, Any] = {
+            "model": model,
+            "messages": api_messages,
+            "temperature": opts["temperature"],
+            "max_tokens": opts["max_tokens"],
+            "tools": anthropic_tools,
+        }
+        if system_content:
+            kwargs["system"] = system_content
+
+        resp = client.messages.create(**kwargs)
+
+        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": round(total_cost, 6),
+            "raw_response": dict(resp),
+            "model_name": model,
+        }
+
+        text = ""
+        tool_calls_out: list[dict[str, Any]] = []
+        for block in resp.content:
+            if block.type == "text":
+                text += block.text
+            elif block.type == "tool_use":
+                tool_calls_out.append({
+                    "id": block.id,
+                    "name": block.name,
+                    "arguments": block.input,
+                })
+
+        return {
+            "text": text,
+            "meta": meta,
+            "tool_calls": tool_calls_out,
+            "stop_reason": resp.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 Anthropic streaming API."""
+        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)
+
+        system_content, api_messages = self._extract_system_and_messages(messages)
+
+        kwargs: dict[str, Any] = {
+            "model": model,
+            "messages": api_messages,
+            "temperature": opts["temperature"],
+            "max_tokens": opts["max_tokens"],
+        }
+        if system_content:
+            kwargs["system"] = system_content
+
+        full_text = ""
+        prompt_tokens = 0
+        completion_tokens = 0
+
+        with client.messages.stream(**kwargs) as stream:
+            for event in stream:
+                if hasattr(event, "type"):
+                    if event.type == "content_block_delta" and hasattr(event, "delta"):
+                        delta_text = getattr(event.delta, "text", "")
+                        if delta_text:
+                            full_text += delta_text
+                            yield {"type": "delta", "text": delta_text}
+                    elif event.type == "message_delta" and hasattr(event, "usage"):
+                        completion_tokens = getattr(event.usage, "output_tokens", 0)
+                    elif event.type == "message_start" and hasattr(event, "message"):
+                        usage = getattr(event.message, "usage", None)
+                        if usage:
+                            prompt_tokens = getattr(usage, "input_tokens", 0)
+
+        total_tokens = prompt_tokens + completion_tokens
+        total_cost = self._calculate_cost("claude", 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,
+            },
         }
-        
-        text = resp.content[0].text
-        return {"text": text, "meta": meta}