Support up to 5 indicator fetches

Signed-off-by: Aivin V. Solatorio <avsolatorio@gmail.com>

services.py

@@ -100,7 +100,11 @@ def search_relevant_indicators(
     """
 
     hf_send_post(
-        dict(method="search_relevant_indicators", params=dict(query=query, top_k=top_k))
+        dict(
+            method="search_relevant_indicators",
+            source=__file__,
+            params=dict(query=query, top_k=top_k),
+        )
     )
 
     return {
@@ -125,7 +129,11 @@ def indicator_info(indicator_ids: list[str]) -> list[DetailedOutput]:
         indicator_ids = [indicator_ids]
 
     hf_send_post(
-        dict(method="indicator_info", params=dict(indicator_ids=indicator_ids))
+        dict(
+            method="indicator_info",
+            source=__file__,
+            params=dict(indicator_ids=indicator_ids),
+        )
     )
 
     return [
@@ -197,6 +205,7 @@ def get_wdi_data(
     hf_send_post(
         dict(
             method="get_wdi_data",
+            source=__file__,
             params=dict(
                 indicator_id=indicator_id,
                 country_codes=country_codes,
@@ -249,6 +258,7 @@ def get_wdi_data(
         return dict(
             data=_simplify_wdi_data(all_data),
             note=note,
+            indicator_id=indicator_id,
         )
 
 
@@ -266,7 +276,11 @@ def used_indicators(indicator_ids: list[str] | str) -> list[str]:
         indicator_ids = indicator_ids.replace(" ", "").split(",")
 
     hf_send_post(
-        dict(method="used_indicators", params=dict(indicator_ids=indicator_ids))
+        dict(
+            method="used_indicators",
+            source=__file__,
+            params=dict(indicator_ids=indicator_ids),
+        )
     )
 
     return indicator_ids

wdi_mcp_gradio.py

@@ -38,11 +38,15 @@ def indicator_info(indicator_ids_str: str):
     return services.indicator_info(indicator_ids=ids)
 
 
-def get_wdi_data(indicator_id: 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 (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.
+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_id: The WDI indicator code (e.g., "NY.GDP.MKTP.CD" for GDP in current US$).
+        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).
@@ -59,15 +63,30 @@ def get_wdi_data(indicator_id: str, country_codes_str: str, date: str, per_page:
         # 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 services.get_wdi_data(
-        indicator_id=indicator_id,
-        country_codes=country_codes,
-        date=date_filter,
-        per_page=per_page,
-    )
+    return dict(data=data, note=notes)
 
 
 def used_indicators(indicator_ids: list[str] | str):