prompture 0.0.36.dev1__py3-none-any.whl → 0.0.37.dev1__py3-none-any.whl

prompture/drivers/async_google_driver.py

@@ -20,6 +20,7 @@ class AsyncGoogleDriver(CostMixin, AsyncDriver):
 
     supports_json_mode = True
     supports_json_schema = True
+    supports_vision = True
 
     MODEL_PRICING = GoogleDriver.MODEL_PRICING
     _PRICING_UNIT = 1_000_000
@@ -50,12 +51,17 @@ class AsyncGoogleDriver(CostMixin, AsyncDriver):
 
     supports_messages = True
 
+    def _prepare_messages(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        from .vision_helpers import _prepare_google_vision_messages
+
+        return _prepare_google_vision_messages(messages)
+
     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)
+        return await self._do_generate(self._prepare_messages(messages), options)
 
     async def _do_generate(
         self, messages: list[dict[str, str]], options: dict[str, Any] | None = None
@@ -90,10 +96,14 @@ class AsyncGoogleDriver(CostMixin, AsyncDriver):
             role = msg.get("role", "user")
             content = msg.get("content", "")
             if role == "system":
-                system_instruction = content
+                system_instruction = content if isinstance(content, str) else str(content)
             else:
                 gemini_role = "model" if role == "assistant" else "user"
-                contents.append({"role": gemini_role, "parts": [content]})
+                if msg.get("_vision_parts"):
+                    # Already converted to Gemini parts by _prepare_messages
+                    contents.append({"role": gemini_role, "parts": content})
+                else:
+                    contents.append({"role": gemini_role, "parts": [content]})
 
         try:
             model_kwargs: dict[str, Any] = {}
@@ -111,7 +121,17 @@ class AsyncGoogleDriver(CostMixin, AsyncDriver):
             if not response.text:
                 raise ValueError("Empty response from model")
 
-            total_prompt_chars = sum(len(msg.get("content", "")) for msg in messages)
+            total_prompt_chars = 0
+            for msg in messages:
+                c = msg.get("content", "")
+                if isinstance(c, str):
+                    total_prompt_chars += len(c)
+                elif isinstance(c, list):
+                    for part in c:
+                        if isinstance(part, str):
+                            total_prompt_chars += len(part)
+                        elif isinstance(part, dict) and "text" in part:
+                            total_prompt_chars += len(part["text"])
             completion_chars = len(response.text)
 
             total_cost = self._calculate_cost_chars(total_prompt_chars, completion_chars)

prompture/drivers/async_grok_driver.py

@@ -14,6 +14,7 @@ from .grok_driver import GrokDriver
 
 class AsyncGrokDriver(CostMixin, AsyncDriver):
     supports_json_mode = True
+    supports_vision = True
 
     MODEL_PRICING = GrokDriver.MODEL_PRICING
     _PRICING_UNIT = 1_000_000
@@ -25,12 +26,17 @@ class AsyncGrokDriver(CostMixin, AsyncDriver):
 
     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)
+
     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)
+        return await self._do_generate(self._prepare_messages(messages), options)
 
     async def _do_generate(self, messages: list[dict[str, str]], options: dict[str, Any]) -> dict[str, Any]:
         if not self.api_key:

prompture/drivers/async_groq_driver.py

@@ -17,6 +17,7 @@ from .groq_driver import GroqDriver
 
 class AsyncGroqDriver(CostMixin, AsyncDriver):
     supports_json_mode = True
+    supports_vision = True
 
     MODEL_PRICING = GroqDriver.MODEL_PRICING
 
@@ -30,12 +31,17 @@ class AsyncGroqDriver(CostMixin, AsyncDriver):
 
     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)
+
     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)
+        return await self._do_generate(self._prepare_messages(messages), options)
 
     async def _do_generate(self, messages: list[dict[str, str]], options: dict[str, Any]) -> dict[str, Any]:
         if self.client is None:

prompture/drivers/async_lmstudio_driver.py

