router-maestro 0.1.2__py3-none-any.whl

router_maestro/stats/heatmap.py

@@ -0,0 +1,154 @@
+"""Terminal heatmap generation for token usage."""
+
+from datetime import datetime, timedelta
+
+import plotext as plt
+from rich.console import Console
+from rich.table import Table
+
+from router_maestro.stats.storage import StatsStorage
+
+console = Console()
+
+
+def generate_heatmap(days: int = 7, provider: str | None = None, model: str | None = None) -> None:
+    """Generate and display a terminal heatmap of token usage.
+
+    Args:
+        days: Number of days to show
+        provider: Filter by provider (optional)
+        model: Filter by model (optional)
+    """
+    storage = StatsStorage()
+    hourly_data = storage.get_usage_by_hour(days=days, provider=provider, model=model)
+
+    if not hourly_data:
+        console.print("[dim]No usage data available for the specified period.[/dim]")
+        return
+
+    # Build a matrix for the heatmap: rows = days, columns = hours
+    # Initialize with zeros
+    today = datetime.now().date()
+    dates = [(today - timedelta(days=i)) for i in range(days - 1, -1, -1)]
+
+    # Create a 2D matrix (days x 24 hours)
+    matrix = [[0 for _ in range(24)] for _ in range(len(dates))]
+    date_to_idx = {d: i for i, d in enumerate(dates)}
+
+    for record in hourly_data:
+        date = datetime.fromisoformat(record["date"]).date()
+        hour = record["hour"]
+        tokens = record["total_tokens"]
+
+        if date in date_to_idx:
+            matrix[date_to_idx[date]][hour] = tokens
+
+    # Display using plotext
+    plt.clear_figure()
+    plt.title(f"Token Usage Heatmap (Last {days} Days)")
+
+    # Create a simple bar chart by day since plotext doesn't have heatmap
+    daily_data = storage.get_usage_by_day(days=days, provider=provider, model=model)
+
+    if daily_data:
+        dates_str = [record["date"] for record in daily_data]
+        tokens = [record["total_tokens"] for record in daily_data]
+
+        plt.bar(dates_str, tokens)
+        plt.xlabel("Date")
+        plt.ylabel("Total Tokens")
+        plt.show()
+
+    # Also show a text-based heatmap using Rich
+    _display_text_heatmap(dates, matrix)
+
+
+def _display_text_heatmap(dates: list, matrix: list[list[int]]) -> None:
+    """Display a text-based heatmap using Rich.
+
+    Args:
+        dates: List of dates
+        matrix: 2D matrix of token counts (days x hours)
+    """
+    console.print("\n[bold]Hourly Activity Heatmap:[/bold]")
+
+    # Find max value for scaling
+    max_val = max(max(row) for row in matrix) if matrix else 1
+    if max_val == 0:
+        max_val = 1
+
+    # Create intensity characters
+    intensity_chars = " ░▒▓█"
+
+    # Build the heatmap
+    hour_labels = "    " + "".join(f"{h:2d}" for h in range(0, 24, 2))
+    console.print(f"[dim]{hour_labels}[/dim]")
+
+    for i, date in enumerate(dates):
+        row_str = f"{date.strftime('%m/%d')} "
+        for h in range(24):
+            value = matrix[i][h]
+            intensity = int((value / max_val) * (len(intensity_chars) - 1))
+            char = intensity_chars[intensity]
+            row_str += char
+        console.print(row_str)
+
+    # Legend
+    console.print(f"\n[dim]Legend: {' '.join(intensity_chars)} (low to high)[/dim]")
+
+
+def display_stats_summary(days: int = 7) -> None:
+    """Display a summary of token usage statistics.
+
+    Args:
+        days: Number of days to summarize
+    """
+    storage = StatsStorage()
+    total = storage.get_total_usage(days=days)
+    by_model = storage.get_usage_by_model(days=days)
+
+    if not total or total.get("total_tokens") is None:
+        console.print("[dim]No usage data available.[/dim]")
+        return
+
+    # Summary table
+    console.print(f"\n[bold]Token Usage Summary (Last {days} Days)[/bold]\n")
+
+    summary_table = Table(show_header=False, box=None)
+    summary_table.add_column("Metric", style="cyan")
+    summary_table.add_column("Value", style="green", justify="right")
+
+    summary_table.add_row("Total Requests", f"{total.get('request_count', 0):,}")
+    summary_table.add_row("Successful", f"{total.get('success_count', 0):,}")
+    summary_table.add_row("Total Tokens", f"{total.get('total_tokens', 0):,}")
+    summary_table.add_row("  Prompt", f"{total.get('prompt_tokens', 0):,}")
+    summary_table.add_row("  Completion", f"{total.get('completion_tokens', 0):,}")
+
+    if total.get("avg_latency_ms"):
+        summary_table.add_row("Avg Latency", f"{total.get('avg_latency_ms', 0):.0f} ms")
+
+    console.print(summary_table)
+
+    # By model table
+    if by_model:
+        console.print("\n[bold]Usage by Model[/bold]\n")
+
+        model_table = Table()
+        model_table.add_column("Model", style="cyan")
+        model_table.add_column("Provider", style="magenta")
+        model_table.add_column("Requests", justify="right")
+        model_table.add_column("Total Tokens", justify="right", style="green")
+        model_table.add_column("Avg Latency", justify="right")
+
+        for record in by_model:
+            avg_latency = record.get("avg_latency_ms")
+            latency = f"{avg_latency:.0f} ms" if avg_latency else "-"
+            model_table.add_row(
+                record["model"],
+                record["provider"],
+                f"{record['request_count']:,}",
+                f"{record['total_tokens']:,}",
+                latency,
+            )
+
+        console.print(model_table)

