From 46b7d79c815b2ddc9f0d523fb74db43704055995 Mon Sep 17 00:00:00 2001 From: Marcelo Soares Date: Mon, 21 Apr 2025 09:05:30 -0300 Subject: [PATCH 1/4] Parse parameter descriptions from docstring into MCP tool --- src/serena/mcp.py | 34 +++++++++- test/serena/test_mcp.py | 139 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 172 insertions(+), 1 deletion(-) create mode 100644 test/serena/test_mcp.py diff --git a/src/serena/mcp.py b/src/serena/mcp.py index dc31e640..8b093c6e 100644 --- a/src/serena/mcp.py +++ b/src/serena/mcp.py @@ -7,7 +7,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from dataclasses import dataclass from logging import Formatter, Logger, StreamHandler -from typing import Literal +from typing import Any, Literal import click # Add click import from mcp.server.fastmcp import server @@ -41,6 +41,36 @@ class SerenaMCPRequestContext: agent: SerenaAgent +def _process_parameter_descriptions( + parameters_properties: dict[str, dict[str, Any]], + func_doc: str, +) -> str: + """ + Move the parameter descriptions from the docstring to the parameters properties + and return the docstring without the parameter descriptions, but keeping the + return description if present. + """ + split_docs = func_doc.split(":param ") + split_docs[0] = split_docs[0].strip().strip(".") + "." + + if ":return:" in split_docs[-1]: + split_docs[-1], return_doc = split_docs[-1].split(":return:") + split_docs[0] = split_docs[0].strip() + " Returns " + return_doc.strip() + + for parameter, properties in parameters_properties.items(): + if ("description" in properties) or (f":param {parameter}:" not in func_doc): + continue + for doc in split_docs: + if doc.startswith(parameter): + param_desc = doc.split(":", 1)[1].strip().strip(".") + "." + param_desc = param_desc[0].upper() + param_desc[1:] + properties["description"] = param_desc + split_docs.remove(doc) + break + + return split_docs[0].strip().strip(".") + "." + + def make_tool( tool: Tool, ) -> MCPTool: @@ -56,6 +86,8 @@ def make_tool( func_arg_metadata = func_metadata(apply_fn) parameters = func_arg_metadata.arg_model.model_json_schema() + func_doc = _process_parameter_descriptions(parameters["properties"], func_doc) + def execute_fn(**kwargs) -> str: # type: ignore return tool.apply_ex(log_call=True, catch_exceptions=True, **kwargs) diff --git a/test/serena/test_mcp.py b/test/serena/test_mcp.py new file mode 100644 index 00000000..5ff77851 --- /dev/null +++ b/test/serena/test_mcp.py @@ -0,0 +1,139 @@ +"""Tests for the mcp.py module in serena.""" + +from serena.mcp import _process_parameter_descriptions + + +def test_basic_parameter_processing() -> None: + """Test basic functionality with parameters and return description.""" + params_props = { + "name": {}, + "age": {}, + } + func_doc = """This is a test function. + + :param name: The person's name + :param age: The person's age + :return: A greeting message + """ + + result_doc = _process_parameter_descriptions(params_props, func_doc) + + assert "param" not in result_doc + assert "Returns A greeting message" in result_doc + assert result_doc == "This is a test function. Returns A greeting message." + assert params_props["name"]["description"] == "The person's name." + assert params_props["age"]["description"] == "The person's age." + + +def test_no_parameters() -> None: + """Test with a docstring that has no parameters.""" + params_props = {} + func_doc = """This is a test function with no parameters. + + :return: A simple result + """ + + result_doc = _process_parameter_descriptions(params_props, func_doc) + + assert result_doc == "This is a test function with no parameters. Returns A simple result." + assert len(params_props) == 0 + + +def test_no_return_description() -> None: + """Test with a docstring that has no return description.""" + params_props = {"name": {}} + func_doc = """This is a test function. + + :param name: The parameter name + """ + + result_doc = _process_parameter_descriptions(params_props, func_doc) + + assert result_doc == "This is a test function." + assert "description" in params_props["name"] + assert params_props["name"]["description"] == "The parameter name." + + +def test_existing_parameter_description() -> None: + """Test when parameters already have descriptions in the properties dict.""" + params_props = { + "name": {"description": "Existing name description"}, + "age": {}, + } + func_doc = """This is a test function. + + :param name: The person's name + :param age: The person's age + """ + + result_doc = _process_parameter_descriptions(params_props, func_doc) + + assert result_doc == "This is a test function." + assert params_props["name"]["description"] == "Existing name description" + assert params_props["age"]["description"] == "The person's age." + + +def test_parameter_not_in_docstring() -> None: + """Test when a parameter in properties is not in the docstring.""" + params_props = { + "name": {}, + "missing_param": {}, + } + func_doc = """This is a test function. + + :param name: The person's name + """ + + result_doc = _process_parameter_descriptions(params_props, func_doc) + + assert result_doc == "This is a test function." + assert params_props["name"]["description"] == "The person's name." + assert "description" not in params_props["missing_param"] + + +def test_multiline_docstring() -> None: + """Test with a complex multi-line docstring.""" + params_props = { + "project_file_path": {}, + "host": {}, + "port": {}, + } + func_doc = """Create an MCP server. + + This function creates and configures a Model Context Protocol server + with the specified settings. + + :param project_file_path: The path to the project file, or None + :param host: The host to bind to + :param port: The port to bind to + :return: A configured FastMCP server instance + """ + + result_doc = _process_parameter_descriptions(params_props, func_doc) + + assert "param" not in result_doc + assert "Returns A configured FastMCP server instance" in result_doc + assert params_props["project_file_path"]["description"] == "The path to the project file, or None." + assert params_props["host"]["description"] == "The host to bind to." + assert params_props["port"]["description"] == "The port to bind to." + + +def test_capitalization_and_periods() -> None: + """Test that the function properly capitalizes descriptions and adds periods.""" + params_props = { + "param1": {}, + "param2": {}, + "param3": {}, + } + func_doc = """Test function. + + :param param1: lowercase description + :param param2: description with period. + :param param3: description with Capitalized word. + """ + + _ = _process_parameter_descriptions(params_props, func_doc) + + assert params_props["param1"]["description"] == "Lowercase description." + assert params_props["param2"]["description"] == "Description with period." + assert params_props["param3"]["description"] == "Description with Capitalized word." From cf029146d17797587d05f7dd0c4ab681b2309fea Mon Sep 17 00:00:00 2001 From: Marcelo Soares Date: Mon, 21 Apr 2025 20:33:00 -0300 Subject: [PATCH 2/4] Refactor test to target the `make_tool` function --- test/serena/test_mcp.py | 269 ++++++++++++++++++++++++---------------- 1 file changed, 165 insertions(+), 104 deletions(-) diff --git a/test/serena/test_mcp.py b/test/serena/test_mcp.py index 5ff77851..6e6a8ae4 100644 --- a/test/serena/test_mcp.py +++ b/test/serena/test_mcp.py @@ -1,139 +1,200 @@ """Tests for the mcp.py module in serena.""" -from serena.mcp import _process_parameter_descriptions +import pytest +from mcp.server.fastmcp.tools.base import Tool as MCPTool + +from serena.agent import Tool +from serena.mcp import make_tool -def test_basic_parameter_processing() -> None: - """Test basic functionality with parameters and return description.""" - params_props = { - "name": {}, - "age": {}, - } - func_doc = """This is a test function. +class BaseMockTool(Tool): + """A mock Tool class for testing.""" - :param name: The person's name - :param age: The person's age - :return: A greeting message - """ - - result_doc = _process_parameter_descriptions(params_props, func_doc) - - assert "param" not in result_doc - assert "Returns A greeting message" in result_doc - assert result_doc == "This is a test function. Returns A greeting message." - assert params_props["name"]["description"] == "The person's name." - assert params_props["age"]["description"] == "The person's age." + def __init__(self) -> None: + pass -def test_no_parameters() -> None: - """Test with a docstring that has no parameters.""" - params_props = {} - func_doc = """This is a test function with no parameters. +class BasicTool(BaseMockTool): + """A mock Tool class for testing.""" - :return: A simple result - """ + def apply(self, name: str, age: int = 0) -> str: + """This is a test function. - result_doc = _process_parameter_descriptions(params_props, func_doc) + :param name: The person's name + :param age: The person's age + :return: A greeting message + """ + return f"Hello {name}, you are {age} years old!" - assert result_doc == "This is a test function with no parameters. Returns A simple result." - assert len(params_props) == 0 + def apply_ex( + self, + log_call: bool = True, + catch_exceptions: bool = True, + **kwargs, + ) -> str: + """Mock implementation of apply_ex.""" + return self.apply(**kwargs) -def test_no_return_description() -> None: - """Test with a docstring that has no return description.""" - params_props = {"name": {}} - func_doc = """This is a test function. +def test_make_tool_basic() -> None: + """Test that make_tool correctly creates an MCP tool from a Tool object.""" + mock_tool = BasicTool() - :param name: The parameter name - """ + mcp_tool = make_tool(mock_tool) - result_doc = _process_parameter_descriptions(params_props, func_doc) + # Test that the MCP tool has the correct properties + assert isinstance(mcp_tool, MCPTool) + assert mcp_tool.name == "basic" + assert "This is a test function. Returns A greeting message." in mcp_tool.description - assert result_doc == "This is a test function." - assert "description" in params_props["name"] - assert params_props["name"]["description"] == "The parameter name." + # Test that the parameters were correctly processed + parameters = mcp_tool.parameters + assert "properties" in parameters + assert "name" in parameters["properties"] + assert "age" in parameters["properties"] + assert parameters["properties"]["name"]["description"] == "The person's name." + assert parameters["properties"]["age"]["description"] == "The person's age." -def test_existing_parameter_description() -> None: - """Test when parameters already have descriptions in the properties dict.""" - params_props = { - "name": {"description": "Existing name description"}, - "age": {}, - } - func_doc = """This is a test function. +def test_make_tool_execution() -> None: + """Test that the execution function created by make_tool works correctly.""" + mock_tool = BasicTool() + mcp_tool = make_tool(mock_tool) - :param name: The person's name - :param age: The person's age - """ + # Execute the MCP tool function + result = mcp_tool.fn(name="Alice", age=30) - result_doc = _process_parameter_descriptions(params_props, func_doc) - - assert result_doc == "This is a test function." - assert params_props["name"]["description"] == "Existing name description" - assert params_props["age"]["description"] == "The person's age." + assert result == "Hello Alice, you are 30 years old!" -def test_parameter_not_in_docstring() -> None: - """Test when a parameter in properties is not in the docstring.""" - params_props = { - "name": {}, - "missing_param": {}, - } - func_doc = """This is a test function. +def test_make_tool_no_params() -> None: + """Test make_tool with a function that has no parameters.""" - :param name: The person's name - """ + class NoParamsTool(BaseMockTool): + def apply(self) -> str: + """This is a test function with no parameters. - result_doc = _process_parameter_descriptions(params_props, func_doc) + :return: A simple result + """ + return "Simple result" - assert result_doc == "This is a test function." - assert params_props["name"]["description"] == "The person's name." - assert "description" not in params_props["missing_param"] + def apply_ex(self, *args, **kwargs) -> str: + return self.apply() + + tool = NoParamsTool() + mcp_tool = make_tool(tool) + + assert mcp_tool.name == "no_params" + assert "This is a test function with no parameters. Returns A simple result." in mcp_tool.description + assert mcp_tool.parameters["properties"] == {} -def test_multiline_docstring() -> None: - """Test with a complex multi-line docstring.""" - params_props = { - "project_file_path": {}, - "host": {}, - "port": {}, - } - func_doc = """Create an MCP server. +def test_make_tool_no_return_description() -> None: + """Test make_tool with a function that has no return description.""" - This function creates and configures a Model Context Protocol server - with the specified settings. + class NoReturnTool(BaseMockTool): + def apply(self, param: str) -> str: + """This is a test function. - :param project_file_path: The path to the project file, or None - :param host: The host to bind to - :param port: The port to bind to - :return: A configured FastMCP server instance - """ + :param param: The parameter + """ + return f"Processed: {param}" - result_doc = _process_parameter_descriptions(params_props, func_doc) + def apply_ex(self, *args, **kwargs) -> str: + return self.apply(**kwargs) - assert "param" not in result_doc - assert "Returns A configured FastMCP server instance" in result_doc - assert params_props["project_file_path"]["description"] == "The path to the project file, or None." - assert params_props["host"]["description"] == "The host to bind to." - assert params_props["port"]["description"] == "The port to bind to." + tool = NoReturnTool() + mcp_tool = make_tool(tool) + + assert mcp_tool.name == "no_return" + assert mcp_tool.description == "This is a test function." + assert mcp_tool.parameters["properties"]["param"]["description"] == "The parameter." -def test_capitalization_and_periods() -> None: - """Test that the function properly capitalizes descriptions and adds periods.""" - params_props = { - "param1": {}, - "param2": {}, - "param3": {}, - } - func_doc = """Test function. +def test_make_tool_parameter_not_in_docstring() -> None: + """Test make_tool when a parameter in properties is not in the docstring.""" - :param param1: lowercase description - :param param2: description with period. - :param param3: description with Capitalized word. - """ + class MissingParamTool(BaseMockTool): + def apply(self, name: str, missing_param: str = "") -> str: + """This is a test function. - _ = _process_parameter_descriptions(params_props, func_doc) + :param name: The person's name + """ + return f"Hello {name}! Missing param: {missing_param}" - assert params_props["param1"]["description"] == "Lowercase description." - assert params_props["param2"]["description"] == "Description with period." - assert params_props["param3"]["description"] == "Description with Capitalized word." + def apply_ex(self, *args, **kwargs) -> str: + return self.apply(**kwargs) + + tool = MissingParamTool() + mcp_tool = make_tool(tool) + + assert "name" in mcp_tool.parameters["properties"] + assert "missing_param" in mcp_tool.parameters["properties"] + assert mcp_tool.parameters["properties"]["name"]["description"] == "The person's name." + assert "description" not in mcp_tool.parameters["properties"]["missing_param"] + + +def test_make_tool_multiline_docstring() -> None: + """Test make_tool with a complex multi-line docstring.""" + + class ComplexDocTool(BaseMockTool): + def apply(self, project_file_path: str, host: str, port: int) -> str: + """Create an MCP server. + + This function creates and configures a Model Context Protocol server + with the specified settings. + + :param project_file_path: The path to the project file, or None + :param host: The host to bind to + :param port: The port to bind to + :return: A configured FastMCP server instance + """ + return f"Server config: {project_file_path}, {host}:{port}" + + def apply_ex(self, *args, **kwargs) -> str: + return self.apply(**kwargs) + + tool = ComplexDocTool() + mcp_tool = make_tool(tool) + + assert "Create an MCP server" in mcp_tool.description + assert "Returns A configured FastMCP server instance" in mcp_tool.description + assert mcp_tool.parameters["properties"]["project_file_path"]["description"] == "The path to the project file, or None." + assert mcp_tool.parameters["properties"]["host"]["description"] == "The host to bind to." + assert mcp_tool.parameters["properties"]["port"]["description"] == "The port to bind to." + + +def test_make_tool_capitalization_and_periods() -> None: + """Test that make_tool properly handles capitalization and periods in descriptions.""" + + class FormatTool(BaseMockTool): + def apply(self, param1: str, param2: str, param3: str) -> str: + """Test function. + + :param param1: lowercase description + :param param2: description with period. + :param param3: description with Capitalized word. + """ + return f"Formatted: {param1}, {param2}, {param3}" + + def apply_ex(self, *args, **kwargs) -> str: + return self.apply(**kwargs) + + tool = FormatTool() + mcp_tool = make_tool(tool) + + assert mcp_tool.parameters["properties"]["param1"]["description"] == "Lowercase description." + assert mcp_tool.parameters["properties"]["param2"]["description"] == "Description with period." + assert mcp_tool.parameters["properties"]["param3"]["description"] == "Description with Capitalized word." + + +def test_make_tool_missing_apply() -> None: + """Test make_tool with a tool that doesn't have an apply method.""" + + class BadTool(BaseMockTool): + pass + + tool = BadTool() + + with pytest.raises(AttributeError): + make_tool(tool) From b007917cf2b73953861b12fcd065f96103edd8ca Mon Sep 17 00:00:00 2001 From: Marcelo Soares Date: Mon, 21 Apr 2025 20:49:47 -0300 Subject: [PATCH 3/4] Refactor the docstring parse to use the `docstring_parser` lib --- pyproject.toml | 1 + src/serena/mcp.py | 48 +++++++++++++++++------------------------------ 2 files changed, 18 insertions(+), 31 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 82aea490..9a21b274 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,7 @@ dependencies = [ "pathspec>=0.12.1", "psutil>=7.0.0", "agno>=1.2.15", + "docstring_parser>=0.16", ] [project.scripts] diff --git a/src/serena/mcp.py b/src/serena/mcp.py index 8b093c6e..ef8b04ca 100644 --- a/src/serena/mcp.py +++ b/src/serena/mcp.py @@ -10,6 +10,7 @@ from logging import Formatter, Logger, StreamHandler from typing import Any, Literal import click # Add click import +import docstring_parser from mcp.server.fastmcp import server from mcp.server.fastmcp.server import FastMCP, Settings from mcp.server.fastmcp.tools.base import Tool as MCPTool @@ -41,36 +42,6 @@ class SerenaMCPRequestContext: agent: SerenaAgent -def _process_parameter_descriptions( - parameters_properties: dict[str, dict[str, Any]], - func_doc: str, -) -> str: - """ - Move the parameter descriptions from the docstring to the parameters properties - and return the docstring without the parameter descriptions, but keeping the - return description if present. - """ - split_docs = func_doc.split(":param ") - split_docs[0] = split_docs[0].strip().strip(".") + "." - - if ":return:" in split_docs[-1]: - split_docs[-1], return_doc = split_docs[-1].split(":return:") - split_docs[0] = split_docs[0].strip() + " Returns " + return_doc.strip() - - for parameter, properties in parameters_properties.items(): - if ("description" in properties) or (f":param {parameter}:" not in func_doc): - continue - for doc in split_docs: - if doc.startswith(parameter): - param_desc = doc.split(":", 1)[1].strip().strip(".") + "." - param_desc = param_desc[0].upper() + param_desc[1:] - properties["description"] = param_desc - split_docs.remove(doc) - break - - return split_docs[0].strip().strip(".") + "." - - def make_tool( tool: Tool, ) -> MCPTool: @@ -86,7 +57,22 @@ def make_tool( func_arg_metadata = func_metadata(apply_fn) parameters = func_arg_metadata.arg_model.model_json_schema() - func_doc = _process_parameter_descriptions(parameters["properties"], func_doc) + docstring = docstring_parser.parse(func_doc) + + # Mount the tool description as a combination of the docstring description and + # the return value description, if it exists. + func_doc = f"{docstring.description.strip().strip('.')}." + if (docstring.returns) and (docstring_returns := docstring.returns.description): + func_doc = f"{func_doc} Returns {docstring_returns.strip().strip('.')}." + + # Parse the parameter descriptions from the docstring and add pass its description + # to the parameters schema. + docstring_params = {param.arg_name: param for param in docstring.params} + parameters_properties: dict[str, dict[str, Any]] = parameters["properties"] + for parameter, properties in parameters_properties.items(): + if (param_doc := docstring_params.get(parameter)) and (param_doc.description): + param_desc = f"{param_doc.description.strip().strip('.') + '.'}" + properties["description"] = param_desc[0].upper() + param_desc[1:] def execute_fn(**kwargs) -> str: # type: ignore return tool.apply_ex(log_call=True, catch_exceptions=True, **kwargs) From 262d05fb7513ce3c8c4a7b7241573c5c4d35ed9b Mon Sep 17 00:00:00 2001 From: Marcelo Soares Date: Mon, 21 Apr 2025 23:53:51 -0300 Subject: [PATCH 4/4] Add missing update to `uv.lock` with new dependency --- uv.lock | 2 ++ 1 file changed, 2 insertions(+) diff --git a/uv.lock b/uv.lock index 456931c3..d8185ded 100644 --- a/uv.lock +++ b/uv.lock @@ -971,6 +971,7 @@ version = "0.1.0" source = { editable = "." } dependencies = [ { name = "agno" }, + { name = "docstring-parser" }, { name = "dotenv" }, { name = "fastmcp" }, { name = "jinja2" }, @@ -1017,6 +1018,7 @@ requires-dist = [ { name = "agno", marker = "extra == 'agno'", specifier = ">=1.2.6" }, { name = "anthropic", marker = "extra == 'anthropic'", specifier = ">=0.49.0" }, { name = "black", extras = ["jupyter"], marker = "extra == 'dev'", specifier = ">=23.7.0" }, + { name = "docstring-parser", specifier = ">=0.16" }, { name = "dotenv", specifier = ">=0.9.9" }, { name = "fastapi", marker = "extra == 'agno'", specifier = ">=0.115.12" }, { name = "fastmcp", specifier = ">=0.4.1" },