prompture 0.0.35__py3-none-any.whl → 0.0.40.dev1__py3-none-any.whl

prompture/image.py

@@ -0,0 +1,180 @@
+"""Image handling utilities for vision-capable LLM drivers."""
+
+from __future__ import annotations
+
+import base64
+import mimetypes
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Union
+
+
+@dataclass(frozen=True)
+class ImageContent:
+    """Normalized image representation for vision-capable drivers.
+
+    Attributes:
+        data: Base64-encoded image data.
+        media_type: MIME type (e.g. ``"image/png"``, ``"image/jpeg"``).
+        source_type: How the image is delivered — ``"base64"`` or ``"url"``.
+        url: Original URL when ``source_type`` is ``"url"``.
+    """
+
+    data: str
+    media_type: str
+    source_type: str = "base64"
+    url: str | None = None
+
+
+# Public type alias accepted by all image-aware APIs.
+ImageInput = Union[bytes, str, Path, ImageContent]
+
+# Known data-URI prefix pattern
+_DATA_URI_RE = re.compile(r"^data:(image/[a-zA-Z0-9.+-]+);base64,(.+)$", re.DOTALL)
+
+# Base64 detection heuristic — must look like pure base64 of reasonable length
+_BASE64_RE = re.compile(r"^[A-Za-z0-9+/\n\r]+=*$")
+
+_MIME_FROM_EXT: dict[str, str] = {
+    ".jpg": "image/jpeg",
+    ".jpeg": "image/jpeg",
+    ".png": "image/png",
+    ".gif": "image/gif",
+    ".webp": "image/webp",
+    ".bmp": "image/bmp",
+    ".svg": "image/svg+xml",
+    ".tiff": "image/tiff",
+    ".tif": "image/tiff",
+}
+
+_MAGIC_BYTES: list[tuple[bytes, str]] = [
+    (b"\x89PNG", "image/png"),
+    (b"\xff\xd8\xff", "image/jpeg"),
+    (b"GIF87a", "image/gif"),
+    (b"GIF89a", "image/gif"),
+    (b"RIFF", "image/webp"),  # WebP starts with RIFF...WEBP
+    (b"BM", "image/bmp"),
+]
+
+
+def _guess_media_type_from_bytes(data: bytes) -> str:
+    """Guess MIME type from the first few bytes of image data."""
+    for magic, mime in _MAGIC_BYTES:
+        if data[: len(magic)] == magic:
+            return mime
+    return "image/png"  # safe fallback
+
+
+def _guess_media_type(path: str) -> str:
+    """Guess MIME type from a file path or URL."""
+    # Strip query strings for URLs
+    clean = path.split("?")[0].split("#")[0]
+    ext = Path(clean).suffix.lower()
+    if ext in _MIME_FROM_EXT:
+        return _MIME_FROM_EXT[ext]
+    guessed = mimetypes.guess_type(clean)[0]
+    return guessed or "image/png"
+
+
+# ------------------------------------------------------------------
+# Constructor functions
+# ------------------------------------------------------------------
+
+
+def image_from_bytes(data: bytes, media_type: str | None = None) -> ImageContent:
+    """Create an :class:`ImageContent` from raw bytes.
+
+    Args:
+        data: Raw image bytes.
+        media_type: MIME type.  Auto-detected from magic bytes when *None*.
+    """
+    if not data:
+        raise ValueError("Image data cannot be empty")
+    b64 = base64.b64encode(data).decode("ascii")
+    mt = media_type or _guess_media_type_from_bytes(data)
+    return ImageContent(data=b64, media_type=mt)
+
+
+def image_from_base64(b64: str, media_type: str = "image/png") -> ImageContent:
+    """Create an :class:`ImageContent` from a base64-encoded string.
+
+    Accepts both raw base64 and ``data:`` URIs.
+    """
+    m = _DATA_URI_RE.match(b64)
+    if m:
+        return ImageContent(data=m.group(2), media_type=m.group(1))
+    return ImageContent(data=b64, media_type=media_type)
+
+
+def image_from_file(path: str | Path, media_type: str | None = None) -> ImageContent:
+    """Create an :class:`ImageContent` by reading a local file.
+
+    Args:
+        path: Path to an image file.
+        media_type: MIME type.  Guessed from extension when *None*.
+    """
+    p = Path(path)
+    if not p.exists():
+        raise FileNotFoundError(f"Image file not found: {p}")
+    raw = p.read_bytes()
+    mt = media_type or _guess_media_type(str(p))
+    return image_from_bytes(raw, mt)
+
+
+def image_from_url(url: str, media_type: str | None = None) -> ImageContent:
+    """Create an :class:`ImageContent` referencing a remote URL.
+
+    The image is **not** downloaded — the URL is stored directly so
+    drivers that accept URL-based images can pass it through.  For
+    drivers that require base64, the URL is embedded as a data URI.
+
+    Args:
+        url: Publicly-accessible image URL.
+        media_type: MIME type.  Guessed from the URL when *None*.
+    """
+    mt = media_type or _guess_media_type(url)
+    return ImageContent(data="", media_type=mt, source_type="url", url=url)
+
+
+# ------------------------------------------------------------------
+# Smart constructor
+# ------------------------------------------------------------------
+
+
+def make_image(source: ImageInput) -> ImageContent:
+    """Auto-detect the source type and return an :class:`ImageContent`.
+
+    Accepts:
+    - ``ImageContent`` — returned as-is.
+    - ``bytes`` — base64-encoded with auto-detected MIME.
+    - ``str`` — tries (in order): data URI, URL, file path, raw base64.
+    - ``pathlib.Path`` — read from disk.
+    """
+    if isinstance(source, ImageContent):
+        return source
+
+    if isinstance(source, bytes):
+        return image_from_bytes(source)
+
+    if isinstance(source, Path):
+        return image_from_file(source)
+
+    if isinstance(source, str):
+        # 1. data URI
+        if source.startswith("data:"):
+            return image_from_base64(source)
+
+        # 2. URL
+        if source.startswith(("http://", "https://")):
+            return image_from_url(source)
+
+        # 3. File path (if exists on disk)
+        p = Path(source)
+        if p.exists():
+            return image_from_file(p)
+
+        # 4. Assume raw base64
+        return image_from_base64(source)
+
+    raise TypeError(f"Unsupported image source type: {type(source).__name__}")

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 ""