router_maestro/stats/storage.py

@@ -0,0 +1,228 @@
+"""SQLite storage for token usage statistics."""
+
+import sqlite3
+from collections.abc import Iterator
+from contextlib import contextmanager
+from datetime import datetime
+from pathlib import Path
+
+from pydantic import BaseModel
+
+from router_maestro.config import STATS_DB_FILE
+
+
+class UsageRecord(BaseModel):
+    """A single token usage record."""
+
+    id: int | None = None
+    timestamp: datetime
+    provider: str
+    model: str
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+    success: bool
+    latency_ms: int | None = None
+
+
+class StatsStorage:
+    """SQLite storage for token usage statistics."""
+
+    def __init__(self, db_path: Path = STATS_DB_FILE) -> None:
+        self.db_path = db_path
+        self._init_db()
+
+    def _init_db(self) -> None:
+        """Initialize the database schema."""
+        self.db_path.parent.mkdir(parents=True, exist_ok=True)
+
+        with self._get_connection() as conn:
+            conn.execute("""
+                CREATE TABLE IF NOT EXISTS usage (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    timestamp TEXT NOT NULL,
+                    provider TEXT NOT NULL,
+                    model TEXT NOT NULL,
+                    prompt_tokens INTEGER NOT NULL,
+                    completion_tokens INTEGER NOT NULL,
+                    total_tokens INTEGER NOT NULL,
+                    success INTEGER NOT NULL,
+                    latency_ms INTEGER
+                )
+            """)
+            conn.execute("""
+                CREATE INDEX IF NOT EXISTS idx_usage_timestamp ON usage(timestamp)
+            """)
+            conn.execute("""
+                CREATE INDEX IF NOT EXISTS idx_usage_provider ON usage(provider)
+            """)
+            conn.execute("""
+                CREATE INDEX IF NOT EXISTS idx_usage_model ON usage(model)
+            """)
+            conn.commit()
+
+    @contextmanager
+    def _get_connection(self) -> Iterator[sqlite3.Connection]:
+        """Get a database connection."""
+        conn = sqlite3.connect(self.db_path)
+        conn.row_factory = sqlite3.Row
+        try:
+            yield conn
+        finally:
+            conn.close()
+
+    def record(self, record: UsageRecord) -> None:
+        """Record a usage event."""
+        with self._get_connection() as conn:
+            conn.execute(
+                """
+                INSERT INTO usage (
+                    timestamp, provider, model,
+                    prompt_tokens, completion_tokens, total_tokens,
+                    success, latency_ms
+                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+                """,
+                (
+                    record.timestamp.isoformat(),
+                    record.provider,
+                    record.model,
+                    record.prompt_tokens,
+                    record.completion_tokens,
+                    record.total_tokens,
+                    1 if record.success else 0,
+                    record.latency_ms,
+                ),
+            )
+            conn.commit()
+
+    def get_usage_by_day(
+        self, days: int = 7, provider: str | None = None, model: str | None = None
+    ) -> list[dict]:
+        """Get usage aggregated by day.
+
+        Args:
+            days: Number of days to look back
+            provider: Filter by provider (optional)
+            model: Filter by model (optional)
+
+        Returns:
+            List of dicts with date, total_tokens, request_count
+        """
+        with self._get_connection() as conn:
+            query = """
+                SELECT
+                    DATE(timestamp) as date,
+                    SUM(total_tokens) as total_tokens,
+                    SUM(prompt_tokens) as prompt_tokens,
+                    SUM(completion_tokens) as completion_tokens,
+                    COUNT(*) as request_count,
+                    SUM(CASE WHEN success = 1 THEN 1 ELSE 0 END) as success_count
+                FROM usage
+                WHERE timestamp >= datetime('now', ?)
+            """
+            params: list = [f"-{days} days"]
+
+            if provider:
+                query += " AND provider = ?"
+                params.append(provider)
+            if model:
+                query += " AND model = ?"
+                params.append(model)
+
+            query += " GROUP BY DATE(timestamp) ORDER BY date"
+
+            cursor = conn.execute(query, params)
+            return [dict(row) for row in cursor.fetchall()]
+
+    def get_usage_by_hour(
+        self, days: int = 7, provider: str | None = None, model: str | None = None
+    ) -> list[dict]:
+        """Get usage aggregated by hour.
+
+        Args:
+            days: Number of days to look back
+            provider: Filter by provider (optional)
+            model: Filter by model (optional)
+
+        Returns:
+            List of dicts with date, hour, total_tokens, request_count
+        """
+        with self._get_connection() as conn:
+            query = """
+                SELECT
+                    DATE(timestamp) as date,
+                    CAST(strftime('%H', timestamp) AS INTEGER) as hour,
+                    SUM(total_tokens) as total_tokens,
+                    COUNT(*) as request_count
+                FROM usage
+                WHERE timestamp >= datetime('now', ?)
+            """
+            params: list = [f"-{days} days"]
+
+            if provider:
+                query += " AND provider = ?"
+                params.append(provider)
+            if model:
+                query += " AND model = ?"
+                params.append(model)
+
+            query += " GROUP BY DATE(timestamp), hour ORDER BY date, hour"
+
+            cursor = conn.execute(query, params)
+            return [dict(row) for row in cursor.fetchall()]
+
+    def get_usage_by_model(self, days: int = 7) -> list[dict]:
+        """Get usage aggregated by model.
+
+        Args:
+            days: Number of days to look back
+
+        Returns:
+            List of dicts with model, provider, total_tokens, request_count
+        """
+        with self._get_connection() as conn:
+            cursor = conn.execute(
+                """
+                SELECT
+                    model,
+                    provider,
+                    SUM(total_tokens) as total_tokens,
+                    SUM(prompt_tokens) as prompt_tokens,
+                    SUM(completion_tokens) as completion_tokens,
+                    COUNT(*) as request_count,
+                    AVG(latency_ms) as avg_latency_ms
+                FROM usage
+                WHERE timestamp >= datetime('now', ?)
+                GROUP BY model, provider
+                ORDER BY total_tokens DESC
+                """,
+                (f"-{days} days",),
+            )
+            return [dict(row) for row in cursor.fetchall()]
+
+    def get_total_usage(self, days: int = 7) -> dict:
+        """Get total usage statistics.
+
+        Args:
+            days: Number of days to look back
+
+        Returns:
+            Dict with total statistics
+        """
+        with self._get_connection() as conn:
+            cursor = conn.execute(
+                """
+                SELECT
+                    SUM(total_tokens) as total_tokens,
+                    SUM(prompt_tokens) as prompt_tokens,
+                    SUM(completion_tokens) as completion_tokens,
+                    COUNT(*) as request_count,
+                    SUM(CASE WHEN success = 1 THEN 1 ELSE 0 END) as success_count,
+                    AVG(latency_ms) as avg_latency_ms
+                FROM usage
+                WHERE timestamp >= datetime('now', ?)
+                """,
+                (f"-{days} days",),
+            )
+            row = cursor.fetchone()
+            return dict(row) if row else {}

