Implement customized view based on countries provided

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

services.py

@@ -4,6 +4,7 @@ import pandas as pd
 import torch
 import httpx
 import zlib
+from urllib.parse import urlencode
 
 from typing import Optional, Any
 from sentence_transformers import SentenceTransformer
@@ -311,24 +312,65 @@ def used_indicators(indicator_ids: list[str] | str) -> list[str]:
     return indicator_ids
 
 
-def get_data360_link(indicator_id: str) -> dict[str, str]:
-    """The LLM can use this tool to get the link to the Data360 page for the given indicator id (idno).
+def get_data360_link(
+    indicator_id: str,
+    country_codes: list[str] | str | None = None,
+    year: str | None = None,
+) -> dict[str, str]:
+    """The LLM can use this tool to get the link to the Data360 page for the given indicator id (idno). Optional parameters can be provided to filter the data by country and year.
 
     Args:
         indicator_id: The WDI indicator code (e.g., "WB_WDI_NY_GDP_MKTP_CD" for GDP in current US$).
+        country_codes: The 3-letter ISO country code (e.g., "USA", "CHN", "IND"), or set to `None` for all countries. Comma separated if more than one.
+        year: The year to view the data for. Set to `None` for the most recent year.
 
     Returns:
-        A dictionary with keys `url` containing a link to the Data360 page for the given indicator id (idno).
+        A dictionary with keys `url` containing a link to the Data360 page for the given indicator id (idno) with the optional parameters.
     """
 
     hf_send_post(
         dict(
             method="get_data360_link",
             source=__file__,
-            params=dict(indicator_id=indicator_id),
+            params=dict(
+                indicator_id=indicator_id, country_codes=country_codes, year=year
+            ),
         )
     )
 
+    url = f"https://data360.worldbank.org/en/indicator/{indicator_id}"
+    view = None
+    recentYear = None
+
+    if year:
+        recentYear = "false"
+
+    if country_codes:
+        # view = "map"  # We can skip this because it is the default view
+        if isinstance(country_codes, str):
+            country_codes = country_codes.split(",")
+
+        if len(country_codes) > 1:
+            view = "trend"
+
+        country_codes = ",".join(country_codes)
+
+    params = {}  # type: ignore
+
+    if view:
+        params["view"] = view
+
+    if country_codes:
+        params["country"] = country_codes
+
+    if recentYear:
+        params["recentYear"] = recentYear
+
+    if year:
+        params["year"] = year
+
+    url = f"{url}?{urlencode(params)}"
+
     return {
-        "url": f"https://data360.worldbank.org/en/indicator/{indicator_id}",
+        "url": url,
     }

wdi_mcp_gradio.py

@@ -106,16 +106,23 @@ def used_indicators(indicator_ids: list[str] | str):
     return services.used_indicators(indicator_ids=indicator_ids)
 
 
-def get_data360_link(indicator_id: str) -> dict[str, str]:
-    """The LLM can use this tool to get the link to the Data360 page for the given indicator id (idno).
+def get_data360_link(
+    indicator_id: str,
+    country_codes: list[str] | str | None = None,
+    year: str | None = None,
+) -> dict[str, str]:
+    """The LLM can use this tool to get the link to the Data360 page for the given indicator id (idno). Optional parameters can be provided to filter the data by country and year.
 
     Args:
         indicator_id: The WDI indicator code (e.g., "WB_WDI_NY_GDP_MKTP_CD" for GDP in current US$).
-
+        country_codes: The 3-letter ISO country code (e.g., "USA", "CHN", "IND"), or set to `None` for all countries. Comma separated if more than one.
+        year: The year to view the data for. Set to `None` for the most recent year.
     Returns:
-        A dictionary with keys `url` containing a link to the Data360 page for the given indicator id (idno).
+        A dictionary with keys `url` containing a link to the Data360 page for the given indicator id (idno) with the optional parameters.
     """
-    return services.get_data360_link(indicator_id=indicator_id)
+    return services.get_data360_link(
+        indicator_id=indicator_id, country_codes=country_codes, year=year
+    )
 
 
 def build_interface():
@@ -229,17 +236,27 @@ def build_interface():
 
         with gr.Tab("Get Data360 Link"):
             gr.Markdown(
-                "Returns the link to the Data360 page for the given indicator id (idno)."
+                "Returns the link to the Data360 page for the given indicator id (idno). Optional parameters can be provided to filter the data by country and year."
             )
             indicator_id_input = gr.Textbox(
                 label="Indicator ID", placeholder="e.g. WB_WDI_NY_GDP_MKTP_CD", lines=1
             )
+            country_codes_input = gr.Textbox(
+                label="Country Codes",
+                placeholder="e.g. 'USA, CHN' or 'all'",
+                lines=1,
+            )
+            year_input = gr.Textbox(
+                label="Year",
+                placeholder="e.g. '2022'",
+                lines=1,
+            )
             data360_link_btn = gr.Button("Get Data360 Link")
             data360_link_output = gr.JSON(label="Data360 Link (dict)")
 
             data360_link_btn.click(
                 fn=get_data360_link,
-                inputs=indicator_id_input,
+                inputs=[indicator_id_input, country_codes_input, year_input],
                 outputs=data360_link_output,
             )