[WIP] Create model context protocol for MP's API

mp_api/client/core/utils.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 import re
 from functools import cache
 from typing import Optional, get_args
@@ -12,6 +13,33 @@
 
 from mp_api.client.core.settings import MAPIClientSettings
 
+_MAPI_SETTINGS = MAPIClientSettings()
+
+
+def validate_api_key(api_key: str | None = None) -> str:
+    """Utility to find and pre-check validity of an API key."""
+    # SETTINGS tries to read API key from ~/.config/.pmgrc.yaml
+    api_key = api_key or os.getenv("MP_API_KEY")
+    if not api_key:
+        from pymatgen.core import SETTINGS
+
+        api_key = SETTINGS.get("PMG_MAPI_KEY")
+
+    if not api_key:
+        raise ValueError(
+            "Please obtain an API key from https://materialsproject.org/api "
+            "and export it as an environment variable `MP_API_KEY`."
+        )
+
+    if api_key and len(api_key) != 32:
+        raise ValueError(
+            "Please use a new API key from https://materialsproject.org/api "
+            "Keys for the new API are 32 characters, whereas keys for the legacy "
+            "API are 16 characters."
+        )
+
+    return api_key
+
 
 def validate_ids(id_list: list[str]):
     """Function to validate material and task IDs.
@@ -25,7 +53,7 @@ def validate_ids(id_list: list[str]):
     Returns:
         id_list: Returns original ID list if everything is formatted correctly.
     """
-    if len(id_list) > MAPIClientSettings().MAX_LIST_LENGTH:
+    if len(id_list) > _MAPI_SETTINGS.MAX_LIST_LENGTH:
         raise ValueError(
             "List of material/molecule IDs provided is too long. Consider removing the ID filter to automatically pull"
             " data for all IDs and filter locally."

mp_api/client/mprester.py

@@ -18,7 +18,7 @@
 from packaging import version
 from pymatgen.analysis.phase_diagram import PhaseDiagram
 from pymatgen.analysis.pourbaix_diagram import IonEntry
-from pymatgen.core import SETTINGS, Composition, Element, Structure
+from pymatgen.core import Composition, Element, Structure
 from pymatgen.core.ion import Ion
 from pymatgen.entries.computed_entries import ComputedStructureEntry
 from pymatgen.io.vasp import Chgcar
@@ -27,7 +27,7 @@
 
 from mp_api.client.core import BaseRester, MPRestError
 from mp_api.client.core.settings import MAPIClientSettings
-from mp_api.client.core.utils import validate_ids
+from mp_api.client.core.utils import validate_api_key, validate_ids
 from mp_api.client.routes import GeneralStoreRester, MessagesRester, UserSettingsRester
 from mp_api.client.routes.materials import (
     AbsorptionRester,
@@ -165,17 +165,7 @@ def __init__(
             mute_progress_bars:  Whether to mute progress bars.
 
         """
-        # SETTINGS tries to read API key from ~/.config/.pmgrc.yaml
-        api_key = api_key or os.getenv("MP_API_KEY") or SETTINGS.get("PMG_MAPI_KEY")
-
-        if api_key and len(api_key) != 32:
-            raise ValueError(
-                "Please use a new API key from https://materialsproject.org/api "
-                "Keys for the new API are 32 characters, whereas keys for the legacy "
-                "API are 16 characters."
-            )
-
-        self.api_key = api_key
+        self.api_key = validate_api_key(api_key)
         self.endpoint = endpoint or os.getenv(
             "MP_API_ENDPOINT", "https://api.materialsproject.org/"
         )

mp_api/mcp/__init__.py

@@ -0,0 +1,6 @@
+"""Get default MCP for Materials Project."""
+from __future__ import annotations
+
+from mp_api.mcp.mp_mcp import MPMcp
+
+__all__ = ["MPMcp"]

mp_api/mcp/__main__.py

@@ -0,0 +1,7 @@
+"""Run MCP."""
+from __future__ import annotations
+
+from mp_api.mcp.mp_mcp import MPMcp
+
+MP_MCP = MPMcp().mcp()
+MP_MCP.run()

mp_api/mcp/mp_mcp.py

@@ -0,0 +1,52 @@
+"""Define custom MCP tools for the Materials Project API."""
+from __future__ import annotations
+
+from typing import Any
+from urllib.parse import urljoin
+
+import httpx
+from fastmcp import FastMCP
+from pydantic import BaseModel, Field, PrivateAttr
+
+from mp_api.client import MPRester
+from mp_api.mcp import tools as mcp_tools
+
+
+class MPMcp(BaseModel):
+    name: str = Field("Materials Project MCP")
+    client_kwargs: dict[str, Any] | None = Field(None)
+    _client: MPRester | None = PrivateAttr(None)
+
+    @property
+    def client(self) -> MPRester:
+        # Always return JSON compliant output for MCP
+        kwargs = {
+            **(self.client_kwargs or {}),
+            "use_document_model": False,
+            "monty_decode": False,
+        }
+        if not self._client:
+            self._client = MPRester(**kwargs)
+        return self._client
+
+    def mcp(self, **kwargs) -> FastMCP:
+        mcp = FastMCP(self.name, **kwargs)
+
+        for attr in {x for x in dir(mcp_tools) if x.startswith("get")}:
+            mcp.tool(getattr(mcp_tools, attr))
+
+        return mcp
+
+    def bootstrap_mcp(self, **kwargs) -> FastMCP:
+        """Bootstrap an MP API MCP only from the OpenAPI spec."""
+        return FastMCP.from_openapi(
+            openapi_spec=httpx.get(
+                urljoin(self.client.endpoint, "openapi.json")
+            ).json(),
+            client=httpx.AsyncClient(
+                base_url=self.client.endpoint,
+                headers={"x-api-key": self.client.api_key},
+            ),
+            name=self.name,
+            **kwargs,
+        )