Skills › Scientific Agent Skills › pymatgen › mp_query.py
mp_query.py
417 lines · skills/pymatgen/scripts/mp_query.py at 5f2bc03
1#!/usr/bin/env python3
2"""Plan or explicitly execute one bounded Materials Project summary query."""
4from __future__ import annotations
6import argparse
7import math
8import os
9import re
10from datetime import datetime, timezone
11from typing import Any
13from _common import (
14 CliError,
15 DEFAULT_MAX_OUTPUT_BYTES,
16 MP_API_VERSION,
17 PYMATGEN_CORE_VERSION,
18 PYMATGEN_VERSION,
19 checked_output_file,
20 emit_json,
21 json_text,
22 package_versions,
23 positive_int,
24 safe_error_message,
25 write_json_new,
26)
29FIELD_PATTERN = re.compile(r"^[A-Za-z][A-Za-z0-9_]{0,79}$")
30MATERIAL_ID_PATTERN = re.compile(r"^(?:mp|mvc)-(?:[0-9]{1,30}|[a-z]{1,30})$")
31TOKEN_PATTERN = re.compile(r"^[A-Za-z0-9.+\-()]+$")
32MAX_RESULTS = 100
33MAX_FIELDS = 20
34MAX_MATERIAL_IDS = 100
37def csv_tokens(value: str | None, label: str) -> list[str] | None:
38 """Parse a bounded, unique comma-separated token list."""
39 if value is None:
40 return None
41 tokens = [token.strip() for token in value.split(",") if token.strip()]
42 if not tokens:
43 raise CliError(f"{label} must contain at least one token")
44 if len(tokens) > 100:
45 raise CliError(f"{label} may contain at most 100 tokens")
46 if any(not TOKEN_PATTERN.fullmatch(token) for token in tokens):
47 raise CliError(f"{label} contains unsupported characters")
48 return list(dict.fromkeys(tokens))
51def parse_fields(value: str) -> list[str]:
52 """Parse explicit response fields."""
53 fields = [field.strip() for field in value.split(",") if field.strip()]
54 fields = list(dict.fromkeys(fields))
55 if not fields:
56 raise CliError("--fields must contain at least one field")
57 if len(fields) > MAX_FIELDS:
58 raise CliError(f"--fields may contain at most {MAX_FIELDS} fields")
59 invalid = [field for field in fields if not FIELD_PATTERN.fullmatch(field)]
60 if invalid:
61 raise CliError(f"invalid field names: {invalid}")
62 return fields
65def checked_range(
66 values: list[float] | None,
67 label: str,
68 *,
69 non_negative: bool = False,
70) -> tuple[float, float] | None:
71 """Validate an inclusive two-number range."""
72 if values is None:
73 return None
74 low, high = values
75 if not all(math.isfinite(value) for value in (low, high)):
76 raise CliError(f"{label} bounds must be finite")
77 if low > high:
78 raise CliError(f"{label} minimum must not exceed maximum")
79 if non_negative and low < 0:
80 raise CliError(f"{label} must be non-negative")
81 return float(low), float(high)
84def query_contract(args: argparse.Namespace) -> dict[str, Any]:
85 """Validate arguments and return the exact public search kwargs."""
86 if len(args.material_id) > MAX_MATERIAL_IDS:
87 raise CliError(
88 f"--material-id may be repeated at most {MAX_MATERIAL_IDS} times"
89 )
90 invalid_ids = [
91 item
92 for item in args.material_id
93 if not MATERIAL_ID_PATTERN.fullmatch(item)
94 ]
95 if invalid_ids:
96 raise CliError(f"invalid Materials Project IDs: {invalid_ids}")
97 for label, value in (("--formula", args.formula), ("--chemsys", args.chemsys)):
98 if value is not None and (
99 len(value) > 200 or not TOKEN_PATTERN.fullmatch(value)
100 ):
101 raise CliError(f"{label} contains unsupported characters")
102 elements = csv_tokens(args.elements, "--elements")
103 exclude_elements = csv_tokens(args.exclude_elements, "--exclude-elements")
104 energy_range = checked_range(
105 args.energy_above_hull,
106 "--energy-above-hull",
107 non_negative=True,
108 )
109 band_gap = checked_range(args.band_gap, "--band-gap", non_negative=True)
110 stable = None
111 if args.is_stable != "any":
112 stable = args.is_stable == "true"
113 fields = parse_fields(args.fields)
114 effective_fields = list(dict.fromkeys(["material_id", *fields]))
115 filters: dict[str, Any] = {
116 "material_ids": args.material_id or None,
117 "formula": args.formula,
118 "chemsys": args.chemsys,
119 "elements": elements,
120 "exclude_elements": exclude_elements,
121 "energy_above_hull": energy_range,
122 "band_gap": band_gap,
123 "is_stable": stable,
124 }
125 if not any(value is not None for value in filters.values()):
126 raise CliError("at least one query filter is required")
127 search_kwargs = {
128 key: value for key, value in filters.items() if value is not None
129 }
130 search_kwargs.update(
131 {
132 "fields": effective_fields,
133 "all_fields": False,
134 "num_chunks": 1,
135 "chunk_size": args.limit,
136 }
137 )
138 return {
139 "filters": filters,
140 "requested_fields": fields,
141 "effective_fields": effective_fields,
142 "search_kwargs": search_kwargs,
143 }
146def plan_payload(
147 args: argparse.Namespace,
148 contract: dict[str, Any],
149) -> dict[str, Any]:
150 """Disclose network, credential, cache, field, limit, and output behavior."""
151 return {
152 "ok": True,
153 "action": "materials_project_summary_query",
154 "execute_requested": bool(args.execute),
155 "network_will_be_accessed": bool(args.execute),
156 "endpoint": "https://api.materialsproject.org/materials/summary/",
157 "client": "mp_api.client.MPRester",
158 "network_operations_when_executed": [
159 "mp-api compatibility/heartbeat metadata checks during MPRester initialization",
160 "one bounded materials.summary.search call",
161 ],
162 "user_agent_with_platform_details": False,
163 "authentication": {
164 "environment_variable": "MP_API_KEY",
165 "read_only_when_execute_is_set": True,
166 "accepted_on_command_line": False,
167 "value_logged_or_serialized": False,
168 },
169 "query": {
170 "filters": contract["filters"],
171 "requested_fields": contract["requested_fields"],
172 "effective_fields": contract["effective_fields"],
173 "limit": args.limit,
174 "num_chunks": 1,
175 },
176 "result_cache": {
177 "enabled": False,
178 "read": False,
179 "written": False,
180 "note": (
181 "This CLI has no implicit result cache. The explicit JSON output "
182 "is the reusable artifact. It does not request mp-api full-dataset "
183 "downloads, so the configured full-dataset cache is not used."
184 ),
185 },
186 "output": {
187 "path_as_provided": args.output,
188 "required_for_execution": True,
189 "maximum_bytes": args.max_output_bytes,
190 "existing_files_overwritten": False,
191 },
192 "rate_and_error_handling": {
193 "custom_retries": False,
194 "client_behavior": (
195 "mp-api 0.46.5 retries 429, 502, and 504 according to its "
196 "configured retry policy and respects Retry-After"
197 ),
198 "numeric_service_quota_assumed": False,
199 "errors_are_redacted_and_bounded": True,
200 },
201 "data_use": {
202 "license": "CC BY 4.0 for Materials Project data; contributed data may differ",
203 "citation_required": True,
204 "computed_data_is_experimental_truth": False,
205 "database_version_recorded_automatically": True,
206 },
207 }
210def document_to_json(document: Any) -> dict[str, Any]:
211 """Convert an mp-api document through its public Pydantic interface."""
212 if isinstance(document, dict):
213 result = document
214 elif hasattr(document, "model_dump"):
215 result = document.model_dump(mode="json")
216 else:
217 raise CliError(
218 f"unsupported mp-api result type: {type(document).__name__}"
219 )
220 if not isinstance(result, dict):
221 raise CliError("mp-api document did not serialize to an object")
222 json_text(result, pretty=False)
223 return result
226def execute_query(
227 args: argparse.Namespace,
228 contract: dict[str, Any],
229 output_path: Any,
230) -> dict[str, Any]:
231 """Execute exactly one bounded summary search."""
232 api_key = os.getenv("MP_API_KEY")
233 if not api_key:
234 raise CliError(
235 "MP_API_KEY is not set; obtain it from the Materials Project dashboard "
236 "and inject only that named secret through your shell or secret manager"
237 )
238 installed = package_versions(("pymatgen", "pymatgen-core", "mp-api"))
239 expected = {
240 "pymatgen": PYMATGEN_VERSION,
241 "pymatgen-core": PYMATGEN_CORE_VERSION,
242 "mp-api": MP_API_VERSION,
243 }
244 if installed != expected:
245 raise CliError(
246 f"network execution requires the verified package snapshot; "
247 f"expected={expected}, installed={installed}"
248 )
249 from mp_api.client import MPRester
250 from mp_api.client.core.exceptions import MPRestError
252 try:
253 with MPRester(
254 api_key=api_key,
255 endpoint="https://api.materialsproject.org/",
256 include_user_agent=False,
257 mute_progress_bars=True,
258 notify_db_version=False,
259 ) as rester:
260 database_version = rester.db_version
261 available = set(rester.materials.summary.available_fields)
262 invalid_fields = [
263 field
264 for field in contract["effective_fields"]
265 if field not in available
266 ]
267 if invalid_fields:
268 raise CliError(
269 f"fields unavailable in this endpoint/client: {invalid_fields}"
270 )
271 documents = rester.materials.summary.search(
272 **contract["search_kwargs"]
273 )
274 except CliError:
275 raise
276 except MPRestError as exc:
277 raise CliError(safe_error_message(exc, secret=api_key)) from exc
278 except Exception as exc:
279 raise CliError(safe_error_message(exc, secret=api_key)) from exc
280 serialized = [document_to_json(document) for document in documents[: args.limit]]
281 result = {
282 "schema_version": "1.0",
283 "provenance": {
284 "retrieved_at_utc": datetime.now(timezone.utc).isoformat(),
285 "endpoint": "https://api.materialsproject.org/materials/summary/",
286 "pymatgen_version": PYMATGEN_VERSION,
287 "pymatgen_core_version": PYMATGEN_CORE_VERSION,
288 "mp_api_version": MP_API_VERSION,
289 "database_version": database_version,
290 "data_license": "CC BY 4.0; contributed data is owned by contributors",
291 "citation": "https://materialsproject.org/about/cite",
292 },
293 "query": {
294 "filters": contract["filters"],
295 "fields": contract["effective_fields"],
296 "limit": args.limit,
297 "num_chunks": 1,
298 },
299 "returned": len(serialized),
300 "returned_equals_requested_limit": len(serialized) == args.limit,
301 "more_results_may_exist": len(serialized) == args.limit,
302 "documents": serialized,
303 "interpretation_limits": [
304 "Materials Project values are computed and method-dependent.",
305 "Aggregated values can change between database releases.",
306 "Missing fields do not imply a measured zero.",
307 "Band gaps and other properties carry documented systematic errors.",
308 ],
309 }
310 write_json_new(output_path, result, max_bytes=args.max_output_bytes)
311 return result
314def build_parser() -> argparse.ArgumentParser:
315 parser = argparse.ArgumentParser(
316 description=(
317 "Plan a bounded Materials Project summary query. No network or "
318 "credential access occurs unless --execute is supplied."
319 )
320 )
321 parser.add_argument("--material-id", action="append", default=[])
322 parser.add_argument("--formula")
323 parser.add_argument("--chemsys")
324 parser.add_argument("--elements", help="Comma-separated required elements")
325 parser.add_argument(
326 "--exclude-elements", help="Comma-separated excluded elements"
327 )
328 parser.add_argument(
329 "--energy-above-hull",
330 nargs=2,
331 type=float,
332 metavar=("MIN", "MAX"),
333 )
334 parser.add_argument("--band-gap", nargs=2, type=float, metavar=("MIN", "MAX"))
335 parser.add_argument(
336 "--is-stable",
337 choices=("any", "true", "false"),
338 default="any",
339 )
340 parser.add_argument(
341 "--fields",
342 required=True,
343 help="Comma-separated response fields; material_id is always added",
344 )
345 parser.add_argument("--limit", type=positive_int, default=25)
346 parser.add_argument("--output", help="New JSON result path (required with --execute)")
347 parser.add_argument("--plan-output", help="New JSON path for the dry-run plan")
348 parser.add_argument(
349 "--execute",
350 action="store_true",
351 help="Explicitly permit the disclosed bounded network query",
352 )
353 parser.add_argument(
354 "--max-output-bytes",
355 type=positive_int,
356 default=DEFAULT_MAX_OUTPUT_BYTES,
357 )
358 return parser
361def main() -> int:
362 args = build_parser().parse_args()
363 try:
364 if args.limit > MAX_RESULTS:
365 raise CliError(f"--limit may not exceed {MAX_RESULTS}")
366 if args.max_output_bytes > 50 * 1024 * 1024:
367 raise CliError("--max-output-bytes may not exceed 50 MiB")
368 contract = query_contract(args)
369 plan = plan_payload(args, contract)
370 if args.execute:
371 if not args.output:
372 raise CliError("--execute requires an explicit new --output path")
373 output_path = checked_output_file(args.output)
374 result = execute_query(args, contract, output_path)
375 emit_json(
376 {
377 "ok": True,
378 "executed": True,
379 "network_accessed": True,
380 "output": output_path.name,
381 "output_bytes": output_path.stat().st_size,
382 "returned": result["returned"],
383 "overwrote_existing": False,
384 "api_key_logged": False,
385 }
386 )
387 elif args.plan_output:
388 plan_path = checked_output_file(args.plan_output)
389 write_json_new(plan_path, plan, max_bytes=args.max_output_bytes)
390 emit_json(
391 {
392 "ok": True,
393 "executed": False,
394 "network_accessed": False,
395 "plan_output": plan_path.name,
396 "overwrote_existing": False,
397 }
398 )
399 else:
400 emit_json(plan)
401 return 0
402 except (CliError, ImportError, OSError, RuntimeError, TypeError, ValueError) as exc:
403 emit_json(
404 {
405 "ok": False,
406 "error": f"{type(exc).__name__}: {exc}"[:1000],
407 "completed": False,
408 "execute_requested": bool(args.execute),
409 "network_may_have_been_accessed": bool(args.execute),
410 "api_key_logged": False,
411 }
412 )
413 return 2
416if __name__ == "__main__":
417 raise SystemExit(main())