prompture 0.0.46__py3-none-any.whl → 0.0.47__py3-none-any.whl

prompture/drivers/async_lmstudio_driver.py

@@ -98,7 +98,12 @@ class AsyncLMStudioDriver(AsyncDriver):
         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"]
+        message = response_data["choices"][0]["message"]
+        text = message.get("content") or ""
+        reasoning_content = message.get("reasoning_content")
+
+        if not text and reasoning_content:
+            text = reasoning_content
 
         usage = response_data.get("usage", {})
         prompt_tokens = usage.get("prompt_tokens", 0)
@@ -114,7 +119,10 @@ class AsyncLMStudioDriver(AsyncDriver):
             "model_name": merged_options.get("model", self.model),
         }
 
-        return {"text": text, "meta": meta}
+        result: dict[str, Any] = {"text": text, "meta": meta}
+        if reasoning_content is not None:
+            result["reasoning_content"] = reasoning_content
+        return result
 
     # -- Model management (LM Studio 0.4.0+) ----------------------------------
 

prompture/drivers/async_moonshot_driver.py

@@ -138,10 +138,11 @@ class AsyncMoonshotDriver(CostMixin, AsyncDriver):
 
         message = resp["choices"][0]["message"]
         text = message.get("content") or ""
+        reasoning_content = message.get("reasoning_content")
 
         # Reasoning models may return content in reasoning_content when content is empty
-        if not text and message.get("reasoning_content"):
-            text = message["reasoning_content"]
+        if not text and reasoning_content:
+            text = reasoning_content
 
         # Structured output fallback: if we used json_schema mode and got an
         # empty response, retry with json_object mode and schema in the prompt.
@@ -184,8 +185,9 @@ class AsyncMoonshotDriver(CostMixin, AsyncDriver):
                     resp = fb_resp
                     fb_message = fb_resp["choices"][0]["message"]
                     text = fb_message.get("content") or ""
-                    if not text and fb_message.get("reasoning_content"):
-                        text = fb_message["reasoning_content"]
+                    reasoning_content = fb_message.get("reasoning_content")
+                    if not text and reasoning_content:
+                        text = reasoning_content
 
         total_cost = self._calculate_cost("moonshot", model, prompt_tokens, completion_tokens)
 
@@ -198,7 +200,10 @@ class AsyncMoonshotDriver(CostMixin, AsyncDriver):
             "model_name": model,
         }
 
-        return {"text": text, "meta": meta}
+        result: dict[str, Any] = {"text": text, "meta": meta}
+        if reasoning_content is not None:
+            result["reasoning_content"] = reasoning_content
+        return result
 
     # ------------------------------------------------------------------
     # Tool use
@@ -271,11 +276,12 @@ class AsyncMoonshotDriver(CostMixin, AsyncDriver):
         }
 
         choice = resp["choices"][0]
-        text = choice["message"].get("content") or ""
+        message = choice["message"]
+        text = message.get("content") or ""
         stop_reason = choice.get("finish_reason")
 
         tool_calls_out: list[dict[str, Any]] = []
-        for tc in choice["message"].get("tool_calls", []):
+        for tc in message.get("tool_calls", []):
             try:
                 args = json.loads(tc["function"]["arguments"])
             except (json.JSONDecodeError, TypeError):
@@ -288,13 +294,21 @@ class AsyncMoonshotDriver(CostMixin, AsyncDriver):
                 }
             )
 
