Start building MCP for WDI

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

wdi_mcp_server.py

@@ -0,0 +1,78 @@
+from mcp.server.fastmcp import FastMCP
+import json
+import sys
+import io
+import time
+from gradio_client import Client
+from pydantic import BaseModel, Field
+
+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")
+
+
+@mcp.tool()
+async def generate_image(prompt: str, width: int = 512, height: int = 512) -> str:
+    """Generate an image using SanaSprint model.
+
+    Args:
+        prompt: Text prompt describing the image to generate
+        width: Image width (default: 512)
+        height: Image height (default: 512)
+    """
+    client = Client("https://ysharma-sanasprint.hf.space/")
+
+    try:
+        result = client.predict(
+            prompt, "0.6B", 0, True, width, height, 4.0, 2, api_name="/infer"
+        )
+
+        if isinstance(result, list) and len(result) >= 1:
+            image_data = result[0]
+            if isinstance(image_data, dict) and "url" in image_data:
+                return json.dumps(
+                    {
+                        "type": "image",
+                        "url": image_data["url"],
+                        "message": f"Generated image for prompt: {prompt}",
+                    }
+                )
+
+        return json.dumps({"type": "error", "message": "Failed to generate image"})
+
+    except Exception as e:
+        return json.dumps(
+            {"type": "error", "message": f"Error generating image: {str(e)}"}
+        )
+
+
+class SearchOutput(BaseModel):
+    idno: str = Field(..., description="The unique identifier of the indicator.")
+    name: str = Field(..., description="The name of the indicator.")
+    definition: str | None = Field(None, description="The indicator definition.")
+
+
+@mcp.tool()
+async def search_relevant_indicators(query: str, top_k: int = 1) -> list[SearchOutput]:
+    """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:
+        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.
+
+    Returns:
+        List of objects with keys indicator name, code
+    """
+
+    return [
+        SearchOutput(
+            idno="NY.GDP.MKTP.KD",
+            name="GDP (constant 2015 US$)",
+            definition="GDP at purchaser's prices is the sum of gross value added by all resident producers in the economy plus any product taxes and minus any subsidies not included in the value of the products. It is calculated without making deductions for depreciation of fabricated assets or for depletion and degradation of natural resources. Data are in constant 2015 prices, expressed in U.S. dollars. Dollar figures for GDP are converted from domestic currencies using 2015 official exchange rates. For a few countries where the official exchange rate does not reflect the rate effectively applied to actual foreign exchange transactions, an alternative conversion factor is used.",
+        )
+    ]
+
+
+if __name__ == "__main__":
+    mcp.run(transport="stdio")