precommit

Author: esoteric-ephemera

mp_api/mcp/_schemas.py

@@ -1,35 +1,37 @@
 """Define auxiliary schemas used by some LLMs."""
+from __future__ import annotations
 
 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
+    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:
+    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}"
+            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([])
+    results: list[OpenAIResult] = Field([])

mp_api/mcp/mp_mcp.py

@@ -17,17 +17,19 @@
 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))
+    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_mcp = FastMCP("Materials_Project_MCP")

mp_api/mcp/server.py

@@ -1,7 +1,7 @@
 """Run MCP."""
 from __future__ import annotations
 
-from mp_api.mcp.mp_mcp import get_mcp, get_openai_compat_mcp
+from mp_api.mcp.mp_mcp import get_openai_compat_mcp
 
 mcp = get_openai_compat_mcp()
 

mp_api/mcp/tools.py

@@ -32,37 +32,35 @@
 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:
+    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: 
+
+        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"]
-                )
+            retults=[
+                OpenAIResult(id=doc["material_id"], text=doc["description"])
                 for doc in self.client.robocrys.search(query.split(","))
             ]
         )
 
-    def fetch(self, idx : str) -> OpenAIResult:
+    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"
@@ -78,33 +76,25 @@ def fetch(self, idx : str) -> OpenAIResult:
         Raises:
             MPRestError: If no identifier is specified
         """
-        if not isinstance(idx,str):
+        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: 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)
+            return OpenAIResult(id=idx)
 
-        metadata : dict[str,str] | None = None
-        if len(
-            summary_docs := self.client.summary.search(
-                material_ids=[idx]
-            )
-        ) > 0:
+        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])
+                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)
             }
@@ -113,16 +103,14 @@ def fetch(self, idx : str) -> OpenAIResult:
             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()
+                        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
-        )
+        return OpenAIResult(id=idx, text=robo_desc, metadata=metadata)
 
 
 class MPMcpTools(_NeedsMPClient):