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 dc31e640..ef8b04ca 100644 --- a/src/serena/mcp.py +++ b/src/serena/mcp.py @@ -7,9 +7,10 @@ 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 +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 @@ -56,6 +57,23 @@ def make_tool( func_arg_metadata = func_metadata(apply_fn) parameters = func_arg_metadata.arg_model.model_json_schema() + 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) diff --git a/test/serena/test_mcp.py b/test/serena/test_mcp.py new file mode 100644 index 00000000..6e6a8ae4 --- /dev/null +++ b/test/serena/test_mcp.py @@ -0,0 +1,200 @@ +"""Tests for the mcp.py module in serena.""" + +import pytest +from mcp.server.fastmcp.tools.base import Tool as MCPTool + +from serena.agent import Tool +from serena.mcp import make_tool + + +class BaseMockTool(Tool): + """A mock Tool class for testing.""" + + def __init__(self) -> None: + pass + + +class BasicTool(BaseMockTool): + """A mock Tool class for testing.""" + + def apply(self, name: str, age: int = 0) -> str: + """This is a test function. + + :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!" + + 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_make_tool_basic() -> None: + """Test that make_tool correctly creates an MCP tool from a Tool object.""" + mock_tool = BasicTool() + + mcp_tool = make_tool(mock_tool) + + # 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 + + # 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_make_tool_execution() -> None: + """Test that the execution function created by make_tool works correctly.""" + mock_tool = BasicTool() + mcp_tool = make_tool(mock_tool) + + # Execute the MCP tool function + result = mcp_tool.fn(name="Alice", age=30) + + assert result == "Hello Alice, you are 30 years old!" + + +def test_make_tool_no_params() -> None: + """Test make_tool with a function that has no parameters.""" + + class NoParamsTool(BaseMockTool): + def apply(self) -> str: + """This is a test function with no parameters. + + :return: A simple result + """ + return "Simple result" + + 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_make_tool_no_return_description() -> None: + """Test make_tool with a function that has no return description.""" + + class NoReturnTool(BaseMockTool): + def apply(self, param: str) -> str: + """This is a test function. + + :param param: The parameter + """ + return f"Processed: {param}" + + def apply_ex(self, *args, **kwargs) -> str: + return self.apply(**kwargs) + + 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_make_tool_parameter_not_in_docstring() -> None: + """Test make_tool when a parameter in properties is not in the docstring.""" + + class MissingParamTool(BaseMockTool): + def apply(self, name: str, missing_param: str = "") -> str: + """This is a test function. + + :param name: The person's name + """ + return f"Hello {name}! Missing param: {missing_param}" + + 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) 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" },