local-deep-research 0.4.3__py3-none-any.whl → 0.5.0__py3-none-any.whl

local_deep_research/metrics/migrate_call_stack_tracking.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+"""
+Database migration script to add call stack tracking columns to token_usage table.
+This adds the Phase 1 call stack tracking functionality.
+"""
+
+import sqlite3
+import sys
+from pathlib import Path
+
+from loguru import logger
+
+
+def migrate_call_stack_tracking(db_path: str):
+    """Add call stack tracking columns to the token_usage table.
+
+    Args:
+        db_path: Path to the SQLite database file
+    """
+    try:
+        conn = sqlite3.connect(db_path)
+        cursor = conn.cursor()
+
+        # Check if columns already exist
+        cursor.execute("PRAGMA table_info(token_usage)")
+        columns = [row[1] for row in cursor.fetchall()]
+
+        # Add call stack tracking columns if they don't exist
+        new_columns = [
+            ("calling_file", "TEXT"),
+            ("calling_function", "TEXT"),
+            ("call_stack", "TEXT"),
+        ]
+
+        for column_name, column_type in new_columns:
+            if column_name not in columns:
+                logger.info(f"Adding column {column_name} to token_usage table")
+                cursor.execute(
+                    f"ALTER TABLE token_usage ADD COLUMN {column_name} {column_type}"
+                )
+            else:
+                logger.info(
+                    f"Column {column_name} already exists in token_usage table"
+                )
+
+        conn.commit()
+        logger.success(
+            "Call stack tracking columns migration completed successfully"
+        )
+
+    except sqlite3.Error as e:
+        logger.error(f"Database error during call stack migration: {e}")
+        raise
+    except Exception as e:
+        logger.error(f"Unexpected error during call stack migration: {e}")
+        raise
+    finally:
+        if conn:
+            conn.close()
+
+
+def find_database_file():
+    """Find the metrics database file."""
+    # Common locations for the database
+    possible_paths = [
+        "data/metrics.db",
+        "../data/metrics.db",
+        "../../data/metrics.db",
+    ]
+
+    for path in possible_paths:
+        db_path = Path(path)
+        if db_path.exists():
+            return str(db_path.absolute())
+
+    return None
+
+
+if __name__ == "__main__":
+    logger.info("Starting call stack tracking migration...")
+
+    # Check if database path provided as argument
+    if len(sys.argv) > 1:
+        db_path = sys.argv[1]
+    else:
+        db_path = find_database_file()
+
+    if not db_path:
+        logger.error("Could not find metrics database file.")
+        logger.info("Please provide the database path as an argument:")
+        logger.info("python migrate_call_stack_tracking.py /path/to/metrics.db")
+        sys.exit(1)
+
+    if not Path(db_path).exists():
+        logger.error(f"Database file does not exist: {db_path}")
+        sys.exit(1)
+
+    logger.info(f"Using database: {db_path}")
+
+    try:
+        migrate_call_stack_tracking(db_path)
+        logger.success("Call stack tracking migration completed!")
+    except Exception as e:
+        logger.error(f"Migration failed: {e}")
+        sys.exit(1)

local_deep_research/metrics/migrate_enhanced_tracking.py

