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

prompture/drivers/async_google_driver.py

@@ -0,0 +1,316 @@
+"""Async Google Generative AI (Gemini) driver."""
+
+from __future__ import annotations
+
+import logging
+import os
+import uuid
+from collections.abc import AsyncIterator
+from typing import Any
+
+import google.generativeai as genai
+
+from ..async_driver import AsyncDriver
+from ..cost_mixin import CostMixin
+from .google_driver import GoogleDriver
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncGoogleDriver(CostMixin, AsyncDriver):
+    """Async driver for Google's Generative AI API (Gemini)."""
+
+    supports_json_mode = True
+    supports_json_schema = True
+    supports_vision = True
+    supports_tool_use = True
+    supports_streaming = True
+
+    MODEL_PRICING = GoogleDriver.MODEL_PRICING
+    _PRICING_UNIT = 1_000_000
+
+    def __init__(self, api_key: str | None = None, model: str = "gemini-1.5-pro"):
+        self.api_key = api_key or os.getenv("GOOGLE_API_KEY")
+        if not self.api_key:
+            raise ValueError("Google API key not found. Set GOOGLE_API_KEY env var or pass api_key to constructor")
+        self.model = model
+        genai.configure(api_key=self.api_key)
+        self.options: dict[str, Any] = {}
+
+    def _calculate_cost_chars(self, prompt_chars: int, completion_chars: int) -> float:
+        """Calculate cost from character counts (same logic as sync GoogleDriver)."""
+        from ..model_rates import get_model_rates
+
+        live_rates = get_model_rates("google", self.model)
+        if live_rates:
+            est_prompt_tokens = prompt_chars / 4
+            est_completion_tokens = completion_chars / 4
+            prompt_cost = (est_prompt_tokens / 1_000_000) * live_rates["input"]
+            completion_cost = (est_completion_tokens / 1_000_000) * live_rates["output"]
+        else:
+            model_pricing = self.MODEL_PRICING.get(self.model, {"prompt": 0, "completion": 0})
+            prompt_cost = (prompt_chars / 1_000_000) * model_pricing["prompt"]
+            completion_cost = (completion_chars / 1_000_000) * model_pricing["completion"]
+        return round(prompt_cost + completion_cost, 6)
+
+    def _extract_usage_metadata(self, response: Any, messages: list[dict[str, Any]]) -> dict[str, Any]:
+        """Extract token counts from response, falling back to character estimation."""
+        usage = getattr(response, "usage_metadata", None)
+        if usage:
+            prompt_tokens = getattr(usage, "prompt_token_count", 0) or 0
+            completion_tokens = getattr(usage, "candidates_token_count", 0) or 0
+            total_tokens = getattr(usage, "total_token_count", 0) or (prompt_tokens + completion_tokens)
+            cost = self._calculate_cost("google", self.model, prompt_tokens, completion_tokens)
+        else:
+            # Fallback: estimate from character counts
+            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) if response.text else 0
+            prompt_tokens = total_prompt_chars // 4
+            completion_tokens = completion_chars // 4
+            total_tokens = prompt_tokens + completion_tokens
+            cost = self._calculate_cost_chars(total_prompt_chars, completion_chars)
+
+        return {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": total_tokens,
+            "cost": round(cost, 6),
+        }
+
+    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 _build_generation_args(
+        self, messages: list[dict[str, Any]], options: dict[str, Any] | None = None
+    ) -> tuple[Any, dict[str, Any], dict[str, Any]]:
+        """Parse messages and options into (gen_input, gen_kwargs, model_kwargs)."""
+        merged_options = self.options.copy()
+        if options:
+            merged_options.update(options)
+
+        generation_config = merged_options.get("generation_config", {})
+        safety_settings = merged_options.get("safety_settings", {})
+
+        if "temperature" in merged_options and "temperature" not in generation_config:
+            generation_config["temperature"] = merged_options["temperature"]
+        if "max_tokens" in merged_options and "max_output_tokens" not in generation_config:
+            generation_config["max_output_tokens"] = merged_options["max_tokens"]
+        if "top_p" in merged_options and "top_p" not in generation_config:
+            generation_config["top_p"] = merged_options["top_p"]
+        if "top_k" in merged_options and "top_k" not in generation_config:
+            generation_config["top_k"] = merged_options["top_k"]
+
+        # Native JSON mode support
+        if merged_options.get("json_mode"):
+            generation_config["response_mime_type"] = "application/json"
+            json_schema = merged_options.get("json_schema")
+            if json_schema:
+                generation_config["response_schema"] = json_schema
+
+        # Convert messages to Gemini format
+        system_instruction = None
+        contents: list[dict[str, Any]] = []
+        for msg in messages:
+            role = msg.get("role", "user")
+            content = msg.get("content", "")
+            if role == "system":
+                system_instruction = content if isinstance(content, str) else str(content)
+            else:
+                gemini_role = "model" if role == "assistant" else "user"
+                if msg.get("_vision_parts"):
+                    contents.append({"role": gemini_role, "parts": content})
+                else:
+                    contents.append({"role": gemini_role, "parts": [content]})
+
+        # For a single message, unwrap only if it has exactly one string part
+        if len(contents) == 1:
+            parts = contents[0]["parts"]
+            if len(parts) == 1 and isinstance(parts[0], str):
+                gen_input = parts[0]
+            else:
+                gen_input = contents
+        else:
+            gen_input = contents
+
+        model_kwargs: dict[str, Any] = {}
+        if system_instruction:
+            model_kwargs["system_instruction"] = system_instruction
+
+        gen_kwargs: dict[str, Any] = {
+            "generation_config": generation_config if generation_config else None,
+            "safety_settings": safety_settings if safety_settings else None,
+        }
+
+        return gen_input, gen_kwargs, model_kwargs
+
+    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]:
+        gen_input, gen_kwargs, model_kwargs = self._build_generation_args(messages, options)
+
+        try:
+            model = genai.GenerativeModel(self.model, **model_kwargs)
+            response = await model.generate_content_async(gen_input, **gen_kwargs)
+
+            if not response.text:
+                raise ValueError("Empty response from model")
+
+            usage_meta = self._extract_usage_metadata(response, messages)
+
+            meta = {
+                **usage_meta,
+                "raw_response": response.prompt_feedback if hasattr(response, "prompt_feedback") else None,
+                "model_name": self.model,
+            }
+
+            return {"text": response.text, "meta": meta}
+
+        except Exception as e:
+            logger.error(f"Google API request failed: {e}")
+            raise RuntimeError(f"Google API request failed: {e}") from e
+
+    # ------------------------------------------------------------------
+    # 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/function calls (async)."""
+        gen_input, gen_kwargs, model_kwargs = self._build_generation_args(
+            self._prepare_messages(messages), options
+        )
+
+        # Convert tools from OpenAI format to Gemini function declarations
+        function_declarations = []
+        for t in tools:
+            if "type" in t and t["type"] == "function":
+                fn = t["function"]
+                decl = {
+                    "name": fn["name"],
+                    "description": fn.get("description", ""),
+                }
+                params = fn.get("parameters")
+                if params:
+                    decl["parameters"] = params
+                function_declarations.append(decl)
+            elif "name" in t:
+                decl = {"name": t["name"], "description": t.get("description", "")}
+                params = t.get("parameters") or t.get("input_schema")
+                if params:
+                    decl["parameters"] = params
+                function_declarations.append(decl)
+
+        try:
+            model = genai.GenerativeModel(self.model, **model_kwargs)
+
+            gemini_tools = [genai.types.Tool(function_declarations=function_declarations)]
+            response = await model.generate_content_async(gen_input, tools=gemini_tools, **gen_kwargs)
+
+            usage_meta = self._extract_usage_metadata(response, messages)
+            meta = {
+                **usage_meta,
+                "raw_response": response.prompt_feedback if hasattr(response, "prompt_feedback") else None,
+                "model_name": self.model,
+            }
+
+            text = ""
+            tool_calls_out: list[dict[str, Any]] = []
+            stop_reason = "stop"
+
+            for candidate in response.candidates:
+                for part in candidate.content.parts:
+                    if hasattr(part, "text") and part.text:
+                        text += part.text
+                    if hasattr(part, "function_call") and part.function_call.name:
+                        fc = part.function_call
+                        tool_calls_out.append({
+                            "id": str(uuid.uuid4()),
+                            "name": fc.name,
+                            "arguments": dict(fc.args) if fc.args else {},
+                        })
+
+                finish_reason = getattr(candidate, "finish_reason", None)
+                if finish_reason is not None:
+                    reason_map = {1: "stop", 2: "max_tokens", 3: "safety", 4: "recitation", 5: "other"}
+                    stop_reason = reason_map.get(finish_reason, "stop")
+
+            if tool_calls_out:
+                stop_reason = "tool_use"
+
+            return {
+                "text": text,
+                "meta": meta,
+                "tool_calls": tool_calls_out,
+                "stop_reason": stop_reason,
+            }
+
+        except Exception as e:
+            logger.error(f"Google API tool call request failed: {e}")
+            raise RuntimeError(f"Google API tool call request failed: {e}") from e
+
+    # ------------------------------------------------------------------
+    # Streaming
+    # ------------------------------------------------------------------
+
+    async def generate_messages_stream(
+        self,
+        messages: list[dict[str, Any]],
+        options: dict[str, Any],
+    ) -> AsyncIterator[dict[str, Any]]:
+        """Yield response chunks via Gemini async streaming API."""
+        gen_input, gen_kwargs, model_kwargs = self._build_generation_args(
+            self._prepare_messages(messages), options
+        )
+
+        try:
+            model = genai.GenerativeModel(self.model, **model_kwargs)
+            response = await model.generate_content_async(gen_input, stream=True, **gen_kwargs)
+
+            full_text = ""
+            async for chunk in response:
+                chunk_text = getattr(chunk, "text", None) or ""
+                if chunk_text:
+                    full_text += chunk_text
+                    yield {"type": "delta", "text": chunk_text}
+
+            # After iteration completes, usage_metadata should be available
+            usage_meta = self._extract_usage_metadata(response, messages)
+
+            yield {
+                "type": "done",
+                "text": full_text,
+                "meta": {
+                    **usage_meta,
+                    "raw_response": {},
+                    "model_name": self.model,
+                },
+            }
+
+        except Exception as e:
+            logger.error(f"Google API streaming request failed: {e}")
+            raise RuntimeError(f"Google API streaming request failed: {e}") from e

prompture/drivers/async_grok_driver.py

@@ -0,0 +1,97 @@
+"""Async xAI Grok 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 .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
+
+    def __init__(self, api_key: str | None = None, model: str = "grok-4-fast-reasoning"):
+        self.api_key = api_key or os.getenv("GROK_API_KEY")
+        self.model = model
+        self.api_base = "https://api.x.ai/v1"
+
+    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 not self.api_key:
+            raise RuntimeError("GROK_API_KEY environment variable is required")
+
+        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}
+
+        payload = {
+            "model": model,
+            "messages": messages,
+        }
+        payload[tokens_param] = opts.get("max_tokens", 512)
+
+        if supports_temperature and "temperature" in opts:
+            payload["temperature"] = opts["temperature"]
+
+        # Native JSON mode support
+        if options.get("json_mode"):
+            payload["response_format"] = {"type": "json_object"}
+
+        headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
+
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.post(
+                    f"{self.api_base}/chat/completions", headers=headers, json=payload, timeout=120
+                )
+                response.raise_for_status()
+                resp = response.json()
+            except httpx.HTTPStatusError as e:
+                raise RuntimeError(f"Grok API request failed: {e!s}") from e
+            except Exception 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": total_cost,
+            "raw_response": resp,
+            "model_name": model,
+        }
+
+        text = resp["choices"][0]["message"]["content"]
+        return {"text": text, "meta": meta}

prompture/drivers/async_groq_driver.py

@@ -0,0 +1,90 @@
+"""Async Groq driver. Requires the ``groq`` package."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+try:
+    import groq
+except Exception:
+    groq = None
+
+from ..async_driver import AsyncDriver
+from ..cost_mixin import CostMixin
+from .groq_driver import GroqDriver
+
+
+class AsyncGroqDriver(CostMixin, AsyncDriver):
+    supports_json_mode = True
+    supports_vision = True
+
+    MODEL_PRICING = GroqDriver.MODEL_PRICING
+
+    def __init__(self, api_key: str | None = None, model: str = "llama2-70b-4096"):
+        self.api_key = api_key or os.getenv("GROQ_API_KEY")
+        self.model = model
+        if groq:
+            self.client = groq.AsyncClient(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("groq package 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": 0.7, "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"):
+            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("groq", 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_hugging_driver.py

@@ -0,0 +1,61 @@
+"""Async Hugging Face driver using httpx."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import httpx
+
+from ..async_driver import AsyncDriver
+
+
+class AsyncHuggingFaceDriver(AsyncDriver):
+    MODEL_PRICING = {"default": {"prompt": 0.0, "completion": 0.0}}
+
+    def __init__(self, endpoint: str | None = None, token: str | None = None, model: str = "bert-base-uncased"):
+        self.endpoint = endpoint or os.getenv("HF_ENDPOINT")
+        self.token = token or os.getenv("HF_TOKEN")
+        self.model = model
+
+        if not self.endpoint:
+            raise ValueError("Hugging Face endpoint is not configured. Set HF_ENDPOINT or pass explicitly.")
+        if not self.token:
+            raise ValueError("Hugging Face token is not configured. Set HF_TOKEN or pass explicitly.")
+
+        self.headers = {"Authorization": f"Bearer {self.token}"}
+
+    async def generate(self, prompt: str, options: dict[str, Any]) -> dict[str, Any]:
+        payload = {
+            "inputs": prompt,
+            "parameters": options,
+        }
+
+        async with httpx.AsyncClient() as client:
+            try:
+                r = await client.post(
+                    self.endpoint, headers=self.headers, json=payload, timeout=options.get("timeout", 60)
+                )
+                r.raise_for_status()
+                response_data = r.json()
+            except Exception as e:
+                raise RuntimeError(f"AsyncHuggingFaceDriver request failed: {e}") from e
+
+        text = None
+        if isinstance(response_data, list) and response_data and "generated_text" in response_data[0]:
+            text = response_data[0]["generated_text"]
+        elif isinstance(response_data, dict) and "generated_text" in response_data:
+            text = response_data["generated_text"]
+        else:
+            text = str(response_data)
+
+        meta = {
+            "prompt_tokens": 0,
+            "completion_tokens": 0,
+            "total_tokens": 0,
+            "cost": 0.0,
+            "raw_response": response_data,
+            "model_name": options.get("model", self.model),
+        }
+
+        return {"text": text, "meta": meta}