learningcircuit--local-deep-research/tests/database/test_metrics_models.py

"""Tests for metrics tracking database models."""

from datetime import datetime, timedelta, timezone

import pytest
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from local_deep_research.database.models import (
    Base,
    ModelUsage,
    ResearchRating,
    SearchCall,
    TokenUsage,
)


class TestMetricsModels:
    """Test suite for metrics tracking models."""

    @pytest.fixture
    def engine(self):
        """Create an in-memory SQLite database for testing."""
        engine = create_engine("sqlite:///:memory:")
        Base.metadata.create_all(engine)
        yield engine
        engine.dispose()

    @pytest.fixture
    def session(self, engine):
        """Create a database session for testing."""
        Session = sessionmaker(bind=engine)
        session = Session()
        yield session
        session.close()

    def test_token_usage_tracking(self, session):
        """Test TokenUsage model for tracking LLM token consumption."""
        usage = TokenUsage(
            research_id="research-123",
            model_provider="openai",
            model_name="gpt-4",
            prompt_tokens=500,
            completion_tokens=150,
            total_tokens=650,
            prompt_cost=0.015,
            completion_cost=0.0045,
            total_cost=0.0195,
            timestamp=datetime.now(timezone.utc),
            operation_type="synthesis",
            operation_details={
                "temperature": 0.7,
                "purpose": "synthesis",
                "request_id": "req_abc123",
            },
        )

        session.add(usage)
        session.commit()

        # Verify the usage record
        saved = session.query(TokenUsage).first()
        assert saved is not None
        assert saved.model_provider == "openai"
        assert saved.model_name == "gpt-4"
        assert saved.total_tokens == 650
        assert saved.total_cost == 0.0195
        assert saved.operation_type == "synthesis"
        assert saved.operation_details["purpose"] == "synthesis"

    def test_model_usage_aggregation(self, session):
        """Test ModelUsage for aggregating model usage statistics."""
        model_usage = ModelUsage(
            model_provider="anthropic",
            model_name="claude-3-opus",
            total_calls=5,
            total_tokens=1450,
            total_cost=0.10,
            avg_response_time_ms=250.5,
            error_count=0,
            success_rate=100.0,
            first_used_at=datetime.now(timezone.utc),
            last_used_at=datetime.now(timezone.utc),
        )

        session.add(model_usage)
        session.commit()

        # Verify aggregated stats
        saved = session.query(ModelUsage).first()
        assert saved is not None
        assert saved.model_provider == "anthropic"
        assert saved.model_name == "claude-3-opus"
        assert saved.total_calls == 5
        assert saved.total_tokens == 1450
        assert saved.total_cost == 0.10
        assert saved.success_rate == 100.0

    def test_research_rating(self, session):
        """Test ResearchRating model for user feedback."""
        rating = ResearchRating(
            research_id="research-456",
            rating=4,
            accuracy=5,
            completeness=4,
            relevance=5,
            readability=3,
            feedback="Great research results, but the summary could be clearer.",
            created_at=datetime.now(timezone.utc),
        )

        session.add(rating)
        session.commit()

        # Verify rating
        saved = session.query(ResearchRating).first()
        assert saved is not None
        assert saved.rating == 4
        assert saved.accuracy == 5
        assert saved.relevance == 5
        assert "summary could be clearer" in saved.feedback

    def test_search_call_tracking(self, session):
        """Test SearchCall model for tracking search engine calls."""
        search = SearchCall(
            research_id="research-789",
            search_engine="google",
            query="quantum computing applications",
            num_results_requested=10,
            num_results_returned=10,
            response_time_ms=150.5,
            success=1,
            error_message=None,
            rate_limited=0,
            timestamp=datetime.now(timezone.utc),
        )

        session.add(search)
        session.commit()

        # Verify search call
        saved = session.query(SearchCall).first()
        assert saved is not None
        assert saved.search_engine == "google"
        assert saved.query == "quantum computing applications"
        assert saved.success == 1
        assert saved.response_time_ms == 150.5

    def test_metrics_relationships(self, session):
        """Test relationships between metrics models."""
        research_id = "research-shared-123"

        # Create related metrics for the same research
        token_usage = TokenUsage(
            research_id=research_id,
            model_provider="openai",
            model_name="gpt-4",
            prompt_tokens=100,
            completion_tokens=50,
            total_tokens=150,
            total_cost=0.0045,
        )

        search_call = SearchCall(
            research_id=research_id,
            search_engine="bing",
            query="test query",
            num_results_returned=5,
        )

        rating = ResearchRating(
            research_id=research_id, rating=5, feedback="Excellent"
        )

        session.add_all([token_usage, search_call, rating])
        session.commit()

        # Query by research_id
        tokens = (
            session.query(TokenUsage).filter_by(research_id=research_id).all()
        )
        searches = (
            session.query(SearchCall).filter_by(research_id=research_id).all()
        )
        ratings = (
            session.query(ResearchRating)
            .filter_by(research_id=research_id)
            .all()
        )

        assert len(tokens) == 1
        assert len(searches) == 1
        assert len(ratings) == 1

    def test_cost_tracking(self, session):
        """Test cost tracking across different models."""
        # Add multiple token usage records
        for i in range(3):
            usage = TokenUsage(
                research_id=f"research-cost-{i}",
                model_provider="openai",
                model_name="gpt-4",
                prompt_tokens=1000,
                completion_tokens=500,
                total_tokens=1500,
                prompt_cost=0.03,
                completion_cost=0.015,
                total_cost=0.045,
            )
            session.add(usage)

        session.commit()

        # Calculate total costs
        from sqlalchemy import func

        total_cost = session.query(func.sum(TokenUsage.total_cost)).scalar()
        assert total_cost == 0.135  # 3 * 0.045

    def test_search_engine_performance(self, session):
        """Test tracking search engine performance metrics."""
        engines = ["google", "bing", "duckduckgo"]

        for engine in engines:
            for i in range(5):
                search = SearchCall(
                    research_id=f"research-perf-{engine}-{i}",
                    search_engine=engine,
                    query=f"test query {i}",
                    num_results_requested=10,
                    num_results_returned=10 if i != 2 else 0,  # One failure
                    response_time_ms=100 + i * 50,
                    success=1 if i != 2 else 0,
                    error_message=None if i != 2 else "Network error",
                )
                session.add(search)

        session.commit()

        # Analyze performance by engine
        from sqlalchemy import func

        engine_stats = (
            session.query(
                SearchCall.search_engine,
                func.count(SearchCall.id).label("total_calls"),
                func.avg(SearchCall.response_time_ms).label(
                    "avg_response_time"
                ),
                func.sum(SearchCall.success).label("successful_calls"),
            )
            .group_by(SearchCall.search_engine)
            .all()
        )

        assert len(engine_stats) == 3
        for stat in engine_stats:
            assert stat.total_calls == 5
            assert stat.successful_calls == 4  # 4 out of 5 successful

    def test_rating_aggregation(self, session):
        """Test aggregating user ratings."""
        # Create multiple ratings
        for i in range(10):
            rating = ResearchRating(
                research_id=f"research-rate-{i}",
                rating=3 + (i % 3),  # Ratings: 3, 4, 5, 3, 4, 5...
                accuracy=4 if i % 2 == 0 else 5,
                completeness=3 + (i % 2),
                relevance=5,
                readability=4,
            )
            session.add(rating)

        session.commit()

        # Calculate average ratings
        from sqlalchemy import func

        avg_rating = session.query(func.avg(ResearchRating.rating)).scalar()
        avg_accuracy = session.query(func.avg(ResearchRating.accuracy)).scalar()

        assert avg_rating > 3.5
        assert avg_accuracy > 4.0

    def test_time_based_metrics(self, session):
        """Test querying metrics by time ranges."""
        now = datetime.now(timezone.utc)

        # Create token usage over different time periods
        for days_ago in range(7):
            for i in range(3):
                usage = TokenUsage(
                    research_id=f"research-time-{days_ago}-{i}",
                    model_provider="anthropic",
                    model_name="claude-3",
                    prompt_tokens=100,
                    completion_tokens=50,
                    total_tokens=150,
                    total_cost=0.005,
                    timestamp=now - timedelta(days=days_ago),
                )
                session.add(usage)

        session.commit()

        # Query last 3 days
        three_days_ago = now - timedelta(days=3)
        recent_usage = (
            session.query(TokenUsage)
            .filter(TokenUsage.timestamp >= three_days_ago)
            .count()
        )

        # Should have 3 days * 3 records per day = 9 records
        assert recent_usage == 12  # days 0, 1, 2, 3 = 4 days * 3 records