Support for displaying token stats in the dashboard

pyproject.toml

@@ -32,6 +32,8 @@ dependencies = [
   "docstring_parser>=0.16",
   "joblib>=1.5.1",
   "tqdm>=4.67.1",
+  "tiktoken>=0.9.0",
+  "anthropic>=0.54.0",
 ]
 
 [project.scripts]

src/serena/agent.py

@@ -21,7 +21,7 @@ from sensai.util.logging import LogTime
 from tqdm import tqdm
 
 from serena import serena_version
-from serena.analytics import record_tool_usage
+from serena.analytics import RegisteredTokenCountEstimator, ToolUsageStats
 from serena.config.context_mode import SerenaAgentContext, SerenaAgentMode
 from serena.config.serena_config import SerenaConfig, ToolSet, get_serena_managed_dir
 from serena.constants import (
@@ -180,11 +180,19 @@ class SerenaAgent:
         if self._gui_log_handler is not None:
             self._gui_log_handler.log_viewer.set_tool_names(tool_names)
 
+        self._tool_usage_stats: ToolUsageStats | None = None
+        if self.serena_config.record_tool_usage_stats:
+            token_count_estimator = RegisteredTokenCountEstimator[self.serena_config.token_count_estimator]
+            log.info(f"Tool usage statistics recording is enabled with token count estimator: {token_count_estimator.name}.")
+            self._tool_usage_stats = ToolUsageStats(token_count_estimator)
+
         # start the dashboard (web frontend), registering its log handler
         if self.serena_config.web_dashboard:
             dashboard_log_handler = MemoryLogHandler(level=serena_log_level)
             Logger.root.addHandler(dashboard_log_handler)
-            self._dashboard_thread, port = SerenaDashboardAPI(dashboard_log_handler, tool_names).run_in_thread()
+            self._dashboard_thread, port = SerenaDashboardAPI(
+                dashboard_log_handler, tool_names, tool_usage_stats=self._tool_usage_stats
+            ).run_in_thread()
             if self.serena_config.web_dashboard_open_on_launch:
                 # open the dashboard URL in the default web browser (using a separate process to control
                 # output redirection)
@@ -241,6 +249,19 @@ class SerenaAgent:
                     "You should now pass either --project <project_name> or --project <project_root>."
                 )
 
+    def record_tool_usage_if_enabled(self, input_kwargs: dict, tool_result: str | dict, tool: Tool) -> None:
+        """
+        Record the usage of a tool with the given input and output strings if tool usage statistics recording is enabled.
+        """
+        tool_name = tool.get_name()
+        if self._tool_usage_stats is not None:
+            input_str = str(input_kwargs)
+            output_str = str(tool_result)
+            log.debug(f"Recording tool usage for tool '{tool_name}'")
+            self._tool_usage_stats.record_tool_usage(tool_name, input_str, output_str)
+        else:
+            log.debug(f"Tool usage statistics recording is disabled, not recording usage of '{tool_name}'.")
+
     @staticmethod
     def _open_dashboard(port: int) -> None:
         # Redirect stdout and stderr file descriptors to /dev/null,

src/serena/analytics.py

@@ -1,39 +1,150 @@
 from __future__ import annotations
 
+import logging
 import threading
+from abc import ABC
 from collections import defaultdict
