diff --git a/pyproject.toml b/pyproject.toml index e0f4794e..7eacf741 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -225,6 +225,7 @@ ignore = [ "B904", # forces use of raise from other_exception "RUF012", # forbids mutable attributes as ClassVar "SIM117", # forbids nested with statements + "C400", # wants to unnecessarily force use of list comprehension ] unfixable = ["F841", "F601", "F602", "B018"] extend-fixable = ["F401", "B905", "W291"] diff --git a/scripts/print_tool_overview.py b/scripts/print_tool_overview.py index 955ee78c..3649a7d5 100644 --- a/scripts/print_tool_overview.py +++ b/scripts/print_tool_overview.py @@ -1,4 +1,4 @@ from serena.agent import ToolRegistry if __name__ == "__main__": - ToolRegistry.print_tool_overview() + ToolRegistry().print_tool_overview() diff --git a/src/serena/agent.py b/src/serena/agent.py index 5a22b2d3..bc5e33ae 100644 --- a/src/serena/agent.py +++ b/src/serena/agent.py @@ -23,7 +23,7 @@ from sensai.util.logging import LogTime from serena import serena_version from serena.config.context_mode import SerenaAgentContext, SerenaAgentMode -from serena.config.serena_config import Project, SerenaConfig, get_serena_managed_dir +from serena.config.serena_config import Project, SerenaConfig, ToolSet, get_serena_managed_dir from serena.constants import ( SERENA_LOG_FORMAT, ) @@ -234,13 +234,13 @@ class SerenaAgent: self._context = context # instantiate all tool classes - self._all_tools: dict[type[Tool], Tool] = {tool_class: tool_class(self) for tool_class in ToolRegistry.get_all_tool_classes()} + self._all_tools: dict[type[Tool], Tool] = {tool_class: tool_class(self) for tool_class in ToolRegistry().get_all_tool_classes()} tool_names = [tool.get_name_from_cls() for tool in self._all_tools.values()] - # determine the set exposed tools (which e.g. the MCP shall see), limited by the context - # (which is fixed for the session) - excluded_tool_classes = set(self._context.get_excluded_tool_classes()) - self._exposed_tools = {tc: t for tc, t in self._all_tools.items() if tc not in excluded_tool_classes} + # determine the set exposed tools (which e.g. the MCP shall see), limited by the Serena config + # as well as the context (which is fixed for the session) + tool_set = ToolSet().apply(self.serena_config, self._context) + self._exposed_tools = {tc: t for tc, t in self._all_tools.items() if tool_set.includes_name(t.get_name())} # If GUI log window is enabled, set the tool names for highlighting if self._gui_log_handler is not None: @@ -408,42 +408,20 @@ class SerenaAgent: def _update_active_tools(self) -> None: """ Update the active tools based on context, modes, and project configuration. - All tool exclusions are merged together. """ - excluded_tool_classes: set[type[Tool]] = set() - # modes - for mode in self._modes: - mode_excluded_tool_classes = mode.get_excluded_tool_classes() - if len(mode_excluded_tool_classes) > 0: - log.info( - f"Mode {mode.name} excluded {len(mode_excluded_tool_classes)} tools: {', '.join([tool.get_name_from_cls() for tool in mode_excluded_tool_classes])}" - ) - excluded_tool_classes.update(mode_excluded_tool_classes) - # context - context_excluded_tool_classes = self._context.get_excluded_tool_classes() - if len(context_excluded_tool_classes) > 0: - log.info( - f"Context {self._context.name} excluded {len(context_excluded_tool_classes)} tools: {', '.join([tool.get_name_from_cls() for tool in context_excluded_tool_classes])}" - ) - excluded_tool_classes.update(context_excluded_tool_classes) - # project config + tool_set = ToolSet().apply(self.serena_config, self._context, *self._modes) if self._active_project is not None: - project_excluded_tool_classes = self._active_project.project_config.get_excluded_tool_classes() - if len(project_excluded_tool_classes) > 0: - log.info( - f"Project {self._active_project.project_name} excluded {len(project_excluded_tool_classes)} tools: {', '.join([tool.get_name_from_cls() for tool in project_excluded_tool_classes])}" - ) - excluded_tool_classes.update(project_excluded_tool_classes) + tool_set.apply(self._active_project.project_config) if self._active_project.project_config.read_only: - for tool_class in self._all_tools: - if tool_class.can_edit(): - excluded_tool_classes.add(tool_class) + tool_set.exclude_editing_tools() self._active_tools = { - tool_class: tool_instance for tool_class, tool_instance in self._all_tools.items() if tool_class not in excluded_tool_classes + tool_class: tool_instance + for tool_class, tool_instance in self._all_tools.items() + if tool_set.includes_name(tool_instance.get_name()) } - log.info(f"Active tools after all exclusions ({len(self._active_tools)}): {', '.join(self.get_active_tool_names())}") + log.info(f"Active tools ({len(self._active_tools)}): {', '.join(self.get_active_tool_names())}") def issue_task(self, task: Callable[[], Any], name: str | None = None) -> Future: """ @@ -648,7 +626,7 @@ class SerenaAgent: return self._all_tools[tool_class] # type: ignore def print_tool_overview(self) -> None: - ToolRegistry.print_tool_overview(self._active_tools.values()) + ToolRegistry().print_tool_overview(self._active_tools.values()) def mark_file_modified(self, relativ_path: str) -> None: assert self.lines_read is not None diff --git a/src/serena/config/context_mode.py b/src/serena/config/context_mode.py index b8cfb9f8..0ee3c264 100644 --- a/src/serena/config/context_mode.py +++ b/src/serena/config/context_mode.py @@ -4,7 +4,7 @@ Context and Mode configuration loader import os from copy import copy -from dataclasses import asdict, dataclass, field +from dataclasses import asdict, dataclass from enum import Enum from pathlib import Path from typing import TYPE_CHECKING, Self @@ -13,16 +13,17 @@ import yaml from sensai.util import logging from sensai.util.string import ToStringMixin +from serena.config.serena_config import ToolInclusionDefinition from serena.constants import CONTEXT_YAMLS_DIR, DEFAULT_CONTEXT, DEFAULT_MODES, MODE_YAMLS_DIR if TYPE_CHECKING: - from serena.agent import Tool + pass log = logging.getLogger(__name__) -@dataclass -class SerenaAgentMode: +@dataclass(kw_only=True) +class SerenaAgentMode(ToolInclusionDefinition): """Represents a mode of operation for the agent, typically read off a YAML file. An agent can be in multiple modes simultaneously as long as they are not mutually exclusive. The modes can be adjusted after the agent is running, for example for switching from planning to editing. @@ -31,7 +32,6 @@ class SerenaAgentMode: name: str prompt: str description: str = "" - excluded_tools: set[str] = field(default_factory=set) def to_json_dict(self) -> dict[str, str | list[str]]: result = asdict(self) @@ -50,12 +50,6 @@ class SerenaAgentMode: if self.excluded_tools: print(" excluded tools:\n " + ", ".join(sorted(self.excluded_tools))) - def get_excluded_tool_classes(self) -> list[type["Tool"]]: - """Get the list of tool classes that are excluded from the mode.""" - from serena.agent import ToolRegistry - - return [ToolRegistry.get_tool_class_by_name(tool_name) for tool_name in self.excluded_tools] - @classmethod def from_yaml(cls, yaml_path: str | Path) -> Self: """Load a mode from a YAML file.""" @@ -93,8 +87,8 @@ class SerenaAgentMode: return cls.from_yaml(name_or_path) -@dataclass -class SerenaAgentContext(ToStringMixin): +@dataclass(kw_only=True) +class SerenaAgentContext(ToolInclusionDefinition, ToStringMixin): """Represents a context where the agent is operating (an IDE, a chat, etc.), typically read off a YAML file. An agent can only be in a single context at a time. The contexts cannot be changed after the agent is running. @@ -103,7 +97,6 @@ class SerenaAgentContext(ToStringMixin): name: str prompt: str description: str = "" - excluded_tools: set[str] = field(default_factory=set) def _tostring_includes(self) -> list[str]: return ["name"] @@ -119,12 +112,6 @@ class SerenaAgentContext(ToStringMixin): data["excluded_tools"] = set(data["excluded_tools"]) return cls(**data) - def get_excluded_tool_classes(self) -> list[type["Tool"]]: - """Get the list of tool classes that are excluded from the context.""" - from serena.agent import ToolRegistry - - return [ToolRegistry.get_tool_class_by_name(tool_name) for tool_name in self.excluded_tools] - @classmethod def from_yaml(cls, yaml_path: str | Path) -> Self: """Load a context from a YAML file.""" diff --git a/src/serena/config/serena_config.py b/src/serena/config/serena_config.py index 883c5931..306393cb 100644 --- a/src/serena/config/serena_config.py +++ b/src/serena/config/serena_config.py @@ -4,6 +4,7 @@ The Serena Model Context Protocol (MCP) Server import os import shutil +from collections.abc import Iterable from copy import deepcopy from dataclasses import dataclass, field from functools import cached_property @@ -22,7 +23,6 @@ from serena.constants import ( SELENA_CONFIG_TEMPLATE_FILE, SERENA_MANAGED_DIR_NAME, ) -from serena.tools import Tool, ToolRegistry from serena.util.general import load_yaml, save_yaml from serena.util.inspection import determine_programming_language_composition from solidlsp.ls_config import Language @@ -35,6 +35,53 @@ T = TypeVar("T") DEFAULT_TOOL_TIMEOUT: float = 240 +class ToolSet: + def __init__(self) -> None: + from serena.tools import ToolRegistry + + self._tool_names = set(ToolRegistry().get_tool_names_default_enabled()) + + def apply(self, *tool_inclusion_definitions: "ToolInclusionDefinition") -> Self: + from serena.tools import ToolRegistry + + registry = ToolRegistry() + for definition in tool_inclusion_definitions: + for included_tool in definition.included_optional_tools: + if not registry.is_valid_tool_name(included_tool): + raise ValueError(f"Invalid tool name '{included_tool}' provided for inclusion") + self._tool_names.add(included_tool) + for excluded_tool in definition.excluded_tools: + if not registry.is_valid_tool_name(excluded_tool): + raise ValueError(f"Invalid tool name '{excluded_tool}' provided for exclusion") + if excluded_tool in self._tool_names: + self._tool_names.remove(excluded_tool) + return self + + def exclude_editing_tools(self) -> Self: + from serena.tools import ToolRegistry + + registry = ToolRegistry() + for tool_name in list(self._tool_names): + if registry.get_tool_class_by_name(tool_name).can_edit(): + self._tool_names.remove(tool_name) + return self + + def get_tool_names(self) -> set[str]: + """ + Returns the names of the tools that are currently included in the tool set. + """ + return self._tool_names + + def includes_name(self, tool_name: str) -> bool: + return tool_name in self._tool_names + + +@dataclass +class ToolInclusionDefinition: + excluded_tools: Iterable[str] = () + included_optional_tools: Iterable[str] = () + + class SerenaConfigError(Exception): pass @@ -56,12 +103,11 @@ def is_running_in_docker() -> bool: return False -@dataclass -class ProjectConfig(ToStringMixin): +@dataclass(kw_only=True) +class ProjectConfig(ToolInclusionDefinition, ToStringMixin): project_name: str language: Language ignored_paths: list[str] = field(default_factory=list) - excluded_tools: set[str] = field(default_factory=set) read_only: bool = False ignore_all_files_in_gitignore: bool = True initial_prompt: str = "" @@ -123,7 +169,8 @@ class ProjectConfig(ToStringMixin): project_name=project_name, language=language, ignored_paths=data.get("ignored_paths", []), - excluded_tools=set(data.get("excluded_tools", [])), + excluded_tools=data.get("excluded_tools", []), + included_optional_tools=data.get("included_optional_tools", []), read_only=data.get("read_only", False), ignore_all_files_in_gitignore=data.get("ignore_all_files_in_gitignore", True), initial_prompt=data.get("initial_prompt", ""), @@ -148,9 +195,6 @@ class ProjectConfig(ToStringMixin): yaml_data["project_name"] = project_root.name return cls._from_dict(yaml_data) - def get_excluded_tool_classes(self) -> set[type["Tool"]]: - return set(ToolRegistry.get_tool_class_by_name(tool_name) for tool_name in self.excluded_tools) - @dataclass class Project: @@ -178,7 +222,7 @@ class Project: @dataclass(kw_only=True) -class SerenaConfig: +class SerenaConfig(ToolInclusionDefinition): """ Holds the Serena agent configuration, which is typically loaded from a YAML configuration file (when instantiated via :method:`from_config_file`), which is updated when projects are added or removed. @@ -287,6 +331,8 @@ class SerenaConfig: instance.web_dashboard_open_on_launch = loaded_commented_yaml.get("web_dashboard_open_on_launch", True) instance.tool_timeout = loaded_commented_yaml.get("tool_timeout", DEFAULT_TOOL_TIMEOUT) instance.trace_lsp_communication = loaded_commented_yaml.get("trace_lsp_communication", False) + instance.excluded_tools = loaded_commented_yaml.get("excluded_tools", []) + instance.included_optional_tools = loaded_commented_yaml.get("included_optional_tools", []) # re-save the configuration file if any migrations were performed if num_project_migrations > 0: diff --git a/src/serena/resources/serena_config.template.yml b/src/serena/resources/serena_config.template.yml index 0013f656..e0a6a098 100644 --- a/src/serena/resources/serena_config.template.yml +++ b/src/serena/resources/serena_config.template.yml @@ -36,6 +36,12 @@ trace_lsp_communication: False tool_timeout: 240 # timeout, in seconds, after which tool executions are terminated +excluded_tools: [] +# list of tools to be globally excluded + +included_optional_tools: [] +# list of optional tools (which are disabled by default) to be included + # 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, diff --git a/src/serena/tools/jetbrains_tools.py b/src/serena/tools/jetbrains_tools.py index 22a40354..751a3123 100644 --- a/src/serena/tools/jetbrains_tools.py +++ b/src/serena/tools/jetbrains_tools.py @@ -1,10 +1,10 @@ import json -from serena.tools import TOOL_DEFAULT_MAX_ANSWER_LENGTH, Tool +from serena.tools import TOOL_DEFAULT_MAX_ANSWER_LENGTH, Tool, ToolMarkerOptional from serena.tools.jetbrains_plugin_client import JetBrainsPluginClient -class JetBrainsFindSymbolTool(Tool): +class JetBrainsFindSymbolTool(Tool, ToolMarkerOptional): """ Performs a global (or local) search for symbols with/containing a given name/substring (optionally filtered by type). """ @@ -70,7 +70,7 @@ class JetBrainsFindSymbolTool(Tool): return self._limit_length(result, max_answer_chars) -class JetBrainsFindSymbolReferencesTool(Tool): +class JetBrainsFindSymbolReferencesTool(Tool, ToolMarkerOptional): """ Finds symbols that reference the given symbol """ diff --git a/src/serena/tools/tools_base.py b/src/serena/tools/tools_base.py index b0a3cd73..8cf53ca6 100644 --- a/src/serena/tools/tools_base.py +++ b/src/serena/tools/tools_base.py @@ -2,8 +2,8 @@ import inspect import os import traceback from abc import ABC, abstractmethod -from collections.abc import Callable, Generator, Iterable -from copy import copy +from collections.abc import Callable, Iterable +from dataclasses import dataclass from types import TracebackType from typing import TYPE_CHECKING, Any, Self, TypeVar @@ -13,6 +13,7 @@ from sensai.util.string import dict_string from serena.prompt_factory import PromptFactory from serena.symbol import SymbolManager +from serena.util.class_decorators import singleton from serena.util.inspection import iter_subclasses from solidlsp import SolidLanguageServer @@ -72,6 +73,12 @@ class ToolMarkerDoesNotRequireActiveProject: pass +class ToolMarkerOptional: + """ + Marker class for optional tools that are disabled by default. + """ + + class ToolInterface(ABC): """Protocol defining the complete interface that make_tool() expects from a tool.""" @@ -311,58 +318,44 @@ class EditedFileContext: # If they do not, we may have to add a call to notify it. +@dataclass(kw_only=True) +class RegisteredTool: + tool_class: type[Tool] + is_optional: bool + tool_name: str + + +@singleton class ToolRegistry: - _tool_dict: dict[str, type[Tool]] | None = None - """maps tool name to the corresponding tool class""" - - @staticmethod - def _iter_tool_classes() -> Generator[type[Tool], None, None]: - """ - Iterate over Tool subclasses. - """ + def __init__(self) -> None: + self._tool_dict: dict[str, RegisteredTool] = {} for cls in iter_subclasses(Tool): - if cls.__module__.startswith("serena.tools"): - yield cls + if not cls.__module__.startswith("serena.tools"): + continue + is_optional = issubclass(cls, ToolMarkerOptional) + name = cls.get_name_from_cls() + if name in self._tool_dict: + raise ValueError(f"Duplicate tool name found: {name}. Tool classes must have unique names.") + self._tool_dict[name] = RegisteredTool(tool_class=cls, is_optional=is_optional, tool_name=name) - @classmethod - def _get_tool_dict(cls) -> dict[str, type[Tool]]: - if cls._tool_dict is None: - cls._tool_dict = {} - for tool_class in cls._iter_tool_classes(): - name = tool_class.get_name_from_cls() - if name in cls._tool_dict: - raise ValueError(f"Duplicate tool name found: {name}. Tool classes must have unique names.") - cls._tool_dict[name] = tool_class - return cls._tool_dict + def get_tool_class_by_name(self, tool_name: str) -> type[Tool]: + return self._tool_dict[tool_name].tool_class - @classmethod - def get_tool_class_by_name(cls, tool_name: str) -> type[Tool]: - try: - return cls._get_tool_dict()[tool_name] - except KeyError as e: - available_tools = "\n".join(ToolRegistry.get_tool_names()) - raise ValueError(f"Tool with name {tool_name} not found. Available tools:\n{available_tools}") from e + def get_all_tool_classes(self) -> list[type[Tool]]: + return list(t.tool_class for t in self._tool_dict.values()) - @classmethod - def get_all_tool_classes(cls) -> list[type[Tool]]: - return list(cls._get_tool_dict().values()) + def get_tool_names_default_enabled(self) -> list[str]: + """ + :return: the list of tool names that are enabled by default (i.e. non-optional tools). + """ + return [t.tool_name for t in self._tool_dict.values() if not t.is_optional] - @classmethod - def get_tool_names(cls) -> list[str]: - return list(cls._get_tool_dict().keys()) - - @classmethod - def tool_dict(cls) -> dict[str, type[Tool]]: - """Maps tool name to the corresponding tool class""" - return copy(cls._get_tool_dict()) - - @classmethod - def print_tool_overview(cls, tools: Iterable[type[Tool] | Tool] | None = None) -> None: + def print_tool_overview(self, tools: Iterable[type[Tool] | Tool] | None = None) -> None: """ Print a summary of the tools. If no tools are passed, a summary of all tools is printed. """ if tools is None: - tools = cls._get_tool_dict().values() + tools = [tool.tool_class for tool in self._tool_dict.values() if not tool.is_optional] tool_dict: dict[str, type[Tool] | Tool] = {} for tool_class in tools: @@ -370,3 +363,6 @@ class ToolRegistry: for tool_name in sorted(tool_dict.keys()): tool_class = tool_dict[tool_name] print(f" * `{tool_name}`: {tool_class.get_tool_description().strip()}") + + def is_valid_tool_name(self, tool_name: str) -> bool: + return tool_name in self._tool_dict diff --git a/test/serena/test_mcp.py b/test/serena/test_mcp.py index 8ffb2c21..37b6eca8 100644 --- a/test/serena/test_mcp.py +++ b/test/serena/test_mcp.py @@ -271,7 +271,7 @@ def is_test_mock_class(tool_class: type) -> bool: ) -@pytest.mark.parametrize("tool_class", ToolRegistry.get_all_tool_classes()) +@pytest.mark.parametrize("tool_class", ToolRegistry().get_all_tool_classes()) def test_make_tool_all_tools(tool_class) -> None: """Test that make_tool works for all tools in the codebase."""