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

local_deep_research/metrics/search_tracker.py

@@ -0,0 +1,380 @@
+"""
+Search call tracking system for metrics collection.
+Similar to token_counter.py but tracks search engine usage.
+"""
+
+from typing import Any, Dict, List, Optional
+
+from loguru import logger
+from sqlalchemy import case, func
+
+from .database import MetricsDatabase
+from .db_models import SearchCall
+from .query_utils import get_research_mode_condition, get_time_filter_condition
+
+
+class SearchTracker:
+    """Track search engine calls and performance metrics."""
+
+    def __init__(self, db: Optional[MetricsDatabase] = None):
+        """Initialize the search tracker."""
+        self.db = db or MetricsDatabase()
+        self.research_context = {}
+
+    def set_research_context(self, context: Dict[str, Any]) -> None:
+        """Set the current research context for search tracking."""
+        self.research_context = context or {}
+        logger.debug(f"Search tracker context updated: {self.research_context}")
+
+    def record_search(
+        self,
+        engine_name: str,
+        query: str,
+        results_count: int = 0,
+        response_time_ms: int = 0,
+        success: bool = True,
+        error_message: Optional[str] = None,
+    ) -> None:
+        """Record a completed search operation directly to database."""
+
+        # Extract research context
+        research_id = self.research_context.get("research_id")
+        research_query = self.research_context.get("research_query")
+        research_mode = self.research_context.get("research_mode", "unknown")
+        research_phase = self.research_context.get("research_phase", "search")
+        search_iteration = self.research_context.get("search_iteration", 0)
+
+        # Determine success status
+        success_status = "success" if success else "error"
+        error_type = None
+        if error_message:
+            error_type = (
+                type(error_message).__name__
+                if isinstance(error_message, Exception)
+                else "unknown_error"
+            )
+
+        # Record search call in database
+        try:
+            with self.db.get_session() as session:
+                # Create search call record
+                search_call = SearchCall(
+                    research_id=research_id,
+                    research_query=research_query,
+                    research_mode=research_mode,
+                    research_phase=research_phase,
+                    search_iteration=search_iteration,
+                    search_engine=engine_name,
+                    query=query,
+                    results_count=results_count,
+                    response_time_ms=response_time_ms,
+                    success_status=success_status,
+                    error_type=error_type,
+                    error_message=str(error_message) if error_message else None,
+                )
+                session.add(search_call)
+                session.commit()
+
+                logger.debug(
+                    f"Search call recorded: {engine_name} - "
+                    f"{results_count} results in {response_time_ms}ms"
+                )
+
+        except Exception as e:
+            logger.error(f"Failed to record search call: {e}")
+
+    def get_search_metrics(
+        self, period: str = "30d", research_mode: str = "all"
+    ) -> Dict[str, Any]:
+        """Get search engine usage metrics."""
+        with self.db.get_session() as session:
+            try:
+                # Build base query with filters
+                query = session.query(SearchCall).filter(
+                    SearchCall.search_engine.isnot(None)
+                )
+
+                # Apply time filter
+                time_condition = get_time_filter_condition(
+                    period, SearchCall.timestamp
+                )
+                if time_condition is not None:
+                    query = query.filter(time_condition)
+
+                # Apply research mode filter
+                mode_condition = get_research_mode_condition(
+                    research_mode, SearchCall.research_mode
+                )
+                if mode_condition is not None:
+                    query = query.filter(mode_condition)
+
+                # Get search engine statistics using ORM aggregation
+                search_stats = session.query(
+                    SearchCall.search_engine,
+                    func.count().label("call_count"),
+                    func.avg(SearchCall.response_time_ms).label(
+                        "avg_response_time"
+                    ),
+                    func.sum(SearchCall.results_count).label("total_results"),
+                    func.avg(SearchCall.results_count).label(
+                        "avg_results_per_call"
+                    ),
+                    func.sum(
+                        case(
+                            (SearchCall.success_status == "success", 1), else_=0
+                        )
+                    ).label("success_count"),
+                    func.sum(
+                        case((SearchCall.success_status == "error", 1), else_=0)
+                    ).label("error_count"),
+                ).filter(SearchCall.search_engine.isnot(None))
+
+                # Apply same filters to stats query
+                if time_condition is not None:
+                    search_stats = search_stats.filter(time_condition)
+                if mode_condition is not None:
+                    search_stats = search_stats.filter(mode_condition)
+
+                search_stats = (
+                    search_stats.group_by(SearchCall.search_engine)
+                    .order_by(func.count().desc())
+                    .all()
+                )
+
+                # Get recent search calls
+                recent_calls_query = session.query(SearchCall)
+                if time_condition is not None:
+                    recent_calls_query = recent_calls_query.filter(
+                        time_condition
+                    )
+                if mode_condition is not None:
+                    recent_calls_query = recent_calls_query.filter(
+                        mode_condition
+                    )
+
+                recent_calls = (
+                    recent_calls_query.order_by(SearchCall.timestamp.desc())
+                    .limit(20)
+                    .all()
+                )
+
+                return {
+                    "search_engine_stats": [
+                        {
+                            "engine": stat.search_engine,
+                            "call_count": stat.call_count,
+                            "avg_response_time": stat.avg_response_time or 0,
+                            "total_results": stat.total_results or 0,
+                            "avg_results_per_call": stat.avg_results_per_call
+                            or 0,
+                            "success_rate": (
+                                (stat.success_count / stat.call_count * 100)
+                                if stat.call_count > 0
+                                else 0
+                            ),
+                            "error_count": stat.error_count or 0,
+                        }
+                        for stat in search_stats
+                    ],
+                    "recent_calls": [
+                        {
+                            "engine": call.search_engine,
+                            "query": (
+                                call.query[:100] + "..."
+                                if len(call.query or "") > 100
+                                else call.query
+                            ),
+                            "results_count": call.results_count,
+                            "response_time_ms": call.response_time_ms,
+                            "success_status": call.success_status,
+                            "timestamp": str(call.timestamp),
+                        }
+                        for call in recent_calls
+                    ],
+                }
+
+            except Exception as e:
+                logger.exception(f"Error getting search metrics: {e}")
+                return {"search_engine_stats": [], "recent_calls": []}
+
+    def get_research_search_metrics(self, research_id: int) -> Dict[str, Any]:
+        """Get search metrics for a specific research session."""
+        with self.db.get_session() as session:
+            try:
+                # Get all search calls for this research
+                search_calls = (
+                    session.query(SearchCall)
+                    .filter(SearchCall.research_id == research_id)
+                    .order_by(SearchCall.timestamp.asc())
+                    .all()
+                )
+
+                # Get search engine stats for this research
+                engine_stats = (
+                    session.query(
+                        SearchCall.search_engine,
+                        func.count().label("call_count"),
+                        func.avg(SearchCall.response_time_ms).label(
+                            "avg_response_time"
+                        ),
+                        func.sum(SearchCall.results_count).label(
+                            "total_results"
+                        ),
+                        func.sum(
+                            case(
+                                (SearchCall.success_status == "success", 1),
+                                else_=0,
+                            )
+                        ).label("success_count"),
+                    )
+                    .filter(SearchCall.research_id == research_id)
+                    .group_by(SearchCall.search_engine)
+                    .order_by(func.count().desc())
+                    .all()
+                )
+
+                # Calculate totals
+                total_searches = len(search_calls)
+                total_results = sum(
+                    call.results_count or 0 for call in search_calls
+                )
+                avg_response_time = (
+                    sum(call.response_time_ms or 0 for call in search_calls)
+                    / total_searches
+                    if total_searches > 0
+                    else 0
+                )
+                successful_searches = sum(
+                    1
+                    for call in search_calls
+                    if call.success_status == "success"
+                )
+                success_rate = (
+                    (successful_searches / total_searches * 100)
+                    if total_searches > 0
+                    else 0
+                )
+
+                return {
+                    "total_searches": total_searches,
+                    "total_results": total_results,
+                    "avg_response_time": round(avg_response_time),
+                    "success_rate": round(success_rate, 1),
+                    "search_calls": [
+                        {
+                            "engine": call.search_engine,
+                            "query": call.query,
+                            "results_count": call.results_count,
+                            "response_time_ms": call.response_time_ms,
+                            "success_status": call.success_status,
+                            "timestamp": str(call.timestamp),
+                        }
+                        for call in search_calls
+                    ],
+                    "engine_stats": [
+                        {
+                            "engine": stat.search_engine,
+                            "call_count": stat.call_count,
+                            "avg_response_time": stat.avg_response_time or 0,
+                            "total_results": stat.total_results or 0,
+                            "success_rate": (
+                                (stat.success_count / stat.call_count * 100)
+                                if stat.call_count > 0
+                                else 0
+                            ),
+                        }
+                        for stat in engine_stats
+                    ],
+                }
+
+            except Exception as e:
+                logger.exception(f"Error getting research search metrics: {e}")
+                return {
+                    "total_searches": 0,
+                    "total_results": 0,
+                    "avg_response_time": 0,
+                    "success_rate": 0,
+                    "search_calls": [],
+                    "engine_stats": [],
+                }
+
+    def get_search_time_series(
+        self, period: str = "30d", research_mode: str = "all"
+    ) -> List[Dict[str, Any]]:
+        """Get search activity time series data for charting.
+
+        Args:
+            period: Time period to filter by ('7d', '30d', '3m', '1y', 'all')
+            research_mode: Research mode to filter by ('quick', 'detailed', 'all')
+
+        Returns:
+            List of time series data points with search engine activity
+        """
+        with self.db.get_session() as session:
+            try:
+                # Build base query
+                query = session.query(SearchCall).filter(
+                    SearchCall.search_engine.isnot(None),
+                    SearchCall.timestamp.isnot(None),
+                )
+
+                # Apply time filter
+                time_condition = get_time_filter_condition(
+                    period, SearchCall.timestamp
+                )
+                if time_condition is not None:
+                    query = query.filter(time_condition)
+
+                # Apply research mode filter
+                mode_condition = get_research_mode_condition(
+                    research_mode, SearchCall.research_mode
+                )
+                if mode_condition is not None:
+                    query = query.filter(mode_condition)
+
+                # Get all search calls ordered by time
+                search_calls = query.order_by(SearchCall.timestamp.asc()).all()
+
+                # Create time series data
+                time_series = []
+                for call in search_calls:
+                    time_series.append(
+                        {
+                            "timestamp": (
+                                str(call.timestamp) if call.timestamp else None
+                            ),
+                            "search_engine": call.search_engine,
+                            "results_count": call.results_count or 0,
+                            "response_time_ms": call.response_time_ms or 0,
+                            "success_status": call.success_status,
+                            "query": (
+                                call.query[:50] + "..."
+                                if call.query and len(call.query) > 50
+                                else call.query
+                            ),
+                        }
+                    )
+
+                return time_series
+
+            except Exception as e:
+                logger.exception(f"Error getting search time series: {e}")
+                return []
+
+
+# Global search tracker instance
+_search_tracker = None
+
+
+def get_search_tracker() -> SearchTracker:
+    """Get the global search tracker instance."""
+    global _search_tracker
+    if _search_tracker is None:
+        _search_tracker = SearchTracker()
+    return _search_tracker
+
+
+def set_search_context(context: Dict[str, Any]) -> None:
+    """Set search context for the global tracker."""
+    tracker = get_search_tracker()
+    tracker.set_research_context(context)