-        return {
+        result: dict[str, Any] = {
             "text": text,
             "meta": meta,
             "tool_calls": tool_calls_out,
             "stop_reason": stop_reason,
         }
 
+        # Preserve reasoning_content for reasoning models so the
+        # conversation loop can include it when sending the assistant
+        # message back (Moonshot requires it on subsequent requests).
+        if message.get("reasoning_content") is not None:
+            result["reasoning_content"] = message["reasoning_content"]
+
+        return result
+
     # ------------------------------------------------------------------
     # Streaming
     # ------------------------------------------------------------------
@@ -325,6 +339,7 @@ class AsyncMoonshotDriver(CostMixin, AsyncDriver):
             data["temperature"] = opts["temperature"]
 
         full_text = ""
+        full_reasoning = ""
         prompt_tokens = 0
         completion_tokens = 0
 
@@ -359,9 +374,11 @@ class AsyncMoonshotDriver(CostMixin, AsyncDriver):
                 if choices:
                     delta = choices[0].get("delta", {})
                     content = delta.get("content") or ""
-                    # Reasoning models stream thinking via reasoning_content
-                    if not content:
-                        content = delta.get("reasoning_content") or ""
+                    reasoning_chunk = delta.get("reasoning_content") or ""
+                    if reasoning_chunk:
+                        full_reasoning += reasoning_chunk
+                    if not content and reasoning_chunk:
+                        content = reasoning_chunk
                     if content:
                         full_text += content
                         yield {"type": "delta", "text": content}
@@ -369,7 +386,7 @@ class AsyncMoonshotDriver(CostMixin, AsyncDriver):
         total_tokens = prompt_tokens + completion_tokens
         total_cost = self._calculate_cost("moonshot", model, prompt_tokens, completion_tokens)
 
-        yield {
+        done_chunk: dict[str, Any] = {
             "type": "done",
             "text": full_text,
             "meta": {
@@ -381,3 +398,6 @@ class AsyncMoonshotDriver(CostMixin, AsyncDriver):
                 "model_name": model,
             },
         }
+        if full_reasoning:
+            done_chunk["reasoning_content"] = full_reasoning
+        yield done_chunk

prompture/drivers/async_ollama_driver.py

@@ -2,8 +2,10 @@
 
 from __future__ import annotations
 
+import json
 import logging
 import os
+import uuid
 from typing import Any
 
 import httpx
@@ -16,6 +18,7 @@ logger = logging.getLogger(__name__)
 class AsyncOllamaDriver(AsyncDriver):
     supports_json_mode = True
     supports_json_schema = True
+    supports_tool_use = True
     supports_vision = True
 
     MODEL_PRICING = {"default": {"prompt": 0.0, "completion": 0.0}}
@@ -80,6 +83,88 @@ class AsyncOllamaDriver(AsyncDriver):
 
         return {"text": response_data.get("response", ""), "meta": meta}
 
+    # ------------------------------------------------------------------
+    # Tool use
+    # ------------------------------------------------------------------
+
+    async 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 via Ollama's /api/chat endpoint."""
+        merged_options = self.options.copy()
+        if options:
+            merged_options.update(options)
+
+        chat_endpoint = self.endpoint.replace("/api/generate", "/api/chat")
+
+        payload: dict[str, Any] = {
+            "model": merged_options.get("model", self.model),
+            "messages": messages,
+            "tools": tools,
+            "stream": False,
+        }
+
+        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 tool use request failed: {e}") from e
+            except Exception as e:
+                raise RuntimeError(f"Ollama tool use 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") or ""
+        stop_reason = response_data.get("done_reason", "stop")
+
+        tool_calls_out: list[dict[str, Any]] = []
+        for tc in message.get("tool_calls", []):
+            func = tc.get("function", {})
+            # Ollama returns arguments as a dict already (no JSON string parsing needed)
+            args = func.get("arguments", {})
+            if isinstance(args, str):
+                try:
+                    args = json.loads(args)
+                except (json.JSONDecodeError, TypeError):
+                    args = {}
+            tool_calls_out.append({
+                # Ollama does not return tool_call IDs — generate one locally
+                "id": f"call_{uuid.uuid4().hex[:24]}",
+                "name": func.get("name", ""),
+                "arguments": args,
+            })
+
+        return {
+            "text": text,
+            "meta": meta,
+            "tool_calls": tool_calls_out,
+            "stop_reason": stop_reason,
+        }
+
     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)

prompture/drivers/async_openrouter_driver.py

@@ -122,8 +122,17 @@ class AsyncOpenRouterDriver(CostMixin, AsyncDriver):
             "model_name": model,
         }
 
-        text = resp["choices"][0]["message"]["content"]
-        return {"text": text, "meta": meta}
+        message = resp["choices"][0]["message"]
+        text = message.get("content") or ""
+        reasoning_content = message.get("reasoning_content")
+
+        if not text and reasoning_content:
+            text = reasoning_content
+
+        result: dict[str, Any] = {"text": text, "meta": meta}
+        if reasoning_content is not None:
+            result["reasoning_content"] = reasoning_content
+        return result
 
     # ------------------------------------------------------------------
     # Tool use