+from copy import copy
+from dataclasses import asdict, dataclass
+from enum import Enum
+
+from anthropic.types import MessageParam, MessageTokensCount
+from dotenv import load_dotenv
+
+log = logging.getLogger(__name__)
 
 
-class ToolStatsEntry:
-    def __init__(self) -> None:
-        self.count = 0
-        self.input_chars = 0
-        self.output_chars = 0
-
-    def to_dict(self) -> dict[str, int]:
-        return {
-            "count": self.count,
-            "input_chars": self.input_chars,
-            "output_chars": self.output_chars,
-        }
-
-_lock = threading.Lock()
-_tool_stats: dict[str, ToolStatsEntry] = defaultdict(ToolStatsEntry)
+class TokenCountEstimator(ABC):
+    def estimate_token_count(self, text: str) -> int:
+        """
+        Estimate the number of tokens in the given text.
+        This is an abstract method that should be implemented by subclasses.
+        """
 
 
-def record_tool_usage(tool_name: str, input_chars: int, output_chars: int) -> None:
-    with _lock:
-        entry = _tool_stats[tool_name]
-        entry.count += 1
-        entry.input_chars += input_chars
-        entry.output_chars += output_chars
+class TiktokenCountEstimator(TokenCountEstimator):
+    """
+    Approximate token count using tiktoken.
+    """
+
+    def __init__(self, model_name: str = "gpt-4o"):
+        """
+        The tokenizer will be downloaded on the first initialization, which may take some time.
+
+        :param model_name: see `tiktoken.model` to see available models.
+        """
+        import tiktoken
+
+        log.info(f"Loading tiktoken encoding for model {model_name}, this may take a while on the first run.")
+        self._encoding = tiktoken.encoding_for_model(model_name)
+
+    def estimate_token_count(self, text: str) -> int:
+        return len(self._encoding.encode(text))
 
 
-def get_tool_stats() -> dict[str, dict[str, int]]:
-    with _lock:
-        return {name: entry.to_dict() for name, entry in _tool_stats.items()}
+class AnthropicTokenCount(TokenCountEstimator):
+    """
+    The exact count using the Anthropic API.
+    Counting is free, but has a rate limit and will require an API key,
+    (typically, set through an env variable).
+    See https://docs.anthropic.com/en/docs/build-with-claude/token-counting
+    """
+
+    def __init__(self, model_name: str = "claude-sonnet-4-20250514", api_key: str | None = None):
+        import anthropic
+
+        self._model_name = model_name
+        if api_key is None:
+            load_dotenv()
+        self._anthropic_client = anthropic.Anthropic(api_key=api_key)
+
+    def _send_count_tokens_request(self, text: str) -> MessageTokensCount:
+        return self._anthropic_client.messages.count_tokens(
+            model=self._model_name,
+            messages=[MessageParam(role="user", content=text)],
+        )
+
+    def estimate_token_count(self, text: str) -> int:
+        return self._send_count_tokens_request(text).input_tokens
 
 
-def clear_tool_stats() -> None:
-    with _lock:
-        _tool_stats.clear()
+_registered_token_estimator_instances_cache: dict[RegisteredTokenCountEstimator, TokenCountEstimator] = {}
+
+
+class RegisteredTokenCountEstimator(Enum):
+    TIKTOKEN_GPT4O = "TIKTOKEN_GPT4O"
+    ANTHROPIC_CLAUDE_SONNET_4 = "ANTHROPIC_CLAUDE_SONNET_4"
+
+    @classmethod
+    def get_valid_names(cls) -> list[str]:
+        """
+        Get a list of all registered token count estimator names.
+        """
+        return [estimator.name for estimator in cls]
+
+    def _create_estimator(self) -> TokenCountEstimator:
+        match self:
+            case self.TIKTOKEN_GPT4O:
+                return TiktokenCountEstimator(model_name="gpt-4o")
+            case self.ANTHROPIC_CLAUDE_SONNET_4:
+                return AnthropicTokenCount(model_name="claude-sonnet-4-20250514")
+            case _:
+                raise ValueError(f"Unknown token count estimator: {self.value}")
+
+    def load_estimator(self) -> TokenCountEstimator:
+        estimator_instance = _registered_token_estimator_instances_cache.get(self)
+        if estimator_instance is None:
+            estimator_instance = self._create_estimator()
+            _registered_token_estimator_instances_cache[self] = estimator_instance
+        return estimator_instance
+
+
+class ToolUsageStats:
+    """
+    A class to record and manage tool usage statistics.
+    """
+
+    def __init__(self, token_count_estimator: RegisteredTokenCountEstimator = RegisteredTokenCountEstimator.TIKTOKEN_GPT4O):
+        self._token_count_estimator = token_count_estimator.load_estimator()
+        self._token_estimator_name = token_count_estimator.value
+        self._tool_stats: dict[str, ToolUsageStats.Entry] = defaultdict(ToolUsageStats.Entry)
+        self._tool_stats_lock = threading.Lock()
+
+    @dataclass(kw_only=True)
+    class Entry:
+        num_times_called: int = 0
+        input_tokens: int = 0
+        output_tokens: int = 0
+
+        def update_on_call(self, input_tokens: int, output_tokens: int) -> None:
+            """
+            Update the entry with the number of tokens used for a single call.
+            """
+            self.num_times_called += 1
+            self.input_tokens += input_tokens
+            self.output_tokens += output_tokens
+
+    def _estimate_token_count(self, text: str) -> int:
+        return self._token_count_estimator.estimate_token_count(text)
+
+    def get_stats(self, tool_name: str) -> ToolUsageStats.Entry:
+        """
+        Get (a copy of) the current usage statistics for a specific tool.
+        """
+        with self._tool_stats_lock:
+            return copy(self._tool_stats[tool_name])
+
+    def record_tool_usage(self, tool_name: str, input_str: str, output_str: str) -> None:
+        input_tokens = self._estimate_token_count(input_str)
+        output_tokens = self._estimate_token_count(output_str)
+        with self._tool_stats_lock:
+            entry = self._tool_stats[tool_name]
+            entry.update_on_call(input_tokens, output_tokens)
+
+    def get_tool_stats_dict(self) -> dict[str, dict[str, int]]:
+        with self._tool_stats_lock:
+            return {name: asdict(entry) for name, entry in self._tool_stats.items()}
+
+    def clear(self) -> None:
+        with self._tool_stats_lock:
+            self._tool_stats.clear()

