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