prompture 0.0.35__py3-none-any.whl → 0.0.40.dev1__py3-none-any.whl

prompture/drivers/grok_driver.py

@@ -13,6 +13,7 @@ from ..driver import Driver
 
 class GrokDriver(CostMixin, Driver):
     supports_json_mode = True
+    supports_vision = True
 
     # Pricing per 1M tokens based on xAI's documentation
     _PRICING_UNIT = 1_000_000
@@ -80,12 +81,17 @@ class GrokDriver(CostMixin, Driver):
 
     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(messages, options)
+        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 not self.api_key:
@@ -93,10 +99,10 @@ class GrokDriver(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("grok", model)
+        tokens_param = model_config["tokens_param"]
+        supports_temperature = model_config["supports_temperature"]
 
         # Defaults
         opts = {"temperature": 1.0, "max_tokens": 512, **options}

prompture/drivers/groq_driver.py

@@ -16,6 +16,7 @@ from ..driver import Driver
 
 class GroqDriver(CostMixin, Driver):
     supports_json_mode = True
+    supports_vision = True
 
     # Approximate pricing per 1K tokens (to be updated with official pricing)
     # Each model entry defines token parameters and temperature support
@@ -50,12 +51,17 @@ class GroqDriver(CostMixin, Driver):
 
     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(messages, options)
+        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:
@@ -63,10 +69,10 @@ class GroqDriver(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("groq", model)
+        tokens_param = model_config["tokens_param"]
+        supports_temperature = model_config["supports_temperature"]
 
         # Base configuration
         opts = {"temperature": 0.7, "max_tokens": 512, **options}

prompture/drivers/lmstudio_driver.py

@@ -12,27 +12,47 @@ logger = logging.getLogger(__name__)
 
 class LMStudioDriver(Driver):
     supports_json_mode = True
+    supports_json_schema = True
+    supports_vision = True
 
     # LM Studio is local – costs are always zero.
     MODEL_PRICING = {"default": {"prompt": 0.0, "completion": 0.0}}
 
-    def __init__(self, endpoint: str | None = None, model: str = "deepseek/deepseek-r1-0528-qwen3-8b"):
+    def __init__(
+        self,
+        endpoint: str | None = None,
+        model: str = "deepseek/deepseek-r1-0528-qwen3-8b",
+        api_key: str | None = None,
+    ):
         # Allow override via env var
         self.endpoint = endpoint or os.getenv("LMSTUDIO_ENDPOINT", "http://127.0.0.1:1234/v1/chat/completions")
         self.model = model
         self.options: dict[str, Any] = {}
 
+        # Derive base_url once for reuse across management endpoints
+        self.base_url = self.endpoint.split("/v1/")[0]
+
+        # API key for LM Studio 0.4.0+ authentication
+        self.api_key = api_key or os.getenv("LMSTUDIO_API_KEY")
+        self._headers = self._build_headers()
+
         # Validate connection to LM Studio server
         self._validate_connection()
 
+    def _build_headers(self) -> dict[str, str]:
+        """Build request headers, including auth if an API key is configured."""
+        headers: dict[str, str] = {"Content-Type": "application/json"}
+        if self.api_key:
+            headers["Authorization"] = f"Bearer {self.api_key}"
+        return headers
+
     def _validate_connection(self):
         """Validate connection to the LM Studio server."""
         try:
-            base_url = self.endpoint.split("/v1/")[0]
-            health_url = f"{base_url}/v1/models"
+            health_url = f"{self.base_url}/v1/models"
 
             logger.debug(f"Validating connection to LM Studio server at: {health_url}")
-            response = requests.get(health_url, timeout=5)
+            response = requests.get(health_url, headers=self._headers, timeout=5)
             response.raise_for_status()
             logger.debug("Connection to LM Studio server validated successfully")
         except requests.exceptions.RequestException as e:
@@ -40,12 +60,17 @@ class LMStudioDriver(Driver):
 
     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: 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)
+        return self._do_generate(self._prepare_messages(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()
@@ -58,15 +83,27 @@ class LMStudioDriver(Driver):
             "temperature": merged_options.get("temperature", 0.7),
         }
 
-        # Native JSON mode support
+        # Native JSON mode support (LM Studio requires json_schema, not json_object)
         if merged_options.get("json_mode"):
-            payload["response_format"] = {"type": "json_object"}
+            json_schema = merged_options.get("json_schema")
+            if json_schema:
+                payload["response_format"] = {
+                    "type": "json_schema",
+                    "json_schema": {
+                        "name": "extraction",
+                        "schema": json_schema,
+                    },
+                }
+            else:
+                # No schema provided — omit response_format entirely;
+                # LM Studio rejects "json_object" type.
+                pass
 
         try:
             logger.debug(f"Sending request to LM Studio endpoint: {self.endpoint}")
             logger.debug(f"Request payload: {payload}")
 
-            r = requests.post(self.endpoint, json=payload, timeout=120)
+            r = requests.post(self.endpoint, json=payload, headers=self._headers, timeout=120)
             r.raise_for_status()
 
             response_data = r.json()
@@ -104,3 +141,31 @@ class LMStudioDriver(Driver):
         }
 
         return {"text": text, "meta": meta}
+
+    # -- Model management (LM Studio 0.4.0+) ----------------------------------
+
+    def list_models(self) -> list[dict[str, Any]]:
+        """List currently loaded models via GET /v1/models (OpenAI-compatible)."""
+        url = f"{self.base_url}/v1/models"
+        r = requests.get(url, headers=self._headers, timeout=10)
+        r.raise_for_status()
+        data = r.json()
+        return data.get("data", [])
+
+    def load_model(self, model: str, context_length: int | None = None) -> dict[str, Any]:
+        """Load a model into LM Studio via POST /api/v1/models/load."""
+        url = f"{self.base_url}/api/v1/models/load"
+        payload: dict[str, Any] = {"model": model}
+        if context_length is not None:
+            payload["context_length"] = context_length
+        r = requests.post(url, json=payload, headers=self._headers, timeout=120)
+        r.raise_for_status()
+        return r.json()
+
+    def unload_model(self, model: str) -> dict[str, Any]:
+        """Unload a model from LM Studio via POST /api/v1/models/unload."""
+        url = f"{self.base_url}/api/v1/models/unload"
+        payload = {"instance_id": model}
+        r = requests.post(url, json=payload, headers=self._headers, timeout=30)
+        r.raise_for_status()
+        return r.json()

prompture/drivers/ollama_driver.py

@@ -13,7 +13,9 @@ logger = logging.getLogger(__name__)
 
 class OllamaDriver(Driver):
     supports_json_mode = True
+    supports_json_schema = True
     supports_streaming = True
+    supports_vision = True
 
     # Ollama is free – costs are always zero.
     MODEL_PRICING = {"default": {"prompt": 0.0, "completion": 0.0}}
@@ -46,6 +48,11 @@ class OllamaDriver(Driver):
 
     supports_messages = True
 
+    def _prepare_messages(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        from .vision_helpers import _prepare_ollama_vision_messages
+
+        return _prepare_ollama_vision_messages(messages)
+
     def generate(self, prompt: str, options: Optional[dict[str, Any]] = None) -> dict[str, Any]:
         # Merge instance options with call-specific options
         merged_options = self.options.copy()
@@ -58,9 +65,10 @@ class OllamaDriver(Driver):
             "stream": False,
         }
 
-        # Native JSON mode support
+        # Native JSON mode / structured output support
         if merged_options.get("json_mode"):
-            payload["format"] = "json"
+            json_schema = merged_options.get("json_schema")
+            payload["format"] = json_schema if json_schema else "json"
 
         # Add any Ollama-specific options from merged_options
         if "temperature" in merged_options:
@@ -146,7 +154,8 @@ class OllamaDriver(Driver):
         }
 
         if merged_options.get("json_mode"):
