Implement simplified data formatting

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

services.py

@@ -126,11 +126,40 @@ def indicator_info(indicator_ids: list[str]) -> list[DetailedOutput]:
     ]
 
 
+def _simplify_wdi_data(data: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Simplifies the WDI data to only include the necessary fields. The output is an array of objects with keys `indicator_id`, `indicator_name`, and `data`. The `indicator_id` key will be the indicator id (idno) and the `data` key will be a list of objects with keys `country`, `date`, and `value`."""
+
+    try:
+        tmp_data = {}
+
+        for item in data:
+            if item["indicator"]["id"] not in tmp_data:
+                tmp_data[item["indicator"]["id"]] = {
+                    "indicator_id": item["indicator"]["id"],
+                    "indicator_name": item["indicator"]["value"],
+                    "data": [],
+                }
+
+            tmp_data[item["indicator"]["id"]]["data"].append(
+                {
+                    "country": item["country"]["value"],
+                    "date": item["date"],
+                    "value": item["value"],
+                }
+            )
+
+        return list(tmp_data.values())
+    except Exception as e:
+        # If the data is not valid, return the original data
+        print(f"ERROR: {e}")
+        return data
+
+
 def get_wdi_data(
     indicator_id: str,
     country_codes: str | list[str],
     date: Optional[str] = None,
-    per_page: Optional[int] = 5,
+    per_page: Optional[int] = 100,
 ) -> dict[str, list[dict[str, Any]] | str]:
     """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.
 
@@ -143,7 +172,7 @@ def get_wdi_data(
     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.
     """
-    MAX_INFO = 20
+    MAX_INFO = 500
     note = ""
 
     if isinstance(country_codes, str):
@@ -188,6 +217,6 @@ def get_wdi_data(
             log.write(json.dumps(dict(all_data=all_data)) + "\n")
 
         return dict(
-            data=all_data,
+            data=_simplify_wdi_data(all_data),
             note=note,
         )