from typing import Any, Callable, Dict, List, Optional, Union
import mcp.types as types
from airflow_client.client.api.dag_stats_api import DagStatsApi
from src.airflow.airflow_client import api_client
dag_stats_api = DagStatsApi(api_client)
def get_all_functions() -> list[tuple[Callable, str, str]]:
return [
(get_dag_stats, "get_dag_stats", "Get DAG stats"),
]
async def get_dag_stats(
dag_ids: Optional[List[str]] = None,
) -> List[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
# Build parameters dictionary
kwargs: Dict[str, Any] = {}
if dag_ids is not None:
kwargs["dag_ids"] = dag_ids
response = dag_stats_api.get_dag_stats(**kwargs)
return [types.TextContent(type="text", text=str(response.to_dict()))]