wdi_mcp_gradio.py

import gradio as gr
import services

# --- Wrapper functions for Gradio (sync versions) ---


def search_relevant_indicators(query: str, top_k: int):
    """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.
    """

    # (If services.search_relevant_indicators is already sync, this just forwards.)
    return services.search_relevant_indicators(query=query, top_k=top_k)


def indicator_info(indicator_ids_str: str):
    """Provides definition information for the given indicator id (idno).

    Args:
        indicator_ids_str: An indicator id or a comma-separated list of indicator ids (idno) that additional information is being requested for.

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

    # Split on commas and strip each ID
    ids = [
        id_.strip()
        for id_ in indicator_ids_str.replace(" ", "").split(",")
        if id_.strip()
    ]
    return services.indicator_info(indicator_ids=ids)


def get_wdi_data(
    indicator_ids: str | list[str], country_codes_str: str, date: str, per_page: int
):
    """After relevant data is identified by using the `search_relevant_indicators`, this tool fetches indicator data for a given indicator id(s) (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.

    IMPORTANT: This tool can only fetch data for at most 5 indicators at a time.

    Args:
        indicator_ids: The WDI indicator code (e.g., "NY.GDP.MKTP.CD" for GDP in current US$). Comma separated if more than one.
        country_codes_str: The 3-letter ISO country code (e.g., "USA", "CHN", "IND"), or "all" for all countries. Comma separated if more than one.
        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.
    """

    # Parse country_codes_str:
    cc_input = country_codes_str.strip()
    if cc_input.lower() == "all":
        country_codes = "all"
    else:
        # Split on commas, uppercase each, strip spaces
        country_codes = [c.strip().upper() for c in cc_input.split(",") if c.strip()]

    if isinstance(indicator_ids, str):
        indicator_ids = indicator_ids.replace(" ", "").split(",")

    if len(indicator_ids) > 5:
        return dict(
            data=[],
            note=f"ERROR: This tool can only fetch data for at most 5 indicators at a time, but you requested {len(indicator_ids)}.",
        )

    # If user left date blank, pass None
    date_filter = date.strip() or None
    data = []
    notes = {}
    for indicator_id in indicator_ids:
        output = services.get_wdi_data(
            indicator_id=indicator_id,
            country_codes=country_codes,
            date=date_filter,
            per_page=per_page,
        )
        data.extend(output["data"])
        notes[output["indicator_id"]] = output["note"]

    return dict(data=data, note=notes)


def used_indicators(indicator_ids: list[str] | str):
    """The LLM can use this tool to let the user know which indicators it has used in generating its response.

    Args:
        indicator_ids: A list or comma-separated list of indicator ids (idno) that have been used by the LLM.

    Returns:
        A list of indicator ids (idno) that have been used by the LLM. This is used to let the user know, in a structured way, which indicators were used.
    """

    return services.used_indicators(indicator_ids=indicator_ids)


def build_interface():
    # --- Build the Gradio interface ---

    with gr.Blocks(title="WDI MCP Gradio") as demo:
        gr.Markdown("## WDI MCP: Gradio Interface")
        gr.Markdown(
            "Use the tabs below to call *search_relevant_indicators*, *indicator_info*, or *get_wdi_data*."
        )

        with gr.Tab("Search Relevant Indicators"):
            gr.Markdown(
                "Search for a shortlist of relevant WDI indicators given a query. "
                "Remember: For best recall, set **Top K ≥ 20**."
            )
            query_input = gr.Textbox(
                label="Query", placeholder="e.g. 'GDP of Asian countries'", lines=1
            )
            top_k_input = gr.Slider(
                label="Top K",
                minimum=1,
                maximum=50,
                step=1,
                value=20,
                info="At least 20 recommended",
            )
            search_btn = gr.Button("Search")
            search_output = gr.JSON(label="Search Results (dict)")

            # When button clicked, call our wrapper and display output in JSON
            search_btn.click(
                fn=search_relevant_indicators,
                inputs=[query_input, top_k_input],
                outputs=search_output,
            )

        with gr.Tab("Indicator Info"):
            gr.Markdown(
                "Provide one or more indicator IDs (comma-separated) to retrieve definitions."
            )
            indicator_ids_input = gr.Textbox(
                label="Indicator IDs",
                placeholder="e.g. NY.GDP.MKTP.CD, SP.POP.TOTL",
                lines=1,
            )
            info_btn = gr.Button("Get Definitions")
            info_output = gr.JSON(label="Indicator Info (list)")

            info_btn.click(
                fn=indicator_info,
                inputs=indicator_ids_input,
                outputs=info_output,
            )

        with gr.Tab("Get WDI Data"):
            gr.Markdown(
                "Fetch actual WDI data for a given indicator and country set. "
                "Set **Country Codes** to ‘all’ or a comma-separated list of 3-letter codes."
            )
            indicator_id_input = gr.Textbox(
                label="Indicator ID", placeholder="e.g. NY.GDP.MKTP.CD", lines=1
            )
            country_codes_input = gr.Textbox(
                label="Country Codes",
                placeholder="e.g. 'USA, CHN' or 'all'",
                lines=1,
            )
            date_input = gr.Textbox(
                label="Date Filter",
                placeholder="Year (e.g. '2022') or range (e.g. '2000:2022') – leave empty for no filter",
                lines=1,
            )
            per_page_input = gr.Number(
                label="Per Page",
                value=5,
                precision=0,
                info="Max allowed is usually 100",
            )
            data_btn = gr.Button("Fetch Data")
            data_output = gr.JSON(label="WDI Data (dict)")

            data_btn.click(
                fn=get_wdi_data,
                inputs=[
                    indicator_id_input,
                    country_codes_input,
                    date_input,
                    per_page_input,
                ],
                outputs=data_output,
            )

        with gr.Tab("Used Indicators"):
            gr.Markdown(
                "Returns the list of indicator ids (idno) that have been used by the LLM."
            )
            indicator_ids_input = gr.Textbox(
                label="Indicator IDs",
                placeholder="e.g. NY.GDP.MKTP.CD, SP.POP.TOTL",
                lines=1,
            )
            used_indicators_btn = gr.Button("Get Used Indicators")
            used_indicators_output = gr.JSON(label="Used Indicators (list)")

            used_indicators_btn.click(
                fn=used_indicators,
                inputs=indicator_ids_input,
                outputs=used_indicators_output,
            )

    return demo


if __name__ == "__main__":
    demo = build_interface()
    demo.launch(mcp_server=True)