@@ -0,0 +1,75 @@
+"""Migration script to add Phase 1 enhanced token tracking fields."""
+
+import sqlite3
+from pathlib import Path
+
+from loguru import logger
+
+from ..utilities.db_utils import DB_PATH
+
+
+def migrate_enhanced_tracking():
+    """Add Phase 1 enhanced tracking columns to existing token_usage table."""
+
+    if not Path(DB_PATH).exists():
+        logger.info("Database doesn't exist yet, skipping migration")
+        return
+
+    try:
+        conn = sqlite3.connect(DB_PATH)
+        cursor = conn.cursor()
+
+        # Check if token_usage table exists
+        cursor.execute(
+            """
+            SELECT name FROM sqlite_master
+            WHERE type='table' AND name='token_usage'
+        """
+        )
+
+        if not cursor.fetchone():
+            logger.info(
+                "token_usage table doesn't exist yet, skipping migration"
+            )
+            conn.close()
+            return
+
+        # Check if enhanced columns already exist
+        cursor.execute("PRAGMA table_info(token_usage)")
+        columns = [column[1] for column in cursor.fetchall()]
+
+        # Define new columns to add
+        new_columns = [
+            ("research_query", "TEXT"),
+            ("research_mode", "TEXT"),
+            ("research_phase", "TEXT"),
+            ("search_iteration", "INTEGER"),
+            ("response_time_ms", "INTEGER"),
+            ("success_status", "TEXT DEFAULT 'success'"),
+            ("error_type", "TEXT"),
+            ("search_engines_planned", "TEXT"),
+            ("search_engine_selected", "TEXT"),
+        ]
+
+        # Add missing columns
+        for column_name, column_type in new_columns:
+            if column_name not in columns:
+                logger.info(f"Adding column {column_name} to token_usage table")
+                cursor.execute(
+                    f"ALTER TABLE token_usage ADD COLUMN {column_name} {column_type}"
+                )
+
+        conn.commit()
+        conn.close()
+
+        logger.info("Enhanced token tracking migration completed successfully")
+
+    except Exception as e:
+        logger.exception(f"Error during enhanced token tracking migration: {e}")
+        if "conn" in locals():
+            conn.close()
+        raise
+
+
+if __name__ == "__main__":
+    migrate_enhanced_tracking()

local_deep_research/metrics/migrate_research_ratings.py

@@ -0,0 +1,31 @@
+#!/usr/bin/env python3
+"""Migration script to add research ratings table."""
+
+import sys
+from pathlib import Path
+
+# Add the project root to Python path
+project_root = Path(__file__).parent.parent.parent.parent
+sys.path.insert(0, str(project_root))
+
+# Import after path modification
+from local_deep_research.metrics.database import MetricsDatabase  # noqa: E402
+from local_deep_research.metrics.db_models import ResearchRating  # noqa: E402
+
+
+def main():
+    """Run the migration to add research ratings table."""
+    print("Creating research ratings table...")
+
+    # Initialize database
+    db = MetricsDatabase()
+
+    # Create the research_ratings table
+    ResearchRating.__table__.create(db.engine, checkfirst=True)
+
+    print("✅ Research ratings table created successfully!")
+    print("Users can now rate their research sessions on a 1-5 star scale.")
+
+
+if __name__ == "__main__":
+    main()

local_deep_research/metrics/models.py

@@ -0,0 +1,61 @@
+"""SQLAlchemy models for metrics."""
+
+from sqlalchemy import Column, DateTime, Integer, String, Text, UniqueConstraint
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.sql import func
+
+Base = declarative_base()
+
+
+class TokenUsage(Base):
+    """Model for tracking individual token usage events."""
+
+    __tablename__ = "token_usage"
+
+    id = Column(Integer, primary_key=True)
+    research_id = Column(Integer, index=True)  # No foreign key for now
+    model_name = Column(String)
+    prompt_tokens = Column(Integer)
+    completion_tokens = Column(Integer)
+    total_tokens = Column(Integer)
+
+    # Phase 1 Enhancement: Research context
+    research_query = Column(Text)
+    research_mode = Column(String)  # 'quick' or 'detailed'
+    research_phase = Column(String)  # 'init', 'iteration_1', etc.
+    search_iteration = Column(Integer)
+
+    # Phase 1 Enhancement: Performance metrics
+    response_time_ms = Column(Integer)
+    success_status = Column(
+        String, default="success"
+    )  # 'success', 'error', 'timeout'
+    error_type = Column(String)
+
+    # Phase 1 Enhancement: Search engine context
+    search_engines_planned = Column(Text)  # JSON array as text
+    search_engine_selected = Column(String)
+
+    # Call stack tracking
+    calling_file = Column(String)  # File that made the LLM call
+    calling_function = Column(String)  # Function that made the LLM call
+    call_stack = Column(Text)  # Full call stack as JSON
+
+    timestamp = Column(DateTime, server_default=func.now())
+
+
+class ModelUsage(Base):
+    """Model for aggregated token usage by model and research."""
+
+    __tablename__ = "model_usage"
+    __table_args__ = (UniqueConstraint("research_id", "model_name"),)
+
+    id = Column(Integer, primary_key=True)
+    research_id = Column(Integer, index=True)  # No foreign key for now
+    model_name = Column(String)
+    provider = Column(String)
+    prompt_tokens = Column(Integer, default=0)
+    completion_tokens = Column(Integer, default=0)
+    total_tokens = Column(Integer, default=0)
+    calls = Column(Integer, default=0)
+    timestamp = Column(DateTime, server_default=func.now())