-            payload["format"] = "json"
+            json_schema = merged_options.get("json_schema")
+            payload["format"] = json_schema if json_schema else "json"
         if "temperature" in merged_options:
             payload["temperature"] = merged_options["temperature"]
         if "top_p" in merged_options:
@@ -190,6 +199,7 @@ class OllamaDriver(Driver):
 
     def generate_messages(self, messages: list[dict[str, Any]], options: dict[str, Any]) -> dict[str, Any]:
         """Use Ollama's /api/chat endpoint for multi-turn conversations."""
+        messages = self._prepare_messages(messages)
         merged_options = self.options.copy()
         if options:
             merged_options.update(options)
@@ -203,9 +213,10 @@ class OllamaDriver(Driver):
             "stream": False,
         }
 
-        # Native JSON mode support
+        # Native JSON mode / structured output support
         if merged_options.get("json_mode"):
-            payload["format"] = "json"
+            json_schema = merged_options.get("json_schema")
+            payload["format"] = json_schema if json_schema else "json"
 
         if "temperature" in merged_options:
             payload["temperature"] = merged_options["temperature"]

prompture/drivers/openai_driver.py

@@ -21,6 +21,7 @@ class OpenAIDriver(CostMixin, Driver):
     supports_json_schema = True
     supports_tool_use = True
     supports_streaming = True
+    supports_vision = True
 
     # Approximate pricing per 1K tokens (keep updated with OpenAI's official pricing)
     # Each model entry also defines which token parameter it supports and
@@ -74,12 +75,17 @@ class OpenAIDriver(CostMixin, Driver):
 
     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, Any]], options: dict[str, Any]) -> dict[str, Any]:
-        return self._do_generate(messages, options)
+        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 self.client is None:
@@ -87,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}
@@ -162,9 +175,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}
 
@@ -233,9 +248,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

@@ -13,6 +13,7 @@ from ..driver import Driver
 
 class OpenRouterDriver(CostMixin, Driver):
     supports_json_mode = True
+    supports_vision = True
 
     # Approximate pricing per 1K tokens based on OpenRouter's pricing
     # https://openrouter.ai/docs#pricing
@@ -66,12 +67,17 @@ class OpenRouterDriver(CostMixin, Driver):
 
     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(messages, options)
+        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 not self.api_key:
@@ -79,10 +85,10 @@ 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"]
 
         # Defaults
         opts = {"temperature": 1.0, "max_tokens": 512, **options}

prompture/drivers/vision_helpers.py

@@ -0,0 +1,153 @@
+"""Shared helpers for converting universal vision message blocks to provider-specific formats."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def _prepare_openai_vision_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Convert universal image blocks to OpenAI-compatible vision format.
+
+    Works for OpenAI, Azure, Groq, Grok, LM Studio, and OpenRouter.
+
+    Universal format::
+
+        {"type": "image", "source": ImageContent(...)}
+
+    OpenAI format::
+
+        {"type": "image_url", "image_url": {"url": "data:mime;base64,..."}}
+    """
+    out: list[dict[str, Any]] = []
+    for msg in messages:
+        content = msg.get("content")
+        if not isinstance(content, list):
+            out.append(msg)
+            continue
+        new_blocks: list[dict[str, Any]] = []
+        for block in content:
+            if isinstance(block, dict) and block.get("type") == "image":
+                source = block["source"]
+                if source.source_type == "url" and source.url:
+                    url = source.url
+                else:
+                    url = f"data:{source.media_type};base64,{source.data}"
+                new_blocks.append(
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": url},
+                    }
+                )
+            else:
+                new_blocks.append(block)
+        out.append({**msg, "content": new_blocks})
+    return out
+
+
+def _prepare_claude_vision_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Convert universal image blocks to Anthropic Claude format.
+
+    Claude format::
+
+        {"type": "image", "source": {"type": "base64", "media_type": "...", "data": "..."}}
+    """
+    out: list[dict[str, Any]] = []
+    for msg in messages:
+        content = msg.get("content")
+        if not isinstance(content, list):
+            out.append(msg)
+            continue
+        new_blocks: list[dict[str, Any]] = []
+        for block in content:
+            if isinstance(block, dict) and block.get("type") == "image":
+                source = block["source"]
+                if source.source_type == "url" and source.url:
+                    new_blocks.append(
+                        {
+                            "type": "image",
+                            "source": {
+                                "type": "url",
+                                "url": source.url,
+                            },
+                        }
+                    )
+                else:
+                    new_blocks.append(
+                        {
+                            "type": "image",
+                            "source": {
+                                "type": "base64",
+                                "media_type": source.media_type,
+                                "data": source.data,
+                            },
+                        }
+                    )
+            else:
+                new_blocks.append(block)
+        out.append({**msg, "content": new_blocks})
+    return out
+
+
+def _prepare_google_vision_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Convert universal image blocks to Google Gemini format.
+
+    Gemini expects ``parts`` arrays containing text and inline_data dicts::
+
+        {"role": "user", "parts": [
+            "text prompt",
+            {"inline_data": {"mime_type": "image/png", "data": "base64..."}},
+        ]}
+    """
+    out: list[dict[str, Any]] = []
+    for msg in messages:
+        content = msg.get("content")
+        if not isinstance(content, list):
+            out.append(msg)
+            continue
+        # Convert content blocks to Gemini parts
+        parts: list[Any] = []
+        for block in content:
+            if isinstance(block, dict) and block.get("type") == "text":
+                parts.append(block["text"])
+            elif isinstance(block, dict) and block.get("type") == "image":
+                source = block["source"]
+                parts.append(
+                    {
+                        "inline_data": {
+                            "mime_type": source.media_type,
+                            "data": source.data,
+                        }
+                    }
+                )
+            else:
+                parts.append(block)
+        out.append({**msg, "content": parts, "_vision_parts": True})
+    return out
+
+
+def _prepare_ollama_vision_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Convert universal image blocks to Ollama format.
+
+    Ollama expects images as a separate field::
+
+        {"role": "user", "content": "text", "images": ["base64..."]}
+    """
+    out: list[dict[str, Any]] = []
+    for msg in messages:
+        content = msg.get("content")
+        if not isinstance(content, list):
+            out.append(msg)
+            continue
+        text_parts: list[str] = []
+        images: list[str] = []
+        for block in content:
+            if isinstance(block, dict) and block.get("type") == "text":
+                text_parts.append(block["text"])
+            elif isinstance(block, dict) and block.get("type") == "image":
+                source = block["source"]
+                images.append(source.data)
+        new_msg = {**msg, "content": " ".join(text_parts)}
+        if images:
+            new_msg["images"] = images
+        out.append(new_msg)
+    return out