headroom-ai 0.2.13__py3-none-any.whl

headroom/storage/base.py

@@ -0,0 +1,125 @@
+"""Base storage interface for Headroom SDK."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Iterator
+from datetime import datetime
+from typing import Any
+
+from ..config import RequestMetrics
+
+
+class Storage(ABC):
+    """Abstract base class for metrics storage."""
+
+    @abstractmethod
+    def save(self, metrics: RequestMetrics) -> None:
+        """
+        Save request metrics.
+
+        Args:
+            metrics: RequestMetrics to save.
+        """
+        pass
+
+    @abstractmethod
+    def get(self, request_id: str) -> RequestMetrics | None:
+        """
+        Get metrics by request ID.
+
+        Args:
+            request_id: The request ID.
+
+        Returns:
+            RequestMetrics or None if not found.
+        """
+        pass
+
+    @abstractmethod
+    def query(
+        self,
+        start_time: datetime | None = None,
+        end_time: datetime | None = None,
+        model: str | None = None,
+        mode: str | None = None,
+        limit: int = 100,
+        offset: int = 0,
+    ) -> list[RequestMetrics]:
+        """
+        Query metrics with filters.
+
+        Args:
+            start_time: Filter by timestamp >= start_time.
+            end_time: Filter by timestamp <= end_time.
+            model: Filter by model name.
+            mode: Filter by mode (audit/optimize).
+            limit: Maximum results to return.
+            offset: Number of results to skip.
+
+        Returns:
+            List of matching RequestMetrics.
+        """
+        pass
+
+    @abstractmethod
+    def count(
+        self,
+        start_time: datetime | None = None,
+        end_time: datetime | None = None,
+        model: str | None = None,
+        mode: str | None = None,
+    ) -> int:
+        """
+        Count metrics matching filters.
+
+        Args:
+            start_time: Filter by timestamp >= start_time.
+            end_time: Filter by timestamp <= end_time.
+            model: Filter by model name.
+            mode: Filter by mode.
+
+        Returns:
+            Count of matching records.
+        """
+        pass
+
+    @abstractmethod
+    def iter_all(self) -> Iterator[RequestMetrics]:
+        """
+        Iterate over all stored metrics.
+
+        Yields:
+            RequestMetrics objects.
+        """
+        pass
+
+    @abstractmethod
+    def get_summary_stats(
+        self,
+        start_time: datetime | None = None,
+        end_time: datetime | None = None,
+    ) -> dict[str, Any]:
+        """
+        Get summary statistics.
+
+        Args:
+            start_time: Filter by timestamp >= start_time.
+            end_time: Filter by timestamp <= end_time.
+
+        Returns:
+            Dict with summary stats (total_requests, tokens_saved, etc.)
+        """
+        pass
+
+    def close(self) -> None:  # noqa: B027
+        """Close storage connection if applicable."""
+        pass
+
+    def __enter__(self) -> Storage:
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Context manager exit."""
+        self.close()

headroom/storage/jsonl.py