router_maestro/stats/tracker.py

@@ -0,0 +1,73 @@
+"""Token usage tracker."""
+
+import time
+from datetime import datetime
+
+from router_maestro.stats.storage import StatsStorage, UsageRecord
+
+
+class UsageTracker:
+    """Tracks token usage for requests."""
+
+    _instance: "UsageTracker | None" = None
+
+    def __init__(self) -> None:
+        self.storage = StatsStorage()
+
+    @classmethod
+    def get_instance(cls) -> "UsageTracker":
+        """Get the singleton instance."""
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    def record(
+        self,
+        provider: str,
+        model: str,
+        prompt_tokens: int,
+        completion_tokens: int,
+        success: bool = True,
+        latency_ms: int | None = None,
+    ) -> None:
+        """Record a usage event.
+
+        Args:
+            provider: Provider name
+            model: Model name
+            prompt_tokens: Number of prompt tokens
+            completion_tokens: Number of completion tokens
+            success: Whether the request was successful
+            latency_ms: Latency in milliseconds
+        """
+        record = UsageRecord(
+            timestamp=datetime.now(),
+            provider=provider,
+            model=model,
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=prompt_tokens + completion_tokens,
+            success=success,
+            latency_ms=latency_ms,
+        )
+        self.storage.record(record)
+
+
+class RequestTimer:
+    """Context manager for timing requests."""
+
+    def __init__(self) -> None:
+        self.start_time: float = 0
+        self.end_time: float = 0
+
+    def __enter__(self) -> "RequestTimer":
+        self.start_time = time.perf_counter()
+        return self
+
+    def __exit__(self, *args) -> None:
+        self.end_time = time.perf_counter()
+
+    @property
+    def elapsed_ms(self) -> int:
+        """Get elapsed time in milliseconds."""
+        return int((self.end_time - self.start_time) * 1000)

