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

prompture/drivers/async_lmstudio_driver.py

@@ -0,0 +1,148 @@
+"""Async LM Studio driver using httpx."""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any
+
+import httpx
+
+from ..async_driver import AsyncDriver
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncLMStudioDriver(AsyncDriver):
+    supports_json_mode = True
+    supports_json_schema = 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",
+        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(self._prepare_messages(messages), options)
+
+    async def _do_generate(
+        self, messages: list[dict[str, str]], options: dict[str, Any] | None = None
+    ) -> dict[str, Any]:
+        merged_options = self.options.copy()
+        if options:
+            merged_options.update(options)
+
+        payload = {
+            "model": merged_options.get("model", self.model),
+            "messages": messages,
+            "temperature": merged_options.get("temperature", 0.7),
+        }
+
+        # Native JSON mode support (LM Studio requires json_schema, not json_object)
+        if merged_options.get("json_mode"):
+            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
+
+        async with httpx.AsyncClient() as client:
+            try:
+                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:
+                raise RuntimeError(f"AsyncLMStudioDriver request failed: {e}") from e
+
+        if "choices" not in response_data or not response_data["choices"]:
+            raise ValueError(f"Unexpected response format: {response_data}")
+
+        text = response_data["choices"][0]["message"]["content"]
+
+        usage = response_data.get("usage", {})
+        prompt_tokens = usage.get("prompt_tokens", 0)
+        completion_tokens = usage.get("completion_tokens", 0)
+        total_tokens = usage.get("total_tokens", prompt_tokens + completion_tokens)
+
+        meta = {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": total_tokens,
+            "cost": 0.0,
+            "raw_response": response_data,
+            "model_name": merged_options.get("model", self.model),
+        }
+
+        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 = {"instance_id": 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_local_http_driver.py

@@ -0,0 +1,44 @@
+"""Async LocalHTTP driver using httpx."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import httpx
+
+from ..async_driver import AsyncDriver
+
+
+class AsyncLocalHTTPDriver(AsyncDriver):
+    MODEL_PRICING = {"default": {"prompt": 0.0, "completion": 0.0}}
+
+    def __init__(self, endpoint: str | None = None, model: str = "local-model"):
+        self.endpoint = endpoint or os.getenv("LOCAL_HTTP_ENDPOINT", "http://localhost:8000/generate")
+        self.model = model
+
+    async def generate(self, prompt: str, options: dict[str, Any]) -> dict[str, Any]:
+        payload = {"prompt": prompt, "options": options}
+
+        async with httpx.AsyncClient() as client:
+            try:
+                r = await client.post(self.endpoint, json=payload, timeout=options.get("timeout", 30))
+                r.raise_for_status()
+                response_data = r.json()
+            except Exception as e:
+                raise RuntimeError(f"AsyncLocalHTTPDriver request failed: {e}") from e
+
+        if "text" in response_data and "meta" in response_data:
+            return response_data
+
+        meta = {
+            "prompt_tokens": response_data.get("prompt_tokens", 0),
+            "completion_tokens": response_data.get("completion_tokens", 0),
+            "total_tokens": response_data.get("total_tokens", 0),
+            "cost": 0.0,
+            "raw_response": response_data,
+            "model_name": options.get("model", self.model),
+        }
+
+        text = response_data.get("text") or response_data.get("response") or str(response_data)
+        return {"text": text, "meta": meta}

prompture/drivers/async_ollama_driver.py

@@ -0,0 +1,135 @@
+"""Async Ollama driver using httpx."""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any
+
+import httpx
+
+from ..async_driver import AsyncDriver
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncOllamaDriver(AsyncDriver):
+    supports_json_mode = True
+    supports_json_schema = True
+    supports_vision = True
+
+    MODEL_PRICING = {"default": {"prompt": 0.0, "completion": 0.0}}
+
+    def __init__(self, endpoint: str | None = None, model: str = "llama3"):
+        self.endpoint = endpoint or os.getenv("OLLAMA_ENDPOINT", "http://localhost:11434/api/generate")
+        self.model = model
+        self.options: dict[str, Any] = {}
+
+    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:
+            merged_options.update(options)
+
+        payload = {
+            "prompt": prompt,
+            "model": merged_options.get("model", self.model),
+            "stream": False,
+        }
+
+        # Native JSON mode / structured output support
+        if merged_options.get("json_mode"):
+            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:
+            payload["top_p"] = merged_options["top_p"]
+        if "top_k" in merged_options:
+            payload["top_k"] = merged_options["top_k"]
+
+        async with httpx.AsyncClient() as client:
+            try:
+                r = await client.post(self.endpoint, json=payload, timeout=120)
+                r.raise_for_status()
+                response_data = r.json()
+            except httpx.HTTPStatusError as e:
+                raise RuntimeError(f"Ollama request failed: {e}") from e
+            except Exception as e:
+                raise RuntimeError(f"Ollama request failed: {e}") from e
+
+        prompt_tokens = response_data.get("prompt_eval_count", 0)
+        completion_tokens = response_data.get("eval_count", 0)
+        total_tokens = prompt_tokens + completion_tokens
+
+        meta = {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": total_tokens,
+            "cost": 0.0,
+            "raw_response": response_data,
+            "model_name": merged_options.get("model", self.model),
+        }
+
+        return {"text": response_data.get("response", ""), "meta": meta}
+
+    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)
+
+        # Derive the chat endpoint from the generate endpoint
+        chat_endpoint = self.endpoint.replace("/api/generate", "/api/chat")
+
+        payload: dict[str, Any] = {
+            "model": merged_options.get("model", self.model),
+            "messages": messages,
+            "stream": False,
+        }
+
+        if merged_options.get("json_mode"):
+            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:
+            payload["top_p"] = merged_options["top_p"]
+        if "top_k" in merged_options:
+            payload["top_k"] = merged_options["top_k"]
+
+        async with httpx.AsyncClient() as client:
+            try:
+                r = await client.post(chat_endpoint, json=payload, timeout=120)
+                r.raise_for_status()
+                response_data = r.json()
+            except httpx.HTTPStatusError as e:
+                raise RuntimeError(f"Ollama chat request failed: {e}") from e
+            except Exception as e:
+                raise RuntimeError(f"Ollama chat request failed: {e}") from e
+
+        prompt_tokens = response_data.get("prompt_eval_count", 0)
+        completion_tokens = response_data.get("eval_count", 0)
+        total_tokens = prompt_tokens + completion_tokens
+
+        meta = {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": total_tokens,
+            "cost": 0.0,
+            "raw_response": response_data,
+            "model_name": merged_options.get("model", self.model),
+        }
+
+        message = response_data.get("message", {})
+        text = message.get("content", "")
+        return {"text": text, "meta": meta}

prompture/drivers/async_openai_driver.py

@@ -0,0 +1,102 @@
+"""Async OpenAI driver. Requires the ``openai`` package (>=1.0.0)."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+try:
+    from openai import AsyncOpenAI
+except Exception:
+    AsyncOpenAI = None
+
+from ..async_driver import AsyncDriver
+from ..cost_mixin import CostMixin
+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
+
+    def __init__(self, api_key: str | None = None, model: str = "gpt-4o-mini"):
+        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
+        self.model = model
+        if AsyncOpenAI:
+            self.client = AsyncOpenAI(api_key=self.api_key)
+        else:
+            self.client = None
+
+    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(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:
+            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)
+
+        opts = {"temperature": 1.0, "max_tokens": 512, **options}
+
+        kwargs = {
+            "model": model,
+            "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 = await self.client.chat.completions.create(**kwargs)
+
+        usage = getattr(resp, "usage", None)
+        prompt_tokens = getattr(usage, "prompt_tokens", 0)
+        completion_tokens = getattr(usage, "completion_tokens", 0)
+        total_tokens = getattr(usage, "total_tokens", 0)
+
+        total_cost = self._calculate_cost("openai", model, prompt_tokens, completion_tokens)
+
+        meta = {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": total_tokens,
+            "cost": total_cost,
+            "raw_response": resp.model_dump(),
+            "model_name": model,
+        }
+
+        text = resp.choices[0].message.content
+        return {"text": text, "meta": meta}

prompture/drivers/async_openrouter_driver.py

@@ -0,0 +1,102 @@
+"""Async OpenRouter driver using httpx."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import httpx
+
+from ..async_driver import AsyncDriver
+from ..cost_mixin import CostMixin
+from .openrouter_driver import OpenRouterDriver
+
+
+class AsyncOpenRouterDriver(CostMixin, AsyncDriver):
+    supports_json_mode = True
+    supports_vision = True
+
+    MODEL_PRICING = OpenRouterDriver.MODEL_PRICING
+
+    def __init__(self, api_key: str | None = None, model: str = "openai/gpt-3.5-turbo"):
+        self.api_key = api_key or os.getenv("OPENROUTER_API_KEY")
+        if not self.api_key:
+            raise ValueError("OpenRouter API key not found. Set OPENROUTER_API_KEY env var.")
+        self.model = model
+        self.base_url = "https://openrouter.ai/api/v1"
+        self.headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "HTTP-Referer": "https://github.com/jhd3197/prompture",
+            "Content-Type": "application/json",
+        }
+
+    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(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)
+
+        model_info = self.MODEL_PRICING.get(model, {})
+        tokens_param = model_info.get("tokens_param", "max_tokens")
+        supports_temperature = model_info.get("supports_temperature", True)
+
+        opts = {"temperature": 1.0, "max_tokens": 512, **options}
+
+        data = {
+            "model": model,
+            "messages": messages,
+        }
+        data[tokens_param] = opts.get("max_tokens", 512)
+
+        if supports_temperature and "temperature" in opts:
+            data["temperature"] = opts["temperature"]
+
+        # Native JSON mode support
+        if options.get("json_mode"):
+            data["response_format"] = {"type": "json_object"}
+
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.post(
+                    f"{self.base_url}/chat/completions",
+                    headers=self.headers,
+                    json=data,
+                    timeout=120,
+                )
+                response.raise_for_status()
+                resp = response.json()
+            except httpx.HTTPStatusError as e:
+                error_msg = f"OpenRouter API request failed: {e!s}"
+                raise RuntimeError(error_msg) from e
+            except Exception 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)
+
+        meta = {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": total_tokens,
+            "cost": total_cost,
+            "raw_response": resp,
+            "model_name": model,
+        }
+
+        text = resp["choices"][0]["message"]["content"]
+        return {"text": text, "meta": meta}

prompture/drivers/async_registry.py

@@ -0,0 +1,133 @@
+"""Async driver registry — mirrors the sync DRIVER_REGISTRY.
+
+This module provides async driver registration and factory functions.
+Custom async drivers can be registered via the ``register_async_driver()``
+function or discovered via entry points.
+
+Entry Point Discovery:
+    Add to your pyproject.toml:
+
+    [project.entry-points."prompture.async_drivers"]
+    my_provider = "my_package.drivers:my_async_driver_factory"
+"""
+
+from __future__ import annotations
+
+from ..settings import settings
+from .async_airllm_driver import AsyncAirLLMDriver
+from .async_azure_driver import AsyncAzureDriver
+from .async_claude_driver import AsyncClaudeDriver
+from .async_google_driver import AsyncGoogleDriver
+from .async_grok_driver import AsyncGrokDriver
+from .async_groq_driver import AsyncGroqDriver
+from .async_lmstudio_driver import AsyncLMStudioDriver
+from .async_local_http_driver import AsyncLocalHTTPDriver
+from .async_ollama_driver import AsyncOllamaDriver
+from .async_openai_driver import AsyncOpenAIDriver
+from .async_openrouter_driver import AsyncOpenRouterDriver
+from .registry import (
+    _get_async_registry,
+    get_async_driver_factory,
+    register_async_driver,
+)
+
+# Register built-in async drivers
+register_async_driver(
+    "openai",
+    lambda model=None: AsyncOpenAIDriver(api_key=settings.openai_api_key, model=model or settings.openai_model),
+    overwrite=True,
+)
+register_async_driver(
+    "ollama",
+    lambda model=None: AsyncOllamaDriver(endpoint=settings.ollama_endpoint, model=model or settings.ollama_model),
+    overwrite=True,
+)
+register_async_driver(
+    "claude",
+    lambda model=None: AsyncClaudeDriver(api_key=settings.claude_api_key, model=model or settings.claude_model),
+    overwrite=True,
+)
+register_async_driver(
+    "lmstudio",
+    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(
+    "azure",
+    lambda model=None: AsyncAzureDriver(
+        api_key=settings.azure_api_key, endpoint=settings.azure_api_endpoint, deployment_id=settings.azure_deployment_id
+    ),
+    overwrite=True,
+)
+register_async_driver(
+    "local_http",
+    lambda model=None: AsyncLocalHTTPDriver(endpoint=settings.local_http_endpoint, model=model),
+    overwrite=True,
+)
+register_async_driver(
+    "google",
+    lambda model=None: AsyncGoogleDriver(api_key=settings.google_api_key, model=model or settings.google_model),
+    overwrite=True,
+)
+register_async_driver(
+    "groq",
+    lambda model=None: AsyncGroqDriver(api_key=settings.groq_api_key, model=model or settings.groq_model),
+    overwrite=True,
+)
+register_async_driver(
+    "openrouter",
+    lambda model=None: AsyncOpenRouterDriver(
+        api_key=settings.openrouter_api_key, model=model or settings.openrouter_model
+    ),
+    overwrite=True,
+)
+register_async_driver(
+    "grok",
+    lambda model=None: AsyncGrokDriver(api_key=settings.grok_api_key, model=model or settings.grok_model),
+    overwrite=True,
+)
+register_async_driver(
+    "airllm",
+    lambda model=None: AsyncAirLLMDriver(
+        model=model or settings.airllm_model,
+        compression=settings.airllm_compression,
+    ),
+    overwrite=True,
+)
+
+# Backwards compatibility: expose registry dict
+ASYNC_DRIVER_REGISTRY = _get_async_registry()
+
+
+def get_async_driver(provider_name: str | None = None):
+    """Factory to get an async driver instance based on the provider name.
+
+    Uses default model from settings if not overridden.
+    """
+    provider = (provider_name or settings.ai_provider or "ollama").strip().lower()
+    factory = get_async_driver_factory(provider)
+    return factory()
+
+
+def get_async_driver_for_model(model_str: str):
+    """Factory to get an async driver instance based on a full model string.
+
+    Format: ``provider/model_id``
+    Example: ``"openai/gpt-4-turbo-preview"``
+    """
+    if not isinstance(model_str, str):
+        raise ValueError("Model string must be a string, got {type(model_str)}")
+
+    if not model_str:
+        raise ValueError("Model string cannot be empty")
+
+    parts = model_str.split("/", 1)
+    provider = parts[0].lower()
+    model_id = parts[1] if len(parts) > 1 else None
+
+    factory = get_async_driver_factory(provider)
+    return factory(model_id)