zntrack mcp

pyproject.toml

@@ -8,9 +8,13 @@ authors = [
 license = "Apache-2.0"
 keywords=["data-science", "data-version-control", "machine-learning", "reproducibility", "collaboration"]
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 dependencies = [
+    "cp2k-input-tools",
     "dvc>=3.60.0",
+    "ipsuite",
+    "massband",
+    "mlipx",
     "pandas>=2.2.3",
     "pyyaml>=6.0.2",
     "tqdm>=4.67.1",
@@ -22,6 +26,7 @@ dependencies = [
 
 [project.scripts]
 zntrack = 'zntrack.cli:app'
+zntrack-mcp = 'zntrack.mcp:mcp.run'
 
 [project.urls]
 Repository = "https://github.com/zincware/zntrack"
@@ -57,6 +62,12 @@ examples = [
     "rdkit2ase>=0.1.4",
 ]
 
+[project.optional-dependencies]
+mcp = [
+    "mcp>=1.12.3",
+    "fastmcp>=2.10.6",
+]
+
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"
@@ -121,5 +132,11 @@ markers = [
 ]
 addopts = "-m 'not benchmark'"
 
+[tool.uv.sources]
+ipsuite = { path = "../ipsuite", editable = true }
+massband = { path = "../massband", editable = true }
+mlipx = { path = "../mlipx", editable = true }
+cp2k-input-tools = { git = "https://github.com/cp2k/cp2k-input-tools.git" }
+
 [project.entry-points.'zntrack.nodes']
 zntrack-examples = 'zntrack.examples:nodes'

zntrack/__init__.py

@@ -2,7 +2,7 @@
 import logging
 import sys
 
-from zntrack import config
+from zntrack import config, entrypoints
 from zntrack.add import add
 from zntrack.apply import apply
 from zntrack.config import NOT_AVAILABLE, FieldTypes

zntrack/entrypoints.py

@@ -0,0 +1,42 @@
+"""Entry points discovery for ZnTrack nodes.
+
+This module provides functionality to discover and load all packages
+that have registered entry points under 'zntrack.nodes'.
+"""
+
+import importlib
+import importlib.metadata
+import logging
+from collections import defaultdict
+
+log = logging.getLogger(__name__)
+
+
+def get_registered_nodes(group: str = "zntrack.nodes") -> dict[str, list[str]]:
+    """Get all packages that registered into [project.entry-points.'zntrack.nodes']."""
+    registered_nodes = defaultdict(list)
+
+    try:
+        # Get all entry points for the 'zntrack.nodes' group
+        entry_points = importlib.metadata.entry_points(group=group)
+
+        for entry_point in entry_points:
+            try:
+                # Load the function registered at this entry point
+                nodes_func = entry_point.load()
+
+                # Call the function to get the dictionary of module -> node names
+                nodes_dict = nodes_func()
+
+                for module_name, node_names in nodes_dict.items():
+                    module_name = module_name.replace("-", "_")  # Normalize module names
+                    registered_nodes[module_name].extend(node_names)
+
+            except Exception as e:
+                log.error(f"Failed to load entry point '{entry_point.name}': {e}")
+                continue
+
+    except Exception as e:
+        log.error(f"Failed to discover entry points: {e}")
+
+    return dict(registered_nodes)

zntrack/examples/__init__.py

@@ -57,7 +57,6 @@
 ]
 
 
-def nodes() -> list[Node]:
+def nodes() -> dict[str, list[str]]:
     """Return the available nodes, grouped into categories."""
-
-    return [globals()[name] for name in __all__]
+    return {"zntrack.examples": __all__}

zntrack/mcp/__init__.py

@@ -0,0 +1,16 @@
+"""ZnTrack Model Context Provider (MCP) integration.
+
+Example
+-------
+To connect the ZnTrack MCP server to claude code at the current working directory, run:
+
+```bash
+bunx @anthropic-ai/claude-code mcp add zntrack-server -- uv run \
+      --project "$(pwd)" zntrack-mcp
+bunx @anthropic-ai/claude-code
+```
+"""
+
+from zntrack.mcp.server import mcp
+
+__all__ = ["mcp"]

zntrack/mcp/resources/graph_getting_started.md

@@ -0,0 +1,13 @@
+```python
+import zntrack
+from package import Node1, Node2
+
+project = zntrack.Project()
+
+with project:
+    node1 = Node1(param="value1")
+    node2 = Node2(input=node1.output)
+
+if __name__ == "__main__":
+    project.build()
+```

zntrack/mcp/resources/graph_with_groups.md

@@ -0,0 +1,18 @@
+```python
+import zntrack
+from package import Node1, Node2
+
+project = zntrack.Project()
+
+with zntrack.group("group1"):
+    node1 = Node1(param="value1")
+
+with zntrack.group("group2"):
+    node2 = Node2(input=node1.output)
+
+with project.group("group1", "nested"):
+    node3 = Node1(param="value2")
+
+if __name__ == "__main__":
+    project.build()
+```

zntrack/mcp/resources/node_getting_started.md

@@ -0,0 +1,25 @@
+```py
+import zntrack
+from pathlib import Path
+
+class MyNode(zntrack.Node):
+    parameter: dict = zntrack.params()
+    output: Path = zntrack.outs_path(zntrack.nwd / "file.txt")
+
+    def run(self) -> None:
+        self.output.parent.mkdir(parents=True, exist_ok=True)
+        self.output.write_text(str(self.parameter))
+```
+
+Options include
+- `zntrack.deps()` for dependencies to other nodes.
+- `zntrack.deps_path` for dependencies to files.
+- `zntrack.params` for parameters.
+- `zntrack.params_path()` for parameter files (yaml / json)
+- `zntrack.outs_path()` for output files.
+- `zntrack.outs()` for outputs that are not files.
+- `zntrack.metrics()` dict output.
+- `zntrack.metrics_path()` for metrics files.
+
+Special directory options include
+- `zntrack.nwd` for the node working directory.