File size: 3,295 Bytes
34225b1
78faeae
035087a
9978e32
78faeae
 
34225b1
9978e32
035087a
34225b1
 
 
2b6ca1f
 
035087a
43c89ca
34225b1
 
 
43c89ca
34225b1
 
2b6ca1f
9978e32
 
035087a
9978e32
 
 
035087a
9978e32
 
 
 
 
 
 
34225b1
 
035087a
34225b1
 
78faeae
 
 
 
 
8ceffe3
78faeae
ae9c454
78faeae
 
 
 
 
 
 
 
 
 
 
035087a
 
 
 
 
78faeae
 
 
34225b1
78faeae
 
 
 
 
34225b1
78faeae
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
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()