update mcp with slightly better docstr + some examples

Author: esoteric-ephemera

.gitignore

@@ -1,5 +1,3 @@
-uv.lock
-
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
@@ -123,3 +121,4 @@ _autosummary
 
 uv.lock
 JANAF_*_data.json
+.gemini

find_stable_si_mcp.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from mp_api.mcp.tools import MPMcpTools
+
+tools = MPMcpTools()
+
+stable_si_docs = tools.get_thermo_data(
+    formula="Si",
+    is_stable=True,
+    fields=["material_id", "formula_pretty", "energy_above_hull"],
+)
+
+if stable_si_docs:
+    print("Found stable silicon structures:")
+    for doc in stable_si_docs:
+        print(
+            f"  Material ID: {doc['material_id']}, "
+            f"Formula: {doc['pretty_formula']}, "
+            f"Energy Above Hull: {doc['energy_above_hull']}"
+        )
+else:
+    print("No stable silicon structures found.")

mp_api/client/core/utils.py

@@ -16,6 +16,7 @@
 
 _MAPI_SETTINGS = MAPIClientSettings()
 
+
 def _compare_emmet_ver(
     ref_version: str, op: Literal["==", ">", ">=", "<", "<="]
 ) -> bool:

mp_api/client/mprester.py

@@ -24,7 +24,12 @@
 from mp_api.client.core import BaseRester, MPRestError
 from mp_api.client.core._oxygen_evolution import OxygenEvolution
 from mp_api.client.core.settings import MAPIClientSettings