@@ -196,18 +205,23 @@ class AsyncOpenRouterDriver(CostMixin, AsyncDriver):
                 args = json.loads(tc["function"]["arguments"])
             except (json.JSONDecodeError, TypeError):
                 args = {}
-            tool_calls_out.append({
-                "id": tc["id"],
-                "name": tc["function"]["name"],
-                "arguments": args,
-            })
+            tool_calls_out.append(
+                {
+                    "id": tc["id"],
+                    "name": tc["function"]["name"],
+                    "arguments": args,
+                }
+            )
 
-        return {
+        result: dict[str, Any] = {
             "text": text,
             "meta": meta,
             "tool_calls": tool_calls_out,
             "stop_reason": stop_reason,
         }
+        if choice["message"].get("reasoning_content") is not None:
+            result["reasoning_content"] = choice["message"]["reasoning_content"]
+        return result
 
     # ------------------------------------------------------------------
     # Streaming
@@ -238,21 +252,25 @@ class AsyncOpenRouterDriver(CostMixin, AsyncDriver):
             data["temperature"] = opts["temperature"]
 
         full_text = ""
+        full_reasoning = ""
         prompt_tokens = 0
         completion_tokens = 0
 
-        async with httpx.AsyncClient() as client, client.stream(
-            "POST",
-            f"{self.base_url}/chat/completions",
-            headers=self.headers,
-            json=data,
-            timeout=120,
-        ) as response:
+        async with (
+            httpx.AsyncClient() as client,
+            client.stream(
+                "POST",
+                f"{self.base_url}/chat/completions",
+                headers=self.headers,
+                json=data,
+                timeout=120,
+            ) as response,
+        ):
             response.raise_for_status()
             async for line in response.aiter_lines():
                 if not line or not line.startswith("data: "):
                     continue
-                payload = line[len("data: "):]
+                payload = line[len("data: ") :]
                 if payload.strip() == "[DONE]":
                     break
                 try:
@@ -270,6 +288,11 @@ class AsyncOpenRouterDriver(CostMixin, AsyncDriver):
                 if choices:
                     delta = choices[0].get("delta", {})
                     content = delta.get("content", "")
+                    reasoning_chunk = delta.get("reasoning_content") or ""
+                    if reasoning_chunk:
+                        full_reasoning += reasoning_chunk
+                    if not content and reasoning_chunk:
+                        content = reasoning_chunk
                     if content:
                         full_text += content
                         yield {"type": "delta", "text": content}
@@ -277,7 +300,7 @@ class AsyncOpenRouterDriver(CostMixin, AsyncDriver):
         total_tokens = prompt_tokens + completion_tokens
         total_cost = self._calculate_cost("openrouter", model, prompt_tokens, completion_tokens)
 
-        yield {
+        done_chunk: dict[str, Any] = {
             "type": "done",
             "text": full_text,
             "meta": {
@@ -289,3 +312,6 @@ class AsyncOpenRouterDriver(CostMixin, AsyncDriver):
                 "model_name": model,
             },
         }
+        if full_reasoning:
+            done_chunk["reasoning_content"] = full_reasoning
+        yield done_chunk

prompture/drivers/azure_driver.py

@@ -2,6 +2,7 @@
 Requires the `openai` package.
 """
 
+import json
 import os
 from typing import Any
 
@@ -17,6 +18,7 @@ from ..driver import Driver
 class AzureDriver(CostMixin, Driver):
     supports_json_mode = True
     supports_json_schema = True
+    supports_tool_use = True
     supports_vision = True
 
     # Pricing per 1K tokens (adjust if your Azure pricing differs from OpenAI defaults)
@@ -164,3 +166,78 @@ class AzureDriver(CostMixin, Driver):
 
         text = resp.choices[0].message.content
         return {"text": text, "meta": meta}
+
+    # ------------------------------------------------------------------
+    # Tool use
+    # ------------------------------------------------------------------
+
+    def generate_messages_with_tools(
+        self,
+        messages: list[dict[str, Any]],
+        tools: list[dict[str, Any]],
+        options: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Generate a response that may include tool calls."""
+        if self.client is None:
+            raise RuntimeError("openai package (>=1.0.0) with AzureOpenAI not installed")
+
+        model = options.get("model", self.model)
+        model_config = self._get_model_config("azure", model)
+        tokens_param = model_config["tokens_param"]
+        supports_temperature = model_config["supports_temperature"]
+
+        self._validate_model_capabilities("azure", model, using_tool_use=True)
+
+        opts = {"temperature": 1.0, "max_tokens": 512, **options}
+
+        kwargs: dict[str, Any] = {
+            "model": self.deployment_id,
+            "messages": messages,
+            "tools": tools,
+        }
+        kwargs[tokens_param] = opts.get("max_tokens", 512)
+
+        if supports_temperature and "temperature" in opts:
+            kwargs["temperature"] = opts["temperature"]
+
+        resp = 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("azure", 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": resp.model_dump(),
+            "model_name": model,
+            "deployment_id": self.deployment_id,
+        }
+
+        choice = resp.choices[0]
+        text = choice.message.content or ""
+        stop_reason = choice.finish_reason
+
+        tool_calls_out: list[dict[str, Any]] = []
+        if choice.message.tool_calls:
+            for tc in choice.message.tool_calls:
+                try:
+                    args = json.loads(tc.function.arguments)
+                except (json.JSONDecodeError, TypeError):
+                    args = {}
+                tool_calls_out.append({
+                    "id": tc.id,
+                    "name": tc.function.name,
+                    "arguments": args,
+                })
+
+        return {
+            "text": text,
+            "meta": meta,
+            "tool_calls": tool_calls_out,
+            "stop_reason": stop_reason,
+        }

prompture/drivers/grok_driver.py

@@ -2,6 +2,7 @@
 Requires the `requests` package. Uses GROK_API_KEY env var.
 """
 
+import json
 import os
 from typing import Any
 
@@ -13,6 +14,7 @@ from ..driver import Driver
 
 class GrokDriver(CostMixin, Driver):
     supports_json_mode = True
+    supports_tool_use = True
     supports_vision = True
 
     # Pricing per 1M tokens based on xAI's documentation
@@ -152,5 +154,102 @@ class GrokDriver(CostMixin, Driver):
             "model_name": model,
         }
 
-        text = resp["choices"][0]["message"]["content"]
-        return {"text": text, "meta": meta}
+        message = resp["choices"][0]["message"]
+        text = message.get("content") or ""
+        reasoning_content = message.get("reasoning_content")
+
+        if not text and reasoning_content:
+            text = reasoning_content
+
+        result: dict[str, Any] = {"text": text, "meta": meta}
+        if reasoning_content is not None:
+            result["reasoning_content"] = reasoning_content
+        return result
+
+    # ------------------------------------------------------------------
+    # Tool use
+    # ------------------------------------------------------------------
+
+    def generate_messages_with_tools(
+        self,
+        messages: list[dict[str, Any]],
+        tools: list[dict[str, Any]],
+        options: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Generate a response that may include tool calls."""
+        if not self.api_key:
+            raise RuntimeError("GROK_API_KEY environment variable is required")
+
+        model = options.get("model", self.model)
+        model_config = self._get_model_config("grok", model)
+        tokens_param = model_config["tokens_param"]
+        supports_temperature = model_config["supports_temperature"]
+
+        self._validate_model_capabilities("grok", model, using_tool_use=True)
+
+        opts = {"temperature": 1.0, "max_tokens": 512, **options}
+
+        payload: dict[str, Any] = {
+            "model": model,
+            "messages": messages,
+            "tools": tools,
+        }
+        payload[tokens_param] = opts.get("max_tokens", 512)
+
+        if supports_temperature and "temperature" in opts:
+            payload["temperature"] = opts["temperature"]
+
+        if "tool_choice" in options:
+            payload["tool_choice"] = options["tool_choice"]
+
+        headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
+
+        try:
+            response = requests.post(f"{self.api_base}/chat/completions", headers=headers, json=payload)
+            response.raise_for_status()
+            resp = response.json()
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"Grok 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("grok", 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": resp,
+            "model_name": model,
+        }
+
+        choice = resp["choices"][0]
+        text = choice["message"].get("content") or ""
+        stop_reason = choice.get("finish_reason")
+
+        tool_calls_out: list[dict[str, Any]] = []
+        for tc in choice["message"].get("tool_calls", []):
+            try:
+                args = json.loads(tc["function"]["arguments"])
+            except (json.JSONDecodeError, TypeError):
+                args = {}
+            tool_calls_out.append(
+                {
+                    "id": tc["id"],
+                    "name": tc["function"]["name"],
+                    "arguments": args,
+                }
+            )
+
+        result: dict[str, Any] = {
+            "text": text,
+            "meta": meta,
+            "tool_calls": tool_calls_out,
+            "stop_reason": stop_reason,
+        }
+        if choice["message"].get("reasoning_content") is not None:
+            result["reasoning_content"] = choice["message"]["reasoning_content"]
+        return result

prompture/drivers/groq_driver.py

@@ -2,6 +2,7 @@
 Requires the `groq` package. Uses GROQ_API_KEY env var.
 """
 
+import json
 import os
 from typing import Any
 
@@ -16,6 +17,7 @@ from ..driver import Driver
 
 class GroqDriver(CostMixin, Driver):
     supports_json_mode = True
+    supports_tool_use = True
     supports_vision = True
 
     # Approximate pricing per 1K tokens (to be updated with official pricing)
@@ -120,5 +122,93 @@ class GroqDriver(CostMixin, Driver):
         }
 
         # Extract generated text
-        text = resp.choices[0].message.content
-        return {"text": text, "meta": meta}
+        text = resp.choices[0].message.content or ""
+        reasoning_content = getattr(resp.choices[0].message, "reasoning_content", None)
+
+        if not text and reasoning_content:
+            text = reasoning_content
+
+        result: dict[str, Any] = {"text": text, "meta": meta}
+        if reasoning_content is not None:
+            result["reasoning_content"] = reasoning_content
+        return result
+
+    # ------------------------------------------------------------------
+    # Tool use
+    # ------------------------------------------------------------------
+
+    def generate_messages_with_tools(
+        self,
+        messages: list[dict[str, Any]],
+        tools: list[dict[str, Any]],
+        options: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Generate a response that may include tool calls."""
+        if self.client is None:
+            raise RuntimeError("groq package is not installed")
+
+        model = options.get("model", self.model)
+        model_config = self._get_model_config("groq", model)
+        tokens_param = model_config["tokens_param"]
+        supports_temperature = model_config["supports_temperature"]
+
+        self._validate_model_capabilities("groq", model, using_tool_use=True)
+
+        opts = {"temperature": 0.7, "max_tokens": 512, **options}
+
+        kwargs: dict[str, Any] = {
+            "model": model,
+            "messages": messages,
+            "tools": tools,
+        }
+        kwargs[tokens_param] = opts.get("max_tokens", 512)
+
+        if supports_temperature and "temperature" in opts:
+            kwargs["temperature"] = opts["temperature"]
+
+        resp = 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("groq", 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": resp.model_dump(),
+            "model_name": model,
+        }
+
+        choice = resp.choices[0]
+        text = choice.message.content or ""
+        stop_reason = choice.finish_reason
+
+        tool_calls_out: list[dict[str, Any]] = []
+        if choice.message.tool_calls:
+            for tc in choice.message.tool_calls:
+                try:
+                    args = json.loads(tc.function.arguments)
+                except (json.JSONDecodeError, TypeError):
+                    args = {}
+                tool_calls_out.append(
+                    {
+                        "id": tc.id,
+                        "name": tc.function.name,
+                        "arguments": args,
+                    }
+                )
+
+        result: dict[str, Any] = {
+            "text": text,
+            "meta": meta,
+            "tool_calls": tool_calls_out,
+            "stop_reason": stop_reason,
+        }
+        reasoning_content = getattr(choice.message, "reasoning_content", None)
+        if reasoning_content is not None:
+            result["reasoning_content"] = reasoning_content
+        return result