local_deep_research/metrics/pricing/__init__.py

@@ -0,0 +1,12 @@
+"""
+LLM Pricing API Module
+
+Provides real-time pricing data for LLM models from various providers.
+Includes caching and cost calculation utilities.
+"""
+
+from .cost_calculator import CostCalculator
+from .pricing_cache import PricingCache
+from .pricing_fetcher import PricingFetcher
+
+__all__ = ["PricingFetcher", "PricingCache", "CostCalculator"]

local_deep_research/metrics/pricing/cost_calculator.py

@@ -0,0 +1,237 @@
+"""
+Cost Calculator
+
+Calculates LLM usage costs based on token usage and pricing data.
+Integrates with pricing fetcher and cache systems.
+"""
+
+import logging
+from typing import Any, Dict, List, Optional
+
+from .pricing_cache import PricingCache
+from .pricing_fetcher import PricingFetcher
+
+logger = logging.getLogger(__name__)
+
+
+class CostCalculator:
+    """Calculates LLM usage costs."""
+
+    def __init__(self, cache_dir: Optional[str] = None):
+        self.cache = PricingCache(cache_dir)
+        self.pricing_fetcher = None
+
+    async def __aenter__(self):
+        self.pricing_fetcher = PricingFetcher()
+        await self.pricing_fetcher.__aenter__()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if self.pricing_fetcher:
+            await self.pricing_fetcher.__aexit__(exc_type, exc_val, exc_tb)
+
+    async def get_model_pricing(
+        self, model_name: str, provider: str = None
+    ) -> Dict[str, float]:
+        """Get pricing for a model and provider (cached or fetched)."""
+        # Create cache key that includes provider
+        cache_key = f"{provider}:{model_name}" if provider else model_name
+
+        # Try cache first
+        cached_pricing = self.cache.get(f"model:{cache_key}")
+        if cached_pricing:
+            return cached_pricing
+
+        # Fetch from API
+        if self.pricing_fetcher:
+            pricing = await self.pricing_fetcher.get_model_pricing(
+                model_name, provider
+            )
+            if pricing:
+                self.cache.set(f"model:{cache_key}", pricing)
+                return pricing
+
+        # No pricing found
+        logger.warning(
+            f"No pricing found for {model_name} (provider: {provider})"
+        )
+        return None
+
+    async def calculate_cost(
+        self,
+        model_name: str,
+        prompt_tokens: int,
+        completion_tokens: int,
+        provider: str = None,
+    ) -> Dict[str, float]:
+        """
+        Calculate cost for a single LLM call.
+
+        Returns:
+            Dict with prompt_cost, completion_cost, total_cost
+        """
+        pricing = await self.get_model_pricing(model_name, provider)
+
+        # If no pricing found, return zero cost
+        if pricing is None:
+            return {
+                "prompt_cost": 0.0,
+                "completion_cost": 0.0,
+                "total_cost": 0.0,
+                "pricing_used": None,
+                "error": "No pricing data available for this model",
+            }
+
+        # Convert tokens to thousands for pricing calculation
+        prompt_cost = (prompt_tokens / 1000) * pricing["prompt"]
+        completion_cost = (completion_tokens / 1000) * pricing["completion"]
+        total_cost = prompt_cost + completion_cost
+
+        return {
+            "prompt_cost": round(prompt_cost, 6),
+            "completion_cost": round(completion_cost, 6),
+            "total_cost": round(total_cost, 6),
+            "pricing_used": pricing,
+        }
+
+    async def calculate_batch_costs(
+        self, usage_records: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """
+        Calculate costs for multiple usage records.
+
+        Expected record format:
+        {
+            "model_name": str,
+            "provider": str (optional),
+            "prompt_tokens": int,
+            "completion_tokens": int,
+            "research_id": int (optional),
+            "timestamp": datetime (optional)
+        }
+        """
+        results = []
+
+        for record in usage_records:
+            try:
+                cost_data = await self.calculate_cost(
+                    record["model_name"],
+                    record["prompt_tokens"],
+                    record["completion_tokens"],
+                    record.get("provider"),
+                )
+
+                result = {**record, **cost_data}
+                results.append(result)
+
+            except Exception as e:
+                logger.error(
+                    f"Failed to calculate cost for record {record}: {e}"
+                )
+                # Add record with zero cost on error
+                results.append(
+                    {
+                        **record,
+                        "prompt_cost": 0.0,
+                        "completion_cost": 0.0,
+                        "total_cost": 0.0,
+                        "error": str(e),
+                    }
+                )
+
+        return results
+
+    def calculate_cost_sync(
+        self, model_name: str, prompt_tokens: int, completion_tokens: int
+    ) -> Dict[str, float]:
+        """
+        Synchronous cost calculation using cached pricing only.
+        Fallback for when async is not available.
+        """
+        # Use cached pricing only
+        pricing = self.cache.get_model_pricing(model_name)
+        if not pricing:
+            # Use static fallback with exact matching only
+            fetcher = PricingFetcher()
+            # Try exact match
+            pricing = fetcher.static_pricing.get(model_name)
+            if not pricing:
+                # Try exact match without provider prefix
+                if "/" in model_name:
+                    model_only = model_name.split("/")[-1]
+                    pricing = fetcher.static_pricing.get(model_only)
+
+        # If no pricing found, return zero cost
+        if not pricing:
+            return {
+                "prompt_cost": 0.0,
+                "completion_cost": 0.0,
+                "total_cost": 0.0,
+                "pricing_used": None,
+                "error": "No pricing data available for this model",
+            }
+
+        prompt_cost = (prompt_tokens / 1000) * pricing["prompt"]
+        completion_cost = (completion_tokens / 1000) * pricing["completion"]
+        total_cost = prompt_cost + completion_cost
+
+        return {
+            "prompt_cost": round(prompt_cost, 6),
+            "completion_cost": round(completion_cost, 6),
+            "total_cost": round(total_cost, 6),
+            "pricing_used": pricing,
+        }
+
+    async def get_research_cost_summary(
+        self, usage_records: List[Dict[str, Any]]
+    ) -> Dict[str, Any]:
+        """
+        Get cost summary for research session(s).
+        """
+        costs = await self.calculate_batch_costs(usage_records)
+
+        total_cost = sum(c["total_cost"] for c in costs)
+        total_prompt_cost = sum(c["prompt_cost"] for c in costs)
+        total_completion_cost = sum(c["completion_cost"] for c in costs)
+
+        total_prompt_tokens = sum(r["prompt_tokens"] for r in usage_records)
+        total_completion_tokens = sum(
+            r["completion_tokens"] for r in usage_records
+        )
+        total_tokens = total_prompt_tokens + total_completion_tokens
+
+        # Model breakdown
+        model_costs = {}
+        for cost in costs:
+            model = cost["model_name"]
+            if model not in model_costs:
+                model_costs[model] = {
+                    "total_cost": 0.0,
+                    "prompt_tokens": 0,
+                    "completion_tokens": 0,
+                    "calls": 0,
+                }
+
+            model_costs[model]["total_cost"] += cost["total_cost"]
+            model_costs[model]["prompt_tokens"] += cost["prompt_tokens"]
+            model_costs[model]["completion_tokens"] += cost["completion_tokens"]
+            model_costs[model]["calls"] += 1
+
+        return {
+            "total_cost": round(total_cost, 6),
+            "prompt_cost": round(total_prompt_cost, 6),
+            "completion_cost": round(total_completion_cost, 6),
+            "total_tokens": total_tokens,
+            "prompt_tokens": total_prompt_tokens,
+            "completion_tokens": total_completion_tokens,
+            "total_calls": len(usage_records),
+            "model_breakdown": model_costs,
+            "avg_cost_per_call": (
+                round(total_cost / len(usage_records), 6)
+                if usage_records
+                else 0.0
+            ),
+            "cost_per_token": (
+                round(total_cost / total_tokens, 8) if total_tokens > 0 else 0.0
+            ),
+        }

local_deep_research/metrics/pricing/pricing_cache.py

@@ -0,0 +1,143 @@
+"""
+Pricing Cache System
+
+Caches pricing data to avoid repeated API calls and improve performance.
+Includes cache expiration and refresh mechanisms.
+"""
+
+import json
+import logging
+import time
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class PricingCache:
+    """Cache for LLM pricing data."""
+
+    def __init__(self, cache_dir: Optional[str] = None, cache_ttl: int = 3600):
+        """
+        Initialize pricing cache.
+
+        Args:
+            cache_dir: Directory to store cache files
+            cache_ttl: Cache time-to-live in seconds (default: 1 hour)
+        """
+        self.cache_ttl = cache_ttl
+
+        if cache_dir:
+            self.cache_dir = Path(cache_dir)
+        else:
+            # Default to data directory
+            self.cache_dir = Path.cwd() / "data" / "cache" / "pricing"
+
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        self.cache_file = self.cache_dir / "pricing_cache.json"
+
+        self._cache = {}
+        self._load_cache()
+
+    def _load_cache(self):
+        """Load cache from disk."""
+        try:
+            if self.cache_file.exists():
+                with open(self.cache_file, "r") as f:
+                    data = json.load(f)
+                    self._cache = data.get("cache", {})
+                    logger.info(
+                        f"Loaded pricing cache with {len(self._cache)} entries"
+                    )
+        except Exception as e:
+            logger.warning(f"Failed to load pricing cache: {e}")
+            self._cache = {}
+
+    def _save_cache(self):
+        """Save cache to disk."""
+        try:
+            cache_data = {
+                "cache": self._cache,
+                "last_updated": datetime.now().isoformat(),
+            }
+            with open(self.cache_file, "w") as f:
+                json.dump(cache_data, f, indent=2)
+        except Exception as e:
+            logger.warning(f"Failed to save pricing cache: {e}")
+
+    def _is_expired(self, timestamp: float) -> bool:
+        """Check if cache entry is expired."""
+        return (time.time() - timestamp) > self.cache_ttl
+
+    def get(self, key: str) -> Optional[Any]:
+        """Get cached pricing data."""
+        if key not in self._cache:
+            return None
+
+        entry = self._cache[key]
+        if self._is_expired(entry["timestamp"]):
+            # Remove expired entry
+            del self._cache[key]
+            self._save_cache()
+            return None
+
+        return entry["data"]
+
+    def set(self, key: str, data: Any):
+        """Set cached pricing data."""
+        self._cache[key] = {"data": data, "timestamp": time.time()}
+        self._save_cache()
+
+    def get_model_pricing(self, model_name: str) -> Optional[Dict[str, float]]:
+        """Get cached pricing for a specific model."""
+        return self.get(f"model:{model_name}")
+
+    def set_model_pricing(self, model_name: str, pricing: Dict[str, float]):
+        """Cache pricing for a specific model."""
+        self.set(f"model:{model_name}", pricing)
+
+    def get_all_pricing(self) -> Optional[Dict[str, Dict[str, float]]]:
+        """Get cached pricing for all models."""
+        return self.get("all_models")
+
+    def set_all_pricing(self, pricing: Dict[str, Dict[str, float]]):
+        """Cache pricing for all models."""
+        self.set("all_models", pricing)
+
+    def clear(self):
+        """Clear all cached data."""
+        self._cache = {}
+        self._save_cache()
+        logger.info("Pricing cache cleared")
+
+    def clear_expired(self):
+        """Remove expired cache entries."""
+        expired_keys = []
+        for key, entry in self._cache.items():
+            if self._is_expired(entry["timestamp"]):
+                expired_keys.append(key)
+
+        for key in expired_keys:
+            del self._cache[key]
+
+        if expired_keys:
+            self._save_cache()
+            logger.info(f"Removed {len(expired_keys)} expired cache entries")
+
+    def get_cache_stats(self) -> Dict[str, Any]:
+        """Get cache statistics."""
+        total_entries = len(self._cache)
+        expired_count = 0
+
+        for entry in self._cache.values():
+            if self._is_expired(entry["timestamp"]):
+                expired_count += 1
+
+        return {
+            "total_entries": total_entries,
+            "expired_entries": expired_count,
+            "valid_entries": total_entries - expired_count,
+            "cache_file": str(self.cache_file),
+            "cache_ttl": self.cache_ttl,
+        }