local-deep-research 0.4.4__py3-none-any.whl → 0.5.2__py3-none-any.whl

local_deep_research/metrics/db_models.py

@@ -0,0 +1,115 @@
+"""SQLAlchemy models for metrics."""
+
+from sqlalchemy import (
+    Column,
+    DateTime,
+    Integer,
+    String,
+    Text,
+    UniqueConstraint,
+    func,
+)
+from sqlalchemy.ext.declarative import declarative_base
+
+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
+    )  # Removed foreign key constraint to fix token tracking
+    model_name = Column(String)
+    provider = Column(
+        String
+    )  # Added provider column for accurate cost tracking
+    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
+    )  # Removed foreign key constraint to fix token tracking
+    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())
+
+
+class ResearchRating(Base):
+    """Model for storing user ratings of research sessions."""
+
+    __tablename__ = "research_ratings"
+
+    id = Column(Integer, primary_key=True)
+    research_id = Column(Integer, unique=True)  # References research session ID
+    rating = Column(Integer)  # 1-5 star rating
+    rated_at = Column(DateTime, server_default=func.now())
+    updated_at = Column(
+        DateTime, server_default=func.now(), onupdate=func.now()
+    )
+
+
+class SearchCall(Base):
+    """Model for individual search engine calls."""
+
+    __tablename__ = "search_calls"
+
+    id = Column(Integer, primary_key=True)
+    research_id = Column(Integer)
+    research_query = Column(Text)
+    research_mode = Column(String)
+    research_phase = Column(String)
+    search_iteration = Column(Integer)
+
+    # Search details
+    search_engine = Column(String)
+    query = Column(Text)
+    results_count = Column(Integer)
+    response_time_ms = Column(Integer)
+
+    # Status tracking
+    success_status = Column(String, default="success")
+    error_type = Column(String)
+    error_message = Column(Text)
+
+    timestamp = Column(DateTime, server_default=func.now())

local_deep_research/metrics/migrate_add_provider_to_token_usage.py

@@ -0,0 +1,148 @@
+"""
+Migration: Add provider column to TokenUsage table
+
+This migration adds the provider column to the TokenUsage table to enable
+accurate cost tracking based on both model and provider information.
+"""
+
+import logging
+from pathlib import Path
+
+from sqlalchemy import text
+
+from .database import get_metrics_db
+
+logger = logging.getLogger(__name__)
+
+
+def add_provider_column_to_token_usage():
+    """Add provider column to TokenUsage table."""
+    try:
+        db = get_metrics_db()
+
+        with db.get_session() as session:
+            # Check if provider column already exists
+            result = session.execute(
+                text(
+                    """
+                SELECT COUNT(*) as count
+                FROM pragma_table_info('token_usage')
+                WHERE name='provider'
+            """
+                )
+            )
+
+            provider_exists = result.fetchone()[0] > 0
+
+            if provider_exists:
+                logger.info(
+                    "Provider column already exists in token_usage table"
+                )
+                return True
+
+            logger.info("Adding provider column to token_usage table...")
+
+            # Add the provider column
+            session.execute(
+                text(
+                    """
+                ALTER TABLE token_usage
+                ADD COLUMN provider VARCHAR
+            """
+                )
+            )
+
+            # Try to populate provider info for existing records based on model name patterns
+            logger.info("Populating provider info for existing records...")
+
+            # Update known local model providers
+            local_model_updates = [
+                (
+                    "ollama",
+                    [
+                        "ollama",
+                        "llama",
+                        "mistral",
+                        "gemma",
+                        "qwen",
+                        "codellama",
+                        "vicuna",
+                        "alpaca",
+                    ],
+                ),
+                ("openai", ["gpt-", "davinci", "curie", "babbage", "ada"]),
+                ("anthropic", ["claude"]),
+                ("google", ["gemini", "bard"]),
+            ]
+
+            for provider, model_patterns in local_model_updates:
+                for pattern in model_patterns:
+                    session.execute(
+                        text(
+                            """
+                        UPDATE token_usage
+                        SET provider = :provider
+                        WHERE provider IS NULL
+                        AND (LOWER(model_name) LIKE :pattern OR LOWER(model_name) LIKE :pattern_percent)
+                    """
+                        ),
+                        {
+                            "provider": provider,
+                            "pattern": pattern,
+                            "pattern_percent": f"%{pattern}%",
+                        },
+                    )
+
+            # Set any remaining NULL providers to 'unknown'
+            session.execute(
+                text(
+                    """
+                UPDATE token_usage
+                SET provider = 'unknown'
+                WHERE provider IS NULL
+            """
+                )
+            )
+
+            session.commit()
+            logger.info(
+                "Successfully added provider column and populated existing data"
+            )
+            return True
+
+    except Exception as e:
+        logger.error(f"Error adding provider column to token_usage: {e}")
+        return False
+
+
+def run_migration():
+    """Run the provider column migration."""
+    logger.info("Starting migration: Add provider column to TokenUsage")
+
+    success = add_provider_column_to_token_usage()
+
+    if success:
+        logger.info("Migration completed successfully")
+    else:
+        logger.error("Migration failed")
+
+    return success
+
+
+if __name__ == "__main__":
+    # Allow running migration directly
+    import sys
+
+    # Add the project root to the path
+    project_root = Path(__file__).parent.parent.parent.parent
+    sys.path.insert(0, str(project_root))
+
+    logging.basicConfig(level=logging.INFO)
+    success = run_migration()
+
+    if success:
+        print("✅ Migration completed successfully")
+        sys.exit(0)
+    else:
+        print("❌ Migration failed")
+        sys.exit(1)

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