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

local_deep_research/metrics/token_counter.py

@@ -0,0 +1,1078 @@
+"""Token counting functionality for LLM usage tracking."""
+
+import inspect
+import json
+import os
+import time
+from typing import Any, Dict, List, Optional
+
+from langchain_core.callbacks import BaseCallbackHandler
+from langchain_core.outputs import LLMResult
+from loguru import logger
+from sqlalchemy import func, text
+
+from .database import get_metrics_db
+from .db_models import ModelUsage, TokenUsage
+from .query_utils import get_research_mode_condition, get_time_filter_condition
+
+
+class TokenCountingCallback(BaseCallbackHandler):
+    """Callback handler for counting tokens across different models."""
+
+    def __init__(
+        self,
+        research_id: Optional[int] = None,
+        research_context: Optional[Dict[str, Any]] = None,
+    ):
+        """Initialize the token counting callback.
+
+        Args:
+            research_id: The ID of the research to track tokens for
+            research_context: Additional research context for enhanced tracking
+        """
+        super().__init__()
+        self.research_id = research_id
+        self.research_context = research_context or {}
+        self.current_model = None
+        self.current_provider = None
+        self.preset_model = None  # Model name set during callback creation
+        self.preset_provider = None  # Provider set during callback creation
+
+        # Phase 1 Enhancement: Track timing and context
+        self.start_time = None
+        self.response_time_ms = None
+        self.success_status = "success"
+        self.error_type = None
+
+        # Call stack tracking
+        self.calling_file = None
+        self.calling_function = None
+        self.call_stack = None
+
+        # Track token counts in memory
+        self.counts = {
+            "total_tokens": 0,
+            "total_prompt_tokens": 0,
+            "total_completion_tokens": 0,
+            "by_model": {},
+        }
+
+    def on_llm_start(
+        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
+    ) -> None:
+        """Called when LLM starts running."""
+        # Phase 1 Enhancement: Start timing
+        self.start_time = time.time()
+
+        # Phase 1 Enhancement: Capture call stack information
+        try:
+            stack = inspect.stack()
+
+            # Skip the first few frames (this method, langchain internals)
+            # Look for the first frame that's in our project directory
+            for frame_info in stack[1:]:
+                file_path = frame_info.filename
+                # Look for any frame containing local_deep_research project
+                if (
+                    "local_deep_research" in file_path
+                    and "site-packages" not in file_path
+                    and "venv" not in file_path
+                ):
+                    # Extract relative path from local_deep_research
+                    if "src/local_deep_research" in file_path:
+                        relative_path = file_path.split(
+                            "src/local_deep_research"
+                        )[-1].lstrip("/")
+                    elif "local_deep_research/src" in file_path:
+                        relative_path = file_path.split(
+                            "local_deep_research/src"
+                        )[-1].lstrip("/")
+                    elif "local_deep_research" in file_path:
+                        # Get everything after local_deep_research
+                        relative_path = file_path.split("local_deep_research")[
+                            -1
+                        ].lstrip("/")
+                    else:
+                        relative_path = os.path.basename(file_path)
+
+                    self.calling_file = relative_path
+                    self.calling_function = frame_info.function
+
+                    # Capture a simplified call stack (just the relevant frames)
+                    call_stack_frames = []
+                    for frame in stack[1:6]:  # Limit to 5 frames
+                        if (
+                            "local_deep_research" in frame.filename
+                            and "site-packages" not in frame.filename
+                            and "venv" not in frame.filename
+                        ):
+                            frame_name = f"{os.path.basename(frame.filename)}:{frame.function}:{frame.lineno}"
+                            call_stack_frames.append(frame_name)
+
+                    self.call_stack = (
+                        " -> ".join(call_stack_frames)
+                        if call_stack_frames
+                        else None
+                    )
+                    break
+        except Exception as e:
+            logger.debug(f"Error capturing call stack: {e}")
+            # Continue without call stack info if there's an error
+
+        # Debug logging
+        logger.debug(f"on_llm_start serialized: {serialized}")
+        logger.debug(f"on_llm_start kwargs: {kwargs}")
+
+        # First, use preset values if available
+        if self.preset_model:
+            self.current_model = self.preset_model
+        else:
+            # Try multiple locations for model name
+            model_name = None
+
+            # First check invocation_params
+            invocation_params = kwargs.get("invocation_params", {})
+            model_name = invocation_params.get(
+                "model"
+            ) or invocation_params.get("model_name")
+
+            # Check kwargs directly
+            if not model_name:
+                model_name = kwargs.get("model") or kwargs.get("model_name")
+
+            # Check serialized data
+            if not model_name and "kwargs" in serialized:
+                model_name = serialized["kwargs"].get("model") or serialized[
+                    "kwargs"
+                ].get("model_name")
+
+            # Check for name in serialized data
+            if not model_name and "name" in serialized:
+                model_name = serialized["name"]
+
+            # If still not found and we have Ollama, try to extract from the instance
+            if (
+                not model_name
+                and "_type" in serialized
+                and "ChatOllama" in serialized["_type"]
+            ):
+                # For Ollama, the model name might be in the serialized kwargs
+                if "kwargs" in serialized and "model" in serialized["kwargs"]:
+                    model_name = serialized["kwargs"]["model"]
+                else:
+                    # Default to the type if we can't find the actual model
+                    model_name = "ollama"
+
+            # Final fallback
+            if not model_name:
+                if "_type" in serialized:
+                    model_name = serialized["_type"]
+                else:
+                    model_name = "unknown"
+
+            self.current_model = model_name
+
+        # Use preset provider if available
+        if self.preset_provider:
+            self.current_provider = self.preset_provider
+        else:
+            # Extract provider from serialized type or kwargs
+            if "_type" in serialized:
+                type_str = serialized["_type"]
+                if "ChatOllama" in type_str:
+                    self.current_provider = "ollama"
+                elif "ChatOpenAI" in type_str:
+                    self.current_provider = "openai"
+                elif "ChatAnthropic" in type_str:
+                    self.current_provider = "anthropic"
+                else:
+                    self.current_provider = kwargs.get("provider", "unknown")
+            else:
+                self.current_provider = kwargs.get("provider", "unknown")
+
+        # Initialize model tracking if needed
+        if self.current_model not in self.counts["by_model"]:
+            self.counts["by_model"][self.current_model] = {
+                "prompt_tokens": 0,
+                "completion_tokens": 0,
+                "total_tokens": 0,
+                "calls": 0,
+                "provider": self.current_provider,
+            }
+
+        # Increment call count
+        self.counts["by_model"][self.current_model]["calls"] += 1
+
+    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        """Called when LLM ends running."""
+        # Phase 1 Enhancement: Calculate response time
+        if self.start_time:
+            self.response_time_ms = int((time.time() - self.start_time) * 1000)
+
+        # Extract token usage from response
+        token_usage = None
+
+        # Check multiple locations for token usage
+        if hasattr(response, "llm_output") and response.llm_output:
+            token_usage = response.llm_output.get(
+                "token_usage"
+            ) or response.llm_output.get("usage", {})
+
+        # Check for usage metadata in generations (Ollama specific)
+        if not token_usage and hasattr(response, "generations"):
+            for generation_list in response.generations:
+                for generation in generation_list:
+                    if hasattr(generation, "message") and hasattr(
+                        generation.message, "usage_metadata"
+                    ):
+                        usage_meta = generation.message.usage_metadata
+                        token_usage = {
+                            "prompt_tokens": usage_meta.get("input_tokens", 0),
+                            "completion_tokens": usage_meta.get(
+                                "output_tokens", 0
+                            ),
+                            "total_tokens": usage_meta.get("total_tokens", 0),
+                        }
+                        break
+                    # Also check response_metadata
+                    elif hasattr(generation, "message") and hasattr(
+                        generation.message, "response_metadata"
+                    ):
+                        resp_meta = generation.message.response_metadata
+                        if resp_meta.get("prompt_eval_count") or resp_meta.get(
+                            "eval_count"
+                        ):
+                            token_usage = {
+                                "prompt_tokens": resp_meta.get(
+                                    "prompt_eval_count", 0
+                                ),
+                                "completion_tokens": resp_meta.get(
+                                    "eval_count", 0
+                                ),
+                                "total_tokens": resp_meta.get(
+                                    "prompt_eval_count", 0
+                                )
+                                + resp_meta.get("eval_count", 0),
+                            }
+                            break
+                if token_usage:
+                    break
+
+        if token_usage:
+            prompt_tokens = token_usage.get("prompt_tokens", 0)
+            completion_tokens = token_usage.get("completion_tokens", 0)
+            total_tokens = token_usage.get(
+                "total_tokens", prompt_tokens + completion_tokens
+            )
+
+            # Update in-memory counts
+            self.counts["total_prompt_tokens"] += prompt_tokens
+            self.counts["total_completion_tokens"] += completion_tokens
+            self.counts["total_tokens"] += total_tokens
+
+            if self.current_model:
+                self.counts["by_model"][self.current_model][
+                    "prompt_tokens"
+                ] += prompt_tokens
+                self.counts["by_model"][self.current_model][
+                    "completion_tokens"
+                ] += completion_tokens
+                self.counts["by_model"][self.current_model]["total_tokens"] += (
+                    total_tokens
+                )
+
+            # Save to database if we have a research_id
+            if self.research_id:
+                self._save_to_db(prompt_tokens, completion_tokens)
+
+    def on_llm_error(self, error, **kwargs: Any) -> None:
+        """Called when LLM encounters an error."""
+        # Phase 1 Enhancement: Track errors
+        if self.start_time:
+            self.response_time_ms = int((time.time() - self.start_time) * 1000)
+
+        self.success_status = "error"
+        self.error_type = str(type(error).__name__)
+
+        # Still save to database to track failed calls
+        if self.research_id:
+            self._save_to_db(0, 0)
+
+    def _save_to_db(self, prompt_tokens: int, completion_tokens: int):
+        """Save token usage to the database."""
+        try:
+            db = get_metrics_db()
+            with db.get_session() as session:
+                # Phase 1 Enhancement: Prepare additional context
+                research_query = self.research_context.get("research_query")
+                research_mode = self.research_context.get("research_mode")
+                research_phase = self.research_context.get("research_phase")
+                search_iteration = self.research_context.get("search_iteration")
+                search_engines_planned = self.research_context.get(
+                    "search_engines_planned"
+                )
+                search_engine_selected = self.research_context.get(
+                    "search_engine_selected"
+                )
+
+                # Debug logging for search engine context
+                if search_engines_planned or search_engine_selected:
+                    logger.info(
+                        f"Token tracking - Search context: planned={search_engines_planned}, selected={search_engine_selected}, phase={research_phase}"
+                    )
+                else:
+                    logger.debug(
+                        f"Token tracking - No search engine context yet, phase={research_phase}"
+                    )
+
+                # Convert list to JSON string if needed
+                if isinstance(search_engines_planned, list):
+                    search_engines_planned = json.dumps(search_engines_planned)
+
+                # Add token usage record with enhanced fields
+                token_usage = TokenUsage(
+                    research_id=self.research_id,
+                    model_name=self.current_model,
+                    provider=self.current_provider,  # Added provider for accurate cost tracking
+                    prompt_tokens=prompt_tokens,
+                    completion_tokens=completion_tokens,
+                    total_tokens=prompt_tokens + completion_tokens,
+                    # Phase 1 Enhancement: Research context
+                    research_query=research_query,
+                    research_mode=research_mode,
+                    research_phase=research_phase,
+                    search_iteration=search_iteration,
+                    # Phase 1 Enhancement: Performance metrics
+                    response_time_ms=self.response_time_ms,
+                    success_status=self.success_status,
+                    error_type=self.error_type,
+                    # Phase 1 Enhancement: Search engine context
+                    search_engines_planned=search_engines_planned,
+                    search_engine_selected=search_engine_selected,
+                    # Phase 1 Enhancement: Call stack tracking
+                    calling_file=self.calling_file,
+                    calling_function=self.calling_function,
+                    call_stack=self.call_stack,
+                )
+                session.add(token_usage)
+
+                # Update or create model usage statistics
+                model_usage = (
+                    session.query(ModelUsage)
+                    .filter_by(
+                        research_id=self.research_id,
+                        model_name=self.current_model,
+                    )
+                    .first()
+                )
+
+                if model_usage:
+                    model_usage.prompt_tokens += prompt_tokens
+                    model_usage.completion_tokens += completion_tokens
+                    model_usage.total_tokens += (
+                        prompt_tokens + completion_tokens
+                    )
+                    model_usage.calls += 1
+                else:
+                    model_usage = ModelUsage(
+                        research_id=self.research_id,
+                        model_name=self.current_model,
+                        provider=self.current_provider,
+                        prompt_tokens=prompt_tokens,
+                        completion_tokens=completion_tokens,
+                        total_tokens=prompt_tokens + completion_tokens,
+                        calls=1,
+                    )
+                    session.add(model_usage)
+
+        except Exception as e:
+            logger.exception(f"Error saving token usage to database: {e}")
+
+    def get_counts(self) -> Dict[str, Any]:
+        """Get the current token counts."""
+        return self.counts
+
+
+class TokenCounter:
+    """Manager class for token counting across the application."""
+
+    def __init__(self):
+        """Initialize the token counter."""
+        self.db = get_metrics_db()
+
+    def create_callback(
+        self,
+        research_id: Optional[int] = None,
+        research_context: Optional[Dict[str, Any]] = None,
+    ) -> TokenCountingCallback:
+        """Create a new token counting callback.
+
+        Args:
+            research_id: The ID of the research to track tokens for
+            research_context: Additional research context for enhanced tracking
+
+        Returns:
+            A new TokenCountingCallback instance
+        """
+        return TokenCountingCallback(
+            research_id=research_id, research_context=research_context
+        )
+
+    def get_research_metrics(self, research_id: int) -> Dict[str, Any]:
+        """Get token metrics for a specific research.
+
+        Args:
+            research_id: The ID of the research
+
+        Returns:
+            Dictionary containing token usage metrics
+        """
+        with self.db.get_session() as session:
+            # Get model usage for this research
+            model_usages = (
+                session.query(ModelUsage)
+                .filter_by(research_id=research_id)
+                .order_by(ModelUsage.total_tokens.desc())
+                .all()
+            )
+
+            model_usage = []
+            total_tokens = 0
+            total_calls = 0
+
+            for usage in model_usages:
+                model_usage.append(
+                    {
+                        "model": usage.model_name,
+                        "provider": usage.provider,
+                        "tokens": usage.total_tokens,
+                        "calls": usage.calls,
+                        "prompt_tokens": usage.prompt_tokens,
+                        "completion_tokens": usage.completion_tokens,
+                    }
+                )
+                total_tokens += usage.total_tokens
+                total_calls += usage.calls
+
+            return {
+                "research_id": research_id,
+                "total_tokens": total_tokens,
+                "total_calls": total_calls,
+                "model_usage": model_usage,
+            }
+
+    def get_overall_metrics(
+        self, period: str = "30d", research_mode: str = "all"
+    ) -> Dict[str, Any]:
+        """Get overall token metrics across all researches.
+
+        Args:
+            period: Time period to filter by ('7d', '30d', '3m', '1y', 'all')
+            research_mode: Research mode to filter by ('quick', 'detailed', 'all')
+
+        Returns:
+            Dictionary containing overall metrics
+        """
+        with self.db.get_session() as session:
+            # Build base query with filters
+            query = session.query(TokenUsage)
+
+            # Apply time filter
+            time_condition = get_time_filter_condition(
+                period, TokenUsage.timestamp
+            )
+            if time_condition is not None:
+                query = query.filter(time_condition)
+
+            # Apply research mode filter
+            mode_condition = get_research_mode_condition(
+                research_mode, TokenUsage.research_mode
+            )
+            if mode_condition is not None:
+                query = query.filter(mode_condition)
+
+            # Total tokens and researches
+            total_tokens = (
+                query.with_entities(func.sum(TokenUsage.total_tokens)).scalar()
+                or 0
+            )
+            total_researches = (
+                query.with_entities(
+                    func.count(func.distinct(TokenUsage.research_id))
+                ).scalar()
+                or 0
+            )
+
+            # Model statistics using ORM aggregation
+            model_stats_query = session.query(
+                TokenUsage.model_name,
+                func.sum(TokenUsage.total_tokens).label("tokens"),
+                func.count().label("calls"),
+                func.sum(TokenUsage.prompt_tokens).label("prompt_tokens"),
+                func.sum(TokenUsage.completion_tokens).label(
+                    "completion_tokens"
+                ),
+            ).filter(TokenUsage.model_name.isnot(None))
+
+            # Apply same filters to model stats
+            if time_condition is not None:
+                model_stats_query = model_stats_query.filter(time_condition)
+            if mode_condition is not None:
+                model_stats_query = model_stats_query.filter(mode_condition)
+
+            model_stats = (
+                model_stats_query.group_by(TokenUsage.model_name)
+                .order_by(func.sum(TokenUsage.total_tokens).desc())
+                .all()
+            )
+
+            # Get provider info from ModelUsage table
+            by_model = []
+            for stat in model_stats:
+                # Try to get provider from ModelUsage table
+                provider_info = (
+                    session.query(ModelUsage.provider)
+                    .filter(ModelUsage.model_name == stat.model_name)
+                    .first()
+                )
+                provider = (
+                    provider_info.provider if provider_info else "unknown"
+                )
+
+                by_model.append(
+                    {
+                        "model": stat.model_name,
+                        "provider": provider,
+                        "tokens": stat.tokens,
+                        "calls": stat.calls,
+                        "prompt_tokens": stat.prompt_tokens,
+                        "completion_tokens": stat.completion_tokens,
+                    }
+                )
+
+            # Get recent researches with token usage
+            # Note: This requires research_history table - for now we'll use available data
+            recent_research_query = session.query(
+                TokenUsage.research_id,
+                func.sum(TokenUsage.total_tokens).label("token_count"),
+                func.max(TokenUsage.timestamp).label("latest_timestamp"),
+            ).filter(TokenUsage.research_id.isnot(None))
+
+            if time_condition is not None:
+                recent_research_query = recent_research_query.filter(
+                    time_condition
+                )
+            if mode_condition is not None:
+                recent_research_query = recent_research_query.filter(
+                    mode_condition
+                )
+
+            recent_research_data = (
+                recent_research_query.group_by(TokenUsage.research_id)
+                .order_by(func.max(TokenUsage.timestamp).desc())
+                .limit(10)
+                .all()
+            )
+
+            recent_researches = []
+            for research_data in recent_research_data:
+                # Get research query from token_usage table if available
+                research_query_data = (
+                    session.query(TokenUsage.research_query)
+                    .filter(
+                        TokenUsage.research_id == research_data.research_id,
+                        TokenUsage.research_query.isnot(None),
+                    )
+                    .first()
+                )
+
+                query_text = (
+                    research_query_data.research_query
+                    if research_query_data
+                    else f"Research {research_data.research_id}"
+                )
+
+                recent_researches.append(
+                    {
+                        "id": research_data.research_id,
+                        "query": query_text,
+                        "tokens": research_data.token_count or 0,
+                        "created_at": research_data.latest_timestamp,
+                    }
+                )
+
+            # Token breakdown statistics
+            breakdown_query = query.with_entities(
+                func.sum(TokenUsage.prompt_tokens).label("total_input_tokens"),
+                func.sum(TokenUsage.completion_tokens).label(
+                    "total_output_tokens"
+                ),
+                func.avg(TokenUsage.prompt_tokens).label("avg_input_tokens"),
+                func.avg(TokenUsage.completion_tokens).label(
+                    "avg_output_tokens"
+                ),
+                func.avg(TokenUsage.total_tokens).label("avg_total_tokens"),
+            )
+            token_breakdown = breakdown_query.first()
+
+            return {
+                "total_tokens": total_tokens,
+                "total_researches": total_researches,
+                "by_model": by_model,
+                "recent_researches": recent_researches,
+                "token_breakdown": {
+                    "total_input_tokens": int(
+                        token_breakdown.total_input_tokens or 0
+                    ),
+                    "total_output_tokens": int(
+                        token_breakdown.total_output_tokens or 0
+                    ),
+                    "avg_input_tokens": int(
+                        token_breakdown.avg_input_tokens or 0
+                    ),
+                    "avg_output_tokens": int(
+                        token_breakdown.avg_output_tokens or 0
+                    ),
+                    "avg_total_tokens": int(
+                        token_breakdown.avg_total_tokens or 0
+                    ),
+                },
+            }
+
+    def get_enhanced_metrics(
+        self, period: str = "30d", research_mode: str = "all"
+    ) -> Dict[str, Any]:
+        """Get enhanced Phase 1 tracking metrics.
+
+        Args:
+            period: Time period to filter by ('7d', '30d', '3m', '1y', 'all')
+            research_mode: Research mode to filter by ('quick', 'detailed', 'all')
+
+        Returns:
+            Dictionary containing enhanced metrics data including time series
+        """
+        with self.db.get_session() as session:
+            # Build base query with filters
+            query = session.query(TokenUsage)
+
+            # Apply time filter
+            time_condition = get_time_filter_condition(
+                period, TokenUsage.timestamp
+            )
+            if time_condition is not None:
+                query = query.filter(time_condition)
+
+            # Apply research mode filter
+            mode_condition = get_research_mode_condition(
+                research_mode, TokenUsage.research_mode
+            )
+            if mode_condition is not None:
+                query = query.filter(mode_condition)
+
+            # Get time series data for the chart - most important for "Token Consumption Over Time"
+            time_series_query = query.filter(
+                TokenUsage.timestamp.isnot(None), TokenUsage.total_tokens > 0
+            ).order_by(TokenUsage.timestamp.asc())
+
+            # Limit to recent data for performance
+            if period != "all":
+                time_series_query = time_series_query.limit(200)
+
+            time_series_data = time_series_query.all()
+
+            # Format time series data with cumulative calculations
+            time_series = []
+            cumulative_tokens = 0
+            cumulative_prompt_tokens = 0
+            cumulative_completion_tokens = 0
+
+            for usage in time_series_data:
+                cumulative_tokens += usage.total_tokens or 0
+                cumulative_prompt_tokens += usage.prompt_tokens or 0
+                cumulative_completion_tokens += usage.completion_tokens or 0
+
+                time_series.append(
+                    {
+                        "timestamp": str(usage.timestamp)
+                        if usage.timestamp
+                        else None,
+                        "tokens": usage.total_tokens or 0,
+                        "prompt_tokens": usage.prompt_tokens or 0,
+                        "completion_tokens": usage.completion_tokens or 0,
+                        "cumulative_tokens": cumulative_tokens,
+                        "cumulative_prompt_tokens": cumulative_prompt_tokens,
+                        "cumulative_completion_tokens": cumulative_completion_tokens,
+                        "research_id": usage.research_id,
+                        "research_query": usage.research_query,
+                    }
+                )
+
+            # Basic performance stats using ORM
+            performance_query = query.filter(
+                TokenUsage.response_time_ms.isnot(None)
+            )
+            total_calls = performance_query.count()
+
+            if total_calls > 0:
+                avg_response_time = (
+                    performance_query.with_entities(
+                        func.avg(TokenUsage.response_time_ms)
+                    ).scalar()
+                    or 0
+                )
+                min_response_time = (
+                    performance_query.with_entities(
+                        func.min(TokenUsage.response_time_ms)
+                    ).scalar()
+                    or 0
+                )
+                max_response_time = (
+                    performance_query.with_entities(
+                        func.max(TokenUsage.response_time_ms)
+                    ).scalar()
+                    or 0
+                )
+                success_count = performance_query.filter(
+                    TokenUsage.success_status == "success"
+                ).count()
+                error_count = performance_query.filter(
+                    TokenUsage.success_status == "error"
+                ).count()
+
+                perf_stats = {
+                    "avg_response_time": round(avg_response_time),
+                    "min_response_time": min_response_time,
+                    "max_response_time": max_response_time,
+                    "success_rate": (
+                        round((success_count / total_calls * 100), 1)
+                        if total_calls > 0
+                        else 0
+                    ),
+                    "error_rate": (
+                        round((error_count / total_calls * 100), 1)
+                        if total_calls > 0
+                        else 0
+                    ),
+                    "total_enhanced_calls": total_calls,
+                }
+            else:
+                perf_stats = {
+                    "avg_response_time": 0,
+                    "min_response_time": 0,
+                    "max_response_time": 0,
+                    "success_rate": 0,
+                    "error_rate": 0,
+                    "total_enhanced_calls": 0,
+                }
+
+            # Research mode breakdown using ORM
+            mode_stats = (
+                query.filter(TokenUsage.research_mode.isnot(None))
+                .with_entities(
+                    TokenUsage.research_mode,
+                    func.count().label("count"),
+                    func.avg(TokenUsage.total_tokens).label("avg_tokens"),
+                    func.avg(TokenUsage.response_time_ms).label(
+                        "avg_response_time"
+                    ),
+                )
+                .group_by(TokenUsage.research_mode)
+                .all()
+            )
+
+            modes = [
+                {
+                    "mode": stat.research_mode,
+                    "count": stat.count,
+                    "avg_tokens": round(stat.avg_tokens or 0),
+                    "avg_response_time": round(stat.avg_response_time or 0),
+                }
+                for stat in mode_stats
+            ]
+
+            # Recent enhanced data (simplified)
+            recent_enhanced_query = (
+                query.filter(TokenUsage.research_query.isnot(None))
+                .order_by(TokenUsage.timestamp.desc())
+                .limit(50)
+            )
+
+            recent_enhanced_data = recent_enhanced_query.all()
+            recent_enhanced = [
+                {
+                    "research_query": usage.research_query,
+                    "research_mode": usage.research_mode,
+                    "research_phase": usage.research_phase,
+                    "search_iteration": usage.search_iteration,
+                    "response_time_ms": usage.response_time_ms,
+                    "success_status": usage.success_status,
+                    "error_type": usage.error_type,
+                    "search_engines_planned": usage.search_engines_planned,
+                    "search_engine_selected": usage.search_engine_selected,
+                    "total_tokens": usage.total_tokens,
+                    "prompt_tokens": usage.prompt_tokens,
+                    "completion_tokens": usage.completion_tokens,
+                    "timestamp": str(usage.timestamp)
+                    if usage.timestamp
+                    else None,
+                    "research_id": usage.research_id,
+                    "calling_file": usage.calling_file,
+                    "calling_function": usage.calling_function,
+                    "call_stack": usage.call_stack,
+                }
+                for usage in recent_enhanced_data
+            ]
+
+            # Search engine breakdown using ORM
+            search_engine_stats = (
+                query.filter(TokenUsage.search_engine_selected.isnot(None))
+                .with_entities(
+                    TokenUsage.search_engine_selected,
+                    func.count().label("count"),
+                    func.avg(TokenUsage.total_tokens).label("avg_tokens"),
+                    func.avg(TokenUsage.response_time_ms).label(
+                        "avg_response_time"
+                    ),
+                )
+                .group_by(TokenUsage.search_engine_selected)
+                .all()
+            )
+
+            search_engines = [
+                {
+                    "search_engine": stat.search_engine_selected,
+                    "count": stat.count,
+                    "avg_tokens": round(stat.avg_tokens or 0),
+                    "avg_response_time": round(stat.avg_response_time or 0),
+                }
+                for stat in search_engine_stats
+            ]
+
+            # Research phase breakdown using ORM
+            phase_stats = (
+                query.filter(TokenUsage.research_phase.isnot(None))
+                .with_entities(
+                    TokenUsage.research_phase,
+                    func.count().label("count"),
+                    func.avg(TokenUsage.total_tokens).label("avg_tokens"),
+                    func.avg(TokenUsage.response_time_ms).label(
+                        "avg_response_time"
+                    ),
+                )
+                .group_by(TokenUsage.research_phase)
+                .all()
+            )
+
+            phases = [
+                {
+                    "phase": stat.research_phase,
+                    "count": stat.count,
+                    "avg_tokens": round(stat.avg_tokens or 0),
+                    "avg_response_time": round(stat.avg_response_time or 0),
+                }
+                for stat in phase_stats
+            ]
+
+            # Call stack analysis using ORM
+            file_stats = (
+                query.filter(TokenUsage.calling_file.isnot(None))
+                .with_entities(
+                    TokenUsage.calling_file,
+                    func.count().label("count"),
+                    func.avg(TokenUsage.total_tokens).label("avg_tokens"),
+                )
+                .group_by(TokenUsage.calling_file)
+                .order_by(func.count().desc())
+                .limit(10)
+                .all()
+            )
+
+            files = [
+                {
+                    "file": stat.calling_file,
+                    "count": stat.count,
+                    "avg_tokens": round(stat.avg_tokens or 0),
+                }
+                for stat in file_stats
+            ]
+
+            function_stats = (
+                query.filter(TokenUsage.calling_function.isnot(None))
+                .with_entities(
+                    TokenUsage.calling_function,
+                    func.count().label("count"),
+                    func.avg(TokenUsage.total_tokens).label("avg_tokens"),
+                )
+                .group_by(TokenUsage.calling_function)
+                .order_by(func.count().desc())
+                .limit(10)
+                .all()
+            )
+
+            functions = [
+                {
+                    "function": stat.calling_function,
+                    "count": stat.count,
+                    "avg_tokens": round(stat.avg_tokens or 0),
+                }
+                for stat in function_stats
+            ]
+
+            return {
+                "recent_enhanced_data": recent_enhanced,
+                "performance_stats": perf_stats,
+                "mode_breakdown": modes,
+                "search_engine_stats": search_engines,
+                "phase_breakdown": phases,
+                "time_series_data": time_series,
+                "call_stack_analysis": {
+                    "by_file": files,
+                    "by_function": functions,
+                },
+            }
+
+    def get_research_timeline_metrics(self, research_id: int) -> Dict[str, Any]:
+        """Get timeline metrics for a specific research.
+
+        Args:
+            research_id: The ID of the research
+
+        Returns:
+            Dictionary containing timeline metrics for the research
+        """
+        with self.db.get_session() as session:
+            # Get all token usage for this research ordered by time including call stack
+            timeline_data = session.execute(
+                text(
+                    """
+                SELECT
+                    timestamp,
+                    total_tokens,
+                    prompt_tokens,
+                    completion_tokens,
+                    response_time_ms,
+                    success_status,
+                    error_type,
+                    research_phase,
+                    search_iteration,
+                    search_engine_selected,
+                    model_name,
+                    calling_file,
+                    calling_function,
+                    call_stack
+                FROM token_usage
+                WHERE research_id = :research_id
+                ORDER BY timestamp ASC
+            """
+                ),
+                {"research_id": research_id},
+            ).fetchall()
+
+            # Format timeline data with cumulative tokens
+            timeline = []
+            cumulative_tokens = 0
+            cumulative_prompt_tokens = 0
+            cumulative_completion_tokens = 0
+
+            for row in timeline_data:
+                cumulative_tokens += row[1] or 0
+                cumulative_prompt_tokens += row[2] or 0
+                cumulative_completion_tokens += row[3] or 0
+
+                timeline.append(
+                    {
+                        "timestamp": str(row[0]) if row[0] else None,
+                        "tokens": row[1] or 0,
+                        "prompt_tokens": row[2] or 0,
+                        "completion_tokens": row[3] or 0,
+                        "cumulative_tokens": cumulative_tokens,
+                        "cumulative_prompt_tokens": cumulative_prompt_tokens,
+                        "cumulative_completion_tokens": cumulative_completion_tokens,
+                        "response_time_ms": row[4],
+                        "success_status": row[5],
+                        "error_type": row[6],
+                        "research_phase": row[7],
+                        "search_iteration": row[8],
+                        "search_engine_selected": row[9],
+                        "model_name": row[10],
+                        "calling_file": row[11],
+                        "calling_function": row[12],
+                        "call_stack": row[13],
+                    }
+                )
+
+            # Get research basic info
+            research_info = session.execute(
+                text(
+                    """
+                SELECT query, mode, status, created_at, completed_at
+                FROM research_history
+                WHERE id = :research_id
+            """
+                ),
+                {"research_id": research_id},
+            ).fetchone()
+
+            research_details = {}
+            if research_info:
+                research_details = {
+                    "query": research_info[0],
+                    "mode": research_info[1],
+                    "status": research_info[2],
+                    "created_at": str(research_info[3])
+                    if research_info[3]
+                    else None,
+                    "completed_at": str(research_info[4])
+                    if research_info[4]
+                    else None,
+                }
+
+            # Calculate summary stats
+            total_calls = len(timeline_data)
+            total_tokens = cumulative_tokens
+            avg_response_time = sum(row[4] or 0 for row in timeline_data) / max(
+                total_calls, 1
+            )
+            success_rate = (
+                sum(1 for row in timeline_data if row[5] == "success")
+                / max(total_calls, 1)
+                * 100
+            )
+
+            # Phase breakdown for this research
+            phase_stats = {}
+            for row in timeline_data:
+                phase = row[7] or "unknown"
+                if phase not in phase_stats:
+                    phase_stats[phase] = {
+                        "count": 0,
+                        "tokens": 0,
+                        "avg_response_time": 0,
+                    }
+                phase_stats[phase]["count"] += 1
+                phase_stats[phase]["tokens"] += row[1] or 0
+                if row[4]:
+                    phase_stats[phase]["avg_response_time"] += row[4]
+
+            # Calculate averages for phases
+            for phase in phase_stats:
+                if phase_stats[phase]["count"] > 0:
+                    phase_stats[phase]["avg_response_time"] = round(
+                        phase_stats[phase]["avg_response_time"]
+                        / phase_stats[phase]["count"]
+                    )
+
+            return {
+                "research_id": research_id,
+                "research_details": research_details,
+                "timeline": timeline,
+                "summary": {
+                    "total_calls": total_calls,
+                    "total_tokens": total_tokens,
+                    "total_prompt_tokens": cumulative_prompt_tokens,
+                    "total_completion_tokens": cumulative_completion_tokens,
+                    "avg_response_time": round(avg_response_time),
+                    "success_rate": round(success_rate, 1),
+                },
+                "phase_stats": phase_stats,
+            }