wdi_mcp_server.py

from mcp.server.fastmcp import FastMCP
from typing import Optional, Any
import services

# sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")
# sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding="utf-8", errors="replace")


mcp = FastMCP("wdi_data_mcp")


@mcp.tool()
async def search_relevant_indicators(
    query: str, top_k: int = 1
) -> dict[str, list[services.SearchOutput] | str]:
    """Search for a shortlist of relevant indicators from the World Development Indicators (WDI) given the query. The search ranking may not be optimal, so the LLM may use this as shortlist and pick the most relevant from the list (if any). You, as an LLM, must always get at least `top_k=20` for better recall.

    Args:
        query: The search query by the user or one formulated by an LLM based on the user's prompt.
        top_k: The number of shortlisted indicators that will be returned that are semantically related to the query. IMPORTANT: You, as an LLM, must ALWAYS set this argument to at least 20.

    Returns:
        A dictionary with keys `indicators` and `note`. The `indicators` key contains a list of indicator objects with keys indicator code/idno and name. The `note` key contains a note about the search.
    """

    return services.search_relevant_indicators(query=query, top_k=top_k)


@mcp.tool()
async def indicator_info(indicator_ids: list[str]) -> list[services.DetailedOutput]:
    """Provides definition information for the given indicator id (idno).

    Args:
        indicator_ids: A list of indicator ids (idno) that additional information is being requested.

    Returns:
        List of objects with keys indicator code/idno, name, and definition.
    """

    return services.indicator_info(indicator_ids=indicator_ids)


@mcp.tool()
async def get_wdi_data(
    indicator_id: str,
    country_codes: str | list[str],
    date: Optional[str] = None,
    per_page: Optional[int] = 5,
) -> dict[str, list[dict[str, Any]] | str]:
    """After relevant data is identified by using the `search_relevant_indicators`, this tool fetches indicator data for a given indicator id (idno) from the World Bank's World Development Indicators (WDI) API. The LLM must exclusively use this tool when the user asks for data. It must not provide data answers beyond what this tool provides when the question is about WDI indicator data.

    Args:
        indicator_id: The WDI indicator code (e.g., "NY.GDP.MKTP.CD" for GDP in current US$).
        country_codes: The 3-letter ISO country code (e.g., "USA", "CHN", "IND"), or "all" for all countries.
        date: A year (e.g., "2022") or a range (e.g., "2000:2022") to filter the results.
        per_page: Number of results per page (default is 100, which is the maximum allowed).

    Returns:
        A dictionary with keys `data` and `note`. The `data` key contains a list of indicator data entries requested. The `note` key contains a note about the data returned.
    """

    return services.get_wdi_data(
        indicator_id=indicator_id,
        country_codes=country_codes,
        date=date,
        per_page=per_page,
    )


if __name__ == "__main__":
    """
    Run the MCP server.

    uv run mcp dev wdi_mcp_server.py
    """
    mcp.run(transport="stdio")
    # mcp.run()