router_maestro/utils/__init__.py

@@ -0,0 +1,16 @@
+"""Utils module for router-maestro."""
+
+from router_maestro.utils.logging import get_logger, setup_logging
+from router_maestro.utils.tokens import (
+    estimate_tokens,
+    estimate_tokens_from_char_count,
+    map_openai_stop_reason_to_anthropic,
+)
+
+__all__ = [
+    "get_logger",
+    "setup_logging",
+    "estimate_tokens",
+    "estimate_tokens_from_char_count",
+    "map_openai_stop_reason_to_anthropic",
+]

router_maestro/utils/logging.py

@@ -0,0 +1,81 @@
+"""Unified logging configuration for router-maestro."""
+
+import logging
+from logging.handlers import RotatingFileHandler
+from pathlib import Path
+
+from rich.logging import RichHandler
+
+from router_maestro.config.paths import get_data_dir
+
+
+def get_log_dir() -> Path:
+    """Get the log directory ~/.local/share/router-maestro/logs/."""
+    log_dir = get_data_dir() / "logs"
+    log_dir.mkdir(parents=True, exist_ok=True)
+    return log_dir
+
+
+def setup_logging(
+    level: str = "INFO",
+    console: bool = True,
+    file: bool = True,
+) -> None:
+    """Configure unified logging with console (Rich) and file output.
+
+    Args:
+        level: Log level (DEBUG, INFO, WARNING, ERROR)
+        console: Enable Rich console handler with colors
+        file: Enable rotating file handler
+    """
+    # Get the root logger for router_maestro
+    logger = logging.getLogger("router_maestro")
+    logger.setLevel(getattr(logging, level.upper(), logging.INFO))
+
+    # Remove existing handlers to avoid duplicates on reconfig
+    logger.handlers.clear()
+
+    # File format (no colors)
+    file_formatter = logging.Formatter(
+        "%(asctime)s | %(levelname)-8s | %(name)s | %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+
+    # Console handler (Rich with colors)
+    if console:
+        console_handler = RichHandler(
+            show_time=True,
+            show_path=False,
+            rich_tracebacks=True,
+            markup=True,
+        )
+        console_handler.setLevel(getattr(logging, level.upper(), logging.INFO))
+        logger.addHandler(console_handler)
+
+    # File handler (rotating, 10MB max, 5 backups)
+    if file:
+        log_file = get_log_dir() / "router-maestro.log"
+        file_handler = RotatingFileHandler(
+            log_file,
+            maxBytes=10 * 1024 * 1024,  # 10MB
+            backupCount=5,
+            encoding="utf-8",
+        )
+        file_handler.setLevel(getattr(logging, level.upper(), logging.INFO))
+        file_handler.setFormatter(file_formatter)
+        logger.addHandler(file_handler)
+
+    # Prevent propagation to root logger
+    logger.propagate = False
+
+
+def get_logger(name: str) -> logging.Logger:
+    """Get a logger for a specific module.
+
+    Args:
+        name: Module name (will be prefixed with 'router_maestro.')
+
+    Returns:
+        Logger instance for the module
+    """
+    return logging.getLogger(f"router_maestro.{name}")

router_maestro/utils/tokens.py

@@ -0,0 +1,51 @@
+"""Token estimation utilities."""
+
+# Approximate characters per token for English text
+CHARS_PER_TOKEN = 4
+
+
+def estimate_tokens(text: str) -> int:
+    """Estimate token count from text.
+
+    Uses a rough approximation of ~4 characters per token for English text.
+    This provides an estimate for context display before actual usage is known.
+
+    Args:
+        text: The text to estimate tokens for
+
+    Returns:
+        Estimated token count
+    """
+    return len(text) // CHARS_PER_TOKEN
+
+
+def estimate_tokens_from_char_count(char_count: int) -> int:
+    """Estimate token count from character count.
+
+    Args:
+        char_count: Number of characters
+
+    Returns:
+        Estimated token count
+    """
+    return char_count // CHARS_PER_TOKEN
+
+
+def map_openai_stop_reason_to_anthropic(openai_reason: str | None) -> str | None:
+    """Map OpenAI finish reason to Anthropic stop reason.
+
+    Args:
+        openai_reason: OpenAI finish reason (stop, length, tool_calls, content_filter)
+
+    Returns:
+        Anthropic stop reason (end_turn, max_tokens, tool_use)
+    """
+    if openai_reason is None:
+        return None
+    mapping = {
+        "stop": "end_turn",
+        "length": "max_tokens",
+        "tool_calls": "tool_use",
+        "content_filter": "end_turn",
+    }
+    return mapping.get(openai_reason, "end_turn")