From 407eb00c1b17adddd34c529f4120364488dfde55 Mon Sep 17 00:00:00 2001 From: Bhavani Ravi Date: Wed, 19 Mar 2025 20:41:38 +0530 Subject: [PATCH] chore: #2 add description to tools --- .../parser/operation_parser.py | 4 +++- src/airflow_mcp_server/tools/base_tools.py | 5 +++-- src/airflow_mcp_server/tools/tool_manager.py | 2 +- tests/parser/test_operation_parser.py | 18 ++++++++++++++++++ 4 files changed, 25 insertions(+), 4 deletions(-) diff --git a/src/airflow_mcp_server/parser/operation_parser.py b/src/airflow_mcp_server/parser/operation_parser.py index 482048d..ac624eb 100644 --- a/src/airflow_mcp_server/parser/operation_parser.py +++ b/src/airflow_mcp_server/parser/operation_parser.py @@ -20,6 +20,7 @@ class OperationDetails: method: str parameters: dict[str, Any] input_model: type[BaseModel] + description: str class OperationParser: @@ -104,6 +105,7 @@ class OperationParser: operation["path"] = path operation["path_item"] = path_item + description = operation.get("description") or operation.get("summary") or operation_id parameters = self.extract_parameters(operation) @@ -119,7 +121,7 @@ class OperationParser: # Create unified input model input_model = self._create_input_model(operation_id, parameters, body_schema) - return OperationDetails(operation_id=operation_id, path=str(path), method=method, parameters=parameters, input_model=input_model) + return OperationDetails(operation_id=operation_id, path=str(path), method=method, parameters=parameters, description=description, input_model=input_model) raise ValueError(f"Operation {operation_id} not found in spec") diff --git a/src/airflow_mcp_server/tools/base_tools.py b/src/airflow_mcp_server/tools/base_tools.py index e154d1d..10567df 100644 --- a/src/airflow_mcp_server/tools/base_tools.py +++ b/src/airflow_mcp_server/tools/base_tools.py @@ -1,9 +1,10 @@ from abc import ABC, abstractmethod from typing import Any + class BaseTools(ABC): """Abstract base class for tools.""" - + @abstractmethod def __init__(self) -> None: """Initialize the tool.""" @@ -12,7 +13,7 @@ class BaseTools(ABC): @abstractmethod def run(self) -> Any: """Execute the tool's main functionality. - + Returns: Any: The result of the tool execution """ diff --git a/src/airflow_mcp_server/tools/tool_manager.py b/src/airflow_mcp_server/tools/tool_manager.py index 3588cc7..ba48392 100644 --- a/src/airflow_mcp_server/tools/tool_manager.py +++ b/src/airflow_mcp_server/tools/tool_manager.py @@ -105,7 +105,7 @@ async def get_airflow_tools(config: AirflowConfig, mode: str = "unsafe") -> list tools.append( Tool( name=operation_id, - description=tool.operation.operation_id, + description=tool.operation.description, inputSchema=schema, ) ) diff --git a/tests/parser/test_operation_parser.py b/tests/parser/test_operation_parser.py index c8cec75..89831d3 100644 --- a/tests/parser/test_operation_parser.py +++ b/tests/parser/test_operation_parser.py @@ -31,6 +31,24 @@ def test_parse_operation_basic(parser: OperationParser) -> None: assert operation.operation_id == "get_dags" assert operation.path == "/dags" assert operation.method == "get" + assert ( + operation.description + == """List DAGs in the database. +`dag_id_pattern` can be set to match dags of a specific pattern +""" + ) + assert isinstance(operation.parameters, dict) + + +def test_parse_operation_with_no_description_but_summary(parser: OperationParser) -> None: + """Test parsing operation with no description but summary.""" + operation = parser.parse_operation("get_connections") + + assert isinstance(operation, OperationDetails) + assert operation.operation_id == "get_connections" + assert operation.path == "/connections" + assert operation.method == "get" + assert operation.description == "List connections" assert isinstance(operation.parameters, dict)