Refactor data

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

mcp_client.py

@@ -162,8 +162,11 @@ class MCPClientWrapper:
         result_messages = []
 
         print(response.content)
+        contents = response.content
+
+        while len(contents) > 0:
+            content = contents.pop(0)
 
-        for content in response.content:
             if content.type == "text":
                 result_messages.append({"role": "assistant", "content": content.text})
                 claude_messages.append({"role": "assistant", "content": content.text})
@@ -278,10 +281,12 @@ class MCPClientWrapper:
 
                 print("next_response", next_response.content)
 
-                if next_response.content and next_response.content[0].type == "text":
-                    result_messages.append(
-                        {"role": "assistant", "content": next_response.content[0].text}
-                    )
+                contents.extend(next_response.content)
+
+                # if next_response.content and next_response.content[0].type == "text":
+                #     result_messages.append(
+                #         {"role": "assistant", "content": next_response.content[0].text}
+                #     )
 
         return result_messages
 

wdi_mcp_server.py

@@ -1,87 +1,18 @@
 from mcp.server.fastmcp import FastMCP
-import json
-
-# import sys
-# import io
-# import time
-# import numpy as np
-import pandas as pd
-import torch
-import httpx
-
-
 from typing import Optional, Any
-from sentence_transformers import SentenceTransformer
-
-# from gradio_client import Client
-from pydantic import BaseModel, Field
-
-
-def get_best_torch_device():
-    if torch.cuda.is_available():
-        return torch.device("cuda")
-    elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
-        return torch.device("mps")
-    else:
-        return torch.device("cpu")
-
-
-device = get_best_torch_device()
+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("huggingface_spaces_wdi_data")
-
-
-# Load the basic WDI metadata and vectors.
-wdi_data_vec_fpath = (
-    "./data/avsolatorio__GIST-small-Embedding-v0__005__indicator_embeddings.json"
-)
-df = pd.read_json(wdi_data_vec_fpath)
-
-# Make it easy to index based on the idno
-df.index = df["idno"]
-
-# Change the IDS naming to metadata standard
-df.rename(columns={"title": "name", "text": "definition"}, inplace=True)
-
-# Extract the vectors into a torch.tensor
-vectors = torch.Tensor(df["embedding"]).to(device)
-
-
-# Load the embedding model
-model_name = "/".join(wdi_data_vec_fpath.split("/")[-1].split("__")[:2])
-embedding_model = SentenceTransformer(model_name, device=device)
-
-
-def get_top_k(query: str, top_k: int = 10, fields: list[str] | None = None):
-    if fields is None:
-        fields = ["idno"]
-
-    # Convert the query to a search vector
-    search_vec = embedding_model.encode([query], convert_to_tensor=True) @ vectors.T
-
-    # Sort by descending similarity score
-    idx = search_vec.argsort(descending=True)[0][:top_k].tolist()
-
-    return df.iloc[idx][fields].to_dict("records")
-
-
-class SearchOutput(BaseModel):
-    idno: str = Field(..., description="The unique identifier of the indicator.")
-    name: str = Field(..., description="The name of the indicator.")
-
-
-class DetailedOutput(SearchOutput):
-    definition: str | None = Field(None, description="The indicator definition.")
+mcp = FastMCP("wdi_data_mcp")
 
 
 @mcp.tool()
 async def search_relevant_indicators(
     query: str, top_k: int = 1
-) -> dict[str, list[SearchOutput] | str]:
+) -> 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).
 
     Args:
@@ -92,17 +23,11 @@ async def search_relevant_indicators(
         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 {
-        "indicators": [
-            SearchOutput(**out)
-            for out in get_top_k(query=query, top_k=top_k, fields=["idno", "name"])
-        ],
-        "note": "IMPORTANT: Let the user know that the search is not exhaustive. The search is based on the semantic similarity of the query to the indicator definitions. It may not be optimal and the LLM may use this as shortlist and pick the most relevant from the list (if any).",
-    }
+    return services.search_relevant_indicators(query=query, top_k=top_k)
 
 
 @mcp.tool()
-async def indicator_info(indicator_ids: list[str]) -> list[DetailedOutput]:
+async def indicator_info(indicator_ids: list[str]) -> list[services.DetailedOutput]:
     """Provides definition information for the given indicator id (idno).
 
     Args:
@@ -111,15 +36,8 @@ async def indicator_info(indicator_ids: list[str]) -> list[DetailedOutput]:
     Returns:
         List of objects with keys indicator code/idno, name, and definition.
     """
-    if isinstance(indicator_ids, str):
-        indicator_ids = [indicator_ids]
 
-    return [
-        DetailedOutput(**out)
-        for out in df.loc[indicator_ids][
-            ["idno", "name", "definition", "time_coverage", "geographic_coverage"]
-        ].to_dict("records")
-    ]
+    return services.indicator_info(indicator_ids=indicator_ids)
 
 
 @mcp.tool()
@@ -140,54 +58,13 @@ async 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
-    note = ""
-
-    if isinstance(country_codes, str):
-        country_codes = [country_codes]
 
-    country_code = ";".join(country_codes)
-    base_url = (
-        f"https://api.worldbank.org/v2/country/{country_code}/indicator/{indicator_id}"
+    return services.get_wdi_data(
+        indicator_id=indicator_id,
+        country_codes=country_codes,
+        date=date,
+        per_page=per_page,
     )
-    params = {"format": "json", "date": date, "per_page": per_page or 100, "page": 1}
-
-    with open("mcp_server.log", "a+") as log:
-        log.write(json.dumps(dict(base_url=base_url, params=params)) + "\n")
-
-    with httpx.Client(timeout=30.0) as client:
-        all_data = []
-        while True:
-            response = client.get(base_url, params=params)
-            if response.status_code != 200:
-                note = f"ERROR: Failed to fetch data: HTTP {response.status_code}"
-                break
-
-            json_response = response.json()
-
-            if not isinstance(json_response, list) or len(json_response) < 2:
-                note = "ERROR: The API response is invalid or empty."
-                break
-
-            metadata, data_page = json_response
-            all_data.extend(data_page)
-
-            if len(all_data) >= MAX_INFO:
-                note = f"IMPORTANT: Let the user know that the data is truncated to the first {MAX_INFO} entries."
-                break
-
-            if params["page"] >= metadata.get("pages", 1):
-                break
-
-            params["page"] += 1
-
-        with open("mcp_server.log", "a+") as log:
-            log.write(json.dumps(dict(all_data=all_data)) + "\n")
-
-        return dict(
-            data=all_data,
-            note=note,
-        )
 
 
 if __name__ == "__main__":