prompture/model_rates.py

@@ -9,6 +9,7 @@ import contextlib
 import json
 import logging
 import threading
+from dataclasses import dataclass
 from datetime import datetime, timezone
 from pathlib import Path
 from typing import Any, Optional
@@ -139,7 +140,12 @@ def _lookup_model(provider: str, model_id: str) -> Optional[dict[str, Any]]:
     if not isinstance(provider_data, dict):
         return None
 
-    return provider_data.get(model_id)
+    # models.dev nests actual models under a "models" key
+    models = provider_data.get("models", provider_data)
+    if not isinstance(models, dict):
+        return None
+
+    return models.get(model_id)
 
 
 # ── Public API ──────────────────────────────────────────────────────────────
@@ -189,7 +195,12 @@ def get_all_provider_models(provider: str) -> list[str]:
     if not isinstance(provider_data, dict):
         return []
 
-    return list(provider_data.keys())
+    # models.dev nests actual models under a "models" key
+    models = provider_data.get("models", provider_data)
+    if not isinstance(models, dict):
+        return []
+
+    return list(models.keys())
 
 
 def refresh_rates_cache(force: bool = False) -> bool:
@@ -215,3 +226,102 @@ def refresh_rates_cache(force: bool = False) -> bool:
             return True
 
         return False
+
+
+# ── Model Capabilities ─────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class ModelCapabilities:
+    """Normalized capability metadata for an LLM model from models.dev.
+
+    All fields default to ``None`` (unknown) so callers can distinguish
+    "the model doesn't support X" from "we have no data about X".
+    """
+
+    supports_temperature: Optional[bool] = None
+    supports_tool_use: Optional[bool] = None
+    supports_structured_output: Optional[bool] = None
+    supports_vision: Optional[bool] = None
+    is_reasoning: Optional[bool] = None
+    context_window: Optional[int] = None
+    max_output_tokens: Optional[int] = None
+    modalities_input: tuple[str, ...] = ()
+    modalities_output: tuple[str, ...] = ()
+
+
+def get_model_capabilities(provider: str, model_id: str) -> Optional[ModelCapabilities]:
+    """Return capability metadata for a model, or ``None`` if unavailable.
+
+    Maps models.dev fields to a :class:`ModelCapabilities` instance:
+
+    - ``temperature`` → ``supports_temperature``
+    - ``tool_call`` → ``supports_tool_use``
+    - ``structured_output`` → ``supports_structured_output``
+    - ``"image" in modalities.input`` → ``supports_vision``
+    - ``reasoning`` → ``is_reasoning``
+    - ``limit.context`` → ``context_window``
+    - ``limit.output`` → ``max_output_tokens``
+    """
+    entry = _lookup_model(provider, model_id)
+    if entry is None:
+        return None
+
+    # Boolean capabilities (True/False/None)
+    supports_temperature: Optional[bool] = None
+    if "temperature" in entry:
+        supports_temperature = bool(entry["temperature"])
+
+    supports_tool_use: Optional[bool] = None
+    if "tool_call" in entry:
+        supports_tool_use = bool(entry["tool_call"])
+
+    supports_structured_output: Optional[bool] = None
+    if "structured_output" in entry:
+        supports_structured_output = bool(entry["structured_output"])
+
+    is_reasoning: Optional[bool] = None
+    if "reasoning" in entry:
+        is_reasoning = bool(entry["reasoning"])
+
+    # Modalities
+    modalities = entry.get("modalities", {})
+    modalities_input: tuple[str, ...] = ()
+    modalities_output: tuple[str, ...] = ()
+    if isinstance(modalities, dict):
+        raw_in = modalities.get("input")
+        if isinstance(raw_in, (list, tuple)):
+            modalities_input = tuple(str(m) for m in raw_in)
+        raw_out = modalities.get("output")
+        if isinstance(raw_out, (list, tuple)):
+            modalities_output = tuple(str(m) for m in raw_out)
+
+    supports_vision: Optional[bool] = None
+    if modalities_input:
+        supports_vision = "image" in modalities_input
+
+    # Limits
+    context_window: Optional[int] = None
+    max_output_tokens: Optional[int] = None
+    limits = entry.get("limit", {})
+    if isinstance(limits, dict):
+        ctx = limits.get("context")
+        if ctx is not None:
+            with contextlib.suppress(TypeError, ValueError):
+                context_window = int(ctx)
+        out = limits.get("output")
+        if out is not None:
+            with contextlib.suppress(TypeError, ValueError):
+                max_output_tokens = int(out)
+
+    return ModelCapabilities(
+        supports_temperature=supports_temperature,
+        supports_tool_use=supports_tool_use,
+        supports_structured_output=supports_structured_output,
+        supports_vision=supports_vision,
+        is_reasoning=is_reasoning,
+        context_window=context_window,
+        max_output_tokens=max_output_tokens,
+        modalities_input=modalities_input,
+        modalities_output=modalities_output,
+    )