-from mp_api.client.core.utils import _compare_emmet_ver, load_json, validate_api_key, validate_ids
+from mp_api.client.core.utils import (
+    _compare_emmet_ver,
+    load_json,
+    validate_api_key,
+    validate_ids,
+)
 from mp_api.client.routes import GeneralStoreRester, MessagesRester, UserSettingsRester
 from mp_api.client.routes.materials import (
     AbsorptionRester,

mp_api/mcp/mp_mcp.py

@@ -16,6 +16,9 @@ def get_mcp() -> FastMCP:
     mcp_tools = MPMcpTools()
     for attr in {x for x in dir(mcp_tools) if x.startswith("get_")}:
         mp_mcp.tool(getattr(mcp_tools, attr))
+
+    # Register tool to set the user's API key
+    mp_mcp.tool(mcp_tools.update_user_api_key)
     return mp_mcp
 
 

mp_api/mcp/tools.py

@@ -2,8 +2,7 @@
 from __future__ import annotations
 
 from datetime import datetime
-import json
-from typing import Literal
+from typing import Literal, Any
 
 import plotly.graph_objects as plotly_go
 
@@ -22,9 +21,9 @@
 from emmet.core.thermo import ThermoType
 from emmet.core.vasp.calc_types import CalcType
 from emmet.core.xas import Edge, Type
-from pymatgen.analysis.phase_diagram import PhaseDiagram
 from pymatgen.analysis.magnetism.analyzer import Ordering
 from pymatgen.core.periodic_table import Element
+from pymatgen.core.composition import Composition
 from pymatgen.core.structure import Structure
 from pymatgen.electronic_structure.core import OrbitalType, Spin
 from pymatgen.entries.computed_entries import ComputedEntry
@@ -35,30 +34,93 @@
 class MPMcpTools(_NeedsMPClient):
     """Define tools needed for the MP MCP client."""
 
+    def get_structure_by_material_id(
+        self,
+        material_id: str,
+        structure_format: Literal["json", "poscar", "cif"],
+    ) -> dict[str, Any] | str:
+        """Find a structure in the Materials Project by its identifier/ID.
+
+        Return type changes based on format:
+            structure_format = "json":
+                Returns the JSON representation of a pymatgen structure object.
+            structure_format = "poscar":
+                Returns a VASP POSCAR-like representation
+            structure_format = "cif":
+                Returns a crystallographic information file (CIF)
+        """
+        struct = self.client.get_structure_by_material_id(
+            material_id=material_id,
+            final=True,
+            conventional_unit_cell=False,
+        )
+        if structure_format == "json":
+            return struct.as_dict() if hasattr(struct, "as_dict") else struct
+
+        if isinstance(struct, dict):
+            struct = Structure.from_dict(struct)
+        return struct.to(fmt=structure_format)
+
     def get_phase_diagram_from_elements(
         self,
         elements: list[str],
         thermo_type: ThermoType | str = "GGA_GGA+U_R2SCAN",
     ) -> plotly_go.Figure:
+        """Find a thermodynamic phase diagram in the Materials Project by specified elements.
+
+        ### Examples:
+        Given elements Na and Cl:
+        ```
+        phase_diagram = MPMcpTools().get_phase_diagram_from_elements(
+            elements = ["Na","Cl"],
+        )
+        ```
+
+        Given a chemical system, "K-P-O":
+        ```
+        phase_diagrasm =  MPMcpTools().get_phase_diagram_from_elements(
+            elements = "K-P-O".split("-"),
+        )
+        ```
+
+        """
         pd = self.client.materials.thermo.get_phase_diagram_from_chemsys(
             "-".join(elements), thermo_type
         )
         return pd.get_plot()  # has to be JSON serializable
 
-    def get_stability(
+    def get_stability_or_energy_above_hull(
         self,
-        composition: dict[str, float],
+        formula: str,
         energy: float,
-        run_type: Literal["GGA", "GGA+U", "R2SCAN"] | None = None,
-        thermo_type: str | ThermoType = "GGA_GGA+U",
-    ) -> list[dict]:
-        data = None
-        if run_type:
-            data = {"run_type": run_type}
-        return self.client.get_stability(
-            entries=ComputedEntry(composition, energy, data=data),
-            thermo_type=thermo_type,
-        )
+        run_type: Literal["GGA", "PBE", "GGA+U", "PBE+U", "R2SCAN"],
+    ) -> float:
+        """Get the stability of a particular material.
+
+        Given a material's formula and energy in eV, tells you the material's
+        stability from the energy above the hull (positive for unstable, 0 for stable)
+        in eV/atom.
+
+        The user must specify a particular functional:
+            - PBE or GGA
+            - PBE+U or GGA+U
+            - R2SCAN
+        this will add necessary corrections to the energy to compare with
+        the Materials Project.
+
+        """
+        run_type = run_type.replace("PBE", "GGA")
+        data = {"run_type": run_type}
+        thermo_type = "GGA_GGA+U" if run_type in {"GGA", "GGA+U"} else "R2SCAN"
+
+        try:
+            stability = self.client.get_stability(
+                entries=[ComputedEntry(Composition(formula), energy, data=data)],
+                thermo_type=thermo_type,
+            )
+            return stability[0]["e_above_hull"]
+        except ValueError:
+            return float("inf")
 
     def get_absorption_data(
         self,
@@ -644,7 +706,7 @@ def get_material_data(
             efermi=efermi,
             elastic_anisotropy=elastic_anisotropy,
             elements=elements,
-            energy_above_hull=tuple(float(x) for x in json.loads(energy_above_hull)),
+            energy_above_hull=energy_above_hull,
             equilibrium_reaction_energy=equilibrium_reaction_energy,
             exclude_elements=exclude_elements,
             formation_energy=formation_energy,
@@ -964,15 +1026,6 @@ def get_entry_by_material_id(
             conventional_unit_cell=conventional_unit_cell,
         )
 
-    def get_structure_by_material_id(
-        self, material_id: str, final: bool = True, conventional_unit_cell: bool = False
-    ) -> list[dict]:
-        return self.client.get_structure_by_material_id(
-            material_id=material_id,
-            final=final,
-            conventional_unit_cell=conventional_unit_cell,
-        )
-
     def get_structures(self, chemsys_formula: str | list[str] = True) -> list[dict]:
         return self.client.get_structures(chemsys_formula=chemsys_formula)
 

mp_api/mcp/utils.py

@@ -10,20 +10,38 @@
     from collections.abc import Callable
     from pathlib import Path
 
+_REQUIRED_CLIENT_KWARGS = {
+    "use_document_model": False,
+    "monty_decode": False,
+    "include_user_agent": True,
+}
+
 
 class _NeedsMPClient:
     def __init__(
         self,
         client_kwargs: dict[str, Any] | None = None,
         client: MPRester | None = None,
     ):
-        self.client = client or MPRester(
-            **{
-                **(client_kwargs or {}),
-                "use_document_model": False,
-                "monty_decode": False,
-            }
-        )
+        self._client_kwargs = {
+            **(client_kwargs or {}),
+            **_REQUIRED_CLIENT_KWARGS,
+        }
+        self.reset_client()
+
+    def reset_client(self) -> None:
+        """Reset the API client."""
+        self.client = MPRester(**self._client_kwargs)
+
+    def update_user_api_key(self, api_key: str) -> None:
+        """Change the API key used in the client.
+
+        Call this method to set the user's API correctly.
+        Ask the user for their API key as plain text,
+        and input the result to this method.
+        """
+        self._client_kwargs["api_key"] = api_key
+        self.reset_client()
 
 
 def get_annotation_signature(