@@ -0,0 +1,220 @@
+"""JSONL file storage implementation for Headroom SDK."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterator
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+from ..config import RequestMetrics
+from ..utils import format_timestamp, parse_timestamp
+from .base import Storage
+
+
+class JSONLStorage(Storage):
+    """JSONL file-based metrics storage."""
+
+    def __init__(self, file_path: str):
+        """
+        Initialize JSONL storage.
+
+        Args:
+            file_path: Path to JSONL file.
+        """
+        self.file_path = file_path
+        self._ensure_file_exists()
+
+    def _ensure_file_exists(self) -> None:
+        """Create file and parent directories if they don't exist."""
+        path = Path(self.file_path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        if not path.exists():
+            path.touch()
+
+    def _metrics_to_dict(self, metrics: RequestMetrics) -> dict[str, Any]:
+        """Convert RequestMetrics to serializable dict."""
+        return {
+            "id": metrics.request_id,
+            "timestamp": format_timestamp(metrics.timestamp),
+            "model": metrics.model,
+            "stream": metrics.stream,
+            "mode": metrics.mode,
+            "tokens_input_before": metrics.tokens_input_before,
+            "tokens_input_after": metrics.tokens_input_after,
+            "tokens_output": metrics.tokens_output,
+            "block_breakdown": metrics.block_breakdown,
+            "waste_signals": metrics.waste_signals,
+            "stable_prefix_hash": metrics.stable_prefix_hash,
+            "cache_alignment_score": metrics.cache_alignment_score,
+            "cached_tokens": metrics.cached_tokens,
+            "transforms_applied": metrics.transforms_applied,
+            "tool_units_dropped": metrics.tool_units_dropped,
+            "turns_dropped": metrics.turns_dropped,
+            "messages_hash": metrics.messages_hash,
+            "error": metrics.error,
+        }
+
+    def _dict_to_metrics(self, data: dict[str, Any]) -> RequestMetrics:
+        """Convert dict to RequestMetrics."""
+        return RequestMetrics(
+            request_id=data["id"],
+            timestamp=parse_timestamp(data["timestamp"]),
+            model=data["model"],
+            stream=data["stream"],
+            mode=data["mode"],
+            tokens_input_before=data["tokens_input_before"],
+            tokens_input_after=data["tokens_input_after"],
+            tokens_output=data.get("tokens_output"),
+            block_breakdown=data.get("block_breakdown", {}),
+            waste_signals=data.get("waste_signals", {}),
+            stable_prefix_hash=data.get("stable_prefix_hash", ""),
+            cache_alignment_score=data.get("cache_alignment_score", 0.0),
+            cached_tokens=data.get("cached_tokens"),
+            transforms_applied=data.get("transforms_applied", []),
+            tool_units_dropped=data.get("tool_units_dropped", 0),
+            turns_dropped=data.get("turns_dropped", 0),
+            messages_hash=data.get("messages_hash", ""),
+            error=data.get("error"),
+        )
+
+    def save(self, metrics: RequestMetrics) -> None:
+        """Save request metrics."""
+        data = self._metrics_to_dict(metrics)
+        with open(self.file_path, "a") as f:
+            f.write(json.dumps(data) + "\n")
+
+    def get(self, request_id: str) -> RequestMetrics | None:
+        """Get metrics by request ID."""
+        for metrics in self.iter_all():
+            if metrics.request_id == request_id:
+                return metrics
+        return None
+
+    def query(
+        self,
+        start_time: datetime | None = None,
+        end_time: datetime | None = None,
+        model: str | None = None,
+        mode: str | None = None,
+        limit: int = 100,
+        offset: int = 0,
+    ) -> list[RequestMetrics]:
+        """Query metrics with filters."""
+        results: list[RequestMetrics] = []
+        skipped = 0
+
+        for metrics in self.iter_all():
+            # Apply filters
+            if start_time is not None and metrics.timestamp < start_time:
+                continue
+            if end_time is not None and metrics.timestamp > end_time:
+                continue
+            if model is not None and metrics.model != model:
+                continue
+            if mode is not None and metrics.mode != mode:
+                continue
+
+            # Handle offset
+            if skipped < offset:
+                skipped += 1
+                continue
+
+            results.append(metrics)
+
+            # Handle limit
+            if len(results) >= limit:
+                break
+
+        # Sort by timestamp descending
+        results.sort(key=lambda m: m.timestamp, reverse=True)
+        return results
+
+    def count(
+        self,
+        start_time: datetime | None = None,
+        end_time: datetime | None = None,
+        model: str | None = None,
+        mode: str | None = None,
+    ) -> int:
+        """Count metrics matching filters."""
+        count = 0
+
+        for metrics in self.iter_all():
+            if start_time is not None and metrics.timestamp < start_time:
+                continue
+            if end_time is not None and metrics.timestamp > end_time:
+                continue
+            if model is not None and metrics.model != model:
+                continue
+            if mode is not None and metrics.mode != mode:
+                continue
+            count += 1
+
+        return count
+
+    def iter_all(self) -> Iterator[RequestMetrics]:
+        """Iterate over all stored metrics."""
+        if not Path(self.file_path).exists():
+            return
+
+        with open(self.file_path) as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    data = json.loads(line)
+                    yield self._dict_to_metrics(data)
+                except json.JSONDecodeError:
+                    # Skip malformed lines
+                    continue
+
+    def get_summary_stats(
+        self,
+        start_time: datetime | None = None,
+        end_time: datetime | None = None,
+    ) -> dict[str, Any]:
+        """Get summary statistics."""
+        total_requests = 0
+        total_tokens_before = 0
+        total_tokens_after = 0
+        total_cache_alignment = 0.0
+        audit_count = 0
+        optimize_count = 0
+
+        for metrics in self.iter_all():
+            if start_time is not None and metrics.timestamp < start_time:
+                continue
+            if end_time is not None and metrics.timestamp > end_time:
+                continue
+
+            total_requests += 1
+            total_tokens_before += metrics.tokens_input_before
+            total_tokens_after += metrics.tokens_input_after
+            total_cache_alignment += metrics.cache_alignment_score
+
+            if metrics.mode == "audit":
+                audit_count += 1
+            elif metrics.mode == "optimize":
+                optimize_count += 1
+
+        total_tokens_saved = total_tokens_before - total_tokens_after
+        avg_tokens_saved = total_tokens_saved / total_requests if total_requests > 0 else 0
+        avg_cache_alignment = total_cache_alignment / total_requests if total_requests > 0 else 0
+
+        return {
+            "total_requests": total_requests,
+            "total_tokens_before": total_tokens_before,
+            "total_tokens_after": total_tokens_after,
+            "total_tokens_saved": total_tokens_saved,
+            "avg_tokens_saved": avg_tokens_saved,
+            "avg_cache_alignment": avg_cache_alignment,
+            "audit_count": audit_count,
+            "optimize_count": optimize_count,
+        }
+
+    def close(self) -> None:
+        """No-op for file storage."""
+        pass

headroom/storage/sqlite.py

@@ -0,0 +1,289 @@
+"""SQLite storage implementation for Headroom SDK."""
+
+from __future__ import annotations
+
+import json
+import sqlite3
+from collections.abc import Iterator
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+from ..config import RequestMetrics
+from ..utils import format_timestamp, parse_timestamp
+from .base import Storage
+
+
+class SQLiteStorage(Storage):
+    """SQLite-based metrics storage."""
+
+    def __init__(self, db_path: str):
+        """
+        Initialize SQLite storage.
+
+        Args:
+            db_path: Path to SQLite database file.
+        """
+        self.db_path = db_path
+        self._ensure_db_exists()
+        self._conn: sqlite3.Connection | None = None
+
+    def _ensure_db_exists(self) -> None:
+        """Create database and tables if they don't exist."""
+        path = Path(self.db_path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        conn = sqlite3.connect(self.db_path)
+        try:
+            cursor = conn.cursor()
+            cursor.execute("""
+                CREATE TABLE IF NOT EXISTS requests (
+                    id TEXT PRIMARY KEY,
+                    timestamp TEXT NOT NULL,
+                    model TEXT NOT NULL,
+                    stream INTEGER NOT NULL,
+                    mode TEXT NOT NULL,
+                    tokens_input_before INTEGER NOT NULL,
+                    tokens_input_after INTEGER NOT NULL,
+                    tokens_output INTEGER,
+                    block_breakdown TEXT NOT NULL,
+                    waste_signals TEXT NOT NULL,
+                    stable_prefix_hash TEXT,
+                    cache_alignment_score REAL,
+                    cached_tokens INTEGER,
+                    transforms_applied TEXT NOT NULL,
+                    tool_units_dropped INTEGER DEFAULT 0,
+                    turns_dropped INTEGER DEFAULT 0,
+                    messages_hash TEXT,
+                    error TEXT
+                )
+            """)
+
+            # Create indices
+            cursor.execute("""
+                CREATE INDEX IF NOT EXISTS idx_timestamp ON requests(timestamp)
+            """)
+            cursor.execute("""
+                CREATE INDEX IF NOT EXISTS idx_model ON requests(model)
+            """)
+            cursor.execute("""
+                CREATE INDEX IF NOT EXISTS idx_mode ON requests(mode)
+            """)
+
+            conn.commit()
+        finally:
+            conn.close()
+
+    def _get_conn(self) -> sqlite3.Connection:
+        """Get or create connection."""
+        if self._conn is None:
+            self._conn = sqlite3.connect(self.db_path)
+            self._conn.row_factory = sqlite3.Row
+        return self._conn
+
+    def save(self, metrics: RequestMetrics) -> None:
+        """Save request metrics."""
+        conn = self._get_conn()
+        cursor = conn.cursor()
+
+        cursor.execute(
+            """
+            INSERT OR REPLACE INTO requests (
+                id, timestamp, model, stream, mode,
+                tokens_input_before, tokens_input_after, tokens_output,
+                block_breakdown, waste_signals,
+                stable_prefix_hash, cache_alignment_score, cached_tokens,
+                transforms_applied, tool_units_dropped, turns_dropped,
+                messages_hash, error
+            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+            """,
+            (
+                metrics.request_id,
+                format_timestamp(metrics.timestamp),
+                metrics.model,
+                1 if metrics.stream else 0,
+                metrics.mode,
+                metrics.tokens_input_before,
+                metrics.tokens_input_after,
+                metrics.tokens_output,
+                json.dumps(metrics.block_breakdown),
+                json.dumps(metrics.waste_signals),
+                metrics.stable_prefix_hash,
+                metrics.cache_alignment_score,
+                metrics.cached_tokens,
+                json.dumps(metrics.transforms_applied),
+                metrics.tool_units_dropped,
+                metrics.turns_dropped,
+                metrics.messages_hash,
+                metrics.error,
+            ),
+        )
+        conn.commit()
+
+    def get(self, request_id: str) -> RequestMetrics | None:
+        """Get metrics by request ID."""
+        conn = self._get_conn()
+        cursor = conn.cursor()
+
+        cursor.execute("SELECT * FROM requests WHERE id = ?", (request_id,))
+        row = cursor.fetchone()
+
+        if row is None:
+            return None
+
+        return self._row_to_metrics(row)
+
+    def query(
+        self,
+        start_time: datetime | None = None,
+        end_time: datetime | None = None,
+        model: str | None = None,
+        mode: str | None = None,
+        limit: int = 100,
+        offset: int = 0,
+    ) -> list[RequestMetrics]:
+        """Query metrics with filters."""
+        conn = self._get_conn()
+        cursor = conn.cursor()
+
+        query = "SELECT * FROM requests WHERE 1=1"
+        params: list[Any] = []
+
+        if start_time is not None:
+            query += " AND timestamp >= ?"
+            params.append(format_timestamp(start_time))
+        if end_time is not None:
+            query += " AND timestamp <= ?"
+            params.append(format_timestamp(end_time))
+        if model is not None:
+            query += " AND model = ?"
+            params.append(model)
+        if mode is not None:
+            query += " AND mode = ?"
+            params.append(mode)
+
+        query += " ORDER BY timestamp DESC LIMIT ? OFFSET ?"
+        params.extend([limit, offset])
+
+        cursor.execute(query, params)
+        rows = cursor.fetchall()
+
+        return [self._row_to_metrics(row) for row in rows]
+
+    def count(
+        self,
+        start_time: datetime | None = None,
+        end_time: datetime | None = None,
+        model: str | None = None,
+        mode: str | None = None,
+    ) -> int:
+        """Count metrics matching filters."""
+        conn = self._get_conn()
+        cursor = conn.cursor()
+
+        query = "SELECT COUNT(*) FROM requests WHERE 1=1"
+        params: list[Any] = []
+
+        if start_time is not None:
+            query += " AND timestamp >= ?"
+            params.append(format_timestamp(start_time))
+        if end_time is not None:
+            query += " AND timestamp <= ?"
+            params.append(format_timestamp(end_time))
+        if model is not None:
+            query += " AND model = ?"
+            params.append(model)
+        if mode is not None:
+            query += " AND mode = ?"
+            params.append(mode)
+
+        cursor.execute(query, params)
+        result = cursor.fetchone()[0]
+        return int(result) if result is not None else 0
+
+    def iter_all(self) -> Iterator[RequestMetrics]:
+        """Iterate over all stored metrics."""
+        conn = self._get_conn()
+        cursor = conn.cursor()
+
+        cursor.execute("SELECT * FROM requests ORDER BY timestamp")
+        for row in cursor:
+            yield self._row_to_metrics(row)
+
+    def get_summary_stats(
+        self,
+        start_time: datetime | None = None,
+        end_time: datetime | None = None,
+    ) -> dict[str, Any]:
+        """Get summary statistics."""
+        conn = self._get_conn()
+        cursor = conn.cursor()
+
+        where_clause = "WHERE 1=1"
+        params: list[Any] = []
+
+        if start_time is not None:
+            where_clause += " AND timestamp >= ?"
+            params.append(format_timestamp(start_time))
+        if end_time is not None:
+            where_clause += " AND timestamp <= ?"
+            params.append(format_timestamp(end_time))
+
+        cursor.execute(
+            f"""
+            SELECT
+                COUNT(*) as total_requests,
+                SUM(tokens_input_before) as total_tokens_before,
+                SUM(tokens_input_after) as total_tokens_after,
+                SUM(tokens_input_before - tokens_input_after) as total_tokens_saved,
+                AVG(tokens_input_before - tokens_input_after) as avg_tokens_saved,
+                AVG(cache_alignment_score) as avg_cache_alignment,
+                SUM(CASE WHEN mode = 'audit' THEN 1 ELSE 0 END) as audit_count,
+                SUM(CASE WHEN mode = 'optimize' THEN 1 ELSE 0 END) as optimize_count
+            FROM requests
+            {where_clause}
+            """,
+            params,
+        )
+
+        row = cursor.fetchone()
+
+        return {
+            "total_requests": row[0] or 0,
+            "total_tokens_before": row[1] or 0,
+            "total_tokens_after": row[2] or 0,
+            "total_tokens_saved": row[3] or 0,
+            "avg_tokens_saved": row[4] or 0,
+            "avg_cache_alignment": row[5] or 0,
+            "audit_count": row[6] or 0,
+            "optimize_count": row[7] or 0,
+        }
+
+    def _row_to_metrics(self, row: sqlite3.Row) -> RequestMetrics:
+        """Convert database row to RequestMetrics."""
+        return RequestMetrics(
+            request_id=row["id"],
+            timestamp=parse_timestamp(row["timestamp"]),
+            model=row["model"],
+            stream=bool(row["stream"]),
+            mode=row["mode"],
+            tokens_input_before=row["tokens_input_before"],
+            tokens_input_after=row["tokens_input_after"],
+            tokens_output=row["tokens_output"],
+            block_breakdown=json.loads(row["block_breakdown"]),
+            waste_signals=json.loads(row["waste_signals"]),
+            stable_prefix_hash=row["stable_prefix_hash"] or "",
+            cache_alignment_score=row["cache_alignment_score"] or 0.0,
+            cached_tokens=row["cached_tokens"],
+            transforms_applied=json.loads(row["transforms_applied"]),
+            tool_units_dropped=row["tool_units_dropped"] or 0,
+            turns_dropped=row["turns_dropped"] or 0,
+            messages_hash=row["messages_hash"] or "",
+            error=row["error"],
+        )
+
+    def close(self) -> None:
+        """Close database connection."""
+        if self._conn is not None:
+            self._conn.close()
+            self._conn = None