draft openai compatible mcp

Author: esoteric-ephemera

mp_api/mcp/_schemas.py

@@ -0,0 +1,35 @@
+"""Define auxiliary schemas used by some LLMs."""
+
+from typing import Any
+from pydantic import BaseModel, Field, model_validator
+
+from mp_api.client.core.utils import validate_ids
+
+class OpenAIResult(BaseModel):
+    """Schematize result for OpenAI support."""
+
+    id : str
+    title : str | None = None
+    text : str | None = None
+    url : str | None = None
+    metadata : dict[str,str] | None = None
+
+    @model_validator(mode="before")
+    def set_url(cls, config : Any) -> Any:
+        """Set default Materials Project URL and title."""
+
+        formatted_mpid = validate_ids([config.get("id")])[0]
+        if not config.get("title"):
+            config["title"] = formatted_mpid
+
+        if not config.get("url"):
+            config["url"]  = (
+                "https://next-gen.materialsproject.org/materials/"
+                f"{formatted_mpid}"
+            )
+        return config
+
+class OpenAISearchOutput(BaseModel):
+    """Schematize data for OpenAI support."""
+
+    results : list[OpenAIResult] = Field([])

mp_api/mcp/mp_mcp.py

@@ -6,9 +6,27 @@
 import httpx
 from fastmcp import FastMCP
 
-from mp_api.mcp.tools import MPMcpTools
+from mp_api.mcp.tools import MPMcpTools, MPOpenAIMcpTools
 from mp_api.mcp.utils import _NeedsMPClient
 
+MCP_SERVER_INSTRUCTIONS = """
+This MCP server defines search and document retrieval capabilities
+for data in the Materials Project.
+Use the search tool to find relevant documents based on materials
+keywords.
+Then use the fetch tool to retrieve complete materials summary information.
+"""
+
+def get_openai_compat_mcp() -> FastMCP:
+    """Create MCP for compatibility with OpenAI models."""
+    mp_mcp = FastMCP(
+        "Materials_Project_MCP",
+        instructions=MCP_SERVER_INSTRUCTIONS,
+    )
+    openai_compat_tools = MPOpenAIMcpTools()
+    for k in {"search","fetch"}:
+        mp_mcp.tool(getattr(openai_compat_tools,k))
+    return mp_mcp
 
 def get_mcp() -> FastMCP:
     """MCP with finer depth of control over tool names."""

mp_api/mcp/server.py

@@ -1,9 +1,9 @@
 """Run MCP."""
 from __future__ import annotations
 
-from mp_api.mcp.mp_mcp import get_mcp
+from mp_api.mcp.mp_mcp import get_mcp, get_openai_compat_mcp
 
-mcp = get_mcp()
+mcp = get_openai_compat_mcp()
 
 if __name__ == "__main__":
     mcp.run()

mp_api/mcp/tools.py

@@ -28,7 +28,101 @@
 from pymatgen.electronic_structure.core import OrbitalType, Spin
 from pymatgen.entries.computed_entries import ComputedEntry
 
+from mp_api.client.core import MPRestError
 from mp_api.mcp.utils import _NeedsMPClient
+from mp_api.mcp._schemas import OpenAISearchOutput, OpenAIResult
+
+class MPOpenAIMcpTools(_NeedsMPClient):
+    """Define OpenAI-specific MCP for the Materials Project API."""
+
+    def search(self, query : str) -> OpenAISearchOutput:
+        """Define OpenAI compatible search.
+        
+        Search through the autogenerated robocrystallographer
+        descriptions of materials to return lists of likely
+        matching materials.
+
+        Args:
+            query (str) : A natural language query of material keywords.
+            It is assumed that the query contains comma-delimited keywords.
+        
+        Returns: 
+            OpenAISearchOutput, a dict of `results` each with structure
+            mp_api.mcp._schemas.OpenAIResult
+        """
+        return OpenAISearchOutput(
+            retults = [
+                OpenAIResult(
+                    id = doc["material_id"],
+                    text = doc["description"]
+                )
+                for doc in self.client.robocrys.search(query.split(","))
+            ]
+        )
+
+    def fetch(self, idx : str) -> OpenAIResult:
+        """Retrieve complete material information by Materials Project ID.
+        
+        Args:
+            idx (str) : A Materials Project ID.
+            Should be an integer prefixed by `mp-`, ex: "mp-149", "mp-13"
+
+        Returns:
+            OpenAIResult : Complete document with id, title, robocrys
+            autogenerated description, URL, and metadata derived from
+            the materials summary collection.
+
+            If no data about the particular id is available, returns a
+            OpenAIResult with only the id field populated.
+
+        Raises:
+            MPRestError: If no identifier is specified
+        """
+        if not isinstance(idx,str):
+            raise MPRestError(
+                f"Unknown {idx=}. Should be an integer prefixed by `mp-`, ex: "
+                "'mp-1', 'mp-1010101'"
+            )
+
+        robo_desc : str | None = None
+        if len(
+            robo_docs := self.client.robocrys.search_docs(
+                material_ids=[idx]
+            )
+        ) > 0:
+            robo_desc = robo_docs[0]["description"]
+
+        if not robo_desc:
+            return OpenAIResult(id = idx)
+
+        metadata : dict[str,str] | None = None
+        if len(
+            summary_docs := self.client.summary.search(
+                material_ids=[idx]
+            )
+        ) > 0:
+            # Try to avoid more nested fields, just provide things with
+            # simple str or numeric type
+            metadata = {
+                k : str(summary_docs[0][k])
+                for k in self.client.summary.document_model.model_fields
+                if isinstance(summary_docs[0][k], str | int | float)
+            }
+
+            # Augment with experimental database id information
+            if summary_docs[0]["database_IDs"]:
+                metadata.update(
+                    {
+                        f"linked_{database}_ids" : ", ".join(matched_ids)
+                        for database, matched_ids in summary_docs[0]["database_IDs"].items()
+                    }
+                )
+
+        return OpenAIResult(
+            id = idx,
+            text = robo_desc,
+            metadata = metadata
+        )
 
 
 class MPMcpTools(_NeedsMPClient):