waldiez 0.5.2__py3-none-any.whl → 0.5.3__py3-none-any.whl

waldiez/running/timeline_processor.py

@@ -0,0 +1,1248 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# pylint: skip-file
+# pyright: reportArgumentType=false,reportUnknownVariableType=false
+# pyright: reportUnknownMemberType=false,reportUnknownArgumentType=false
+# flake8: noqa: C901
+"""
+Timeline Analysis Data Processor.
+
+Processes CSV files and outputs JSON structure for timeline visualization
+"""
+
+import argparse
+import json
+import os
+import re
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import pandas as pd
+
+from waldiez.logger import WaldiezLogger
+
+if TYPE_CHECKING:
+    Series = pd.Series[Any]
+else:
+    Series = pd.Series
+
+# Color palettes
+AGENT_COLORS = [
+    "#FF6B35",
+    "#4A90E2",
+    "#7ED321",
+    "#9013FE",
+    "#FF9500",
+    "#FF3B30",
+    "#007AFF",
+    "#34C759",
+    "#AF52DE",
+    "#FF9F0A",
+    "#FF2D92",
+    "#5AC8FA",
+    "#30D158",
+    "#BF5AF2",
+    "#FFD60A",
+    "#FF453A",
+    "#64D2FF",
+    "#32D74B",
+    "#DA70D6",
+    "#FFD23F",
+]
+
+ACTIVITY_COLORS = {
+    "human_input_waiting": "#FF8C00",
+    "user_thinking": "#87CEEB",
+    "agent_transition": "#FF7043",
+    "tool_call": "#4CAF50",
+    "function_call": "#9C27B0",
+    "processing": "#BDBDBD",
+    "session": "#8B5CF6",
+}
+
+DEFAULT_AGENT_COLOR = "#E5E7EB"
+
+LOG = WaldiezLogger()
+
+
+class TimelineProcessor:
+    """Class to process timeline data from CSV files."""
+
+    agents_data: pd.DataFrame | None
+    chat_data: pd.DataFrame | None
+    events_data: pd.DataFrame | None
+    functions_data: pd.DataFrame | None
+
+    def __init__(self) -> None:
+        """Initialize the TimelineProcessor with empty data attributes."""
+        self.agents_data = None
+        self.chat_data = None
+        self.events_data = None
+        self.functions_data = None
+
+    def is_missing_or_nan(self, value: Any) -> bool:
+        """Check if a value is missing, NaN, or empty.
+
+        Parameters
+        ----------
+        value : Any
+            The value to check.
+
+        Returns
+        -------
+        bool
+            True if the value is missing, NaN, or empty; False otherwise.
+        """
+        if pd.isna(value):  # pyright: ignore
+            return True
+        if isinstance(value, str) and (
+            value.strip() == "" or value.lower() == "nan"
+        ):
+            return True
+        return False
+
+    def fill_missing_agent_names(
+        self, data: pd.DataFrame | None, name_column: str = "source_name"
+    ) -> pd.DataFrame | None:
+        """Fill missing agent names with the previous valid name.
+
+        Parameters
+        ----------
+        data : pd.DataFrame | None
+            DataFrame containing agent names.
+        name_column : str, optional
+            The column name containing agent names, by default "source_name".
+
+        Returns
+        -------
+        pd.DataFrame | None
+            DataFrame with missing agent names filled.
+        """
+        if data is None or data.empty:
+            return data
+
+        data = data.copy()
+        last_valid_name: str | None = None
+
+        for idx in range(len(data)):
+            current_name = data.iloc[idx][name_column]
+
+            if self.is_missing_or_nan(current_name):
+                if last_valid_name is not None:
+                    column = data.columns.get_loc(name_column)
+                    data.iloc[idx, column] = last_valid_name  # type: ignore[index]
+                    LOG.debug(
+                        "Row %d: Replaced missing agent name with '%s'",
+                        idx,
+                        last_valid_name,
+                    )
+                else:
+                    # If no previous valid name, use a default
+                    default_name = "unknown_agent"
+                    column = data.columns.get_loc(name_column)
+                    data.iloc[idx, column] = default_name  # type: ignore[index]
+                    last_valid_name = default_name
+                    LOG.debug(
+                        "Row %d: Used default agent name '%s'",
+                        idx,
+                        default_name,
+                    )
+            else:
+                last_valid_name = current_name
+
+        return data
+
+    def fill_missing_agent_data(self) -> None:
+        """Fill missing agent names in agents_data."""
+        if self.agents_data is None:
+            return
+
+        self.agents_data = self.fill_missing_agent_names(
+            self.agents_data, "name"
+        )
+
+    def load_csv_files(
+        self,
+        agents_file: str | None = None,
+        chat_file: str | None = None,
+        events_file: str | None = None,
+        functions_file: str | None = None,
+    ) -> None:
+        """Load CSV files into pandas DataFrames.
+
+        Parameters
+        ----------
+        agents_file : str | None
+            Path to the agents CSV file.
+        chat_file : str | None
+            Path to the chat CSV file.
+        events_file : str | None
+            Path to the events CSV file.
+        functions_file : str | None
+            Path to the functions CSV file.
+        """
+        if agents_file:
+            self.agents_data = pd.read_csv(agents_file)
+            LOG.info("Loaded agents data: %d rows", len(self.agents_data))
+            # Fill missing agent names
+            self.fill_missing_agent_data()
+
+        if chat_file:
+            self.chat_data = pd.read_csv(chat_file)
+            LOG.info("Loaded chat data: %d rows", len(self.chat_data))
+            # Fill missing agent names in chat data
+            self.chat_data = self.fill_missing_agent_names(
+                self.chat_data, "source_name"
+            )
+
+        if events_file:
+            self.events_data = pd.read_csv(events_file)
+            LOG.info("Loaded events data: %d rows", len(self.events_data))
+
+        if functions_file:
+            self.functions_data = pd.read_csv(functions_file)
+            LOG.info("Loaded functions data: %d rows", len(self.functions_data))
+
+    def parse_date(self, date_str: str) -> pd.Timestamp:
+        """Parse date string to datetime.
+
+        Parameters
+        ----------
+        date_str : str
+            The date string to parse.
+
+        Returns
+        -------
+        pd.Timestamp
+            The parsed datetime.
+        """
+        try:
+            return pd.to_datetime(date_str)
+        except Exception:
+            coerced = pd.to_datetime(date_str, errors="coerce")
+            if isinstance(coerced, pd.Timestamp):
+                return coerced
+            return pd.Timestamp("1970-01-01")
+
+    def generate_agent_colors(self, agent_names: list[str]) -> dict[str, str]:
+        """Generate color mapping for agents.
+
+        Parameters
+        ----------
+        agent_names : list[str]
+            List of agent names.
+
+        Returns
+        -------
+        dict[str, str]
+            Mapping of agent names to their assigned colors.
+        """
+        colors = {}
+        for i, agent in enumerate(agent_names):
+            colors[agent] = AGENT_COLORS[i % len(AGENT_COLORS)]
+        return colors
+
+    def extract_token_info(
+        self,
+        request_str: Any,
+        response_str: Any,
+    ) -> dict[str, int]:
+        """Extract token information from request/response strings.
+
+        Parameters
+        ----------
+        request_str : Any
+            The request string containing token usage information.
+        response_str : Any
+            The response string containing token usage information.
+
+        Returns
+        -------
+        dict[str, int]
+            A dictionary containing the extracted token information.
+        """
+        prompt_tokens = 0
+        completion_tokens = 0
+        total_tokens = 0
+        try:
+            # Try to parse as JSON first
+            if (
+                request_str
+                and isinstance(request_str, str)
+                and request_str.strip().startswith("{")
+            ):
+                request_data = json.loads(request_str)
+                if "usage" in request_data:
+                    prompt_tokens = request_data["usage"].get(
+                        "prompt_tokens", 0
+                    )
+                elif "prompt_tokens" in request_data:
+                    prompt_tokens = request_data["prompt_tokens"]
+                elif "messages" in request_data:
+                    # Estimate tokens from content length
+                    content_length = sum(
+                        len(msg.get("content", ""))
+                        for msg in request_data["messages"]
+                        if "content" in msg and msg["content"]
+                    )
+                    prompt_tokens = max(1, content_length // 4)
+
+            if (
+                response_str
+                and isinstance(response_str, str)
+                and response_str.strip().startswith("{")
+            ):
+                response_data = json.loads(response_str)
+                if "usage" in response_data:
+                    prompt_tokens = response_data["usage"].get(
+                        "prompt_tokens", prompt_tokens
+                    )
+                    completion_tokens = response_data["usage"].get(
+                        "completion_tokens", 0
+                    )
+                    total_tokens = response_data["usage"].get(
+                        "total_tokens", prompt_tokens + completion_tokens
+                    )
+        except json.JSONDecodeError:
+            # Fallback to regex patterns if JSON parsing fails
+            pass
+
+        if total_tokens == 0 and (prompt_tokens > 0 or completion_tokens > 0):
+            total_tokens = prompt_tokens + completion_tokens
+
+        return {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": total_tokens,
+        }
+
+    def extract_llm_model(
+        self, agent_name: str, request_str: Any = None
+    ) -> str:
+        """Extract LLM model from agent data or request.
+
+        Parameters
+        ----------
+        agent_name : str
+            The name of the agent.
+        request_str : Any, optional
+            The request string containing token usage information.
+
+        Returns
+        -------
+        str
+            The extracted LLM model name.
+        """
+        # Handle missing/nan agent names
+        if self.is_missing_or_nan(agent_name):
+            agent_name = "unknown_agent"
+
+        # First try to extract from request_str (chat_completions.csv)
+        if request_str:
+            model = self._extract_model_from_text(str(request_str))
+            if model != "Unknown":
+                return model
+
+        # Then try to extract from agents data
+        if self.agents_data is not None:
+            agent_row = self.agents_data[self.agents_data["name"] == agent_name]
+            if not agent_row.empty and "init_args" in agent_row.columns:
+                init_args = str(agent_row.iloc[0]["init_args"])
+                model = self._extract_model_from_text(init_args)
+                if model != "Unknown":
+                    return model
+
+        return "Unknown"
+
+    def _extract_model_from_text(self, text: Any) -> str:
+        """Extract model name from text using dynamic parsing.
+
+        Parameters
+        ----------
+        text : Any
+            The text to extract the model name from.
+
+        Returns
+        -------
+        str
+            The extracted model name.
+        """
+        if not text or not isinstance(text, str):
+            return "Unknown"
+
+        try:
+            # Try JSON parsing first
+            if text.strip().startswith("{"):
+                model = self._extract_model_from_json(text)
+                if model != "Unknown":
+                    return model
+        except json.JSONDecodeError:
+            pass
+
+        # Use dynamic regex patterns to catch any model-like strings
+        model = self._extract_model_with_regex(text)
+        if model != "Unknown":
+            return model
+
+        return "Unknown"
+
+    def _extract_model_from_json(self, text: str) -> str:
+        """Extract model from JSON text using comprehensive key search.
+
+        Parameters
+        ----------
+        text : str
+            The JSON text to extract the model name from.
+
+        Returns
+        -------
+        str
+            The extracted model name.
+        """
+        try:
+            parsed = json.loads(text)
+
+            # Direct model keys
+            model_keys = [
+                "model",
+                "llm_model",
+                "engine",
+                "model_name",
+                "model_id",
+            ]
+            for key in model_keys:
+                if key in parsed and isinstance(parsed[key], str):
+                    return parsed[key]
+
+            # Nested searches for different structures
+            # Structure 1: config_list array (from agents.csv)
+            if "config_list" in parsed and isinstance(
+                parsed["config_list"], list
+            ):
+                for config in parsed["config_list"]:
+                    if isinstance(config, dict) and "model" in config:
+                        return config["model"]
+
+            # Structure 2: llm_config._model.config_list (from agents.csv)
+            if "llm_config" in parsed:
+                llm_config = parsed["llm_config"]
+                if isinstance(llm_config, dict):
+                    if "_model" in llm_config and isinstance(
+                        llm_config["_model"], dict
+                    ):
+                        model_config = llm_config["_model"]
+                        if "config_list" in model_config and isinstance(
+                            model_config["config_list"], list
+                        ):
+                            for config in model_config["config_list"]:
+                                if (
+                                    isinstance(config, dict)
+                                    and "model" in config
+                                ):
+                                    return config["model"]
+                        # Also check direct model keys in _model
+                        for key in model_keys:
+                            if key in model_config and isinstance(
+                                model_config[key], str
+                            ):
+                                return model_config[key]
+
+                    # Check llm_config level for model keys
+                    for key in model_keys:
+                        if key in llm_config and isinstance(
+                            llm_config[key], str
+                        ):
+                            return llm_config[key]
+
+            # Structure 3: Recursive search for any model key in nested objects
+
+            model = recursive_search(parsed, model_keys)
+            if model != "Unknown":
+                return model
+
+        except (json.JSONDecodeError, AttributeError, TypeError):
+            pass
+
+        return "Unknown"
+
+    def _extract_model_with_regex(self, text: str) -> str:
+        """Extract model using flexible regex patterns.
+
+        Parameters
+        ----------
+        text : str
+            The input text from which to extract the model name.
+
+        Returns
+        -------
+        str
+            The extracted model name or "Unknown" if not found.
+        """
+        # Dynamic patterns that can catch various model names
+        patterns = [
+            # OpenAI models - flexible to catch versions like gpt-4.1, gpt-4o...
+            r"\bgpt-[0-9]+(?:\.[0-9]+)?[a-zA-Z]*(?:-[a-zA-Z0-9]+)*\b",
+            # Claude models - flexible for various versions
+            r"\bclaude-[0-9]+(?:\.[0-9]+)?(?:-[a-zA-Z0-9]+)*\b",
+            # Gemini models
+            r"\bgemini-[a-zA-Z0-9]+(?:-[a-zA-Z0-9]+)*\b",
+            # Generic model patterns
+            r"\b[a-zA-Z]+-[0-9]+(?:\.[0-9]+)?[a-zA-Z]*(?:-[a-zA-Z0-9]+)*\b",
+            # Anthropic models
+            r"\b(?:anthropic|claude)[/_-][a-zA-Z0-9]+(?:[._-][a-zA-Z0-9]+)*\b",
+            # Other common patterns
+            r"\b(?:llama|mistral|falcon|vicuna|alpaca)[/_-]?[0-9]+[a-zA-Z]*(?:[._-][a-zA-Z0-9]+)*\b",
+            # Cohere models
+            r"\bcommand[/_-]?[a-zA-Z0-9]*\b",
+            # Generic AI model patterns
+            r"\b[a-zA-Z]+(?:ai|ml|model)[/_-]?[0-9]+[a-zA-Z]*\b",
+        ]
+
+        for pattern in patterns:
+            matches = re.findall(pattern, text, re.IGNORECASE)
+            if matches:
+                # Return the first match, but prefer longer matches
+                best_match = max(matches, key=len)
+                return best_match
+
+        # Last resort: look for any word that might be a model name
+        # This catches custom or unknown models
+        model_indicators = [
+            r'model["\']?\s*[:=]\s*["\']?([a-zA-Z0-9._-]+)',
+            r'engine["\']?\s*[:=]\s*["\']?([a-zA-Z0-9._-]+)',
+            r'"model"\s*:\s*"([^"]+)"',
+            r"'model'\s*:\s*'([^']+)'",
+        ]
+
+        for pattern in model_indicators:
+            matches = re.findall(pattern, text, re.IGNORECASE)
+            if matches:
+                return matches[0]
+
+        return "Unknown"
+
+    def is_human_input_waiting_period(
+        self,
+        prev_session: Series,
+        current_session: Series,
+        gap_duration: float,
+    ) -> bool:
+        """Detect if gap represents human input waiting.
+
+        Parameters
+        ----------
+        prev_session : Series
+            The previous session data.
+        current_session : Series
+            The current session data.
+        gap_duration : float
+            The duration of the gap to analyze.
+
+        Returns
+        -------
+        bool
+            True if gap likely represents human input waiting, False otherwise.
+        """
+        if gap_duration < 1.0:  # Reduced threshold for better detection
+            return False
+
+        if self.events_data is None:
+            return False
+
+        # Get events around the gap period
+        prev_end = self.parse_date(prev_session["end_time"])
+        current_start = self.parse_date(current_session["start_time"])
+
+        # Look for user message events right after the gap (within 1 second)
+        after_gap_window = current_start + pd.Timedelta(seconds=1)
+
+        user_events_after_gap = self.events_data[
+            (pd.to_datetime(self.events_data["timestamp"]) >= current_start)
+            & (
+                pd.to_datetime(self.events_data["timestamp"])
+                <= after_gap_window
+            )
+            & (self.events_data["event_name"] == "received_message")
+        ]
+
+        # Check if any of these events contain user messages
+        user_message_found = False
+        for _, event in user_events_after_gap.iterrows():
+            if self.is_user_message_event(event):
+                user_message_found = True
+                break
+
+        if user_message_found:
+            return True
+
+        # Alternative check: look for gaps that are longer and likely represent
+        # user thinking time
+        # This catches cases where user input detection might be missed
+        if gap_duration > 5.0:  # Longer gaps are more likely to be user input
+            # Check if there are any user messages in the broader timeline
+            # around this gap
+            broader_window_start = prev_end - pd.Timedelta(seconds=2)
+            broader_window_end = current_start + pd.Timedelta(seconds=5)
+
+            broader_events = self.events_data[
+                (
+                    pd.to_datetime(self.events_data["timestamp"])
+                    >= broader_window_start
+                )
+                & (
+                    pd.to_datetime(self.events_data["timestamp"])
+                    <= broader_window_end
+                )
+            ]
+
+            # Look for user message patterns
+            for _, event in broader_events.iterrows():
+                if self.is_user_message_event(event):
+                    return True
+
+        return False
+
+    def is_user_message_event(self, event: Series) -> bool:
+        """Check if an event represents a user message.
+
+        Parameters
+        ----------
+        event : Series
+            The event data to check.
+
+        Returns
+        -------
+        bool
+            True if the event represents a user message, False otherwise.
+        """
+        json_state = event.get("json_state", "")
+        if not json_state or not isinstance(json_state, str):
+            return False
+
+        try:
+            parsed = json.loads(json_state)
+            # Check for user role in message
+            if parsed.get("message", {}).get("role") == "user":
+                return True
+            # Check for customer sender (another indicator of user input)
+            if parsed.get("sender") == "customer":
+                return True
+        except (json.JSONDecodeError, AttributeError, TypeError):
+            pass
+
+        return False
+
+    def categorize_gap_activity(
+        self,
+        prev_session: Series,
+        current_session: Series,
+        gap_duration: float,
+    ) -> dict[str, Any]:
+        """Categorize what happened during a gap.
+
+        Parameters
+        ----------
+        prev_session : Series
+            The previous session data.
+        current_session : Series
+            The current session data.
+        gap_duration : float
+            The duration of the gap in seconds.
+
+        Returns
+        -------
+        dict[str, Any]
+            A dictionary categorizing the gap activity.
+        """
+        # First check for human input waiting period
+        if self.is_human_input_waiting_period(
+            prev_session, current_session, gap_duration
+        ):
+            return {
+                "type": "human_input_waiting",
+                "label": "👤 Human Input",
+                "detail": f"Waiting for user ({gap_duration:.1f}s)",
+            }
+
+        # Check for function calls during gap
+        if self.functions_data is not None:
+            prev_end = self.parse_date(prev_session["end_time"])
+            current_start = self.parse_date(current_session["start_time"])
+
+            gap_functions = self.functions_data[
+                (pd.to_datetime(self.functions_data["timestamp"]) >= prev_end)
+                & (
+                    pd.to_datetime(self.functions_data["timestamp"])
+                    <= current_start
+                )
+            ]
+
+            if not gap_functions.empty:
+                primary_function = gap_functions.iloc[0]["function_name"]
+
+                if (
+                    "transfer" in primary_function.lower()
+                    or "switch" in primary_function.lower()
+                ):
+                    detail = (
+                        f"{primary_function} → {current_session['source_name']}"
+                    )
+                    return {
+                        "type": "agent_transition",
+                        "label": "🔄 Transfer",
+                        "detail": detail,
+                    }
+                else:
+                    return {
+                        "type": "tool_call",
+                        "label": f"🛠️ {primary_function.replace('_', ' ')}",
+                        "detail": "Tool execution",
+                    }
+
+        # Check if agent changed
+        if prev_session["source_name"] != current_session["source_name"]:
+            detail = (
+                f"{prev_session['source_name']} → "
+                f"{current_session['source_name']}"
+            )
+            return {
+                "type": "agent_transition",
+                "label": "🔄 Agent Switch",
+                "detail": detail,
+            }
+
+        # For longer gaps without clear indicators,
+        # they might still be user input
+        # This provides a fallback for cases
+        # where the user input detection might miss
+        if gap_duration > 8.0:  # Longer gaps are more likely to be user input
+            return {
+                "type": "human_input_waiting",
+                "label": "👤 Likely User Input",
+                "detail": f"Probable user input ({gap_duration:.1f}s)",
+            }
+
+        return {
+            "type": "processing",
+            "label": "⚙️ Processing",
+            "detail": f"Processing ({gap_duration:.1f}s)",
+        }
+
+    def compress_timeline(
+        self,
+    ) -> tuple[list[dict[str, Any]], list[dict[str, Any]], float, float]:
+        """Create compressed timeline from chat data.
+
+        Processes chat data and generates a compressed timeline with gaps,
+        sessions, and cost information.
+
+        Returns
+        -------
+        tuple[list[dict[str, Any]], list[dict[str, Any]], float, float]
+            A tuple containing:
+            - Compressed timeline as a list of dictionaries.
+            - Cost timeline as a list of dictionaries.
+            - Total compressed time.
+            - Cumulative cost.
+
+        Raises
+        ------
+        ValueError
+            If chat data is not provided.
+        """
+        if self.chat_data is None:
+            raise ValueError("Chat data is required")
+
+        LOG.info("Starting timeline compression...")
+
+        # Sort by start time and calculate durations
+        chat_sorted = self.chat_data.copy()
+        chat_sorted["start_time"] = pd.to_datetime(chat_sorted["start_time"])
+        chat_sorted["end_time"] = pd.to_datetime(chat_sorted["end_time"])
+        chat_sorted = chat_sorted.sort_values("start_time")
+        chat_sorted["duration"] = (
+            chat_sorted["end_time"] - chat_sorted["start_time"]
+        ).dt.total_seconds()
+
+        LOG.info(
+            "Sorted chat data by start time. Total sessions: %d",
+            len(chat_sorted),
+        )
+
+        timeline: list[dict[str, Any]] = []
+        cost_timeline: list[dict[str, Any]] = []
+        current_compressed_time = 0.0
+        cumulative_cost = 0.0
+        session_id = 1
+
+        for _i, (_idx, row) in enumerate(chat_sorted.iterrows()):
+            try:
+                # Get agent name and handle missing values
+                agent_name = row["source_name"]
+                if self.is_missing_or_nan(agent_name):
+                    agent_name = "unknown_agent"
+
+                LOG.debug(
+                    "Processing session %d: %s",
+                    session_id,
+                    agent_name,
+                )
+                start_compressed = current_compressed_time
+                gap_before = 0
+                gap_activity = None
+
+                if session_id > 1:  # Not the first session
+                    prev_row = chat_sorted.iloc[session_id - 2]  # Previous row
+                    gap_duration = (
+                        row["start_time"] - prev_row["end_time"]
+                    ).total_seconds()
+
+                    gap_activity = self.categorize_gap_activity(
+                        prev_row, row, gap_duration
+                    )
+
+                    # Determine compressed gap duration
+                    if gap_activity["type"] == "human_input_waiting":
+                        compressed_gap = 1.0
+                        gap_before = gap_duration
+                    elif gap_duration > 2.0 and gap_activity["type"] in [
+                        "processing",
+                        "user_thinking",
+                    ]:
+                        compressed_gap = 2.0
+                        gap_before = gap_duration
+                    else:
+                        compressed_gap = gap_duration
+                        gap_before = gap_duration
+
+                    # Add gap to timeline if significant
+                    if gap_before > 0.1:
+                        gap_start = current_compressed_time
+                        gap_end = gap_start + compressed_gap
+
+                        timeline.append(
+                            {
+                                "id": f"gap_{session_id - 1}",
+                                "type": "gap",
+                                "gap_type": gap_activity["type"],
+                                "start": gap_start,
+                                "end": gap_end,
+                                "duration": compressed_gap,
+                                "value": compressed_gap,
+                                "real_duration": gap_before,
+                                "compressed": gap_activity["type"]
+                                == "human_input_waiting"
+                                or (
+                                    gap_duration > 2.0
+                                    and gap_activity["type"]
+                                    in ["processing", "user_thinking"]
+                                ),
+                                "color": ACTIVITY_COLORS.get(
+                                    gap_activity["type"],
+                                    ACTIVITY_COLORS["processing"],
+                                ),
+                                "label": (
+                                    gap_activity["label"]
+                                    + f" ({gap_before:.1f}s)"
+                                    if gap_before != compressed_gap
+                                    else gap_activity["label"]
+                                ),
+                                "y_position": session_id - 0.5,
+                            }
+                        )
+
+                    current_compressed_time += compressed_gap
+                    start_compressed = current_compressed_time
+
+                end_compressed = start_compressed + row["duration"]
+
+                # Extract token info with error handling
+                try:
+                    token_info = self.extract_token_info(
+                        row.get("request", ""), row.get("response", "")
+                    )
+                except Exception as e:
+                    LOG.error(
+                        "Error extracting token info for session %s: %s",
+                        session_id,
+                        e,
+                    )
+                    token_info = {
+                        "prompt_tokens": 0,
+                        "completion_tokens": 0,
+                        "total_tokens": 0,
+                    }
+
+                # Add session to timeline
+                timeline.append(
+                    {
+                        "id": f"session_{session_id}",
+                        "type": "session",
+                        "start": start_compressed,
+                        "end": end_compressed,
+                        "duration": row["duration"],
+                        "value": row["duration"],
+                        "agent": agent_name,
+                        # Will be updated if agents data available
+                        "agent_class": agent_name,
+                        "cost": row.get("cost", 0),
+                        "tokens": token_info["total_tokens"],
+                        "prompt_tokens": token_info["prompt_tokens"],
+                        "completion_tokens": token_info["completion_tokens"],
+                        "events": 1,  # Placeholder
+                        "color": DEFAULT_AGENT_COLOR,  # Will be updated later
+                        "label": f"S{session_id}: {agent_name}",
+                        "is_cached": bool(row.get("is_cached", False)),
+                        "y_position": session_id,
+                        "llm_model": self.extract_llm_model(
+                            agent_name, row.get("request", "")
+                        ),
+                        "session_id": row.get(
+                            "session_id", f"session_{session_id}"
+                        ),
+                        "real_start_time": row["start_time"].strftime(
+                            "%H:%M:%S"
+                        ),
+                        "request": row.get("request", ""),
+                        "response": row.get("response", ""),
+                    }
+                )
+
+                # Add to cost timeline
+                cumulative_cost += row.get("cost", 0)
+                cost_timeline.append(
+                    {
+                        "time": start_compressed + row["duration"] / 2,
+                        "cumulative_cost": cumulative_cost,
+                        "session_cost": row.get("cost", 0),
+                        "session_id": session_id,
+                    }
+                )
+
+                current_compressed_time = end_compressed
+                session_id += 1
+
+            except Exception as e:
+                LOG.error(
+                    "Error processing session %d: %s",
+                    session_id,
+                    e,
+                )
+                LOG.error("Row data: %s", dict(row))
+                raise
+
+        # Finalize timeline
+        if not timeline:
+            LOG.warning("No valid sessions found in chat data.")
+            return [], [], 0.0, 0.0
+        LOG.info(
+            "Timeline compression complete. Generated %d items.",
+            len(timeline),
+        )
+        return timeline, cost_timeline, current_compressed_time, cumulative_cost
+
+    def process_timeline(self) -> dict[str, Any]:
+        """Timeline processing function.
+
+        Processes chat data and generates a timeline with summary statistics.
+
+        Returns
+        -------
+        dict
+            A dictionary containing the processed timeline, cost timeline,
+        """
+        if self.chat_data is None:
+            raise ValueError("Chat data is required for processing")
+
+        timeline, cost_timeline, total_time, total_cost = (
+            self.compress_timeline()
+        )
+
+        # Get unique agents and assign colors
+        # (filter out any remaining NaN values)
+        agents_in_timeline = list(
+            {
+                item["agent"]
+                for item in timeline
+                if item["type"] == "session"
+                and not self.is_missing_or_nan(item["agent"])
+            }
+        )
+        agent_colors = self.generate_agent_colors(agents_in_timeline)
+
+        # Update timeline with colors and agent classes
+        for item in timeline:
+            if item["type"] == "session":
+                agent_name = item["agent"]
+                if self.is_missing_or_nan(agent_name):
+                    agent_name = "unknown_agent"
+                    item["agent"] = agent_name
+
+                item["color"] = agent_colors.get(
+                    agent_name, DEFAULT_AGENT_COLOR
+                )
+
+                # Update agent class if agents data available
+                if self.agents_data is not None:
+                    agent_row = self.agents_data[
+                        self.agents_data["name"] == agent_name
+                    ]
+                    if not agent_row.empty and "class" in agent_row.columns:
+                        agent_class = agent_row.iloc[0]["class"]
+                        if not self.is_missing_or_nan(agent_class):
+                            item["agent_class"] = agent_class
+
+        # Create agents list
+        agents: list[dict[str, Any]] = []
+        for agent_name in agents_in_timeline:
+            if self.is_missing_or_nan(agent_name):
+                continue
+
+            agent_class = agent_name  # Default
+            if self.agents_data is not None:
+                agent_row = self.agents_data[
+                    self.agents_data["name"] == agent_name
+                ]
+                if not agent_row.empty and "class" in agent_row.columns:
+                    agent_class_value = agent_row.iloc[0]["class"]
+                    if not self.is_missing_or_nan(agent_class_value):
+                        agent_class = agent_class_value
+
+            agents.append(
+                {
+                    "name": agent_name,
+                    "class": agent_class,
+                    "color": agent_colors.get(agent_name, DEFAULT_AGENT_COLOR),
+                }
+            )
+
+        # Calculate summary statistics
+        sessions = [item for item in timeline if item["type"] == "session"]
+        gaps = [item for item in timeline if item["type"] == "gap"]
+
+        total_tokens = sum(session["tokens"] for session in sessions)
+        gaps_compressed = sum(1 for gap in gaps if gap["compressed"])
+        time_saved = sum(
+            gap["real_duration"] - gap["duration"]
+            for gap in gaps
+            if gap["compressed"]
+        )
+
+        # Get model statistics
+        model_stats = {}
+        for session in sessions:
+            model = session.get("llm_model", "Unknown")
+            if model not in model_stats:
+                model_stats[model] = {"count": 0, "tokens": 0, "cost": 0}
+            model_stats[model]["count"] += 1
+            model_stats[model]["tokens"] += session.get("tokens", 0)
+            model_stats[model]["cost"] += session.get("cost", 0)
+
+        summary = {
+            "total_sessions": len(sessions),
+            "total_time": total_time,
+            "total_cost": total_cost,
+            "total_agents": len(agents_in_timeline),
+            "total_events": sum(session["events"] for session in sessions),
+            "total_tokens": total_tokens,
+            "avg_cost_per_session": (
+                total_cost / len(sessions) if sessions else 0
+            ),
+            "compression_info": {
+                "gaps_compressed": gaps_compressed,
+                "time_saved": time_saved,
+            },
+            "model_stats": model_stats,
+        }
+
+        # Create metadata
+        max_time = max([item["end"] for item in timeline]) if timeline else 0
+        max_cost = (
+            max([point["cumulative_cost"] for point in cost_timeline])
+            if cost_timeline
+            else 0
+        )
+
+        metadata = {
+            "time_range": [0, max_time * 1.1],
+            "cost_range": [0, max_cost * 1.1],
+            "colors": {
+                "human_input": ACTIVITY_COLORS["human_input_waiting"],
+                "processing": ACTIVITY_COLORS["processing"],
+                "agent_transition": ACTIVITY_COLORS["agent_transition"],
+                "cost_line": "#E91E63",
+            },
+        }
+
+        return {
+            "timeline": timeline,
+            "cost_timeline": cost_timeline,
+            "summary": summary,
+            "metadata": metadata,
+            "agents": agents,
+        }
+
+    @staticmethod
+    def get_short_results(results: dict[str, Any]) -> dict[str, Any]:
+        """Remove request/response from the timeline entries.
+
+        Parameters
+        ----------
+        results : dict[str, Any]
+            The original results dictionary.
+
+        Returns
+        -------
+        dict[str, Any]
+            The modified results dictionary with shortened timeline.
+        """
+        new_results = results.copy()
+        new_results["timeline"] = []
+        for item in results["timeline"]:
+            new_item = item.copy()
+            # Remove request and response fields
+            new_item.pop("request", None)
+            new_item.pop("response", None)
+            new_results["timeline"].append(new_item)
+        return new_results
+
+    @staticmethod
+    def get_files(logs_dir: Path | str) -> dict[str, str | None]:
+        """Get all CSV files in the specified directory.
+
+        Parameters
+        ----------
+        logs_dir : Path | str
+            The directory to search for CSV files.
+
+        Returns
+        -------
+        dict[str, str | None]
+            A dictionary mapping CSV file names to their paths
+            or None if not found.
+        """
+        agents_file = os.path.join(logs_dir, "agents.csv")
+        chat_file = os.path.join(logs_dir, "chat_completions.csv")
+        events_file = os.path.join(logs_dir, "events.csv")
+        functions_file = os.path.join(logs_dir, "function_calls.csv")
+
+        return {
+            "agents": agents_file if os.path.exists(agents_file) else None,
+            "chat": chat_file if os.path.exists(chat_file) else None,
+            "events": events_file if os.path.exists(events_file) else None,
+            "functions": (
+                functions_file if os.path.exists(functions_file) else None
+            ),
+        }
+
+
+def recursive_search(obj: Any, keys_to_find: list[str]) -> str:
+    """Recursively search for keys in a nested structure.
+
+    Parameters
+    ----------
+    obj : Any
+        The object to search within.
+    keys_to_find : list[str]
+        The keys to search for.
+
+    Returns
+    -------
+    str
+        The found value or "Unknown" if not found.
+    """
+    if isinstance(obj, dict):
+        for key in keys_to_find:
+            if key in obj and isinstance(obj[key], str) and obj[key].strip():
+                return obj[key]
+        for value in obj.values():
+            result = recursive_search(value, keys_to_find)
+            if result != "Unknown":
+                return result
+    elif isinstance(obj, list):
+        for item in obj:
+            result = recursive_search(item, keys_to_find)
+            if result != "Unknown":
+                return result
+    return "Unknown"
+
+
+def main() -> int:
+    """Run the timeline processor."""
+    parser = argparse.ArgumentParser(
+        description="Process timeline CSV files and output JSON"
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        required=True,
+        help="Directory to get the csv files from and save the output JSON",
+    )
+
+    args = parser.parse_args()
+
+    output_dir = Path(args.output_dir).resolve()
+    if not output_dir.is_dir():
+        LOG.error("Output directory does not exist: %s", output_dir)
+        return 1
+    LOG.info("Using output directory: %s", output_dir)
+    files = TimelineProcessor.get_files(output_dir)
+    agents_file = files["agents"]
+    chat_file = files["chat"]
+    events_file = files["events"]
+    functions_file = files["functions"]
+    if not any(files.values()):
+        LOG.error(
+            "No CSV files found in the output directory: %s",
+            output_dir,
+        )
+        return 1
+    output_file = os.path.join(output_dir, "timeline.json")
+    if not agents_file or not os.path.exists(agents_file):
+        agents_file = None
+    if not chat_file or not os.path.exists(chat_file):
+        chat_file = None
+    if not events_file or not os.path.exists(events_file):
+        events_file = None
+    if not functions_file or not os.path.exists(functions_file):
+        functions_file = None
+
+    processor = TimelineProcessor()
+
+    try:
+        processor.load_csv_files(
+            agents_file=agents_file,
+            chat_file=chat_file,
+            events_file=events_file,
+            functions_file=functions_file,
+        )
+
+        result = processor.process_timeline()
+
+        with open(output_file, "w", encoding="utf-8") as f:
+            json.dump(result, f, indent=2, default=str)
+
+        print(f"Timeline data saved to {output_file}")
+        print(
+            f"Summary: {result['summary']['total_sessions']} sessions, "
+            f"${result['summary']['total_cost']:.6f} total cost, "
+            f"{result['summary']['total_time']:.1f}s duration"
+        )
+
+        # Print model statistics
+        print("\nModel Statistics:")
+        for model, stats in result["summary"]["model_stats"].items():
+            print(
+                f"  {model}: {stats['count']} sessions, "
+                f"{stats['tokens']} tokens, ${stats['cost']:.6f}"
+            )
+
+    except Exception as e:
+        print(f"Error: {e}")
+        return 1
+
+    return 0
+
+
+if __name__ == "__main__":
+    exit(main())