513 lines
19 KiB
Python
513 lines
19 KiB
Python
# materials_mp_cif_mcp.py
|
||
#
|
||
# 说明：
|
||
# - 这是一个基于 FastMCP 的“Materials Project CIF”服务器，提供：
|
||
# 1) search_cifs：使用 Materials Project 的 mp-api 进行检索（两步：summary 过滤 + materials 获取结构）
|
||
# 2) get_cif：按材料 ID 获取结构并导出 CIF 文本（pymatgen）
|
||
# 3) filter_cifs：在本地对候选进行硬筛选
|
||
# 4) download_cif_bulk：批量下载 CIF，带并发与进度
|
||
# 5) 资源：cif://mp/{id} 直接读取 CIF 文本
|
||
#
|
||
# - 依赖：pip install "mcp[cli]" mp_api pymatgen
|
||
# - API Key：
|
||
# 默认从环境变量 MP_API_KEY 读取；若未设置则回退为你提供的 Key（仅用于演示，不建议在生产中硬编码）[3]
|
||
# - MCP/Context/Lifespan 的使用方法参考 SDK 文档（工具、结构化输出、进度等）[1]
|
||
#
|
||
# 引用：
|
||
# - MPRester/MP_API_KEY 的用法与推荐的会话管理方式见 Materials Project 文档 [3]
|
||
# - summary.search 与 materials.search 的查询与 fields 限制示例见文档 [4]
|
||
# - FastMCP 的工具、资源、Context、结构化输出见 README 示例 [1]
|
||
|
||
import os
|
||
import asyncio
|
||
from dataclasses import dataclass
|
||
from typing import Any, Literal
|
||
from collections.abc import AsyncIterator
|
||
from contextlib import asynccontextmanager
|
||
import re
|
||
|
||
from pydantic import BaseModel, Field
|
||
from mp_api.client import MPRester # [3]
|
||
from pymatgen.core import Structure # 用于转 CIF
|
||
from mcp.server.fastmcp import FastMCP, Context # [1]
|
||
from mcp.server.session import ServerSession # [1]
|
||
|
||
# ========= 配置 =========
|
||
|
||
# 生产中请通过环境变量 MP_API_KEY 配置；此处为演示而提供回退到你给的 API Key（不建议硬编码）[3]
|
||
DEFAULT_MP_API_KEY = os.getenv("MP_API_KEY", "X0NqhxkFeupGy7k09HRmLS6PI4VmvTBW")
|
||
|
||
# ========= 数据模型 =========
|
||
|
||
class LatticeParameters(BaseModel):
|
||
a: float | None = Field(default=None, description="a 轴长度 (Å)")
|
||
b: float | None = Field(default=None, description="b 轴长度 (Å)")
|
||
c: float | None = Field(default=None, description="c 轴长度 (Å)")
|
||
alpha: float | None = Field(default=None, description="α 角 (度)")
|
||
beta: float | None = Field(default=None, description="β 角 (度)")
|
||
gamma: float | None = Field(default=None, description="γ 角 (度)")
|
||
|
||
|
||
class CIFCandidate(BaseModel):
|
||
id: str = Field(description="Materials Project ID，如 mp-149")
|
||
source: Literal["mp"] = Field(default="mp", description="来源库")
|
||
chemical_formula: str | None = Field(default=None, description="化学式（pretty）")
|
||
space_group: str | None = Field(default=None, description="空间群符号")
|
||
cell_parameters: LatticeParameters | None = Field(default=None, description="晶格参数")
|
||
bandgap: float | None = Field(default=None, description="带隙 (eV)")
|
||
stability_energy: float | None = Field(default=None, description="E_above_hull (eV)")
|
||
magnetic_ordering: str | None = Field(default=None, description="磁性信息")
|
||
has_cif: bool = Field(default=True, description="是否存在可导出的 CIF（只要有结构即可）")
|
||
extras: dict[str, Any] | None = Field(default=None, description="原始文档片段（健壮解析）")
|
||
|
||
|
||
class SearchQuery(BaseModel):
|
||
# 基础化学约束
|
||
include_elements: list[str] | None = Field(default=None, description="包含元素集合（如 ['Si','O']）") # [4]
|
||
exclude_elements: list[str] | None = Field(default=None, description="排除元素集合（如 ['Li']）")
|
||
formula: str | None = Field(default=None, description="化学式片段（本地二次过滤）")
|
||
|
||
# 结构与性能约束（summary 可过滤带隙与 E_above_hull）[4]
|
||
bandgap_min: float | None = None
|
||
bandgap_max: float | None = None
|
||
e_above_hull_max: float | None = None
|
||
|
||
# 空间群、晶格范围（在本地二次过滤）
|
||
space_group: str | None = None
|
||
cell_a_min: float | None = None
|
||
cell_a_max: float | None = None
|
||
cell_b_min: float | None = None
|
||
cell_b_max: float | None = None
|
||
cell_c_min: float | None = None
|
||
cell_c_max: float | None = None
|
||
|
||
magnetic_ordering: str | None = None # 本地二次过滤
|
||
|
||
page: int = Field(default=1, ge=1, description="页码（仅本地分页；MP 端点请自行扩展）")
|
||
per_page: int = Field(default=50, ge=1, le=200, description="每页数量（本地分页）")
|
||
sort_by: str | None = Field(default=None, description="排序字段（bandgap/e_above_hull等，本地排序）")
|
||
sort_order: Literal["asc", "desc"] | None = Field(default=None, description="排序方向")
|
||
|
||
|
||
class CIFText(BaseModel):
|
||
id: str
|
||
source: Literal["mp"]
|
||
text: str = Field(description="CIF 文件文本内容")
|
||
metadata: dict[str, Any] | None = Field(default=None, description="附加元数据（如结构来源）")
|
||
|
||
|
||
class FilterCriteria(BaseModel):
|
||
include_elements: list[str] | None = None
|
||
exclude_elements: list[str] | None = None
|
||
formula: str | None = None
|
||
space_group: str | None = None
|
||
cell_a_min: float | None = None
|
||
cell_a_max: float | None = None
|
||
cell_b_min: float | None = None
|
||
cell_b_max: float | None = None
|
||
cell_c_min: float | None = None
|
||
cell_c_max: float | None = None
|
||
bandgap_min: float | None = None
|
||
bandgap_max: float | None = None
|
||
e_above_hull_max: float | None = None
|
||
magnetic_ordering: str | None = None
|
||
|
||
|
||
class BulkDownloadItem(BaseModel):
|
||
id: str
|
||
source: Literal["mp"] = "mp"
|
||
success: bool
|
||
text: str | None = None
|
||
error: str | None = None
|
||
|
||
|
||
class BulkDownloadResult(BaseModel):
|
||
items: list[BulkDownloadItem]
|
||
total: int
|
||
succeeded: int
|
||
failed: int
|
||
|
||
|
||
# ========= Lifespan =========
|
||
|
||
@dataclass
|
||
class AppContext:
|
||
mpr: MPRester
|
||
api_key: str
|
||
|
||
|
||
@asynccontextmanager
|
||
async def lifespan(_server: FastMCP) -> AsyncIterator[AppContext]:
|
||
# 推荐使用 MP_API_KEY 环境变量；此处回退到提供的 Key 字符串 [3]
|
||
api_key = DEFAULT_MP_API_KEY
|
||
mpr = MPRester(api_key) # 会话管理在 mp-api 内部 [3]
|
||
try:
|
||
yield AppContext(mpr=mpr, api_key=api_key)
|
||
finally:
|
||
# mp-api 未强制要求手动关闭；若未来提供关闭方法，可在此清理 [3]
|
||
pass
|
||
|
||
|
||
# ========= 服务器 =========
|
||
|
||
mcp = FastMCP(
|
||
name="Materials Project CIF",
|
||
instructions="通过 Materials Project API 检索并导出 CIF 文件。",
|
||
lifespan=lifespan,
|
||
streamable_http_path="/", # 便于在 Starlette 下挂载到子路径 [1]
|
||
)
|
||
|
||
|
||
# ========= 工具实现 =========
|
||
|
||
def _match_elements_set(text: str | None, required: list[str] | None, excluded: list[str] | None) -> bool:
|
||
if text is None:
|
||
return not required and not excluded
|
||
s = text.upper()
|
||
if required:
|
||
for el in required:
|
||
if el.upper() not in s:
|
||
return False
|
||
if excluded:
|
||
for el in excluded:
|
||
if el.upper() in s:
|
||
return False
|
||
return True
|
||
|
||
|
||
def _lattice_from_structure(struct: Structure | None) -> LatticeParameters | None:
|
||
if struct is None:
|
||
return None
|
||
lat = struct.lattice
|
||
return LatticeParameters(
|
||
a=float(lat.a),
|
||
b=float(lat.b),
|
||
c=float(lat.c),
|
||
alpha=float(lat.alpha),
|
||
beta=float(lat.beta),
|
||
gamma=float(lat.gamma),
|
||
)
|
||
|
||
|
||
def _sg_symbol_from_doc(doc: Any) -> str | None:
|
||
"""
|
||
尝试从 materials.search 返回的文档中提取空间群符号。
|
||
MP 文档中 symmetry 通常包含 spacegroup 信息（不同版本字段路径可能变化）[4]
|
||
"""
|
||
try:
|
||
sym = getattr(doc, "symmetry", None)
|
||
if sym and getattr(sym, "spacegroup", None):
|
||
sg = sym.spacegroup
|
||
# 常见字段：symbol 或 international
|
||
return getattr(sg, "symbol", None) or getattr(sg, "international", None)
|
||
except Exception:
|
||
pass
|
||
return None
|
||
|
||
|
||
async def _to_thread(fn, *args, **kwargs):
|
||
return await asyncio.to_thread(fn, *args, **kwargs)
|
||
|
||
|
||
def _parse_elements_from_formula(formula: str) -> list[str]:
|
||
"""从化学式中解析元素符号（简易版）。"""
|
||
if not formula:
|
||
return []
|
||
elems = re.findall(r"[A-Z][a-z]?", formula)
|
||
seen, ordered = set(), []
|
||
for e in elems:
|
||
if e not in seen:
|
||
seen.add(e)
|
||
ordered.append(e)
|
||
return ordered
|
||
|
||
@mcp.tool()
|
||
async def search_cifs(
|
||
query: SearchQuery,
|
||
ctx: Context[ServerSession, AppContext],
|
||
) -> list[CIFCandidate]:
|
||
"""
|
||
通过 summary 做初筛（限制返回数量），再用 materials 拉取当前页的结构与对称性。
|
||
修复点：
|
||
- 仅请求当前页（chunk_size/per_page + num_chunks=1），避免全库下载超时
|
||
- 绝不在返回值中包含 mp-api 文档对象或 Structure，extras 只放可序列化轻量字段
|
||
"""
|
||
app = ctx.request_context.lifespan_context
|
||
mpr = app.mpr
|
||
|
||
# 将 formula 自动转为元素集合（若用户未显式提供）
|
||
include_elements = query.include_elements or _parse_elements_from_formula(query.formula or "")
|
||
|
||
# 至少要有一个约束，避免全库扫描
|
||
if not include_elements and not query.exclude_elements and query.bandgap_min is None \
|
||
and query.bandgap_max is None and query.e_above_hull_max is None:
|
||
raise ValueError("请至少提供一个过滤条件，例如 include_elements、bandgap 或 energy_above_hull 上限，以避免超时。")
|
||
|
||
# 1) summary 初筛：只取“当前页”的文档（轻字段）
|
||
summary_kwargs: dict[str, Any] = {}
|
||
if include_elements:
|
||
summary_kwargs["elements"] = include_elements # MP 支持按元素集合检索
|
||
if query.exclude_elements:
|
||
summary_kwargs["exclude_elements"] = query.exclude_elements
|
||
if query.bandgap_min is not None or query.bandgap_max is not None:
|
||
summary_kwargs["band_gap"] = (
|
||
float(query.bandgap_min) if query.bandgap_min is not None else None,
|
||
float(query.bandgap_max) if query.bandgap_max is not None else None,
|
||
)
|
||
if query.e_above_hull_max is not None:
|
||
summary_kwargs["energy_above_hull"] = (None, float(query.e_above_hull_max))
|
||
|
||
fields_summary = [
|
||
"material_id",
|
||
"formula_pretty",
|
||
"band_gap",
|
||
"energy_above_hull",
|
||
"ordering",
|
||
]
|
||
|
||
# 关键：限制 summary 返回规模
|
||
try:
|
||
summary_docs = await asyncio.to_thread(
|
||
mpr.materials.summary.search,
|
||
fields=fields_summary,
|
||
chunk_size=query.per_page, # 每块（页）大小
|
||
num_chunks=1, # 只取一块 => 只取当前页数量
|
||
**summary_kwargs,
|
||
)
|
||
except TypeError:
|
||
# 如果你的 mp-api 版本不支持上述参数，提示升级，避免一次性抓回十几万条导致超时
|
||
raise RuntimeError(
|
||
"当前 mp-api 版本不支持 chunk_size/num_chunks，请升级 mp-api（pip install -U mp_api）。"
|
||
)
|
||
|
||
if not summary_docs:
|
||
await ctx.debug("summary 检索结果为空")
|
||
return []
|
||
|
||
# 本地排序（如需）
|
||
results_all: list[CIFCandidate] = []
|
||
for sdoc in summary_docs:
|
||
e_hull = getattr(sdoc, "energy_above_hull", None)
|
||
bandgap = getattr(sdoc, "band_gap", None)
|
||
ordering = getattr(sdoc, "ordering", None)
|
||
|
||
# 暂不放 structure/doc 到 extras，避免序列化错误
|
||
results_all.append(
|
||
CIFCandidate(
|
||
id=sdoc.material_id,
|
||
source="mp",
|
||
chemical_formula=getattr(sdoc, "formula_pretty", None),
|
||
space_group=None,
|
||
cell_parameters=None,
|
||
bandgap=float(bandgap) if bandgap is not None else None,
|
||
stability_energy=float(e_hull) if e_hull is not None else None,
|
||
magnetic_ordering=ordering,
|
||
has_cif=True,
|
||
extras={
|
||
"summary": {
|
||
"material_id": sdoc.material_id,
|
||
"formula_pretty": getattr(sdoc, "formula_pretty", None),
|
||
"band_gap": float(bandgap) if bandgap is not None else None,
|
||
"energy_above_hull": float(e_hull) if e_hull is not None else None,
|
||
"ordering": ordering,
|
||
}
|
||
},
|
||
)
|
||
)
|
||
|
||
# 本地分页：summary 已限制返回条数，这里只做安全截断和排序
|
||
if query.sort_by:
|
||
reverse = query.sort_order == "desc"
|
||
key_map = {
|
||
"bandgap": lambda c: (c.bandgap if c.bandgap is not None else float("inf")),
|
||
"e_above_hull": lambda c: (c.stability_energy if c.stability_energy is not None else float("inf")),
|
||
"energy_above_hull": lambda c: (c.stability_energy if c.stability_energy is not None else float("inf")),
|
||
}
|
||
if query.sort_by in key_map:
|
||
results_all.sort(key=key_map[query.sort_by], reverse=reverse)
|
||
|
||
page_slice = results_all[: query.per_page]
|
||
page_ids = [c.id for c in page_slice]
|
||
|
||
# 2) 只为当前页材料获取结构与对称性
|
||
if not page_ids:
|
||
return []
|
||
|
||
fields_materials = ["material_id", "structure", "symmetry"]
|
||
materials_docs = await asyncio.to_thread(
|
||
mpr.materials.search,
|
||
material_ids=page_ids,
|
||
fields=fields_materials,
|
||
)
|
||
mat_by_id = {doc.material_id: doc for doc in materials_docs}
|
||
|
||
# 合并结构/空间群，仍然不返回不可序列化对象
|
||
merged: list[CIFCandidate] = []
|
||
for c in page_slice:
|
||
mdoc = mat_by_id.get(c.id)
|
||
struct = getattr(mdoc, "structure", None) if mdoc else None
|
||
|
||
# 空间群符号
|
||
sg = None
|
||
try:
|
||
sym = getattr(mdoc, "symmetry", None)
|
||
if sym and getattr(sym, "spacegroup", None):
|
||
sgobj = sym.spacegroup
|
||
sg = getattr(sgobj, "symbol", None) or getattr(sgobj, "international", None)
|
||
except Exception:
|
||
pass
|
||
|
||
c.space_group = sg
|
||
c.cell_parameters = _lattice_from_structure(struct)
|
||
c.has_cif = struct is not None
|
||
# 保持 extras 为轻量 JSON
|
||
c.extras = {
|
||
**(c.extras or {}),
|
||
"materials": {"has_structure": struct is not None},
|
||
}
|
||
merged.append(c)
|
||
|
||
await ctx.debug(f"返回 {len(merged)} 条（已限制每页 {query.per_page} 条）")
|
||
return merged
|
||
|
||
@mcp.tool()
|
||
async def get_cif(
|
||
id: str,
|
||
ctx: Context[ServerSession, AppContext],
|
||
) -> CIFText:
|
||
"""
|
||
使用 materials.search 获取结构并导出 CIF 文本。
|
||
- fields=["structure"] 以提升速度 [4]
|
||
- 结构到 CIF 通过 pymatgen 的 Structure.to(fmt="cif")
|
||
"""
|
||
app = ctx.request_context.lifespan_context
|
||
mpr = app.mpr
|
||
|
||
await ctx.info(f"获取结构并导出 CIF：{id}")
|
||
docs = await _to_thread(mpr.materials.search, material_ids=[id], fields=["material_id", "structure"]) # [4]
|
||
if not docs:
|
||
raise ValueError(f"未找到材料：{id}")
|
||
|
||
doc = docs[0]
|
||
struct: Structure | None = getattr(doc, "structure", None)
|
||
if struct is None:
|
||
raise ValueError(f"材料 {id} 无结构数据，无法导出 CIF")
|
||
|
||
cif_text = struct.to(fmt="cif") # 直接输出 CIF 文本
|
||
return CIFText(id=id, source="mp", text=cif_text, metadata={"from": "materials.search"})
|
||
|
||
|
||
@mcp.tool()
|
||
async def filter_cifs(
|
||
candidates: list[CIFCandidate],
|
||
criteria: FilterCriteria,
|
||
ctx: Context[ServerSession, AppContext],
|
||
) -> list[CIFCandidate]:
|
||
def in_range(v: float | None, vmin: float | None, vmax: float | None) -> bool:
|
||
if v is None:
|
||
return vmin is None and vmax is None
|
||
if vmin is not None and v < vmin:
|
||
return False
|
||
if vmax is not None and v > vmax:
|
||
return False
|
||
return True
|
||
|
||
filtered: list[CIFCandidate] = []
|
||
for c in candidates:
|
||
# 元素集合近似匹配（基于化学式字符串）
|
||
if not _match_elements_set(c.chemical_formula, criteria.include_elements, criteria.exclude_elements):
|
||
continue
|
||
# 化学式包含判断
|
||
if criteria.formula:
|
||
if not c.chemical_formula or criteria.formula.upper() not in c.chemical_formula.upper():
|
||
continue
|
||
# 空间群
|
||
if criteria.space_group:
|
||
spg = (c.space_group or "").lower()
|
||
if criteria.space_group.lower() not in spg:
|
||
continue
|
||
# 晶格范围
|
||
cp = c.cell_parameters or LatticeParameters()
|
||
if not in_range(cp.a, criteria.cell_a_min, criteria.cell_a_max):
|
||
continue
|
||
if not in_range(cp.b, criteria.cell_b_min, criteria.cell_b_max):
|
||
continue
|
||
if not in_range(cp.c, criteria.cell_c_min, criteria.cell_c_max):
|
||
continue
|
||
# 带隙与稳定性
|
||
if not in_range(c.bandgap, criteria.bandgap_min, criteria.bandgap_max):
|
||
continue
|
||
if criteria.e_above_hull_max is not None:
|
||
if c.stability_energy is None or c.stability_energy > criteria.e_above_hull_max:
|
||
continue
|
||
# 磁性
|
||
if criteria.magnetic_ordering:
|
||
mo = "" # 当前未填充，保留兼容
|
||
if criteria.magnetic_ordering.upper() not in mo.upper():
|
||
continue
|
||
|
||
filtered.append(c)
|
||
|
||
await ctx.debug(f"筛选后：{len(filtered)} / 原始 {len(candidates)}")
|
||
return filtered
|
||
|
||
|
||
@mcp.tool()
|
||
async def download_cif_bulk(
|
||
ids: list[str],
|
||
concurrency: int = 4,
|
||
ctx: Context[ServerSession, AppContext] | None = None,
|
||
) -> BulkDownloadResult:
|
||
semaphore = asyncio.Semaphore(max(1, concurrency))
|
||
items: list[BulkDownloadItem] = []
|
||
|
||
async def worker(mid: str, idx: int) -> None:
|
||
nonlocal items
|
||
try:
|
||
async with semaphore:
|
||
# 复用当前 ctx
|
||
assert ctx is not None
|
||
cif = await get_cif(id=mid, ctx=ctx)
|
||
items.append(BulkDownloadItem(id=mid, success=True, text=cif.text))
|
||
except Exception as e:
|
||
items.append(BulkDownloadItem(id=mid, success=False, error=str(e)))
|
||
if ctx is not None:
|
||
progress = (idx + 1) / len(ids) if len(ids) > 0 else 1.0
|
||
await ctx.report_progress(progress=progress, total=1.0, message=f"{idx + 1}/{len(ids)}") # [1]
|
||
|
||
tasks = [worker(mid, i) for i, mid in enumerate(ids)]
|
||
await asyncio.gather(*tasks)
|
||
|
||
succeeded = sum(1 for it in items if it.success)
|
||
failed = len(items) - succeeded
|
||
return BulkDownloadResult(items=items, total=len(items), succeeded=succeeded, failed=failed)
|
||
|
||
|
||
# ========= 资源：cif://mp/{id} =========
|
||
|
||
@mcp.resource("cif://mp/{id}")
|
||
async def cif_resource(id: str) -> str:
|
||
"""
|
||
基于 URI 的 CIF 内容读取。资源函数无法直接拿到 Context，
|
||
这里简单地新的 MPRester 会话读取结构并转 CIF（仅资源路径用途）。
|
||
"""
|
||
# 资源中使用一次性会话（不复用 lifespan），适合偶发读取 [1]
|
||
mpr = MPRester(DEFAULT_MP_API_KEY) # [3]
|
||
try:
|
||
docs = mpr.materials.search(material_ids=[id], fields=["material_id", "structure"]) # [4]
|
||
if not docs or getattr(docs[0], "structure", None) is None:
|
||
raise ValueError(f"材料 {id} 无结构数据")
|
||
struct: Structure = docs[0].structure
|
||
return struct.to(fmt="cif")
|
||
finally:
|
||
pass
|
||
|
||
|
||
def create_materials_mcp() -> FastMCP:
|
||
"""供 Starlette 挂载的工厂函数。"""
|
||
return mcp
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 独立运行（非挂载模式），使用 Streamable HTTP 传输 [1]
|
||
mcp.run(transport="streamable-http") |