@@ -15,22 +15,47 @@ logger = logging.getLogger(__name__)
 
 class AsyncLMStudioDriver(AsyncDriver):
     supports_json_mode = True
+    supports_vision = True
 
     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,
+    ):
         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()
+
     supports_messages = True
 
+    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 _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)
+
     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)
+        return await self._do_generate(self._prepare_messages(messages), options)
 
     async def _do_generate(
         self, messages: list[dict[str, str]], options: dict[str, Any] | None = None
@@ -51,7 +76,7 @@ class AsyncLMStudioDriver(AsyncDriver):
 
         async with httpx.AsyncClient() as client:
             try:
-                r = await client.post(self.endpoint, json=payload, timeout=120)
+                r = await client.post(self.endpoint, json=payload, headers=self._headers, timeout=120)
                 r.raise_for_status()
                 response_data = r.json()
             except Exception as e:
@@ -77,3 +102,34 @@ class AsyncLMStudioDriver(AsyncDriver):
         }
 
         return {"text": text, "meta": meta}
+
+    # -- Model management (LM Studio 0.4.0+) ----------------------------------
+
+    async 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"
+        async with httpx.AsyncClient() as client:
+            r = await client.get(url, headers=self._headers, timeout=10)
+            r.raise_for_status()
+            data = r.json()
+        return data.get("data", [])
+
+    async 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
+        async with httpx.AsyncClient() as client:
+            r = await client.post(url, json=payload, headers=self._headers, timeout=120)
+            r.raise_for_status()
+        return r.json()
+
+    async 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 = {"model": model}
+        async with httpx.AsyncClient() as client:
+            r = await client.post(url, json=payload, headers=self._headers, timeout=30)
+            r.raise_for_status()
+        return r.json()

prompture/drivers/async_ollama_driver.py

@@ -15,6 +15,7 @@ logger = logging.getLogger(__name__)
 
 class AsyncOllamaDriver(AsyncDriver):
     supports_json_mode = True
+    supports_vision = True
 
     MODEL_PRICING = {"default": {"prompt": 0.0, "completion": 0.0}}
 
@@ -25,6 +26,11 @@ class AsyncOllamaDriver(AsyncDriver):
 
     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)
+
     async def generate(self, prompt: str, options: dict[str, Any] | None = None) -> dict[str, Any]:
         merged_options = self.options.copy()
         if options:
@@ -74,6 +80,7 @@ class AsyncOllamaDriver(AsyncDriver):
 
     async def generate_messages(self, messages: list[dict[str, str]], 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)

prompture/drivers/async_openai_driver.py

@@ -18,6 +18,7 @@ from .openai_driver import OpenAIDriver
 class AsyncOpenAIDriver(CostMixin, AsyncDriver):
     supports_json_mode = True
     supports_json_schema = True
+    supports_vision = True
 
     MODEL_PRICING = OpenAIDriver.MODEL_PRICING
 
@@ -31,12 +32,17 @@ class AsyncOpenAIDriver(CostMixin, AsyncDriver):
 
     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)
+
     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)
+        return await self._do_generate(self._prepare_messages(messages), options)
 
     async def _do_generate(self, messages: list[dict[str, str]], options: dict[str, Any]) -> dict[str, Any]:
         if self.client is None:

prompture/drivers/async_openrouter_driver.py

@@ -14,6 +14,7 @@ from .openrouter_driver import OpenRouterDriver
 
 class AsyncOpenRouterDriver(CostMixin, AsyncDriver):
     supports_json_mode = True
+    supports_vision = True
 
     MODEL_PRICING = OpenRouterDriver.MODEL_PRICING
 
@@ -31,12 +32,17 @@ class AsyncOpenRouterDriver(CostMixin, AsyncDriver):
 
     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)
+
     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)
+        return await self._do_generate(self._prepare_messages(messages), options)
 
     async def _do_generate(self, messages: list[dict[str, str]], options: dict[str, Any]) -> dict[str, Any]:
         model = options.get("model", self.model)

