chore: #2 add description to tools

This commit is contained in:
Bhavani Ravi
2025-03-19 20:41:38 +05:30
parent 707f3747d7
commit 407eb00c1b
4 changed files with 25 additions and 4 deletions

View File

@@ -20,6 +20,7 @@ class OperationDetails:
method: str method: str
parameters: dict[str, Any] parameters: dict[str, Any]
input_model: type[BaseModel] input_model: type[BaseModel]
description: str
class OperationParser: class OperationParser:
@@ -104,6 +105,7 @@ class OperationParser:
operation["path"] = path operation["path"] = path
operation["path_item"] = path_item operation["path_item"] = path_item
description = operation.get("description") or operation.get("summary") or operation_id
parameters = self.extract_parameters(operation) parameters = self.extract_parameters(operation)
@@ -119,7 +121,7 @@ class OperationParser:
# Create unified input model # Create unified input model
input_model = self._create_input_model(operation_id, parameters, body_schema) 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") raise ValueError(f"Operation {operation_id} not found in spec")

View File

@@ -1,6 +1,7 @@
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from typing import Any from typing import Any
class BaseTools(ABC): class BaseTools(ABC):
"""Abstract base class for tools.""" """Abstract base class for tools."""

View File

@@ -105,7 +105,7 @@ async def get_airflow_tools(config: AirflowConfig, mode: str = "unsafe") -> list
tools.append( tools.append(
Tool( Tool(
name=operation_id, name=operation_id,
description=tool.operation.operation_id, description=tool.operation.description,
inputSchema=schema, inputSchema=schema,
) )
) )

View File

@@ -31,6 +31,24 @@ def test_parse_operation_basic(parser: OperationParser) -> None:
assert operation.operation_id == "get_dags" assert operation.operation_id == "get_dags"
assert operation.path == "/dags" assert operation.path == "/dags"
assert operation.method == "get" 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) assert isinstance(operation.parameters, dict)