|
| 1 | +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +"""Particle physics MCP tool wrapping the PDG particle database. |
| 16 | +
|
| 17 | +Provides lookup of particle properties (mass, width, charge, spin, lifetime) |
| 18 | +and search by name. Uses the ``particle`` pip package. No API key required. |
| 19 | +
|
| 20 | +Prerequisites: |
| 21 | + pip install particle |
| 22 | +
|
| 23 | +Usage: |
| 24 | + ++tool_modules=[nemo_skills.mcp.servers.particle_tool::ParticleTool] |
| 25 | +""" |
| 26 | + |
| 27 | +import logging |
| 28 | +from typing import Annotated, Any |
| 29 | + |
| 30 | +from pydantic import Field |
| 31 | + |
| 32 | +from nemo_skills.mcp.tool_manager import Tool |
| 33 | + |
| 34 | +logger = logging.getLogger(__name__) |
| 35 | + |
| 36 | +MAX_SEARCH_RESULTS = 10 |
| 37 | + |
| 38 | + |
| 39 | +def _format_particle(p) -> str: |
| 40 | + lines = [f"**{p.name}**"] |
| 41 | + lines.append(f"PDG ID: {p.pdgid}") |
| 42 | + if p.latex_name: |
| 43 | + lines.append(f"LaTeX: {p.latex_name}") |
| 44 | + if p.mass is not None: |
| 45 | + lines.append(f"Mass: {p.mass} MeV/c^2") |
| 46 | + if p.width is not None: |
| 47 | + lines.append(f"Width: {p.width} MeV") |
| 48 | + if p.charge is not None: |
| 49 | + lines.append(f"Charge: {p.charge} e") |
| 50 | + if p.J is not None: |
| 51 | + lines.append(f"Spin (J): {p.J}") |
| 52 | + if p.lifetime is not None: |
| 53 | + lines.append(f"Lifetime: {p.lifetime:.6e} ns") |
| 54 | + if p.anti_flag: |
| 55 | + lines.append(f"Anti-flag: {p.anti_flag.name}") |
| 56 | + if p.P is not None: |
| 57 | + lines.append(f"Parity (P): {'+' if p.P == 1 else '-'}") |
| 58 | + if p.C is not None: |
| 59 | + lines.append(f"C-parity: {'+' if p.C == 1 else '-'}") |
| 60 | + return "\n".join(lines) |
| 61 | + |
| 62 | + |
| 63 | +def particle_lookup( |
| 64 | + name_or_id: Annotated[ |
| 65 | + str, |
| 66 | + Field(description="Particle name (e.g. 'pi+', 'K0', 'J/psi(1S)') or PDG ID as a string (e.g. '211')."), |
| 67 | + ], |
| 68 | +) -> str: |
| 69 | + """Look up a particle by name or PDG ID. Returns mass, width, charge, spin, lifetime, and quantum numbers.""" |
| 70 | + from particle import InvalidParticle, Particle, ParticleNotFound |
| 71 | + |
| 72 | + try: |
| 73 | + p = Particle.from_pdgid(int(name_or_id)) |
| 74 | + return _format_particle(p) |
| 75 | + except (ValueError, ParticleNotFound, InvalidParticle): |
| 76 | + pass |
| 77 | + |
| 78 | + try: |
| 79 | + p = Particle.from_name(name_or_id) |
| 80 | + return _format_particle(p) |
| 81 | + except (ParticleNotFound, InvalidParticle): |
| 82 | + pass |
| 83 | + |
| 84 | + matches = Particle.findall(name_or_id) |
| 85 | + if matches: |
| 86 | + return _format_particle(matches[0]) |
| 87 | + |
| 88 | + return f"Particle '{name_or_id}' not found. Try a standard name like 'pi+', 'K-', 'D0', or a PDG ID." |
| 89 | + |
| 90 | + |
| 91 | +def particle_search( |
| 92 | + query: Annotated[ |
| 93 | + str, Field(description="Search query - substring match on particle names (e.g. 'K', 'charm', 'omega').") |
| 94 | + ], |
| 95 | +) -> str: |
| 96 | + """Search for particles by name. Returns a list of matching particles with key properties.""" |
| 97 | + from particle import Particle |
| 98 | + |
| 99 | + matches = Particle.findall(query) |
| 100 | + if not matches: |
| 101 | + return f"No particles found matching '{query}'." |
| 102 | + |
| 103 | + matches = matches[:MAX_SEARCH_RESULTS] |
| 104 | + results = [] |
| 105 | + for p in matches: |
| 106 | + mass_str = f"{p.mass} MeV/c^2" if p.mass is not None else "n/a" |
| 107 | + results.append(f"**{p.name}** (PDG {p.pdgid}) - mass: {mass_str}, charge: {p.charge}") |
| 108 | + header = f"Found {len(results)} particle(s) matching '{query}':\n" |
| 109 | + return header + "\n".join(results) |
| 110 | + |
| 111 | + |
| 112 | +class ParticleTool(Tool): |
| 113 | + def __init__(self) -> None: |
| 114 | + self._config: dict[str, Any] = {} |
| 115 | + |
| 116 | + def default_config(self) -> dict[str, Any]: |
| 117 | + return dict(self._config) |
| 118 | + |
| 119 | + def configure(self, overrides: dict[str, Any] | None = None, context: dict[str, Any] | None = None) -> None: |
| 120 | + if overrides: |
| 121 | + unknown = ", ".join(sorted(overrides)) |
| 122 | + raise ValueError(f"ParticleTool does not support overrides: {unknown}") |
| 123 | + |
| 124 | + async def list_tools(self) -> list[dict[str, Any]]: |
| 125 | + return [ |
| 126 | + { |
| 127 | + "name": "particle-lookup", |
| 128 | + "description": "Look up a particle by name or PDG ID.", |
| 129 | + "input_schema": { |
| 130 | + "type": "object", |
| 131 | + "properties": {"name_or_id": {"type": "string", "description": "Particle name or PDG ID."}}, |
| 132 | + "required": ["name_or_id"], |
| 133 | + }, |
| 134 | + }, |
| 135 | + { |
| 136 | + "name": "particle-search", |
| 137 | + "description": "Search for particles by name.", |
| 138 | + "input_schema": { |
| 139 | + "type": "object", |
| 140 | + "properties": {"query": {"type": "string", "description": "Substring query on particle names."}}, |
| 141 | + "required": ["query"], |
| 142 | + }, |
| 143 | + }, |
| 144 | + ] |
| 145 | + |
| 146 | + async def execute(self, tool_name: str, arguments: dict[str, Any], extra_args: dict[str, Any] | None = None): |
| 147 | + if extra_args: |
| 148 | + unknown = ", ".join(sorted(map(str, extra_args))) |
| 149 | + raise ValueError(f"ParticleTool does not support extra_args: {unknown}") |
| 150 | + arguments = dict(arguments or {}) |
| 151 | + if tool_name == "particle-lookup": |
| 152 | + return particle_lookup(**arguments) |
| 153 | + if tool_name == "particle-search": |
| 154 | + return particle_search(**arguments) |
| 155 | + return f"Error: unknown tool '{tool_name}'" |
0 commit comments