prompture/drivers/async_registry.py

@@ -49,7 +49,11 @@ register_async_driver(
 )
 register_async_driver(
     "lmstudio",
-    lambda model=None: AsyncLMStudioDriver(endpoint=settings.lmstudio_endpoint, model=model or settings.lmstudio_model),
+    lambda model=None: AsyncLMStudioDriver(
+        endpoint=settings.lmstudio_endpoint,
+        model=model or settings.lmstudio_model,
+        api_key=settings.lmstudio_api_key,
+    ),
     overwrite=True,
 )
 register_async_driver(

prompture/drivers/azure_driver.py

@@ -17,6 +17,7 @@ from ..driver import 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 = {
@@ -90,12 +91,17 @@ class AzureDriver(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:

prompture/drivers/claude_driver.py

@@ -21,6 +21,7 @@ class ClaudeDriver(CostMixin, Driver):
     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 = {
@@ -57,12 +58,17 @@ class ClaudeDriver(CostMixin, Driver):
 
     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(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 anthropic is None:

prompture/drivers/google_driver.py

@@ -15,6 +15,7 @@ class GoogleDriver(CostMixin, Driver):
 
     supports_json_mode = True
     supports_json_schema = True
+    supports_vision = True
 
     # Based on current Gemini pricing (as of 2025)
     # Source: https://cloud.google.com/vertex-ai/pricing#gemini_models
@@ -107,12 +108,17 @@ class GoogleDriver(CostMixin, Driver):
 
     supports_messages = True
 
+    def _prepare_messages(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        from .vision_helpers import _prepare_google_vision_messages
+
+        return _prepare_google_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()
@@ -147,11 +153,15 @@ class GoogleDriver(CostMixin, Driver):
             role = msg.get("role", "user")
             content = msg.get("content", "")
             if role == "system":
-                system_instruction = content
+                system_instruction = content if isinstance(content, str) else str(content)
             else:
                 # Gemini uses "model" for assistant role
                 gemini_role = "model" if role == "assistant" else "user"
-                contents.append({"role": gemini_role, "parts": [content]})
+                if msg.get("_vision_parts"):
+                    # Already converted to Gemini parts by _prepare_messages
+                    contents.append({"role": gemini_role, "parts": content})
+                else:
+                    contents.append({"role": gemini_role, "parts": [content]})
 
         try:
             logger.debug(f"Initializing {self.model} for generation")
@@ -174,7 +184,17 @@ class GoogleDriver(CostMixin, Driver):
                 raise ValueError("Empty response from model")
 
             # Calculate token usage and cost
-            total_prompt_chars = sum(len(msg.get("content", "")) for msg in messages)
+            total_prompt_chars = 0
+            for msg in messages:
+                c = msg.get("content", "")
+                if isinstance(c, str):
+                    total_prompt_chars += len(c)
+                elif isinstance(c, list):
+                    for part in c:
+                        if isinstance(part, str):
+                            total_prompt_chars += len(part)
+                        elif isinstance(part, dict) and "text" in part:
+                            total_prompt_chars += len(part["text"])
             completion_chars = len(response.text)
 
             # Google uses character-based cost estimation

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:

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:

prompture/drivers/lmstudio_driver.py

@@ -12,27 +12,46 @@ logger = logging.getLogger(__name__)
 
 class LMStudioDriver(Driver):
     supports_json_mode = 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 +59,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()
@@ -66,7 +90,7 @@ class LMStudioDriver(Driver):
             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 +128,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 = {"model": model}
+        r = requests.post(url, json=payload, headers=self._headers, timeout=30)
+        r.raise_for_status()
+        return r.json()

prompture/drivers/ollama_driver.py

@@ -14,6 +14,7 @@ logger = logging.getLogger(__name__)
 class OllamaDriver(Driver):
     supports_json_mode = 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 +47,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()
@@ -190,6 +196,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)

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: