prompture 0.0.38.dev2__py3-none-any.whl → 0.0.42__py3-none-any.whl

prompture/drivers/zai_driver.py

@@ -0,0 +1,318 @@
+"""Z.ai (Zhipu AI) driver implementation.
+Requires the `requests` package. Uses ZHIPU_API_KEY env var.
+
+The Z.ai API is fully OpenAI-compatible (/chat/completions).
+All pricing comes from models.dev (provider: "zai") — no hardcoded pricing.
+"""
+
+import json
+import os
+from collections.abc import Iterator
+from typing import Any
+
+import requests
+
+from ..cost_mixin import CostMixin, prepare_strict_schema
+from ..driver import Driver
+
+
+class ZaiDriver(CostMixin, Driver):
+    supports_json_mode = True
+    supports_json_schema = True
+    supports_tool_use = True
+    supports_streaming = True
+    supports_vision = True
+
+    # All pricing resolved live from models.dev (provider: "zai")
+    MODEL_PRICING: dict[str, dict[str, Any]] = {}
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        model: str = "glm-4.7",
+        endpoint: str = "https://api.z.ai/api/paas/v4",
+    ):
+        """Initialize Z.ai driver.
+
+        Args:
+            api_key: Zhipu API key. If not provided, will look for ZHIPU_API_KEY env var.
+            model: Model to use. Defaults to glm-4.7.
+            endpoint: API base URL. Defaults to https://api.z.ai/api/paas/v4.
+        """
+        self.api_key = api_key or os.getenv("ZHIPU_API_KEY")
+        if not self.api_key:
+            raise ValueError("Zhipu API key not found. Set ZHIPU_API_KEY env var.")
+
+        self.model = model
+        self.base_url = endpoint.rstrip("/")
+
+        self.headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "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)
+
+    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(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:
+            raise RuntimeError("Zhipu API key not found")
+
+        model = options.get("model", self.model)
+
+        model_config = self._get_model_config("zai", model)
+        tokens_param = model_config["tokens_param"]
+        supports_temperature = model_config["supports_temperature"]
+
+        self._validate_model_capabilities(
+            "zai",
+            model,
+            using_json_schema=bool(options.get("json_schema")),
+        )
+
+        opts = {"temperature": 1.0, "max_tokens": 512, **options}
+
+        data: dict[str, Any] = {
+            "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"):
+            json_schema = options.get("json_schema")
+            if json_schema:
+                schema_copy = prepare_strict_schema(json_schema)
+                data["response_format"] = {
+                    "type": "json_schema",
+                    "json_schema": {
+                        "name": "extraction",
+                        "strict": True,
+                        "schema": schema_copy,
+                    },
+                }
+            else:
+                data["response_format"] = {"type": "json_object"}
+
+        try:
+            response = requests.post(
+                f"{self.base_url}/chat/completions",
+                headers=self.headers,
+                json=data,
+                timeout=120,
+            )
+            response.raise_for_status()
+            resp = response.json()
+        except requests.exceptions.HTTPError as e:
+            error_msg = f"Z.ai API request failed: {e!s}"
+            raise RuntimeError(error_msg) from e
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"Z.ai 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("zai", 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,
+        }
+
+        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 not self.api_key:
+            raise RuntimeError("Zhipu API key not found")
+
+        model = options.get("model", self.model)
+        model_config = self._get_model_config("zai", model)
+        tokens_param = model_config["tokens_param"]
+        supports_temperature = model_config["supports_temperature"]
+
+        self._validate_model_capabilities("zai", model, using_tool_use=True)
+
+        opts = {"temperature": 1.0, "max_tokens": 512, **options}
+
+        data: dict[str, Any] = {
+            "model": model,
+            "messages": messages,
+            "tools": tools,
+        }
+        data[tokens_param] = opts.get("max_tokens", 512)
+
+        if supports_temperature and "temperature" in opts:
+            data["temperature"] = opts["temperature"]
+
+        if "tool_choice" in options:
+            data["tool_choice"] = options["tool_choice"]
+
+        try:
+            response = requests.post(
+                f"{self.base_url}/chat/completions",
+                headers=self.headers,
+                json=data,
+                timeout=120,
+            )
+            response.raise_for_status()
+            resp = response.json()
+        except requests.exceptions.HTTPError as e:
+            error_msg = f"Z.ai API request failed: {e!s}"
+            raise RuntimeError(error_msg) from e
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"Z.ai 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("zai", 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,
+                }
+            )
+
+        return {
+            "text": text,
+            "meta": meta,
+            "tool_calls": tool_calls_out,
+            "stop_reason": stop_reason,
+        }
+
+    # ------------------------------------------------------------------
+    # Streaming
+    # ------------------------------------------------------------------
+
+    def generate_messages_stream(
+        self,
+        messages: list[dict[str, Any]],
+        options: dict[str, Any],
+    ) -> Iterator[dict[str, Any]]:
+        """Yield response chunks via Z.ai streaming API."""
+        if not self.api_key:
+            raise RuntimeError("Zhipu API key not found")
+
+        model = options.get("model", self.model)
+        model_config = self._get_model_config("zai", model)
+        tokens_param = model_config["tokens_param"]
+        supports_temperature = model_config["supports_temperature"]
+
+        opts = {"temperature": 1.0, "max_tokens": 512, **options}
+
+        data: dict[str, Any] = {
+            "model": model,
+            "messages": messages,
+            "stream": True,
+            "stream_options": {"include_usage": True},
+        }
+        data[tokens_param] = opts.get("max_tokens", 512)
+
+        if supports_temperature and "temperature" in opts:
+            data["temperature"] = opts["temperature"]
+
+        response = requests.post(
+            f"{self.base_url}/chat/completions",
+            headers=self.headers,
+            json=data,
+            stream=True,
+            timeout=120,
+        )
+        response.raise_for_status()
+
+        full_text = ""
+        prompt_tokens = 0
+        completion_tokens = 0
+
+        for line in response.iter_lines(decode_unicode=True):
+            if not line or not line.startswith("data: "):
+                continue
+            payload = line[len("data: ") :]
+            if payload.strip() == "[DONE]":
+                break
+            try:
+                chunk = json.loads(payload)
+            except json.JSONDecodeError:
+                continue
+
+            usage = chunk.get("usage")
+            if usage:
+                prompt_tokens = usage.get("prompt_tokens", 0)
+                completion_tokens = usage.get("completion_tokens", 0)
+
+            choices = chunk.get("choices", [])
+            if choices:
+                delta = choices[0].get("delta", {})
+                content = delta.get("content", "")
+                if content:
+                    full_text += content
+                    yield {"type": "delta", "text": content}
+
+        total_tokens = prompt_tokens + completion_tokens
+        total_cost = self._calculate_cost("zai", model, prompt_tokens, completion_tokens)
+
+        yield {
+            "type": "done",
+            "text": full_text,
+            "meta": {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": total_tokens,
+                "cost": round(total_cost, 6),
+                "raw_response": {},
+                "model_name": model,
+            },
+        }

prompture/groups.py

@@ -114,6 +114,27 @@ class SequentialGroup:
         """Request graceful shutdown after the current agent finishes."""
         self._stop_requested = True
 
+    @property
+    def shared_state(self) -> dict[str, Any]:
+        """Return a copy of the current shared execution state."""
+        return dict(self._state)
+
+    def inject_state(self, state: dict[str, Any], *, recursive: bool = False) -> None:
+        """Merge external key-value pairs into this group's shared state.
+
+        Existing keys are NOT overwritten (uses setdefault semantics).
+
+        Args:
+            state: Key-value pairs to inject.
+            recursive: If True, also inject into nested sub-groups.
+        """
+        for k, v in state.items():
+            self._state.setdefault(k, v)
+        if recursive:
+            for agent, _ in self._agents:
+                if hasattr(agent, "inject_state"):
+                    agent.inject_state(state, recursive=True)
+
     def save(self, path: str) -> None:
         """Run and save result to file. Convenience wrapper."""
         result = self.run()
@@ -267,6 +288,27 @@ class LoopGroup:
         """Request graceful shutdown."""
         self._stop_requested = True
 
+    @property
+    def shared_state(self) -> dict[str, Any]:
+        """Return a copy of the current shared execution state."""
+        return dict(self._state)
+
+    def inject_state(self, state: dict[str, Any], *, recursive: bool = False) -> None:
+        """Merge external key-value pairs into this group's shared state.
+
+        Existing keys are NOT overwritten (uses setdefault semantics).
+
+        Args:
+            state: Key-value pairs to inject.
+            recursive: If True, also inject into nested sub-groups.
+        """
+        for k, v in state.items():
+            self._state.setdefault(k, v)
+        if recursive:
+            for agent, _ in self._agents:
+                if hasattr(agent, "inject_state"):
+                    agent.inject_state(state, recursive=True)
+
     def run(self, prompt: str = "") -> GroupResult:
         """Execute the loop."""
         self._stop_requested = False

prompture/ledger.py

@@ -0,0 +1,252 @@
+"""Persistent model usage ledger — tracks which LLM models have been used.
+
+Stores per-model usage stats (call count, tokens, cost, timestamps) in a
+SQLite database at ``~/.prompture/usage/model_ledger.db``.  The public
+convenience functions are fire-and-forget: they never raise exceptions so
+they cannot break existing extraction/conversation flows.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+import sqlite3
+import threading
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+logger = logging.getLogger("prompture.ledger")
+
+_DEFAULT_DB_DIR = Path.home() / ".prompture" / "usage"
+_DEFAULT_DB_PATH = _DEFAULT_DB_DIR / "model_ledger.db"
+
+_SCHEMA_SQL = """
+CREATE TABLE IF NOT EXISTS model_usage (
+    model_name  TEXT NOT NULL,
+    api_key_hash TEXT NOT NULL,
+    use_count   INTEGER NOT NULL DEFAULT 1,
+    total_tokens INTEGER NOT NULL DEFAULT 0,
+    total_cost  REAL NOT NULL DEFAULT 0.0,
+    first_used  TEXT NOT NULL,
+    last_used   TEXT NOT NULL,
+    last_status TEXT NOT NULL DEFAULT 'success',
+    PRIMARY KEY (model_name, api_key_hash)
+);
+"""
+
+
+class ModelUsageLedger:
+    """SQLite-backed model usage tracker.
+
+    Thread-safe via an internal :class:`threading.Lock`.
+
+    Args:
+        db_path: Path to the SQLite database file.  Defaults to
+            ``~/.prompture/usage/model_ledger.db``.
+    """
+
+    def __init__(self, db_path: str | Path | None = None) -> None:
+        self._db_path = Path(db_path) if db_path else _DEFAULT_DB_PATH
+        self._db_path.parent.mkdir(parents=True, exist_ok=True)
+        self._lock = threading.Lock()
+        self._init_db()
+
+    def _init_db(self) -> None:
+        with self._lock:
+            conn = sqlite3.connect(str(self._db_path))
+            try:
+                conn.executescript(_SCHEMA_SQL)
+                conn.commit()
+            finally:
+                conn.close()
+
+    def _connect(self) -> sqlite3.Connection:
+        conn = sqlite3.connect(str(self._db_path))
+        conn.row_factory = sqlite3.Row
+        return conn
+
+    # ------------------------------------------------------------------ #
+    # Recording
+    # ------------------------------------------------------------------ #
+
+    def record_usage(
+        self,
+        model_name: str,
+        *,
+        api_key_hash: str = "",
+        tokens: int = 0,
+        cost: float = 0.0,
+        status: str = "success",
+    ) -> None:
+        """Record a model usage event (upsert).
+
+        On conflict the row's counters are incremented and ``last_used``
+        is updated.
+        """
+        now = datetime.now(timezone.utc).isoformat()
+        with self._lock:
+            conn = self._connect()
+            try:
+                conn.execute(
+                    """
+                    INSERT INTO model_usage
+                        (model_name, api_key_hash, use_count, total_tokens, total_cost,
+                         first_used, last_used, last_status)
+                    VALUES (?, ?, 1, ?, ?, ?, ?, ?)
+                    ON CONFLICT(model_name, api_key_hash) DO UPDATE SET
+                        use_count    = use_count + 1,
+                        total_tokens = total_tokens + excluded.total_tokens,
+                        total_cost   = total_cost + excluded.total_cost,
+                        last_used    = excluded.last_used,
+                        last_status  = excluded.last_status
+                    """,
+                    (model_name, api_key_hash, tokens, cost, now, now, status),
+                )
+                conn.commit()
+            finally:
+                conn.close()
+
+    # ------------------------------------------------------------------ #
+    # Queries
+    # ------------------------------------------------------------------ #
+
+    def get_model_stats(self, model_name: str, api_key_hash: str = "") -> dict[str, Any] | None:
+        """Return stats for a specific model + key combination, or ``None``."""
+        with self._lock:
+            conn = self._connect()
+            try:
+                row = conn.execute(
+                    "SELECT * FROM model_usage WHERE model_name = ? AND api_key_hash = ?",
+                    (model_name, api_key_hash),
+                ).fetchone()
+                if row is None:
+                    return None
+                return dict(row)
+            finally:
+                conn.close()
+
+    def get_verified_models(self) -> set[str]:
+        """Return model names that have at least one successful usage."""
+        with self._lock:
+            conn = self._connect()
+            try:
+                rows = conn.execute(
+                    "SELECT DISTINCT model_name FROM model_usage WHERE last_status = 'success'"
+                ).fetchall()
+                return {r["model_name"] for r in rows}
+            finally:
+                conn.close()
+
+    def get_recently_used(self, limit: int = 10) -> list[dict[str, Any]]:
+        """Return recent model usage rows ordered by ``last_used`` descending."""
+        with self._lock:
+            conn = self._connect()
+            try:
+                rows = conn.execute(
+                    "SELECT * FROM model_usage ORDER BY last_used DESC LIMIT ?",
+                    (limit,),
+                ).fetchall()
+                return [dict(r) for r in rows]
+            finally:
+                conn.close()
+
+    def get_all_stats(self) -> list[dict[str, Any]]:
+        """Return all usage rows."""
+        with self._lock:
+            conn = self._connect()
+            try:
+                rows = conn.execute("SELECT * FROM model_usage ORDER BY last_used DESC").fetchall()
+                return [dict(r) for r in rows]
+            finally:
+                conn.close()
+
+
+# ------------------------------------------------------------------
+# Module-level singleton
+# ------------------------------------------------------------------
+
+_ledger: ModelUsageLedger | None = None
+_ledger_lock = threading.Lock()
+
+
+def _get_ledger() -> ModelUsageLedger:
+    """Return (and lazily create) the module-level singleton ledger."""
+    global _ledger
+    if _ledger is None:
+        with _ledger_lock:
+            if _ledger is None:
+                _ledger = ModelUsageLedger()
+    return _ledger
+
+
+# ------------------------------------------------------------------
+# Public convenience functions (fire-and-forget)
+# ------------------------------------------------------------------
+
+
+def record_model_usage(
+    model_name: str,
+    *,
+    api_key_hash: str = "",
+    tokens: int = 0,
+    cost: float = 0.0,
+    status: str = "success",
+) -> None:
+    """Record a model usage event.  Never raises — all exceptions are swallowed."""
+    try:
+        _get_ledger().record_usage(
+            model_name,
+            api_key_hash=api_key_hash,
+            tokens=tokens,
+            cost=cost,
+            status=status,
+        )
+    except Exception:
+        logger.debug("Failed to record model usage for %s", model_name, exc_info=True)
+
+
+def get_recently_used_models(limit: int = 10) -> list[dict[str, Any]]:
+    """Return recently used models.  Returns empty list on error."""
+    try:
+        return _get_ledger().get_recently_used(limit)
+    except Exception:
+        logger.debug("Failed to get recently used models", exc_info=True)
+        return []
+
+
+# ------------------------------------------------------------------
+# API key hash helper
+# ------------------------------------------------------------------
+
+_LOCAL_PROVIDERS = frozenset({"ollama", "lmstudio", "local_http", "airllm"})
+
+
+def _resolve_api_key_hash(model_name: str) -> str:
+    """Derive an 8-char hex hash of the API key for the given model's provider.
+
+    Local providers (ollama, lmstudio, etc.) return ``""``.
+    """
+    try:
+        provider = model_name.split("/", 1)[0].lower() if "/" in model_name else model_name.lower()
+        if provider in _LOCAL_PROVIDERS:
+            return ""
+
+        from .settings import settings
+
+        key_map: dict[str, str | None] = {
+            "openai": settings.openai_api_key,
+            "claude": settings.claude_api_key,
+            "google": settings.google_api_key,
+            "groq": settings.groq_api_key,
+            "grok": settings.grok_api_key,
+            "openrouter": settings.openrouter_api_key,
+            "azure": settings.azure_api_key,
+            "huggingface": settings.hf_token,
+        }
+        api_key = key_map.get(provider)
+        if not api_key:
+            return ""
+        return hashlib.sha256(api_key.encode()).hexdigest()[:8]
+    except Exception:
+        return ""