services.py

import json
import os
import pandas as pd
import torch
import httpx


from typing import Optional, Any
from sentence_transformers import SentenceTransformer

from pydantic import BaseModel, Field
from urllib.request import urlretrieve


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()

# sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")
# sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding="utf-8", errors="replace")


# Load the basic WDI metadata and vectors.

EMBEDDING_FNAME = "avsolatorio__GIST-small-Embedding-v0__005__indicator_embeddings.json"
EMBEDDING_SOURCE = (
    f"https://raw.githubusercontent.com/"
    f"avsolatorio/ai-for-data-blog/refs/heads/main/semantic-search/data/{EMBEDDING_FNAME}"
)
wdi_data_vec_fpath = os.path.join("data", EMBEDDING_FNAME)

os.makedirs(os.path.dirname(wdi_data_vec_fpath), exist_ok=True)

if not os.path.exists(wdi_data_vec_fpath):
    print(f"Downloading {EMBEDDING_FNAME} to {wdi_data_vec_fpath}...")
    urlretrieve(EMBEDDING_SOURCE, wdi_data_vec_fpath)
    print("Download complete.")
else:
    print(f"File already exists at {wdi_data_vec_fpath}.")

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.")


def search_relevant_indicators(
    query: str, top_k: int = 1
) -> dict[str, list[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). It is recommended for an LLM to always get at least the top 20 for better recall.

    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:
        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).",
    }


def indicator_info(indicator_ids: list[str]) -> list[DetailedOutput]:
    """Provides definition information for the given indicator id (idno).

    Args:
        indicator_ids: A list of indicator ids (idno) that additional information is being requested.

    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")
    ]


def get_wdi_data(
    indicator_id: str,
    country_codes: str | list[str],
    date: Optional[str] = None,
    per_page: Optional[int] = 5,
) -> 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.

    Args:
        indicator_id: The WDI indicator code (e.g., "NY.GDP.MKTP.CD" for GDP in current US$).
        country_codes: The 3-letter ISO country code (e.g., "USA", "CHN", "IND"), or "all" for all countries.
        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).

    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}"
    )
    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,
        )