mcp-gen-material
This commit is contained in:
29
mcp/main.py
29
mcp/main.py
@@ -1,31 +1,42 @@
|
||||
# main.py
|
||||
# 将 Materials CIF MCP 与 System Tools MCP 一起挂载到 Starlette。
|
||||
# 关键点：
|
||||
# - 在 lifespan 中启动每个 MCP 的 session_manager.run()（参考 SDK README 的 Starlette 挂载示例与 streamable_http_app 用法 [1]）
|
||||
# - 通过 Mount 指定各自的子路径（如 /system 与 /materials）
|
||||
|
||||
import contextlib
|
||||
from starlette.applications import Starlette
|
||||
from starlette.routing import Mount
|
||||
|
||||
from test_tools import create_test_mcp
|
||||
from system_tools import create_system_mcp # 如果暂时不用可先不挂
|
||||
from system_tools import create_system_mcp
|
||||
from materialproject_mcp import create_materials_mcp
|
||||
|
||||
# 先创建 MCP 实例
|
||||
test_mcp = create_test_mcp()
|
||||
# 创建 MCP 实例
|
||||
system_mcp = create_system_mcp()
|
||||
materials_mcp = create_materials_mcp()
|
||||
|
||||
# 关键：在 Starlette 的 lifespan 中启动 MCP 的 session manager
|
||||
# 在 Starlette 的 lifespan 中启动 MCP 的 session manager
|
||||
@contextlib.asynccontextmanager
|
||||
async def lifespan(app: Starlette):
|
||||
async with contextlib.AsyncExitStack() as stack:
|
||||
# await stack.enter_async_context(test_mcp.session_manager.run())
|
||||
await stack.enter_async_context(system_mcp.session_manager.run())
|
||||
await stack.enter_async_context(materials_mcp.session_manager.run())
|
||||
yield # 服务器运行期间
|
||||
# 退出时自动清理
|
||||
|
||||
# 挂载两个 MCP 的 Streamable HTTP App
|
||||
app = Starlette(
|
||||
lifespan=lifespan,
|
||||
routes=[
|
||||
# Mount("/test", app=test_mcp.streamable_http_app()),
|
||||
Mount("/system", app=system_mcp.streamable_http_app()),
|
||||
Mount("/materials", app=materials_mcp.streamable_http_app()),
|
||||
],
|
||||
)
|
||||
|
||||
# 启动代码为uvicorn main:app --host 0.0.0.0 --port 8000
|
||||
# url为http://localhost:8000/system
|
||||
# 启动命令（终端执行）：
|
||||
# uvicorn main:app --host 0.0.0.0 --port 8000
|
||||
# 访问：
|
||||
# http://localhost:8000/system
|
||||
# http://localhost:8000/materials
|
||||
#
|
||||
# 如果需要浏览器客户端访问（CORS 暴露 Mcp-Session-Id），请参考 README 中的 CORS 配置示例 [1]
|
||||
513
mcp/materialproject_mcp.py
Normal file
513
mcp/materialproject_mcp.py
Normal file
@@ -0,0 +1,513 @@
|
||||
# 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")
|
||||
Reference in New Issue
Block a user