src/serena/config/serena_config.py

@@ -27,6 +27,8 @@ from serena.util.general import load_yaml, save_yaml
 from serena.util.inspection import determine_programming_language_composition
 from solidlsp.ls_config import Language
 
+from ..analytics import RegisteredTokenCountEstimator
+
 if TYPE_CHECKING:
     from ..project import Project
 
@@ -254,6 +256,17 @@ class SerenaConfig(ToolInclusionDefinition, ToStringMixin):
     """
     whether to apply JetBrains mode
     """
+    record_tool_usage_stats: bool = False
+    """Whether to record tool usage statistics, they will be shown in the web dashboard if recording is active. 
+    """
+    token_count_estimator: str = RegisteredTokenCountEstimator.TIKTOKEN_GPT4O.name
+    """Only relevant if `record_tool_usage` is True; the name of the token count estimator to use for tool usage statistics.
+    See the `RegisteredTokenCountEstimator` enum for available options.
+    
+    Note: some token estimators (like tiktoken) may require downloading data files
+    on the first run, which can take some time and require internet access. Others, like the Anthropic ones, may require an API key
+    and rate limits may apply.
+    """
 
     CONFIG_FILE = "serena_config.yml"
     CONFIG_FILE_DOCKER = "serena_config.docker.yml"  # Docker-specific config file; auto-generated if missing, mounted via docker-compose for user customization
@@ -351,6 +364,10 @@ class SerenaConfig(ToolInclusionDefinition, ToStringMixin):
         instance.excluded_tools = loaded_commented_yaml.get("excluded_tools", [])
         instance.included_optional_tools = loaded_commented_yaml.get("included_optional_tools", [])
         instance.jetbrains = loaded_commented_yaml.get("jetbrains", False)
+        instance.record_tool_usage_stats = loaded_commented_yaml.get("record_tool_usage_stats", False)
+        instance.token_count_estimator = loaded_commented_yaml.get(
+            "token_count_estimator", RegisteredTokenCountEstimator.TIKTOKEN_GPT4O.name
+        )
 
         # re-save the configuration file if any migrations were performed
         if num_project_migrations > 0:

src/serena/dashboard.py

@@ -9,7 +9,7 @@ from flask import Flask, Response, request, send_from_directory
 from pydantic import BaseModel
 from sensai.util import logging
 
-from serena.analytics import clear_tool_stats, get_tool_stats
+from serena.analytics import ToolUsageStats
 from serena.constants import SERENA_DASHBOARD_DIR, SERENA_LOG_FORMAT
 
 log = logging.getLogger(__name__)
@@ -76,12 +76,17 @@ class SerenaDashboardAPI:
     log = logging.getLogger(__qualname__)
 
     def __init__(
-        self, memory_log_handler: MemoryLogHandler, tool_names: list[str], shutdown_callback: Callable[[], None] | None = None
+        self,
+        memory_log_handler: MemoryLogHandler,
+        tool_names: list[str],
+        shutdown_callback: Callable[[], None] | None = None,
+        tool_usage_stats: ToolUsageStats | None = None,
     ) -> None:
         self._memory_log_handler = memory_log_handler
         self._tool_names = tool_names
         self._shutdown_callback = shutdown_callback
         self._app = Flask(__name__)
+        self._tool_usage_stats = tool_usage_stats
         self._setup_routes()
 
     @property
@@ -139,11 +144,13 @@ class SerenaDashboardAPI:
         return ResponseToolNames(tool_names=self._tool_names)
 
     def _get_tool_stats(self) -> ResponseToolStats:
-        return ResponseToolStats(stats=get_tool_stats())
+        if self._tool_usage_stats is not None:
+            return ResponseToolStats(stats=self._tool_usage_stats.get_tool_stats_dict())
+        else:
+            return ResponseToolStats(stats={})
 
-    @staticmethod
-    def _clear_tool_stats() -> None:
-        clear_tool_stats()
+    def _clear_tool_stats(self) -> None:
+        self._tool_usage_stats.clear()
 
     def _shutdown(self) -> None:
         log.info("Shutting down Serena")

src/serena/resources/dashboard/dashboard.js

@@ -226,9 +226,9 @@ class Dashboard {
 
     displayStats(stats) {
         const names = Object.keys(stats);
-        const counts = names.map(n => stats[n].count);
-        const inputChars = names.map(n => stats[n].input_chars);
-        const outputChars = names.map(n => stats[n].output_chars);
+        const counts = names.map(n => stats[n].num_times_called);
+        const inputTokens = names.map(n => stats[n].input_tokens);
+        const outputTokens = names.map(n => stats[n].output_tokens);
 
         const countCtx = document.getElementById('count-chart');
         const inputCtx = document.getElementById('input-chart');
@@ -255,13 +255,13 @@ class Dashboard {
 
         this.inputChart = new Chart(inputCtx, {
             type: 'bar',
-            data: { labels: names, datasets: [{ label: 'Input Chars', data: inputChars }] },
+            data: { labels: names, datasets: [{ label: 'Input Tokens', data: inputTokens }] },
             options: { scales: { y: { beginAtZero: true } } }
         });
 
         this.outputChart = new Chart(outputCtx, {
             type: 'bar',
-            data: { labels: names, datasets: [{ label: 'Output Chars', data: outputChars }] },
+            data: { labels: names, datasets: [{ label: 'Output Tokens', data: outputTokens }] },
             options: { scales: { y: { beginAtZero: true } } }
         });
     }

src/serena/resources/serena_config.template.yml

@@ -48,6 +48,18 @@ jetbrains: False
 # NOTE: The plugin is yet unreleased. This is for Serena developers only.
 
 
+record_tool_usage_stats:  False
+# whether to record tool usage statistics, they will be shown in the web dashboard if recording is active.
+
+token_count_estimator: TIKTOKEN_GPT4O
+# Only relevant if `record_tool_usage` is True; the name of the token count estimator to use for tool usage statistics.
+# See the `RegisteredTokenCountEstimator` enum for available options.
+#
+# Note: some token estimators (like tiktoken) may require downloading data files
+# on the first run, which can take some time and require internet access. Others, like the Anthropic ones, may require an API key
+# and rate limits may apply.
+
+
 # MANAGED BY SERENA, KEEP AT THE BOTTOM OF THE YAML AND DON'T EDIT WITHOUT NEED
 # The list of registered projects.
 # To add a project, within a chat, simply ask Serena to "activate the project /path/to/project" or,

src/serena/tools/tools_base.py

@@ -5,13 +5,13 @@ from abc import ABC
 from collections.abc import Iterable
 from dataclasses import dataclass
 from types import TracebackType
-from typing import Any, Protocol, Self, TYPE_CHECKING, TypeVar
+from typing import TYPE_CHECKING, Any, Protocol, Self, TypeVar
 
 from mcp.server.fastmcp.utilities.func_metadata import FuncMetadata, func_metadata
 from sensai.util import logging
 from sensai.util.string import dict_string
-from serena.project import Project
 
+from serena.project import Project
 from serena.prompt_factory import PromptFactory
 from serena.symbol import LanguageServerSymbolRetriever
 from serena.util.class_decorators import singleton
@@ -92,7 +92,8 @@ class ToolMarkerOptional:
 class ApplyMethodProtocol(Protocol):
     """Callable protocol for the apply method of a tool."""
 
-    def __call__(self, *args: Any, **kwargs: Any) -> str: pass
+    def __call__(self, *args: Any, **kwargs: Any) -> str:
+        pass
 
 
 class Tool(Component):
@@ -237,6 +238,7 @@ class Tool(Component):
 
                 # apply the actual tool
                 result = apply_fn(**kwargs)
+                self.agent.record_tool_usage_if_enabled(kwargs, result, self)
 
             except Exception as e:
                 if not catch_exceptions: