From c5a31d50fb125c1a5d832e802ce57cd75a2f990c Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 06:30:20 +0200 Subject: [PATCH 01/37] feat(sn): extend StandardName schema with rich fields and fix persistence bugs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 1: - Add 12 rich fields to StandardName schema (documentation, kind, tags, links, ids_paths, validity_domain, constraints, subject, component, coordinate, position, process) - Add StandardNameKind enum (scalar/vector/metadata) - Fix StandardNameReviewStatus: rename candidate→drafted, add published - Fix MEASURES→HAS_STANDARD_NAME in schema doc and signals.py query Phase 4: - Fix coalesce bug in write_standard_names — all fields use coalesce(new, existing) to prevent data loss on re-runs - Write all rich fields to graph - Create CANONICAL_UNITS relationship per schema range convention - Wire embedding generation in persist_worker Tests: - Add tests/sn/test_graph_ops.py (12 tests) covering coalesce, DD/signal/unit relationships, and query filtering - Add tests/sn/conftest.py with shared fixtures - Update test_publish.py for candidate→drafted rename --- imas_codex/schemas/standard_name.yaml | 56 ++++- imas_codex/sn/graph_ops.py | 81 ++++--- imas_codex/sn/models.py | 2 +- imas_codex/sn/publish.py | 2 +- imas_codex/sn/sources/signals.py | 2 +- imas_codex/sn/workers.py | 43 +++- tests/sn/conftest.py | 56 +++++ tests/sn/test_graph_ops.py | 315 ++++++++++++++++++++++++++ tests/sn/test_publish.py | 6 +- 9 files changed, 523 insertions(+), 40 deletions(-) create mode 100644 tests/sn/conftest.py create mode 100644 tests/sn/test_graph_ops.py diff --git a/imas_codex/schemas/standard_name.yaml b/imas_codex/schemas/standard_name.yaml index 43511b580..6ec09de2e 100644 --- a/imas_codex/schemas/standard_name.yaml +++ b/imas_codex/schemas/standard_name.yaml @@ -54,15 +54,27 @@ enums: StandardNameReviewStatus: description: Review lifecycle for standard names permissible_values: - candidate: - description: Generated by LLM, awaiting review + drafted: + description: LLM-generated, awaiting review + published: + description: Exported to catalog PR for review accepted: - description: Reviewed and accepted into vocabulary + description: Imported from merged catalog entry rejected: description: Reviewed and rejected skipped: description: Skipped during review (e.g., low confidence) + StandardNameKind: + description: Entry kind for standard names + permissible_values: + scalar: + description: Scalar quantity + vector: + description: Vector quantity (R,Z or multi-component) + metadata: + description: Non-measurable concept or classification + # ============================================================================= # Classes # ============================================================================= @@ -79,7 +91,7 @@ classes: Inbound relationships: - (IMASNode)-[:HAS_STANDARD_NAME]->(StandardName) - - (FacilitySignal)-[:MEASURES]->(StandardName) + - (FacilitySignal)-[:HAS_STANDARD_NAME]->(StandardName) Example: plasma_current, electron_density_core, major_radius class_uri: sn:StandardName @@ -139,3 +151,39 @@ classes: embedded_at: description: When the embedding was last computed range: datetime + documentation: + description: >- + Rich documentation with LaTeX equations, governing physics, + measurement methods, typical values, sign conventions. + Uses [name](#name) inline links to other standard names. + kind: + description: Entry kind — scalar, vector, or metadata + range: StandardNameKind + tags: + description: Classification tags from controlled vocabulary + multivalued: true + range: string + links: + description: Internal cross-references to related standard names (name only) + multivalued: true + range: string + ids_paths: + description: IMAS DD paths mapped to this standard name + multivalued: true + range: string + validity_domain: + description: Physical region where this quantity is defined + constraints: + description: Physical/mathematical constraints (e.g., T_e > 0) + multivalued: true + range: string + subject: + description: Particle species (electron, ion, deuterium, etc.) + component: + description: Vector component (radial, toroidal, vertical, etc.) + coordinate: + description: Coordinate qualifier + position: + description: Spatial location qualifier (magnetic_axis, midplane, etc.) + process: + description: Physical process qualifier (ohmic, bootstrap, etc.) diff --git a/imas_codex/sn/graph_ops.py b/imas_codex/sn/graph_ops.py index 6c74b63e2..0df5a13a4 100644 --- a/imas_codex/sn/graph_ops.py +++ b/imas_codex/sn/graph_ops.py @@ -13,6 +13,8 @@ import logging from typing import Any +from imas_codex.graph.client import GraphClient + logger = logging.getLogger(__name__) @@ -31,8 +33,6 @@ def get_extraction_candidates_dd( Returns dynamic leaf nodes that have been enriched (status=embedded), optionally filtered by IDS or physics domain. """ - from imas_codex.graph.client import GraphClient - with GraphClient() as gc: params: dict[str, Any] = {"limit": limit} where_clauses = [ @@ -76,8 +76,6 @@ def get_extraction_candidates_signals( Returns signals that have been enriched, optionally filtered by physics domain. """ - from imas_codex.graph.client import GraphClient - with GraphClient() as gc: params: dict[str, Any] = {"facility": facility, "limit": limit} where_clauses = ["s.status = 'enriched'"] @@ -112,8 +110,6 @@ def get_extraction_candidates_signals( def get_existing_standard_names() -> set[str]: """Return the set of existing StandardName node IDs for deduplication.""" - from imas_codex.graph.client import GraphClient - with GraphClient() as gc: results = gc.query("MATCH (sn:StandardName) RETURN sn.id AS id") return {r["id"] for r in results} @@ -125,8 +121,6 @@ def get_named_source_ids() -> set[str]: Used for resumability: extract skips sources that already have a standard name unless --force is specified. """ - from imas_codex.graph.client import GraphClient - with GraphClient() as gc: results = gc.query(""" MATCH (src)-[:HAS_STANDARD_NAME]->(sn:StandardName) @@ -153,47 +147,63 @@ def write_standard_names(names: list[dict[str, Any]]) -> int: - ``source_id``: the originating path / signal ID Optional fields: ``physical_base``, ``subject``, ``component``, - ``coordinate``, ``position``, ``units``, ``description``, - ``model``, ``review_status``, ``generated_at``, ``confidence``. + ``coordinate``, ``position``, ``process``, ``units``, ``description``, + ``documentation``, ``kind``, ``tags``, ``links``, ``ids_paths``, + ``validity_domain``, ``constraints``, ``model``, ``review_status``, + ``generated_at``, ``confidence``. Returns the number of nodes written. """ - from imas_codex.graph.client import GraphClient - if not names: return 0 with GraphClient() as gc: - # MERGE StandardName nodes with provenance + # MERGE StandardName nodes with provenance — coalesce to preserve existing data gc.query( """ UNWIND $batch AS b MERGE (sn:StandardName {id: b.id}) - SET sn.source_type = b.source_type, - sn.physical_base = b.physical_base, - sn.subject = b.subject, - sn.component = b.component, - sn.coordinate = b.coordinate, - sn.position = b.position, - sn.units = b.units, - sn.description = b.description, - sn.model = b.model, - sn.review_status = b.review_status, - sn.generated_at = b.generated_at, - sn.confidence = b.confidence, + SET sn.source_type = coalesce(b.source_type, sn.source_type), + sn.physical_base = coalesce(b.physical_base, sn.physical_base), + sn.subject = coalesce(b.subject, sn.subject), + sn.component = coalesce(b.component, sn.component), + sn.coordinate = coalesce(b.coordinate, sn.coordinate), + sn.position = coalesce(b.position, sn.position), + sn.process = coalesce(b.process, sn.process), + sn.description = coalesce(b.description, sn.description), + sn.documentation = coalesce(b.documentation, sn.documentation), + sn.kind = coalesce(b.kind, sn.kind), + sn.tags = coalesce(b.tags, sn.tags), + sn.links = coalesce(b.links, sn.links), + sn.ids_paths = coalesce(b.ids_paths, sn.ids_paths), + sn.validity_domain = coalesce(b.validity_domain, sn.validity_domain), + sn.constraints = coalesce(b.constraints, sn.constraints), + sn.canonical_units = coalesce(b.units, sn.canonical_units), + sn.model = coalesce(b.model, sn.model), + sn.review_status = coalesce(b.review_status, sn.review_status), + sn.generated_at = coalesce(b.generated_at, sn.generated_at), + sn.confidence = coalesce(b.confidence, sn.confidence), sn.created_at = coalesce(sn.created_at, datetime()) """, batch=[ { "id": n["id"], - "source_type": n.get("source_type", ""), + "source_type": n.get("source_type") or None, "physical_base": n.get("physical_base"), "subject": n.get("subject"), "component": n.get("component"), "coordinate": n.get("coordinate"), "position": n.get("position"), - "units": n.get("units"), + "process": n.get("process"), "description": n.get("description"), + "documentation": n.get("documentation"), + "kind": n.get("kind"), + "tags": n.get("tags") or None, + "links": n.get("links") or None, + "ids_paths": n.get("ids_paths") or None, + "validity_domain": n.get("validity_domain"), + "constraints": n.get("constraints") or None, + "units": n.get("units"), "model": n.get("model"), "review_status": n.get("review_status"), "generated_at": n.get("generated_at"), @@ -236,6 +246,21 @@ def write_standard_names(names: list[dict[str, Any]]) -> int: ], ) + # Create CANONICAL_UNITS relationships: StandardName → Unit + units_batch = [ + {"id": n["id"], "unit": n["units"]} for n in names if n.get("units") + ] + if units_batch: + gc.query( + """ + UNWIND $batch AS b + MATCH (sn:StandardName {id: b.id}) + MERGE (u:Unit {id: b.unit}) + MERGE (sn)-[:CANONICAL_UNITS]->(u) + """, + batch=units_batch, + ) + written = len(names) logger.info("Wrote %d StandardName nodes", written) return written @@ -272,8 +297,6 @@ def get_validated_standard_names( list of dicts with keys: name, description, source, source_path, canonical_units, confidence, ids_name. """ - from imas_codex.graph.client import GraphClient - with GraphClient() as gc: params: dict[str, Any] = {"confidence_min": confidence_min} diff --git a/imas_codex/sn/models.py b/imas_codex/sn/models.py index 412a64678..091858fe6 100644 --- a/imas_codex/sn/models.py +++ b/imas_codex/sn/models.py @@ -52,7 +52,7 @@ class SNPublishEntry(BaseModel): ) unit: str | None = Field(default=None, description="SI unit string") tags: list[str] = Field(default_factory=list, description="Classification tags") - status: str = Field(default="candidate", description="Entry status") + status: str = Field(default="drafted", description="Entry status") description: str = Field(default="", description="Human-readable description") provenance: SNProvenance = Field(description="Generation provenance") diff --git a/imas_codex/sn/publish.py b/imas_codex/sn/publish.py index 97823c678..745e5f9ea 100644 --- a/imas_codex/sn/publish.py +++ b/imas_codex/sn/publish.py @@ -275,7 +275,7 @@ def graph_records_to_entries( kind="physical", unit=unit, tags=tags, - status="candidate", + status="drafted", description=description[:500] if description else "", provenance=provenance, ) diff --git a/imas_codex/sn/sources/signals.py b/imas_codex/sn/sources/signals.py index e36a516e6..63c4085ba 100644 --- a/imas_codex/sn/sources/signals.py +++ b/imas_codex/sn/sources/signals.py @@ -53,7 +53,7 @@ def extract_signal_candidates( f""" MATCH (s:FacilitySignal) WHERE {where_clause} - OPTIONAL MATCH (s)-[:MEASURES]->(sn:StandardName) + OPTIONAL MATCH (s)-[:HAS_STANDARD_NAME]->(sn:StandardName) RETURN s.id AS signal_id, s.description AS description, s.physics_domain AS physics_domain, s.canonical_units AS units, diff --git a/imas_codex/sn/workers.py b/imas_codex/sn/workers.py index b81d8d33a..3e6f9f10b 100644 --- a/imas_codex/sn/workers.py +++ b/imas_codex/sn/workers.py @@ -657,7 +657,7 @@ async def persist_worker(state: SNBuildState, **_kwargs) -> None: # Enrich with provenance for entry in state.validated: entry.setdefault("model", model) - entry.setdefault("review_status", "skipped") + entry.setdefault("review_status", "drafted") entry.setdefault("generated_at", now) # confidence comes from LLM output — never default to 1.0 @@ -681,6 +681,47 @@ async def persist_worker(state: SNBuildState, **_kwargs) -> None: written = await asyncio.to_thread(write_standard_names, state.validated) + # Embed descriptions for vector search + if written > 0: + try: + from imas_codex.embeddings.description import embed_descriptions_batch + + embed_items = [ + {"id": e["id"], "description": e.get("description", "")} + for e in state.validated + if e.get("description") + ] + if embed_items: + enriched = await asyncio.to_thread( + embed_descriptions_batch, embed_items + ) + # Write embeddings back to graph + from imas_codex.graph.client import GraphClient + + def _write_embeddings(): + with GraphClient() as gc: + gc.query( + """ + UNWIND $batch AS b + MATCH (sn:StandardName {id: b.id}) + SET sn.embedding = b.embedding, + sn.embedded_at = datetime() + """, + batch=[ + {"id": e["id"], "embedding": e["embedding"]} + for e in enriched + if e.get("embedding") + ], + ) + + await asyncio.to_thread(_write_embeddings) + wlog.info("Embedded %d StandardName descriptions", len(embed_items)) + except Exception: + wlog.warning( + "Embedding generation failed — names persisted without embeddings", + exc_info=True, + ) + state.persist_stats.processed = written state.persist_stats.record_batch(written) state.stats["persist_written"] = written diff --git a/tests/sn/conftest.py b/tests/sn/conftest.py new file mode 100644 index 000000000..aa5f65052 --- /dev/null +++ b/tests/sn/conftest.py @@ -0,0 +1,56 @@ +"""Shared fixtures for standard name tests.""" + +from __future__ import annotations + +import pytest + + +@pytest.fixture() +def sample_standard_names() -> list[dict]: + """Sample standard name dicts for write_standard_names testing.""" + return [ + { + "id": "electron_temperature", + "source_type": "dd", + "source_id": "core_profiles/profiles_1d/electrons/temperature", + "physical_base": "temperature", + "subject": "electron", + "description": "Electron temperature profile", + "documentation": "The electron temperature $T_e$ is measured by Thomson scattering.", + "kind": "scalar", + "tags": ["core_profiles", "kinetics"], + "links": ["ion_temperature", "electron_density"], + "ids_paths": ["core_profiles/profiles_1d/electrons/temperature"], + "validity_domain": "core plasma", + "constraints": ["T_e > 0"], + "units": "eV", + "model": "test/model", + "review_status": "drafted", + "generated_at": "2024-01-01T00:00:00Z", + "confidence": 0.95, + }, + { + "id": "plasma_current", + "source_type": "signal", + "source_id": "tcv:ip/measured", + "physical_base": "current", + "description": "Plasma current", + "units": "A", + "kind": "scalar", + "tags": ["magnetics"], + "model": "test/model", + "review_status": "drafted", + "generated_at": "2024-01-01T00:00:00Z", + "confidence": 0.88, + }, + ] + + +@pytest.fixture() +def mock_graph_client(): + """A mock GraphClient that records query calls.""" + from unittest.mock import MagicMock + + client = MagicMock() + client.query = MagicMock(return_value=[]) + return client diff --git a/tests/sn/test_graph_ops.py b/tests/sn/test_graph_ops.py new file mode 100644 index 000000000..f1f64b044 --- /dev/null +++ b/tests/sn/test_graph_ops.py @@ -0,0 +1,315 @@ +"""Tests for standard name graph operations. + +Tests write_standard_names coalesce behavior, relationship creation, +and get_validated_standard_names filtering — all mocked, no live Neo4j. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + + +class TestWriteStandardNames: + """Test write_standard_names Cypher generation and coalesce behavior.""" + + def _call_write(self, names: list[dict], mock_gc: MagicMock) -> int: + """Call write_standard_names with a mocked GraphClient.""" + with patch("imas_codex.sn.graph_ops.GraphClient") as MockGC: + MockGC.return_value.__enter__ = MagicMock(return_value=mock_gc) + MockGC.return_value.__exit__ = MagicMock(return_value=False) + from imas_codex.sn.graph_ops import write_standard_names + + return write_standard_names(names) + + def test_write_coalesce_preserves_existing( + self, sample_standard_names: list[dict] + ) -> None: + """Re-running write with None fields should NOT overwrite existing data. + + The Cypher must use coalesce(b.field, sn.field) so that passing + None for a field preserves whatever is already in the graph. + """ + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + # First write: all fields populated + self._call_write(sample_standard_names, mock_gc) + + # Verify MERGE query uses coalesce + merge_call = mock_gc.query.call_args_list[0] + cypher = merge_call[0][0] + + # Every field SET should use coalesce pattern + assert "coalesce(b.source_type, sn.source_type)" in cypher + assert "coalesce(b.description, sn.description)" in cypher + assert "coalesce(b.documentation, sn.documentation)" in cypher + assert "coalesce(b.kind, sn.kind)" in cypher + assert "coalesce(b.tags, sn.tags)" in cypher + assert "coalesce(b.links, sn.links)" in cypher + assert "coalesce(b.ids_paths, sn.ids_paths)" in cypher + assert "coalesce(b.validity_domain, sn.validity_domain)" in cypher + assert "coalesce(b.constraints, sn.constraints)" in cypher + assert "coalesce(b.confidence, sn.confidence)" in cypher + assert "coalesce(b.process, sn.process)" in cypher + + # created_at should use coalesce(sn.created_at, datetime()) — preserve existing + assert "coalesce(sn.created_at, datetime())" in cypher + + def test_write_empty_returns_zero(self) -> None: + """Empty list should return 0 without touching the graph.""" + from imas_codex.sn.graph_ops import write_standard_names + + result = write_standard_names([]) + assert result == 0 + + def test_dd_relationship_created(self, sample_standard_names: list[dict]) -> None: + """DD-sourced names should create (IMASNode)-[:HAS_STANDARD_NAME]->(StandardName).""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + self._call_write(sample_standard_names, mock_gc) + + # Find the DD relationship query + dd_calls = [ + call for call in mock_gc.query.call_args_list if "IMASNode" in str(call) + ] + assert len(dd_calls) >= 1, "Should create DD HAS_STANDARD_NAME relationship" + + dd_cypher = dd_calls[0][0][0] + assert "HAS_STANDARD_NAME" in dd_cypher + assert "MEASURES" not in dd_cypher # Old relationship name must not appear + assert "IMASNode" in dd_cypher + + def test_signal_relationship_created( + self, sample_standard_names: list[dict] + ) -> None: + """Signal-sourced names should create (FacilitySignal)-[:HAS_STANDARD_NAME]->(StandardName).""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + self._call_write(sample_standard_names, mock_gc) + + # Find the signal relationship query + signal_calls = [ + call + for call in mock_gc.query.call_args_list + if "FacilitySignal" in str(call) + ] + assert len(signal_calls) >= 1, ( + "Should create signal HAS_STANDARD_NAME relationship" + ) + + signal_cypher = signal_calls[0][0][0] + assert "HAS_STANDARD_NAME" in signal_cypher + assert "MEASURES" not in signal_cypher + assert "FacilitySignal" in signal_cypher + + def test_unit_relationship_created(self, sample_standard_names: list[dict]) -> None: + """Names with units should create (StandardName)-[:CANONICAL_UNITS]->(Unit).""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + self._call_write(sample_standard_names, mock_gc) + + # Find the CANONICAL_UNITS relationship query + unit_calls = [ + call + for call in mock_gc.query.call_args_list + if "CANONICAL_UNITS" in str(call) + ] + assert len(unit_calls) >= 1, "Should create CANONICAL_UNITS relationship" + + unit_cypher = unit_calls[0][0][0] + assert "Unit" in unit_cypher + assert "MERGE (u:Unit" in unit_cypher + assert "MERGE (sn)-[:CANONICAL_UNITS]->(u)" in unit_cypher + + def test_no_unit_relationship_when_no_units(self) -> None: + """Names without units should NOT create CANONICAL_UNITS relationships.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + names = [ + { + "id": "test_name", + "source_type": "dd", + "source_id": "some/path", + } + ] + self._call_write(names, mock_gc) + + unit_calls = [ + call + for call in mock_gc.query.call_args_list + if "CANONICAL_UNITS" in str(call) + ] + assert len(unit_calls) == 0, "Should NOT create CANONICAL_UNITS when no units" + + def test_rich_fields_in_batch(self, sample_standard_names: list[dict]) -> None: + """All rich fields should be included in the batch parameter.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + self._call_write(sample_standard_names, mock_gc) + + merge_call = mock_gc.query.call_args_list[0] + batch = merge_call[1]["batch"] + + # First entry should have all rich fields + first = batch[0] + assert first["id"] == "electron_temperature" + assert first["documentation"] is not None + assert first["kind"] == "scalar" + assert ( + "core_profiles" in (first.get("tags") or []) + or first.get("tags") is not None + ) + assert first["validity_domain"] == "core plasma" + + def test_empty_lists_become_none(self) -> None: + """Empty list fields should be converted to None for coalesce to work.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + names = [ + { + "id": "test_name", + "source_type": "dd", + "source_id": "some/path", + "tags": [], + "links": [], + "ids_paths": [], + "constraints": [], + } + ] + self._call_write(names, mock_gc) + + merge_call = mock_gc.query.call_args_list[0] + batch = merge_call[1]["batch"] + first = batch[0] + + # Empty lists should become None so coalesce preserves existing + assert first["tags"] is None + assert first["links"] is None + assert first["ids_paths"] is None + assert first["constraints"] is None + + +class TestGetValidatedStandardNames: + """Test get_validated_standard_names query filtering.""" + + def test_confidence_filter(self) -> None: + """Should filter by minimum confidence.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock( + return_value=[ + { + "name": "electron_temperature", + "description": "Te", + "source": "dd", + "source_path": "core_profiles/profiles_1d/electrons/temperature", + "canonical_units": "eV", + "confidence": 0.95, + "ids_name": "core_profiles", + } + ] + ) + + with patch("imas_codex.sn.graph_ops.GraphClient") as MockGC: + MockGC.return_value.__enter__ = MagicMock(return_value=mock_gc) + MockGC.return_value.__exit__ = MagicMock(return_value=False) + + from imas_codex.sn.graph_ops import get_validated_standard_names + + results = get_validated_standard_names(confidence_min=0.9) + + # Verify confidence_min was passed to query + call_kwargs = mock_gc.query.call_args[1] + assert call_kwargs["confidence_min"] == 0.9 + assert len(results) == 1 + + def test_ids_filter(self) -> None: + """Should filter by IDS name.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + with patch("imas_codex.sn.graph_ops.GraphClient") as MockGC: + MockGC.return_value.__enter__ = MagicMock(return_value=mock_gc) + MockGC.return_value.__exit__ = MagicMock(return_value=False) + + from imas_codex.sn.graph_ops import get_validated_standard_names + + get_validated_standard_names(ids_filter="equilibrium") + + # Verify ids_filter was passed to query + call_kwargs = mock_gc.query.call_args[1] + assert call_kwargs["ids_filter"] == "equilibrium" + + # Verify the Cypher includes the IDS filter clause + cypher = mock_gc.query.call_args[0][0] + assert "ids_filter" in cypher + + def test_no_filters(self) -> None: + """With no filters, should return all standard names.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock( + return_value=[ + { + "name": "a", + "description": "", + "source": "dd", + "source_path": "x", + "canonical_units": None, + "confidence": 1.0, + "ids_name": None, + }, + { + "name": "b", + "description": "", + "source": "dd", + "source_path": "y", + "canonical_units": None, + "confidence": 1.0, + "ids_name": None, + }, + ] + ) + + with patch("imas_codex.sn.graph_ops.GraphClient") as MockGC: + MockGC.return_value.__enter__ = MagicMock(return_value=mock_gc) + MockGC.return_value.__exit__ = MagicMock(return_value=False) + + from imas_codex.sn.graph_ops import get_validated_standard_names + + results = get_validated_standard_names() + + assert len(results) == 2 + + +class TestGetExistingStandardNames: + """Test deduplication query.""" + + def test_returns_set_of_ids(self) -> None: + """Should return a set of standard name IDs.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock( + return_value=[ + {"id": "electron_temperature"}, + {"id": "plasma_current"}, + ] + ) + + with patch("imas_codex.sn.graph_ops.GraphClient") as MockGC: + MockGC.return_value.__enter__ = MagicMock(return_value=mock_gc) + MockGC.return_value.__exit__ = MagicMock(return_value=False) + + from imas_codex.sn.graph_ops import get_existing_standard_names + + result = get_existing_standard_names() + + assert isinstance(result, set) + assert "electron_temperature" in result + assert "plasma_current" in result + assert len(result) == 2 diff --git a/tests/sn/test_publish.py b/tests/sn/test_publish.py index 73be92b30..a416d1f02 100644 --- a/tests/sn/test_publish.py +++ b/tests/sn/test_publish.py @@ -46,7 +46,7 @@ def sample_entry(sample_provenance: SNProvenance) -> SNPublishEntry: kind="physical", unit="eV", tags=["equilibrium", "core_profiles"], - status="candidate", + status="drafted", description="Electron temperature profile", provenance=sample_provenance, ) @@ -144,7 +144,7 @@ def test_defaults(self, sample_provenance: SNProvenance) -> None: provenance=sample_provenance, ) assert entry.kind == "physical" - assert entry.status == "candidate" + assert entry.status == "drafted" assert entry.tags == [] assert entry.unit is None @@ -201,7 +201,7 @@ def test_format(self, sample_entry: SNPublishEntry) -> None: assert doc["name"] == "electron_temperature" assert doc["kind"] == "physical" assert doc["unit"] == "eV" - assert doc["status"] == "candidate" + assert doc["status"] == "drafted" assert doc["description"] == "Electron temperature profile" assert doc["provenance"]["source"] == "dd" assert doc["provenance"]["confidence"] == 0.95 From a2b28a361e6afc3dc756da3d62e4b204ccf847f3 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 06:42:06 +0200 Subject: [PATCH 02/37] refactor(sn): rename ids_paths to imas_paths and add physics_domain for schema consistency --- imas_codex/schemas/standard_name.yaml | 4 +++- imas_codex/sn/graph_ops.py | 6 +++--- tests/sn/conftest.py | 2 +- tests/sn/test_graph_ops.py | 6 +++--- 4 files changed, 10 insertions(+), 8 deletions(-) diff --git a/imas_codex/schemas/standard_name.yaml b/imas_codex/schemas/standard_name.yaml index 6ec09de2e..269d8e15d 100644 --- a/imas_codex/schemas/standard_name.yaml +++ b/imas_codex/schemas/standard_name.yaml @@ -167,7 +167,7 @@ classes: description: Internal cross-references to related standard names (name only) multivalued: true range: string - ids_paths: + imas_paths: description: IMAS DD paths mapped to this standard name multivalued: true range: string @@ -185,5 +185,7 @@ classes: description: Coordinate qualifier position: description: Spatial location qualifier (magnetic_axis, midplane, etc.) + physics_domain: + description: Physics domain classification (equilibrium, transport, etc.) process: description: Physical process qualifier (ohmic, bootstrap, etc.) diff --git a/imas_codex/sn/graph_ops.py b/imas_codex/sn/graph_ops.py index 0df5a13a4..96d1fb5cf 100644 --- a/imas_codex/sn/graph_ops.py +++ b/imas_codex/sn/graph_ops.py @@ -148,7 +148,7 @@ def write_standard_names(names: list[dict[str, Any]]) -> int: Optional fields: ``physical_base``, ``subject``, ``component``, ``coordinate``, ``position``, ``process``, ``units``, ``description``, - ``documentation``, ``kind``, ``tags``, ``links``, ``ids_paths``, + ``documentation``, ``kind``, ``tags``, ``links``, ``imas_paths``, ``validity_domain``, ``constraints``, ``model``, ``review_status``, ``generated_at``, ``confidence``. @@ -175,7 +175,7 @@ def write_standard_names(names: list[dict[str, Any]]) -> int: sn.kind = coalesce(b.kind, sn.kind), sn.tags = coalesce(b.tags, sn.tags), sn.links = coalesce(b.links, sn.links), - sn.ids_paths = coalesce(b.ids_paths, sn.ids_paths), + sn.imas_paths = coalesce(b.imas_paths, sn.imas_paths), sn.validity_domain = coalesce(b.validity_domain, sn.validity_domain), sn.constraints = coalesce(b.constraints, sn.constraints), sn.canonical_units = coalesce(b.units, sn.canonical_units), @@ -200,7 +200,7 @@ def write_standard_names(names: list[dict[str, Any]]) -> int: "kind": n.get("kind"), "tags": n.get("tags") or None, "links": n.get("links") or None, - "ids_paths": n.get("ids_paths") or None, + "imas_paths": n.get("imas_paths") or None, "validity_domain": n.get("validity_domain"), "constraints": n.get("constraints") or None, "units": n.get("units"), diff --git a/tests/sn/conftest.py b/tests/sn/conftest.py index aa5f65052..e15ebb9cb 100644 --- a/tests/sn/conftest.py +++ b/tests/sn/conftest.py @@ -20,7 +20,7 @@ def sample_standard_names() -> list[dict]: "kind": "scalar", "tags": ["core_profiles", "kinetics"], "links": ["ion_temperature", "electron_density"], - "ids_paths": ["core_profiles/profiles_1d/electrons/temperature"], + "imas_paths": ["core_profiles/profiles_1d/electrons/temperature"], "validity_domain": "core plasma", "constraints": ["T_e > 0"], "units": "eV", diff --git a/tests/sn/test_graph_ops.py b/tests/sn/test_graph_ops.py index f1f64b044..6504c6498 100644 --- a/tests/sn/test_graph_ops.py +++ b/tests/sn/test_graph_ops.py @@ -48,7 +48,7 @@ def test_write_coalesce_preserves_existing( assert "coalesce(b.kind, sn.kind)" in cypher assert "coalesce(b.tags, sn.tags)" in cypher assert "coalesce(b.links, sn.links)" in cypher - assert "coalesce(b.ids_paths, sn.ids_paths)" in cypher + assert "coalesce(b.imas_paths, sn.imas_paths)" in cypher assert "coalesce(b.validity_domain, sn.validity_domain)" in cypher assert "coalesce(b.constraints, sn.constraints)" in cypher assert "coalesce(b.confidence, sn.confidence)" in cypher @@ -180,7 +180,7 @@ def test_empty_lists_become_none(self) -> None: "source_id": "some/path", "tags": [], "links": [], - "ids_paths": [], + "imas_paths": [], "constraints": [], } ] @@ -193,7 +193,7 @@ def test_empty_lists_become_none(self) -> None: # Empty lists should become None so coalesce preserves existing assert first["tags"] is None assert first["links"] is None - assert first["ids_paths"] is None + assert first["imas_paths"] is None assert first["constraints"] is None From 952f9ecf36a32ab29b5852d4902e435a6d923731 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 07:17:51 +0200 Subject: [PATCH 03/37] docs: plan to import PhysicsDomain from imas-standard-names --- .../15-import-physics-domain.md | 209 ++++++++++++++++++ 1 file changed, 209 insertions(+) create mode 100644 plans/features/standard-names/15-import-physics-domain.md diff --git a/plans/features/standard-names/15-import-physics-domain.md b/plans/features/standard-names/15-import-physics-domain.md new file mode 100644 index 000000000..372e15c99 --- /dev/null +++ b/plans/features/standard-names/15-import-physics-domain.md @@ -0,0 +1,209 @@ +# Import PhysicsDomain from imas-standard-names + +> **Repo**: imas-codex +> **Status**: planned +> **Depends on**: imas-standard-names `01-rename-tags-to-physics-domain.md` (Phase 6 RC release) + +## Problem + +imas-codex maintains its own `PhysicsDomain` enum (22 values) codegen'd from +a 250-line LinkML schema (`physics_domains.yaml`). This duplicates the physics +domain vocabulary that imas-standard-names now owns as the canonical source. + +After imas-standard-names publishes `PhysicsDomain` as a `str, Enum` (31 values), +imas-codex should import it rather than maintain a parallel definition. + +## Current Architecture + +``` +imas_codex/schemas/physics_domains.yaml ← 250-line LinkML schema (SOURCE) + ↓ gen_physics_domains.py (codegen) +imas_codex/core/physics_domain.py ← PhysicsDomain(str, Enum) 22 values + ↓ imported by +imas_codex/schemas/common.yaml ← `imports: - physics_domains` + ↓ used by +facility.yaml, standard_name.yaml ← `range: PhysicsDomain` + ↓ used by +15+ Python modules ← runtime enum validation +``` + +## Target Architecture + +``` +imas-standard-names (PyPI) + ↓ exports +PhysicsDomain(str, Enum) 31 values + ↓ imported by +imas_codex/core/physics_domain.py ← re-export (one-line change) + ↓ imported by (unchanged) +15+ Python modules ← same import path, no changes +``` + +## Phase 1: Replace codegen with import + +### 1a. Update `imas_codex/core/physics_domain.py` + +Replace the entire codegen'd file with a re-export: + +```python +"""Physics domain enum — canonical source is imas-standard-names. + +This module re-exports PhysicsDomain from imas-standard-names so that +all imas-codex code continues to import from the same path: + from imas_codex.core.physics_domain import PhysicsDomain +""" + +from imas_standard_names.grammar.tag_types import PhysicsDomain + +__all__ = ["PhysicsDomain"] +``` + +This file is currently auto-generated and gitignored. Change it to a +hand-written re-export and **remove from .gitignore**. + +### 1b. Delete `imas_codex/schemas/physics_domains.yaml` + +The 250-line LinkML schema is no longer needed. The enum source of truth +is now in imas-standard-names. + +### 1c. Delete `scripts/gen_physics_domains.py` + +The codegen script is no longer needed. + +### 1d. Update `scripts/build_models.py` + +Remove the physics_domains codegen step. The build hook currently: +1. Generates physics_domain.py from physics_domains.yaml +2. Generates graph models from facility.yaml/common.yaml/etc + +Remove step 1. Step 2 continues unchanged. + +### 1e. Update `hatch_build_hooks.py` + +Remove physics_domain from the build hook's generated file list. + +### 1f. Update `.gitignore` + +Remove `imas_codex/core/physics_domain.py` from gitignore — it's now a +hand-written file that should be tracked. + +### 1g. Update `imas_codex/schemas/common.yaml` + +Remove `imports: - physics_domains` from the imports list. The LinkML schema +no longer needs to reference the physics_domains schema since the enum now +comes from Python at runtime. The `range: PhysicsDomain` annotations in +facility.yaml and standard_name.yaml become string-validated at the LinkML +level but enum-validated at the Python level. + +**Option A** (clean): Remove `range: PhysicsDomain` from LinkML, use +`range: string` — runtime Python code does enum validation via the imported +enum. This avoids needing a phantom LinkML enum. + +**Option B** (keep schema validation): Create a minimal LinkML enum stub +that lists just the values (no descriptions/categories) and is auto-synced +from imas-standard-names. This preserves `range: PhysicsDomain` in schema. + +**Decision**: Option A is cleaner. The PhysicsDomain values appear in the +generated schema_context_data.py regardless (built from Python enum, not +LinkML). Schema compliance tests validate against the Python enum. + +### 1h. Update `pyproject.toml` + +Add `imas-standard-names >= 0.8.0` to dependencies (the version with +PhysicsDomain export). + +## Phase 2: Add new enum values to graph + +The unified enum has 9 new values not currently in the graph. These are +additive — no existing data needs migration. + +New values: `core_plasma_physics`, `fast_particles`, `runaway_electrons`, +`waves`, `fueling`, `plasma_initiation`, `spectroscopy`, `neutronics`, +`gyrokinetics`. + +No graph migration needed — new values become available for new data +automatically. Existing classification prompts will start using them +as they appear in the enum. + +## Phase 3: Update tests + +### 3a. `tests/core/test_physics_categorization.py` + +Update test expectations for the new 31-value enum. Add tests for +new values. + +### 3b. `tests/graph/test_schema_compliance.py` + +Verify schema compliance tests pass with the new enum source. +The tests should work unchanged since they validate against the +Python enum, not the LinkML schema directly. + +### 3c. Run full test suite + +```bash +uv run build-models --force +uv run pytest tests/ -x -q +``` + +## Phase 4: Clean up and commit + +```bash +# Lint +uv run ruff check --fix . +uv run ruff format . + +# Stage changes (NOT auto-generated files) +git add imas_codex/core/physics_domain.py # Now hand-written, tracked +git add -u # Stage deletions and modifications +# DO NOT stage: graph/models.py, graph/dd_models.py, config/models.py + +uv run git commit -m "refactor: import PhysicsDomain from imas-standard-names + +BREAKING CHANGE: PhysicsDomain enum now has 31 values (was 22). +Nine new values added: core_plasma_physics, fast_particles, +runaway_electrons, waves, fueling, plasma_initiation, spectroscopy, +neutronics, gyrokinetics. + +Removed physics_domains.yaml LinkML schema and gen_physics_domains.py +codegen script. PhysicsDomain is now imported from imas-standard-names +and re-exported from imas_codex.core.physics_domain." + +git pull --no-rebase origin develop +git push origin develop +``` + +## Documentation Updates + +| Target | Changes | +|--------|---------| +| `AGENTS.md` | Update PhysicsDomain section: note it's imported from imas-standard-names | +| `AGENTS.md` | Remove physics_domains.yaml from schema files list | +| `AGENTS.md` | Remove gen_physics_domains.py from build pipeline | + +## Files Changed + +| Action | File | Notes | +|--------|------|-------| +| REWRITE | `imas_codex/core/physics_domain.py` | Codegen → hand-written re-export | +| DELETE | `imas_codex/schemas/physics_domains.yaml` | No longer source of truth | +| DELETE | `scripts/gen_physics_domains.py` | No longer needed | +| MODIFY | `scripts/build_models.py` | Remove physics_domains step | +| MODIFY | `hatch_build_hooks.py` | Remove from generated files | +| MODIFY | `.gitignore` | Un-ignore physics_domain.py | +| MODIFY | `imas_codex/schemas/common.yaml` | Remove physics_domains import | +| MODIFY | `pyproject.toml` | Add imas-standard-names >= 0.8.0 dep | +| MODIFY | `tests/core/test_physics_categorization.py` | Update for 31 values | +| MODIFY | `AGENTS.md` | Update documentation | + +## Risks + +- **Version pinning**: If imas-standard-names adds/removes PhysicsDomain + values in a future release, imas-codex graph data could become inconsistent. + Mitigation: pin to `>= 0.8.0, < 1.0.0` and review on major bumps. +- **Build order**: `uv sync` must install imas-standard-names before the + build hook runs. Since physics_domain.py is now hand-written (not codegen'd), + this is only a runtime concern, not a build-time concern. +- **LinkML validation gap**: With Option A, LinkML schemas lose + `range: PhysicsDomain` validation. Schema compliance tests still work + because they validate against the Python enum. The gap is only in the + LinkML schema itself (used for documentation, not runtime). From b67bd2c95628f4963e6ec92252454fe7a7a64e7e Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 09:02:13 +0200 Subject: [PATCH 04/37] fix: hardcode ACR image name so fork RC builds reach Azure Fork CI builds used github.repository for ACR image path, producing simon-mcintosh/imas-codex instead of iterorganization/imas-codex. Azure watches the upstream path only, so fork RCs were invisible. Split into IMAGE_NAME (per-fork, for GHCR) and ACR_IMAGE_NAME (hardcoded upstream path, for ACR). GHCR stays per-fork since each fork has its own container registry namespace. --- .github/workflows/docker-build-push.yml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/docker-build-push.yml b/.github/workflows/docker-build-push.yml index c059166f4..68f4be742 100644 --- a/.github/workflows/docker-build-push.yml +++ b/.github/workflows/docker-build-push.yml @@ -15,6 +15,9 @@ env: GHCR_REGISTRY: ghcr.io ACR_REGISTRY: crcommonallfrc.azurecr.io IMAGE_NAME: ${{ github.repository }} + # ACR image name is always the upstream path so Azure picks up fork RC + # builds. GHCR uses IMAGE_NAME (per-fork) — each fork has its own GHCR. + ACR_IMAGE_NAME: iterorganization/imas-codex IMAS_DD_VERSION: "4.1.0" # Graph OCI artifacts: try the repo owner's registry first (where # `imas-codex release` pushes them), fall back to the upstream org @@ -529,7 +532,7 @@ jobs: id: meta-acr uses: docker/metadata-action@v5 with: - images: ${{ env.ACR_REGISTRY }}/${{ env.IMAGE_NAME }} + images: ${{ env.ACR_REGISTRY }}/${{ env.ACR_IMAGE_NAME }} tags: | type=semver,pattern=v{{version}}${{ matrix.graph_variant.suffix }} type=semver,pattern=v{{major}}.{{minor}}${{ matrix.graph_variant.suffix }} From 5bf6247a433c9bbf0e563c8b5fee4857bb70d611 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 09:06:40 +0200 Subject: [PATCH 05/37] refactor: import PhysicsDomain from imas-standard-names Replace codegen'd PhysicsDomain enum (22 values from LinkML schema) with import from imas-standard-names PyPI package (32 values, StrEnum). Deleted: physics_domains.yaml, gen_physics_domains.py, domains.yaml Removed: codegen steps from build_models.py and hatch_build_hooks.py Added: imas_codex/core/physics_domain.py as tracked re-export file BREAKING CHANGE: PhysicsDomain enum now has 32 values (was 22). New values: core_plasma_physics, fast_particles, runaway_electrons, waves, fueling, plasma_initiation, spectroscopy, neutronics, gyrokinetics, plasma_measurement_diagnostics. --- .gitignore | 1 - AGENTS.md | 10 +- hatch_build_hooks.py | 52 --- imas_codex/core/physics_domain.py | 10 + imas_codex/definitions/physics/domains.yaml | 345 ------------------ imas_codex/schemas/common.yaml | 7 +- imas_codex/schemas/facility.yaml | 4 +- imas_codex/schemas/imas_dd.yaml | 6 +- imas_codex/schemas/physics_domains.yaml | 379 -------------------- pyproject.toml | 5 +- scripts/build_models.py | 66 +--- scripts/gen_physics_domains.py | 148 -------- uv.lock | 57 +-- 13 files changed, 34 insertions(+), 1056 deletions(-) create mode 100644 imas_codex/core/physics_domain.py delete mode 100644 imas_codex/definitions/physics/domains.yaml delete mode 100644 imas_codex/schemas/physics_domains.yaml delete mode 100644 scripts/gen_physics_domains.py diff --git a/.gitignore b/.gitignore index b3772b5d8..1177388b1 100644 --- a/.gitignore +++ b/.gitignore @@ -8,7 +8,6 @@ __pycache__/ imas_codex/graph/models.py imas_codex/graph/schema_context_data.py imas_codex/config/models.py -imas_codex/core/physics_domain.py agents/schema-reference.md # Generated index files diff --git a/AGENTS.md b/AGENTS.md index a688e127d..b23ac5b43 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -74,19 +74,21 @@ All graph node types, relationships, and properties are defined in LinkML schema **Schema files:** - `imas_codex/schemas/facility.yaml` - Facility graph: SourceFile, SignalNode, CodeChunk, etc. - `imas_codex/schemas/imas_dd.yaml` - DD graph: IMASNode, DDVersion, Unit, IMASCoordinateSpec -- `imas_codex/schemas/common.yaml` - Shared: status enums, PhysicsDomain +- `imas_codex/schemas/common.yaml` - Shared: status enums **Build pipeline:** - Models auto-generated during `uv sync` via hatch build hook - Regenerate manually: `uv run build-models --force` -- Output: `imas_codex/graph/models.py`, `imas_codex/graph/dd_models.py`, `imas_codex/config/models.py` +- Output: `imas_codex/graph/models.py`, `imas_codex/graph/dd_models.py`, `imas_codex/config/models.py`, `agents/schema-reference.md`, `imas_codex/graph/schema_context_data.py` **CRITICAL: Never commit auto-generated files.** These are gitignored and rebuilt on `uv sync`. If `git status` shows a generated model file as untracked or modified, do NOT stage it. Generated files: - `imas_codex/graph/models.py` - `imas_codex/graph/dd_models.py` - `imas_codex/config/models.py` -- `imas_codex/core/physics_domain.py` - `agents/schema-reference.md` +- `imas_codex/graph/schema_context_data.py` + +**PhysicsDomain enum**: Imported from the `imas-standard-names` PyPI package and re-exported from `imas_codex.core.physics_domain`. The canonical vocabulary is maintained in the imas-standard-names project. Contains 32 physics domain values. `imas_codex/core/physics_domain.py` is a hand-written one-line re-export — it IS committed and should NOT be treated as auto-generated. Always import enums and classes from generated models. Never hardcode status values: @@ -743,7 +745,7 @@ git pull --no-rebase origin # Merge fork changes first git push origin # Push to fork (NEVER upstream) ``` -**Never stage:** auto-generated files (models.py, dd_models.py, physics_domain.py), gitignored files, `*_private.yaml` files. +**Never stage:** auto-generated files (models.py, dd_models.py, schema_context_data.py), gitignored files, `*_private.yaml` files. | Type | Purpose | |------|---------| diff --git a/hatch_build_hooks.py b/hatch_build_hooks.py index d770363fc..0cd4cfb29 100644 --- a/hatch_build_hooks.py +++ b/hatch_build_hooks.py @@ -85,50 +85,6 @@ def _generate_graph_models(self, package_root: Path) -> None: finally: sys.path[:] = original_path - def _check_physics_domain_exists(self) -> bool: - """Check if physics domain enum file exists and is up to date.""" - package_root = Path(__file__).parent - generated_file = package_root / "imas_codex" / "core" / "physics_domain.py" - schema_file = ( - package_root / "imas_codex" / "definitions" / "physics" / "domains.yaml" - ) - - if not generated_file.exists(): - return False - - if not schema_file.exists(): - return True # No schema, nothing to generate - - # Check if schema is newer than generated file - return generated_file.stat().st_mtime >= schema_file.stat().st_mtime - - def _generate_physics_domain(self, package_root: Path) -> None: - """Generate physics domain enum from LinkML schema.""" - schema_file = ( - package_root / "imas_codex" / "definitions" / "physics" / "domains.yaml" - ) - output_file = package_root / "imas_codex" / "core" / "physics_domain.py" - - if not schema_file.exists(): - self._trace(f"Schema file not found: {schema_file}") - return - - # Import generator function - original_path = sys.path[:] - if str(package_root) not in sys.path: - sys.path.insert(0, str(package_root)) - - try: - from scripts.gen_physics_domains import generate_enum_code - - code = generate_enum_code(schema_file) - output_file.write_text(code) - self._trace(f"Generated {output_file}") - except Exception as e: - self._trace(f"Failed to generate physics domain: {e}") - finally: - sys.path[:] = original_path - def _generate_schema_reference(self, package_root: Path) -> None: """Generate agents/schema-reference.md from LinkML schemas.""" output_file = package_root / "agents" / "schema-reference.md" @@ -269,14 +225,6 @@ def initialize(self, version: str, build_data: dict[str, Any]) -> None: self._trace(f"Using DD version: {resolved_dd_version}") - # Check if physics domain enum needs generation - physics_domain_exists = self._check_physics_domain_exists() - self._trace(f"physics_domain_exists={physics_domain_exists}") - - if not physics_domain_exists: - self._trace("Generating physics domain enum from LinkML schema...") - self._generate_physics_domain(package_root) - # Check if graph models need generation graph_models_exist = self._check_graph_models_exist() self._trace(f"graph_models_exist={graph_models_exist}") diff --git a/imas_codex/core/physics_domain.py b/imas_codex/core/physics_domain.py new file mode 100644 index 000000000..62a58eb7f --- /dev/null +++ b/imas_codex/core/physics_domain.py @@ -0,0 +1,10 @@ +"""Physics domain enum — canonical source is imas-standard-names. + +This module re-exports PhysicsDomain so that all imas-codex code +continues to import from the same path: + from imas_codex.core.physics_domain import PhysicsDomain +""" + +from imas_standard_names.grammar.tag_types import PhysicsDomain + +__all__ = ["PhysicsDomain"] diff --git a/imas_codex/definitions/physics/domains.yaml b/imas_codex/definitions/physics/domains.yaml deleted file mode 100644 index b666bc8f6..000000000 --- a/imas_codex/definitions/physics/domains.yaml +++ /dev/null @@ -1,345 +0,0 @@ -id: https://imas.iter.org/codex/physics-domains -name: physics_domains -title: IMAS Physics Domain Definitions -description: >- - LinkML schema defining physics domains for categorizing IMAS IDS entries. - This schema is the source of truth for the PhysicsDomain enum used throughout - the IMAS Codex system. Version and license inherited from project pyproject.toml. - -prefixes: - linkml: https://w3id.org/linkml/ - imas: https://imas.iter.org/codex/ - -imports: - - linkml:types - -default_range: string - -types: - DomainIdentifier: - typeof: string - description: A valid physics domain identifier - -slots: - domain_name: - range: PhysicsDomain - description: The physics domain identifier - - category: - range: DomainCategory - description: High-level category grouping related domains - - characteristics: - range: string - multivalued: true - description: Key characteristics and phenomena associated with this domain - - related_domains: - range: PhysicsDomain - multivalued: true - description: Other physics domains that frequently interact with this domain - - domain_description: - range: string - description: Human-readable description of the domain - -classes: - PhysicsDomainDefinition: - description: Complete definition of a physics domain with its relationships and characteristics - slots: - - domain_name - - domain_description - - category - - characteristics - - related_domains - slot_usage: - domain_name: - required: true - category: - required: true - -enums: - DomainCategory: - description: >- - High-level categories for grouping physics domains. Categories provide - a coarse classification for filtering and organizing domains. - permissible_values: - core_plasma_physics: - description: Fundamental plasma physics phenomena including equilibrium, transport, and instabilities - meaning: imas:core_plasma_physics - heating_and_current_drive: - description: Auxiliary heating systems and non-inductive current drive methods - meaning: imas:heating_and_current_drive - plasma_material_interactions: - description: Physics at the plasma boundary including wall, divertor, and edge phenomena - meaning: imas:plasma_material_interactions - diagnostics: - description: Measurement and analysis systems for plasma and machine parameters - meaning: imas:diagnostics - control_and_operations: - description: Plasma control, feedback systems, and operational parameters - meaning: imas:control_and_operations - engineering_systems: - description: Machine components, structural elements, and plant systems - meaning: imas:engineering_systems - data_and_workflow: - description: Data organization, metadata management, and computational workflows - meaning: imas:data_and_workflow - uncategorized: - description: General purpose or uncategorized data structures - meaning: imas:uncategorized - - PhysicsDomain: - description: >- - Physics domains for categorizing IMAS Interface Data Structures (IDS). - Each domain represents a distinct area of fusion plasma physics or - tokamak engineering. - permissible_values: - # Core Plasma Physics - equilibrium: - description: Magnetohydrodynamic equilibrium and magnetic field configuration - meaning: imas:equilibrium - annotations: - category: core_plasma_physics - characteristics: >- - Magnetic flux surfaces and geometry, - Pressure and current density profiles, - Shafranov shift and elongation - related_domains: magnetohydrodynamics, transport, magnetic_field_systems - - transport: - description: Particle, energy, and momentum transport processes - meaning: imas:transport - annotations: - category: core_plasma_physics - characteristics: >- - Diffusion coefficients, - Heat and particle fluxes, - Confinement time scaling - related_domains: turbulence, equilibrium, auxiliary_heating - - magnetohydrodynamics: - description: Magnetohydrodynamic instabilities and plasma modes - meaning: imas:magnetohydrodynamics - annotations: - category: core_plasma_physics - characteristics: >- - Tearing modes and islands, - Sawteeth and edge localized modes, - Resistive wall modes - related_domains: equilibrium, plasma_control, magnetic_field_diagnostics - - turbulence: - description: Microscopic turbulence and anomalous transport phenomena - meaning: imas:turbulence - annotations: - category: core_plasma_physics - characteristics: >- - Ion temperature gradient and trapped electron modes, - Zonal flows and geodesic acoustic modes, - Fluctuation measurements - related_domains: transport, electromagnetic_wave_diagnostics - - # Heating and Current Drive - auxiliary_heating: - description: Auxiliary heating systems including neutral beam injection and radiofrequency heating - meaning: imas:auxiliary_heating - annotations: - category: heating_and_current_drive - characteristics: >- - Power deposition profiles, - Heating efficiency, - Fast particle generation - related_domains: current_drive, transport, particle_measurement_diagnostics - - current_drive: - description: Non-inductive current drive methods - meaning: imas:current_drive - annotations: - category: heating_and_current_drive - characteristics: >- - Driven current profiles, - Current drive efficiency, - Bootstrap current - related_domains: auxiliary_heating, equilibrium, plasma_control - - # Plasma-Material Interactions - plasma_wall_interactions: - description: Plasma-wall interactions and first wall components - meaning: imas:plasma_wall_interactions - annotations: - category: plasma_material_interactions - characteristics: >- - Heat loads and erosion, - Material migration, - Recycling and retention - related_domains: divertor_physics, edge_plasma_physics, radiation_measurement_diagnostics - - divertor_physics: - description: Divertor physics and power exhaust mechanisms - meaning: imas:divertor_physics - annotations: - category: plasma_material_interactions - characteristics: >- - Target heat flux, - Detachment and radiation, - Neutral dynamics - related_domains: plasma_wall_interactions, edge_plasma_physics, particle_measurement_diagnostics - - edge_plasma_physics: - description: Edge plasma and scrape-off layer physics - meaning: imas:edge_plasma_physics - annotations: - category: plasma_material_interactions - characteristics: >- - Scrape-off layer width and decay lengths, - Pedestal structure, - Edge localized mode dynamics - related_domains: divertor_physics, plasma_wall_interactions, magnetohydrodynamics - - # Diagnostics - particle_measurement_diagnostics: - description: Particle measurement and analysis diagnostic systems - meaning: imas:particle_measurement_diagnostics - annotations: - category: diagnostics - characteristics: >- - Neutral particle analyzers, - Mass spectrometry, - Thomson scattering particle measurements - related_domains: transport, auxiliary_heating, edge_plasma_physics - - electromagnetic_wave_diagnostics: - description: Electromagnetic wave and field diagnostic systems - meaning: imas:electromagnetic_wave_diagnostics - annotations: - category: diagnostics - characteristics: >- - Reflectometry, - Electron cyclotron emission and microwave diagnostics, - Interferometry - related_domains: turbulence, equilibrium, magnetohydrodynamics - - radiation_measurement_diagnostics: - description: Radiation-based diagnostic systems - meaning: imas:radiation_measurement_diagnostics - annotations: - category: diagnostics - characteristics: >- - Bolometry and radiometry, - X-ray diagnostics, - Spectroscopy - related_domains: transport, plasma_wall_interactions, edge_plasma_physics - - magnetic_field_diagnostics: - description: Magnetic field measurement diagnostic systems - meaning: imas:magnetic_field_diagnostics - annotations: - category: diagnostics - characteristics: >- - Magnetic probes and flux loops, - Rogowski coils, - Motional Stark effect - related_domains: equilibrium, magnetohydrodynamics, plasma_control - - mechanical_measurement_diagnostics: - description: Mechanical and pressure measurement diagnostic systems - meaning: imas:mechanical_measurement_diagnostics - annotations: - category: diagnostics - characteristics: >- - Pressure gauges, - Strain and force sensors, - Vibration monitoring - related_domains: structural_components, plasma_wall_interactions, machine_operations - - # Control and Operation - plasma_control: - description: Plasma control and feedback systems - meaning: imas:plasma_control - annotations: - category: control_and_operations - characteristics: >- - Shape and position control, - Instability suppression, - Scenario development - related_domains: equilibrium, magnetohydrodynamics, magnetic_field_systems - - machine_operations: - description: Operational parameters and machine status monitoring - meaning: imas:machine_operations - annotations: - category: control_and_operations - characteristics: >- - Pulse scheduling, - Interlocks and limits, - Machine state - related_domains: plasma_control, plant_systems, data_management - - # System Components - magnetic_field_systems: - description: Magnetic field coil systems and field generation equipment - meaning: imas:magnetic_field_systems - annotations: - category: engineering_systems - characteristics: >- - Poloidal and toroidal field coils, - Superconducting magnets, - Power supplies - related_domains: equilibrium, plasma_control, structural_components - - structural_components: - description: Structural components and mechanical systems - meaning: imas:structural_components - annotations: - category: engineering_systems - characteristics: >- - Vacuum vessel, - Support structures, - Thermal shields - related_domains: magnetic_field_systems, plasma_wall_interactions, mechanical_measurement_diagnostics - - plant_systems: - description: Engineering plant systems and auxiliary components - meaning: imas:plant_systems - annotations: - category: engineering_systems - characteristics: >- - Cryogenics, - Vacuum systems, - Cooling systems - related_domains: structural_components, machine_operations, plasma_control - - # Data and Workflow - data_management: - description: Data organization, metadata, and information management - meaning: imas:data_management - annotations: - category: data_and_workflow - characteristics: >- - Pulse databases, - Data provenance, - Signal definitions - related_domains: computational_workflow, machine_operations - - computational_workflow: - description: Computational workflows and process management - meaning: imas:computational_workflow - annotations: - category: data_and_workflow - characteristics: >- - Simulation pipelines, - Analysis chains, - Reproducibility - related_domains: data_management - - # Fallback - general: - description: General purpose or uncategorized data structures - meaning: imas:general - annotations: - category: uncategorized - characteristics: >- - Common utilities, - Generic structures - related_domains: "" diff --git a/imas_codex/schemas/common.yaml b/imas_codex/schemas/common.yaml index af6e1269a..52792831f 100644 --- a/imas_codex/schemas/common.yaml +++ b/imas_codex/schemas/common.yaml @@ -25,7 +25,6 @@ default_range: string imports: - linkml:types - - physics_domains # ============================================================================= # Status Enums - Unified Lifecycle Terminology @@ -615,7 +614,7 @@ classes: multivalued: true physics_domain: description: Primary physics domain inferred from content - range: PhysicsDomain + range: string score_cost: description: LLM/VLM cost in USD for scoring (batch cost / batch size) range: float @@ -709,10 +708,6 @@ classes: Higher when validated against physics quantities. range: float - # NOTE: PhysicsDomain is now an ENUM imported from physics_domains.yaml - # The class was removed to avoid conflict with the enum. Use the enum - # for categorizing FacilitySignal, IMASPath, DataNode, etc. - Unit: description: >- A physical unit used in both IMAS Data Dictionary and facility-specific diff --git a/imas_codex/schemas/facility.yaml b/imas_codex/schemas/facility.yaml index ec5722cef..bc8f88aa6 100644 --- a/imas_codex/schemas/facility.yaml +++ b/imas_codex/schemas/facility.yaml @@ -1851,7 +1851,7 @@ classes: multivalued: true physics_domain: description: Physics domain classification - range: PhysicsDomain + range: string status: description: Lifecycle status range: SignalSourceStatus @@ -2174,7 +2174,7 @@ classes: relationship_type: AT_FACILITY physics_domain: description: Physics domain for categorizing this signal - range: PhysicsDomain + range: string annotations: required_after: enriched name: diff --git a/imas_codex/schemas/imas_dd.yaml b/imas_codex/schemas/imas_dd.yaml index f07da1393..50e8b6fa5 100644 --- a/imas_codex/schemas/imas_dd.yaml +++ b/imas_codex/schemas/imas_dd.yaml @@ -446,7 +446,7 @@ classes: description: IDS description from DD documentation (raw XML text) physics_domain: description: Primary physics domain - range: PhysicsDomain + range: string path_count: description: Number of paths in this IDS (current version) range: integer @@ -588,7 +588,7 @@ classes: range: DDNodeType physics_domain: description: Physics domain (derived from path/IDS) - range: PhysicsDomain + range: string maxoccur: description: Maximum occurrences (for struct arrays) range: integer @@ -874,7 +874,7 @@ classes: description: Extended description of the cluster concept physics_domain: description: Primary physics domain of cluster members - range: PhysicsDomain + range: string path_count: description: Number of paths in this cluster range: integer diff --git a/imas_codex/schemas/physics_domains.yaml b/imas_codex/schemas/physics_domains.yaml deleted file mode 100644 index 01c672f03..000000000 --- a/imas_codex/schemas/physics_domains.yaml +++ /dev/null @@ -1,379 +0,0 @@ -id: https://imas.iter.org/codex/physics-domains -name: physics_domains -title: IMAS Physics Domain Definitions -description: >- - LinkML schema defining physics domains for categorizing IMAS IDS entries. - This schema is the source of truth for the PhysicsDomain enum used throughout - the IMAS Codex system. Version and license inherited from project pyproject.toml. - -prefixes: - linkml: https://w3id.org/linkml/ - imas: https://imas.iter.org/codex/ - -imports: - - linkml:types - -default_range: string - -types: - DomainIdentifier: - typeof: string - description: A valid physics domain identifier - -slots: - domain_name: - range: PhysicsDomain - description: The physics domain identifier - - category: - range: DomainCategory - description: High-level category grouping related domains - - characteristics: - range: string - multivalued: true - description: Key characteristics and phenomena associated with this domain - - related_domains: - range: PhysicsDomain - multivalued: true - description: Other physics domains that frequently interact with this domain - - domain_description: - range: string - description: Human-readable description of the domain - -classes: - PhysicsDomainDefinition: - description: Complete definition of a physics domain with its relationships and characteristics - slots: - - domain_name - - domain_description - - category - - characteristics - - related_domains - slot_usage: - domain_name: - required: true - category: - required: true - -enums: - DomainCategory: - description: >- - High-level categories for grouping physics domains. Categories provide - a coarse classification for filtering and organizing domains. - permissible_values: - core_plasma_physics: - description: Fundamental plasma physics phenomena including equilibrium, transport, and instabilities - meaning: imas:core_plasma_physics - heating_and_current_drive: - description: Auxiliary heating systems and non-inductive current drive methods - meaning: imas:heating_and_current_drive - plasma_material_interactions: - description: Physics at the plasma boundary including wall, divertor, and edge phenomena - meaning: imas:plasma_material_interactions - diagnostics: - description: Measurement and analysis systems for plasma and machine parameters - meaning: imas:diagnostics - control_and_operations: - description: Plasma control, feedback systems, and operational parameters - meaning: imas:control_and_operations - engineering_systems: - description: Machine components, structural elements, and plant systems - meaning: imas:engineering_systems - data_and_workflow: - description: Data organization, metadata management, and computational workflows - meaning: imas:data_and_workflow - uncategorized: - description: General purpose or uncategorized data structures - meaning: imas:uncategorized - - PhysicsDomain: - description: >- - Physics domains for categorizing IMAS Interface Data Structures (IDS). - Each domain represents a distinct area of fusion plasma physics or - tokamak engineering. - permissible_values: - # Core Plasma Physics - equilibrium: - description: Magnetohydrodynamic equilibrium and magnetic field configuration - meaning: imas:equilibrium - annotations: - category: core_plasma_physics - characteristics: >- - Magnetic flux surfaces and geometry, - Pressure and current density profiles, - Shafranov shift and elongation - related_domains: magnetohydrodynamics, transport, magnetic_field_systems - - transport: - description: Particle, energy, and momentum transport processes - meaning: imas:transport - annotations: - category: core_plasma_physics - characteristics: >- - Diffusion coefficients, - Heat and particle fluxes, - Confinement time scaling - related_domains: turbulence, equilibrium, auxiliary_heating - - magnetohydrodynamics: - description: Magnetohydrodynamic instabilities and plasma modes - meaning: imas:magnetohydrodynamics - annotations: - category: core_plasma_physics - characteristics: >- - Tearing modes and islands, - Sawteeth and edge localized modes, - Resistive wall modes - related_domains: equilibrium, plasma_control, magnetic_field_diagnostics - - turbulence: - description: Microscopic turbulence and anomalous transport phenomena - meaning: imas:turbulence - annotations: - category: core_plasma_physics - characteristics: >- - Ion temperature gradient and trapped electron modes, - Zonal flows and geodesic acoustic modes, - Fluctuation measurements - related_domains: transport, electromagnetic_wave_diagnostics - - # Heating and Current Drive - auxiliary_heating: - description: Auxiliary heating systems including neutral beam injection and radiofrequency heating - meaning: imas:auxiliary_heating - annotations: - category: heating_and_current_drive - characteristics: >- - Power deposition profiles, - Heating efficiency, - Fast particle generation - related_domains: current_drive, transport, particle_measurement_diagnostics - - current_drive: - description: Non-inductive current drive methods - meaning: imas:current_drive - annotations: - category: heating_and_current_drive - characteristics: >- - Driven current profiles, - Current drive efficiency, - Bootstrap current - related_domains: auxiliary_heating, equilibrium, plasma_control - - # Plasma-Material Interactions - plasma_wall_interactions: - description: Plasma-wall interactions and first wall components - meaning: imas:plasma_wall_interactions - annotations: - category: plasma_material_interactions - characteristics: >- - Heat loads and erosion, - Material migration, - Recycling and retention - related_domains: divertor_physics, edge_plasma_physics, radiation_measurement_diagnostics - - divertor_physics: - description: Divertor physics and power exhaust mechanisms - meaning: imas:divertor_physics - annotations: - category: plasma_material_interactions - characteristics: >- - Target heat flux, - Detachment and radiation, - Neutral dynamics - related_domains: plasma_wall_interactions, edge_plasma_physics, particle_measurement_diagnostics - - edge_plasma_physics: - description: Edge plasma and scrape-off layer physics - meaning: imas:edge_plasma_physics - annotations: - category: plasma_material_interactions - characteristics: >- - Scrape-off layer width and decay lengths, - Pedestal structure, - Edge localized mode dynamics - related_domains: divertor_physics, plasma_wall_interactions, magnetohydrodynamics - - # Diagnostics - particle_measurement_diagnostics: - description: Particle measurement and analysis diagnostic systems - meaning: imas:particle_measurement_diagnostics - annotations: - category: diagnostics - characteristics: >- - Neutral particle analyzers, - Mass spectrometry, - Thomson scattering particle measurements - related_domains: transport, auxiliary_heating, edge_plasma_physics - - plasma_measurement_diagnostics: - description: Plasma measurement diagnostic systems covering multiple measurement techniques - meaning: imas:plasma_measurement_diagnostics - annotations: - category: diagnostics - characteristics: >- - Combined plasma diagnostics, - Multi-technique measurement systems, - Integrated plasma monitoring - related_domains: particle_measurement_diagnostics, electromagnetic_wave_diagnostics, radiation_measurement_diagnostics - - electromagnetic_wave_diagnostics: - description: Electromagnetic wave and field diagnostic systems - meaning: imas:electromagnetic_wave_diagnostics - annotations: - category: diagnostics - characteristics: >- - Reflectometry, - Electron cyclotron emission and microwave diagnostics, - Interferometry - related_domains: turbulence, equilibrium, magnetohydrodynamics - - radiation_measurement_diagnostics: - description: Radiation-based diagnostic systems - meaning: imas:radiation_measurement_diagnostics - annotations: - category: diagnostics - characteristics: >- - Bolometry and radiometry, - X-ray diagnostics, - Spectroscopy - related_domains: transport, plasma_wall_interactions, edge_plasma_physics - - magnetic_field_diagnostics: - description: Magnetic field measurement diagnostic systems - meaning: imas:magnetic_field_diagnostics - annotations: - category: diagnostics - characteristics: >- - Magnetic probes and flux loops, - Rogowski coils, - Motional Stark effect - related_domains: equilibrium, magnetohydrodynamics, plasma_control - - mechanical_measurement_diagnostics: - description: Mechanical and pressure measurement diagnostic systems - meaning: imas:mechanical_measurement_diagnostics - annotations: - category: diagnostics - characteristics: >- - Pressure gauges, - Strain and force sensors, - Vibration monitoring - related_domains: structural_components, plasma_wall_interactions, machine_operations - - # Control and Operation - plasma_control: - description: Plasma control and feedback systems - meaning: imas:plasma_control - annotations: - category: control_and_operations - characteristics: >- - Shape and position control, - Instability suppression, - Scenario development - related_domains: equilibrium, magnetohydrodynamics, magnetic_field_systems - - machine_operations: - description: Operational parameters and machine status monitoring - meaning: imas:machine_operations - annotations: - category: control_and_operations - characteristics: >- - Pulse scheduling, - Interlocks and limits, - Machine state - related_domains: plasma_control, plant_systems, data_management - - # System Components - magnetic_field_systems: - description: Magnetic field coil systems and field generation equipment - meaning: imas:magnetic_field_systems - annotations: - category: engineering_systems - characteristics: >- - Poloidal and toroidal field coils, - Superconducting magnets, - Power supplies - related_domains: equilibrium, plasma_control, structural_components - - structural_components: - description: Structural components and mechanical systems - meaning: imas:structural_components - annotations: - category: engineering_systems - characteristics: >- - Vacuum vessel, - Support structures, - Thermal shields - related_domains: magnetic_field_systems, plasma_wall_interactions, mechanical_measurement_diagnostics - - plant_systems: - description: Engineering plant systems and auxiliary components - meaning: imas:plant_systems - annotations: - category: engineering_systems - characteristics: >- - Cryogenics, - Vacuum systems, - Cooling systems - related_domains: structural_components, machine_operations, plasma_control - - # Data and Workflow - data_management: - description: Data organization, metadata, and information management - meaning: imas:data_management - annotations: - category: data_and_workflow - characteristics: >- - Pulse databases, - Data provenance, - Signal definitions - related_domains: computational_workflow, machine_operations - - computational_workflow: - description: Computational workflows and process management - meaning: imas:computational_workflow - annotations: - category: data_and_workflow - characteristics: >- - Simulation pipelines, - Analysis chains, - Reproducibility - related_domains: data_management - - # IDS-level Domains (matching IMAS IDS names directly) - magnetics: - description: Magnetic measurement systems and data from the magnetics IDS - meaning: imas:magnetics - annotations: - category: diagnostics - characteristics: >- - Magnetic probes and flux loops, - Plasma current measurements, - Equilibrium reconstruction inputs - related_domains: magnetic_field_diagnostics, equilibrium, plasma_control - - gyrokinetics: - description: Gyrokinetic simulation data including wavevectors and eigenmodes - meaning: imas:gyrokinetics - annotations: - category: core_plasma_physics - characteristics: >- - Linear and nonlinear gyrokinetic modes, - Wavevector spectra, - Growth rates and frequencies - related_domains: turbulence, transport - - # Fallback - general: - description: General purpose or uncategorized data structures - meaning: imas:general - annotations: - category: uncategorized - characteristics: >- - Common utilities, - Generic structures - related_domains: "" diff --git a/pyproject.toml b/pyproject.toml index ddc2ebc31..4c502ea2b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -154,7 +154,7 @@ dev = [ # --- Serve (embedding + LLM proxy) --- "fastapi>=0.115.0", "uvicorn>=0.31.1", - "imas-standard-names", + "imas-standard-names>=0.7.0rc2", ] [project.urls] @@ -431,9 +431,10 @@ torch = [ { index = "pytorch-cpu", extra = "test" }, { index = "pytorch-gpu", extra = "gpu" }, ] -imas-standard-names = { path = "../imas-standard-names", editable = true } + [tool.uv] +prerelease = "if-necessary-or-explicit" # cpu, gpu, and test extras for torch are mutually exclusive with gpu conflicts = [ [ diff --git a/scripts/build_models.py b/scripts/build_models.py index 51fe87b96..c80397333 100644 --- a/scripts/build_models.py +++ b/scripts/build_models.py @@ -31,67 +31,11 @@ def get_graph_dir() -> Path: return get_project_root() / "imas_codex" / "graph" -def get_core_dir() -> Path: - """Get the core module directory.""" - return get_project_root() / "imas_codex" / "core" - - -def get_definitions_dir() -> Path: - """Get the definitions directory.""" - return get_project_root() / "imas_codex" / "definitions" - - def get_config_module_dir() -> Path: """Get the config module directory for generated models.""" return get_project_root() / "imas_codex" / "config" -def _generate_physics_domain( - logger: logging.Logger, - force: bool, - dry_run: bool, -) -> int: - """Generate physics domain enum from LinkML schema. - - Returns: - 0 on success, non-zero on failure. - """ - from scripts.gen_physics_domains import generate_enum_code - - definitions_dir = get_definitions_dir() - core_dir = get_core_dir() - - schema_file = definitions_dir / "physics" / "domains.yaml" - output_file = core_dir / "physics_domain.py" - - if not schema_file.exists(): - logger.error(f"Physics domain schema not found: {schema_file}") - return 1 - - # Check if output already exists and is up to date - if output_file.exists() and not force: - if schema_file.stat().st_mtime <= output_file.stat().st_mtime: - logger.info(f"Physics domain up to date at {output_file}") - return 0 - logger.info("Physics schema newer than enum, regenerating...") - - logger.info(f"Generating physics domain enum from {schema_file}") - - if dry_run: - click.echo(f"Would generate: {output_file}") - return 0 - - try: - code = generate_enum_code(schema_file) - output_file.write_text(code, encoding="utf-8") - logger.info(f"Generated physics domain written to {output_file}") - click.echo(f"Generated: {output_file}") - return 0 - except Exception as e: - logger.error(f"Failed to generate physics domain: {e}") - return 1 - - def _generate_schema_reference( logger: logging.Logger, force: bool, @@ -190,9 +134,8 @@ def build_models( """Generate Pydantic models from LinkML schemas. This command generates: - 1. Physics domain enum from definitions/physics/domains.yaml - 2. Graph Pydantic models from schemas/facility.yaml - 3. IMAS DD models from schemas/imas_dd.yaml + 1. Graph Pydantic models from schemas/facility.yaml + 2. IMAS DD models from schemas/imas_dd.yaml Examples: build-models # Generate all models @@ -214,11 +157,6 @@ def build_models( logger = logging.getLogger(__name__) try: - # Generate physics domain enum first (required by other modules) - physics_result = _generate_physics_domain(logger, force, dry_run) - if physics_result != 0: - return physics_result - # Generate graph models schemas_dir = get_schemas_dir() graph_dir = get_graph_dir() diff --git a/scripts/gen_physics_domains.py b/scripts/gen_physics_domains.py deleted file mode 100644 index 586c35d82..000000000 --- a/scripts/gen_physics_domains.py +++ /dev/null @@ -1,148 +0,0 @@ -#!/usr/bin/env python3 -""" -Generate PhysicsDomain enum from LinkML schema. - -This script reads the LinkML domains.yaml schema and generates a Python -enum compatible with Pydantic (str, Enum). The generated enum replaces -the manually maintained PhysicsDomain in data_model.py. -""" - -import sys -from datetime import UTC, datetime -from pathlib import Path - -import click -from linkml_runtime.utils.schemaview import SchemaView - - -def generate_enum_code(schema_path: Path) -> str: - """Generate Python enum code from LinkML schema. - - Args: - schema_path: Path to the LinkML YAML schema file. - - Returns: - Python source code for the enum. - """ - sv = SchemaView(str(schema_path)) - schema = sv.schema - - # Get the PhysicsDomain enum - enum_def = sv.get_enum("PhysicsDomain") - if not enum_def: - raise ValueError("PhysicsDomain enum not found in schema") - - # Build enum members - members = [] - for pv_name, pv in enum_def.permissible_values.items(): - # Create enum member name (uppercase with underscores) - member_name = pv_name.upper() - - # Add member with description as comment - description = pv.description or "" - members.append(f' {member_name} = "{pv_name}" # {description}') - - # Generate the code - code = f'''""" -Physics domain enum generated from LinkML schema. - -DO NOT EDIT THIS FILE DIRECTLY. -Edit imas_codex/definitions/physics/domains.yaml and run: - uv run gen-physics-domains - -Generated: {datetime.now(UTC).isoformat()} -Schema: {schema.name} -""" - -from enum import Enum - - -class PhysicsDomain(str, Enum): - """Physics domains for categorizing IMAS Interface Data Structures (IDS). - - Each domain represents a distinct area of fusion plasma physics or - tokamak engineering. This enum is generated from the LinkML schema - at definitions/physics/domains.yaml. - """ - -{chr(10).join(members)} -''' - - return code - - -@click.command() -@click.option( - "--schema", - type=click.Path(exists=True, path_type=Path), - default=Path(__file__).parent.parent - / "imas_codex/definitions/physics/domains.yaml", - help="Path to LinkML schema file", -) -@click.option( - "--output", - type=click.Path(path_type=Path), - default=None, - help="Output file path (default: stdout)", -) -@click.option( - "--check", - is_flag=True, - help="Check if generated code matches existing file", -) -def gen_physics_domains(schema: Path, output: Path | None, check: bool) -> int: - """Generate PhysicsDomain enum from LinkML schema. - - This command reads the domains.yaml LinkML schema and generates a - Python enum that can be used with Pydantic models. - - Examples: - gen-physics-domains # Print to stdout - gen-physics-domains --output src/enum.py # Write to file - gen-physics-domains --check # Verify file is up to date - """ - try: - code = generate_enum_code(schema) - - if check: - if not output: - click.echo("Error: --check requires --output", err=True) - return 1 - if not output.exists(): - click.echo(f"Error: {output} does not exist", err=True) - return 1 - existing = output.read_text() - # Compare ignoring the timestamp line - existing_lines = [ - line - for line in existing.splitlines() - if not line.startswith("Generated:") - ] - new_lines = [ - line for line in code.splitlines() if not line.startswith("Generated:") - ] - if existing_lines != new_lines: - click.echo( - f"Error: {output} is out of date. Run gen-physics-domains to update.", - err=True, - ) - return 1 - click.echo(f"OK: {output} is up to date") - return 0 - - if output: - output.parent.mkdir(parents=True, exist_ok=True) - output.write_text(code) - click.echo(f"Generated {output}") - else: - click.echo(code) - - return 0 - - except Exception as e: - click.echo(f"Error: {e}", err=True) - return 1 - - -if __name__ == "__main__": - sys.exit(gen_physics_domains()) diff --git a/uv.lock b/uv.lock index 7f470523a..ce808559c 100644 --- a/uv.lock +++ b/uv.lock @@ -1399,7 +1399,7 @@ dev = [ { name = "fastapi", specifier = ">=0.115.0" }, { name = "hdbscan", specifier = ">=0.8.41" }, { name = "imas-python", specifier = ">=2.0.1" }, - { name = "imas-standard-names", editable = "../imas-standard-names" }, + { name = "imas-standard-names", specifier = ">=0.7.0rc2" }, { name = "ipykernel", specifier = ">=6.29.5" }, { name = "ipython", specifier = ">=9.2.0" }, { name = "jellyfish", specifier = ">=1.2.1" }, @@ -1478,12 +1478,14 @@ wheels = [ [[package]] name = "imas-standard-names" -source = { editable = "../imas-standard-names" } +version = "0.7.0rc2" +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, { name = "dotenv" }, { name = "fastmcp" }, { name = "markdown" }, + { name = "nest-asyncio" }, { name = "pint" }, { name = "pydantic" }, { name = "pydantic-ai" }, @@ -1492,54 +1494,9 @@ dependencies = [ { name = "strictyaml" }, { name = "textual" }, ] - -[package.metadata] -requires-dist = [ - { name = "click", specifier = ">=8.1.8,<9.0.0" }, - { name = "dotenv", specifier = ">=0.9.9,<0.10.0" }, - { name = "en-core-web-sm", marker = "extra == 'quality'", url = "https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl" }, - { name = "fastmcp", specifier = "~=3.2.0" }, - { name = "markdown", specifier = ">=3.8,<4.0" }, - { name = "mike", marker = "extra == 'docs'", specifier = ">=2.1.3,<3.0.0" }, - { name = "mkdocs", marker = "extra == 'docs'", specifier = ">=1.6.1,<2.0.0" }, - { name = "mkdocs-data-plugin", marker = "extra == 'docs'", specifier = ">=0.2.0,<0.3.0" }, - { name = "mkdocs-include-markdown-plugin", marker = "extra == 'docs'", specifier = ">=7.0.0,<8.0.0" }, - { name = "mkdocs-macros-plugin", marker = "extra == 'docs'", specifier = ">=1.0.4,<2.0.0" }, - { name = "mkdocs-material", marker = "extra == 'docs'", specifier = ">=9.6.5,<10.0.0" }, - { name = "mkdocs-table-reader-plugin", marker = "extra == 'docs'", specifier = ">=3.1.0,<4.0.0" }, - { name = "pint", specifier = ">=0.24.4,<0.25.0" }, - { name = "proselint", marker = "extra == 'quality'", specifier = ">=0.14.0,<0.15.0" }, - { name = "pydantic", specifier = ">=2.10.6,<3.0.0" }, - { name = "pydantic-ai", specifier = ">=1.56.0" }, - { name = "pytest", marker = "extra == 'test'", specifier = ">=8.3.4,<9.0.0" }, - { name = "pytest-cov", marker = "extra == 'test'", specifier = ">=4.1.0,<5.0.0" }, - { name = "pytest-html", marker = "extra == 'test'", specifier = ">=4.1.1,<5.0.0" }, - { name = "pyyaml", specifier = ">=6.0.2,<7.0.0" }, - { name = "requests", specifier = ">=2.33.0,<3.0.0" }, - { name = "ruff", marker = "extra == 'test'", specifier = ">=0.9.8,<1.0.0" }, - { name = "spacy", marker = "extra == 'quality'", specifier = ">=3.8.0,<4.0.0" }, - { name = "strictyaml", specifier = ">=1.7.3,<2.0.0" }, - { name = "textual", specifier = ">=6.1.0" }, -] -provides-extras = ["docs", "quality", "test"] - -[package.metadata.requires-dev] -dev = [ - { name = "en-core-web-sm", url = "https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl" }, - { name = "ipykernel", specifier = ">=6.29.5,<7.0.0" }, - { name = "logfire", specifier = ">=4.16.0,<5.0.0" }, - { name = "mcp-cli", specifier = ">=0.1.0,<1.0.0" }, - { name = "pandas", specifier = ">=2.2.3,<3.0.0" }, - { name = "pandas-stubs", specifier = ">=2.2.3.250308,<3.0.0" }, - { name = "pre-commit", specifier = ">=4.1.0,<5.0.0" }, - { name = "proselint", specifier = ">=0.14.0,<0.15.0" }, - { name = "pytest", specifier = ">=8.3.4,<9.0.0" }, - { name = "pytest-cov", specifier = ">=4.1.0,<5.0.0" }, - { name = "pytest-html", specifier = ">=4.1.1,<5.0.0" }, - { name = "ruff", specifier = ">=0.11.10,<1.0.0" }, - { name = "spacy", specifier = ">=3.8.0,<4.0.0" }, - { name = "textual-dev", specifier = ">=1.7.0" }, - { name = "types-pyyaml", specifier = ">=6.0.12.20241230,<7.0.0" }, +sdist = { url = "https://files.pythonhosted.org/packages/26/d8/adcf860a3586aa213cc0d9e6ed4c73392920a940b91de1fa9bc452da14fd/imas_standard_names-0.7.0rc2.tar.gz", hash = "sha256:828aefe86bcd1822bae309529b5ab43ad3b1132da090ef8a9cb1f8023c38f13a", size = 657595, upload-time = "2026-04-10T06:37:42.5Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/64/6c/f6e7cf8abd22ce908f37964997d76e36e1f8ffeaf40c8f94f5d4d07cf5fe/imas_standard_names-0.7.0rc2-py3-none-any.whl", hash = "sha256:453ae5f96c552578ef449ad632c768b56b44af30b75be7f7989fd19b2dac8027", size = 310171, upload-time = "2026-04-10T06:37:40.828Z" }, ] [[package]] From b61b36171ed5085291ec9e5ba038d4d0d7ac0ae5 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 09:08:43 +0200 Subject: [PATCH 06/37] fix: remove dead IMAS_DD_VERSION from container build workflow MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The IMAS_DD_VERSION env var and build-arg were passed to docker build but the Dockerfile never declared a matching ARG — it was dead code. The container gets its DD version from the graph data loaded from GHCR, not from a build argument. The DD version single source of truth is pyproject.toml under [tool.imas-codex.data-dictionary].version, read by get_dd_version() at runtime. The test workflow correctly uses IMAS_DD_VERSION as an env var override for multi-version matrix testing. --- .github/workflows/docker-build-push.yml | 3 --- 1 file changed, 3 deletions(-) diff --git a/.github/workflows/docker-build-push.yml b/.github/workflows/docker-build-push.yml index 68f4be742..5dfbc77e6 100644 --- a/.github/workflows/docker-build-push.yml +++ b/.github/workflows/docker-build-push.yml @@ -18,7 +18,6 @@ env: # ACR image name is always the upstream path so Azure picks up fork RC # builds. GHCR uses IMAGE_NAME (per-fork) — each fork has its own GHCR. ACR_IMAGE_NAME: iterorganization/imas-codex - IMAS_DD_VERSION: "4.1.0" # Graph OCI artifacts: try the repo owner's registry first (where # `imas-codex release` pushes them), fall back to the upstream org # registry for forks that haven't pushed their own graph yet. @@ -559,7 +558,6 @@ jobs: ${{ steps.meta-ghcr-rc.outputs.labels }} build-args: | IDS_FILTER= - IMAS_DD_VERSION=${{ env.IMAS_DD_VERSION }} GHCR_REGISTRY=${{ needs.graph-quality.outputs.graph-registry }} GRAPH_TAG=${{ steps.variant.outputs.graph-tag }} GRAPH_PACKAGE=${{ matrix.graph_variant.package }} @@ -587,7 +585,6 @@ jobs: ${{ steps.meta-acr.outputs.labels }} build-args: | IDS_FILTER= - IMAS_DD_VERSION=${{ env.IMAS_DD_VERSION }} GHCR_REGISTRY=${{ needs.graph-quality.outputs.graph-registry }} GRAPH_TAG=${{ steps.variant.outputs.graph-tag }} GRAPH_PACKAGE=${{ matrix.graph_variant.package }} From 32bc0ca85ace45b79518e67fc64acbdde3704de1 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 09:17:37 +0200 Subject: [PATCH 07/37] docs: update AGENTS.md with fork/main workflow MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace multi-branch workflow with fork-based main-only workflow. All development happens on fork's main branch — no feature branches. Document ACR deployment path and Azure test URL. Add rule against pushing same tag to both origin and upstream. --- AGENTS.md | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index b23ac5b43..b6870446a 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -2,10 +2,10 @@ Use terminal for direct operations (`rg`, `fd`, `git`), MCP `repl()` for chained processing and graph queries, `uv run` for git/tests/CLI. Conventional commits. **CRITICAL: Always commit and push when files have been modified — no confirmation, no asking, just do it. This is non-negotiable. Every response that modifies files MUST end with `git add`, `git commit`, and `git push`.** **Never use `vscode_askQuestions` or any interactive VS Code popup/dialog tools — present all questions inline in the chat response so the user can answer them in one message.** -**Git sync discipline (multi-instance workflow):** This repo is edited from multiple machines and by multiple agents concurrently. Always **merge** on pull — never rebase. -1. **Session start:** `git pull origin` before any work (pulls current branch from fork). -2. **Before push:** `git pull origin && git push origin` — never push without pulling first. Push to `origin` (fork), **never directly to `upstream`**. -3. **Stay on current branch:** Push to whatever branch you're on. If the branch is `develop`, push to `origin develop`. If `main`, push to `origin main`. **Never merge branches or switch to `main` without explicit user approval.** +**Git sync discipline (fork-based workflow):** All development happens on the fork's `main` branch. Always **merge** on pull — never rebase. Never use feature branches (`develop`, `feature/*`) — they add merge overhead and break the release CLI which requires `main`. +1. **Session start:** `git pull origin main` before any work. +2. **Before push:** `git pull origin main && git push origin main` — never push without pulling first. Push to `origin` (fork), **never directly to `upstream`**. +3. **Always work on `main`** — the release CLI requires `main` branch. Never create or switch to feature branches without explicit user approval. 4. **Dirty worktree:** Commit or stash your own files before pulling. Never stash everything (`git stash`) — only your files: `git stash push -- file1 file2`. 5. **Conflict resolution:** If merge conflicts, resolve and commit. Never force-push without user approval. 6. **Repo-local config:** Each clone must run the setup commands below to override any global/system rebase defaults. @@ -691,14 +691,15 @@ Azure Web App has continuous deployment enabled on ACR. When a new image appears **Fork-based development workflow:** -1. **Develop on fork's `main`** — all work happens on `origin` (your fork) -2. **RC releases → fork** — `imas-codex release -m "..."` pushes graph + tag to origin, fork CI validates and deploys to Azure test URL -3. **Verify RC** — exercise tools on test deployment, run A/B tests -4. **PR to upstream** — when RC is confirmed working, PR fork/main → upstream/main -5. **Final release → upstream** — after PR merges: `imas-codex release --final -m "..."` tags upstream, production CI deploys +1. **All work on fork's `main`** — no feature branches. Multiple agents use the same `main` branch with merge discipline. +2. **RC releases → fork CI → Azure test** — `imas-codex release -m "..."` pushes graph + tag to origin. Fork CI builds and pushes to ACR (hardcoded `iterorganization/` path). Azure auto-deploys the `latest-rc` tag. +3. **Verify RC** — exercise tools on test deployment at `https://app-imas-mcp-server-test-frc.azurewebsites.net/health`. Run A/B tests against all MCP tools. +4. **PR to upstream** — when RC is confirmed working, PR fork/main → upstream/main. +5. **Final release → upstream** — after PR merges: `imas-codex release --final -m "..."` tags upstream, production CI deploys with `latest-stable` tag. **Rules:** - **Never push directly to `upstream/main`** — always PR. Use `git push origin main` for day-to-day work. +- **Never push the same tag to both origin and upstream** — RC tags go to origin only, final tags to upstream only. Duplicate tags cause ACR race conditions. - RC tags on fork are disposable — iterate freely - Graph push runs from the ITER machine where Neo4j runs — CI cannot build graph data - The release CLI handles everything — do not manually push graphs or tags separately @@ -741,8 +742,8 @@ uv run ruff check --fix . # Lint (Python only) uv run ruff format . # Format git add ... # Stage specific files (never git add -A) uv run git commit -m "type: concise summary" # Conventional format -git pull --no-rebase origin # Merge fork changes first -git push origin # Push to fork (NEVER upstream) +git pull --no-rebase origin main # Merge fork changes first +git push origin main # Push to fork (NEVER upstream) ``` **Never stage:** auto-generated files (models.py, dd_models.py, schema_context_data.py), gitignored files, `*_private.yaml` files. From a21b3bb49ef0d09d93d42500f990eb9ce69c2b9b Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 09:57:47 +0200 Subject: [PATCH 08/37] chore: update GitHub Actions to Node.js 24 compatible versions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit actions/checkout v4→v6, actions/upload-artifact v4→v7, actions/cache v4→v5, astral-sh/setup-uv v5→v8, codecov/codecov-action v4→v6, softprops/action-gh-release v1→v2, actions/attest-build-provenance v1→v4. Remove FORCE_JAVASCRIPT_ACTIONS_TO_NODE24 env var — no longer needed with native Node 24 actions. --- .github/workflows/benchmark.yml | 5 ++--- .github/workflows/container-cleanup.yml | 1 - .github/workflows/docker-build-push.yml | 13 ++++++------- .github/workflows/graph-quality.yml | 7 +++---- .github/workflows/release.yml | 7 +++---- .github/workflows/test.yml | 13 ++++++------- 6 files changed, 20 insertions(+), 26 deletions(-) diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 8cfbb70e5..992de7724 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -1,7 +1,6 @@ name: Benchmark env: - FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true on: push: @@ -54,7 +53,7 @@ jobs: steps: # ── Setup ────────────────────────────────────────────────────────── - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: fetch-depth: 0 @@ -94,7 +93,7 @@ jobs: df -h - name: Install UV + ASV - uses: astral-sh/setup-uv@v5 + uses: astral-sh/setup-uv@v8 with: github-token: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/container-cleanup.yml b/.github/workflows/container-cleanup.yml index 4f9b105fa..67e0f9a09 100644 --- a/.github/workflows/container-cleanup.yml +++ b/.github/workflows/container-cleanup.yml @@ -16,7 +16,6 @@ on: type: boolean env: - FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true ACR_REGISTRY: crcommonallfrc.azurecr.io IMAGE_NAME: iterorganization/imas-codex diff --git a/.github/workflows/docker-build-push.yml b/.github/workflows/docker-build-push.yml index 8c9d4136c..15ed79c48 100644 --- a/.github/workflows/docker-build-push.yml +++ b/.github/workflows/docker-build-push.yml @@ -11,7 +11,6 @@ concurrency: cancel-in-progress: true env: - FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true GHCR_REGISTRY: ghcr.io ACR_REGISTRY: crcommonallfrc.azurecr.io IMAGE_NAME: ${{ github.repository }} @@ -54,7 +53,7 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Install oras CLI run: | @@ -209,7 +208,7 @@ jobs: - name: Install uv if: steps.graph-tag.outputs.imas-tag != 'none' - uses: astral-sh/setup-uv@v5 + uses: astral-sh/setup-uv@v8 with: enable-cache: true github-token: ${{ secrets.GITHUB_TOKEN }} @@ -259,7 +258,7 @@ jobs: - name: Upload test results if: always() && steps.graph-tag.outputs.imas-tag != 'none' - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v7 with: name: graph-quality-results path: graph-quality-results.xml @@ -272,7 +271,7 @@ jobs: if: startsWith(github.ref, 'refs/tags/') || github.event_name == 'workflow_dispatch' steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: fetch-depth: 0 @@ -438,7 +437,7 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: fetch-depth: 0 # Full git history for dynamic versioning # Ensure all refs are available @@ -641,7 +640,7 @@ jobs: - name: Generate artifact attestation for GHCR (releases only) if: github.event_name != 'pull_request' && startsWith(github.ref, 'refs/tags/v') && !contains(github.ref, '-rc') - uses: actions/attest-build-provenance@v1 + uses: actions/attest-build-provenance@v4 with: subject-name: ${{ env.GHCR_REGISTRY }}/${{ env.IMAGE_NAME }} subject-digest: ${{ steps.build-release.outputs.digest }} diff --git a/.github/workflows/graph-quality.yml b/.github/workflows/graph-quality.yml index e2954efca..79a93e91d 100644 --- a/.github/workflows/graph-quality.yml +++ b/.github/workflows/graph-quality.yml @@ -29,7 +29,6 @@ concurrency: cancel-in-progress: true env: - FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true NEO4J_PASSWORD: imas-codex NEO4J_URI: bolt://localhost:7687 @@ -58,7 +57,7 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Install oras CLI run: | @@ -186,7 +185,7 @@ jobs: "MATCH (n) RETURN count(n) AS nodes, labels(n)[0] AS label ORDER BY nodes DESC LIMIT 10" - name: Install uv - uses: astral-sh/setup-uv@v5 + uses: astral-sh/setup-uv@v8 with: enable-cache: true github-token: ${{ secrets.GITHUB_TOKEN }} @@ -230,7 +229,7 @@ jobs: - name: Upload test results if: always() - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v7 with: name: graph-quality-results path: graph-quality-results.xml diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 14b1855d7..7218b9742 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -1,7 +1,6 @@ name: Release env: - FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true on: push: @@ -23,12 +22,12 @@ jobs: steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: fetch-depth: 0 - name: Install uv - uses: astral-sh/setup-uv@v5 + uses: astral-sh/setup-uv@v8 with: github-token: ${{ secrets.GITHUB_TOKEN }} @@ -90,7 +89,7 @@ jobs: echo "tag_name=${TAG_NAME}" >> $GITHUB_OUTPUT - name: Create Release - uses: softprops/action-gh-release@v1 + uses: softprops/action-gh-release@v2 with: tag_name: ${{ steps.changelog.outputs.tag_name }} name: Release ${{ steps.changelog.outputs.tag_name }} diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index c8becbf92..b9367da8a 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -1,7 +1,6 @@ name: Test env: - FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true on: push: @@ -23,7 +22,7 @@ jobs: imas-dd-version: ["3.42.2", "4.1.0"] steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v6 - name: Set up Node.js uses: actions/setup-node@v4 @@ -37,7 +36,7 @@ jobs: npx --version - name: Cache npm packages - uses: actions/cache@v4 + uses: actions/cache@v5 with: path: ~/.npm key: npm-${{ runner.os }}-${{ hashFiles('.github/workflows/test.yml') }} @@ -51,7 +50,7 @@ jobs: run: playwright install --with-deps chromium - name: Install uv - uses: astral-sh/setup-uv@v5 + uses: astral-sh/setup-uv@v8 with: enable-cache: true github-token: ${{ secrets.GITHUB_TOKEN }} @@ -85,7 +84,7 @@ jobs: IMAS_DD_VERSION: ${{ matrix.imas-dd-version }} - name: Upload coverage reports to Codecov - uses: codecov/codecov-action@v4 + uses: codecov/codecov-action@v6 if: ${{ !cancelled() }} with: file: ./coverage.xml @@ -116,10 +115,10 @@ jobs: --health-start-period 30s steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v6 - name: Install uv - uses: astral-sh/setup-uv@v5 + uses: astral-sh/setup-uv@v8 with: enable-cache: true github-token: ${{ secrets.GITHUB_TOKEN }} From 1a5d91041214a7bbb7a5b370b0e40ed10a05e510 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 10:03:45 +0200 Subject: [PATCH 09/37] feat(sn): extend SNCandidate with rich fields and update compose/validate workers Phase 2: Add description, documentation, unit, kind, tags, links, ids_paths, validity_domain, constraints to SNCandidate model. Update compose_worker to pass all fields through to state.composed. Phase 3: Enhance compose_system prompt with rich output format, documentation template, tags vocabulary, kind rules, links guidance. Phase 5: Add soft validation checks for description length, doc length, unit validity, kind enum, tags vocabulary, links references. --- imas_codex/llm/prompts/sn/compose_system.md | 50 +++++++- imas_codex/sn/models.py | 28 ++++- imas_codex/sn/workers.py | 121 ++++++++++++++++++++ 3 files changed, 195 insertions(+), 4 deletions(-) diff --git a/imas_codex/llm/prompts/sn/compose_system.md b/imas_codex/llm/prompts/sn/compose_system.md index 5752fcf96..aff82352d 100644 --- a/imas_codex/llm/prompts/sn/compose_system.md +++ b/imas_codex/llm/prompts/sn/compose_system.md @@ -99,12 +99,58 @@ Physics: B_T={{ machine.physics.get('toroidal_magnetic_field', {}).get('value', ## Output Format Return a JSON object with: -- `candidates`: array of standard name compositions +- `candidates`: array of standard name compositions (see schema below) - `skipped`: array of source_ids that are not distinct physics quantities -Each candidate has: +### Candidate Schema + +Each candidate MUST include: - `source_id`: full DD path (e.g., "equilibrium/time_slice/profiles_1d/psi") - `standard_name`: the composed name in snake_case +- `description`: one-sentence summary, **under 120 characters** (e.g., "Electron temperature profile on the poloidal flux grid") +- `documentation`: rich documentation paragraph (200-500 chars) — see template below +- `unit`: SI unit string (`eV`, `m`, `A`, `T`, `Pa`, `W`, `m^-3`, `s`, `V`, `kg`, `rad`, `K`, `Wb`, `ohm`, `Hz`, `J`) or `null` for dimensionless +- `kind`: one of `"scalar"`, `"vector"`, `"metadata"` — see classification rules +- `tags`: array of 1-2 primary + 0-3 secondary tags from the controlled vocabulary +- `links`: array of 4-8 related standard names from the existing_names list +- `ids_paths`: array of IMAS DD paths this name maps to (include the source_id at minimum) - `fields`: dict of grammar fields used (only non-null fields) - `confidence`: float 0.0-1.0 - `reason`: brief justification +- `validity_domain`: physical region where this quantity is meaningful (e.g., "core plasma", "scrape-off layer", "entire plasma", "pedestal region") or `null` +- `constraints`: array of physical constraints (e.g., `["T_e > 0"]`, `["0 ≤ ρ ≤ 1"]`) + +### Documentation Template + +Write documentation following this structure (200-500 characters): + +1. **Opening statement** — what the quantity is and where it appears +2. **Governing physics** — equations or relationships using LaTeX ($T_e$, $\psi$, $n_e$) +3. **Physical significance** — why this quantity matters for plasma performance +4. **Measurement context** — how it is typically measured or computed +5. **Typical values** — use ranges from the tokamak parameter data above +6. **Sign conventions** — note any COCOS dependencies if applicable +7. **Cross-references** — use `[name](#name)` format to link related quantities + +Example documentation: +> The electron temperature $T_e$ is a fundamental kinetic quantity representing the thermal energy of the plasma electron population. Measured primarily by Thomson scattering and electron cyclotron emission (ECE) diagnostics. Typical values range from ~100 eV at the edge to 1-20 keV in the core depending on heating power and confinement regime. Related to [electron_density](#electron_density) via the electron pressure $p_e = n_e T_e$. + +### Tags — Controlled Vocabulary + +**Primary tags** (include 1-2): fundamental, equilibrium, core-physics, transport, edge-physics, mhd, nbi, ec-heating, ic-heating, lh-heating, waves, fast-particles, runaway-electrons, magnetics, thomson-scattering, interferometry, reflectometry, spectroscopy, radiation-diagnostics, imaging, neutronics, coils-and-control, fueling, wall-and-structures, pulse-management, data-products, utilities, turbulence, plasma-initiation + +**Secondary tags** (include 0-3): time-dependent, steady-state, spatial-profile, flux-surface-average, volume-average, line-integrated, local-measurement, global-quantity, measured, reconstructed, simulated, derived, validated, equilibrium-reconstruction, transport-modeling, mhd-stability-analysis, heating-deposition, calibrated, real-time, post-shot-analysis, benchmark-quantity, performance-metric + +### Kind Classification Rules + +- **scalar**: single value per spatial point or time — temperature, density, current, pressure, energy, power, frequency, flux, beta, safety factor +- **vector**: has R/Z or multi-component structure — magnetic field, velocity field, gradient, current density vector, force density +- **metadata**: non-measurable concepts, technique names, classifications, indices, status flags — confinement mode label, scenario identifier + +### Links Guidance + +Reference 4-8 related standard names from the `existing_names` list. Only include names that actually exist — do NOT invent new names for links. Prefer names that are: +- Same physical quantity in a different context (electron_temperature ↔ ion_temperature) +- Derived or input quantities (pressure ↔ temperature + density) +- Measured by the same diagnostic +- Commonly plotted together diff --git a/imas_codex/sn/models.py b/imas_codex/sn/models.py index 091858fe6..99dc344bc 100644 --- a/imas_codex/sn/models.py +++ b/imas_codex/sn/models.py @@ -3,6 +3,7 @@ from __future__ import annotations from enum import StrEnum +from typing import Literal from pydantic import BaseModel, Field @@ -11,10 +12,33 @@ class SNCandidate(BaseModel): """A single standard name candidate from LLM composition.""" source_id: str = Field(description="Source entity ID (DD path or signal ID)") - standard_name: str = Field(description="Composed standard name") - fields: dict[str, str] = Field(description="Grammar fields used") + standard_name: str = Field(description="Composed standard name in snake_case") + description: str = Field(default="", description="One sentence, <120 chars") + documentation: str = Field( + default="", description="Rich docs with LaTeX, links, typical values" + ) + unit: str | None = Field( + default=None, description="SI unit string (eV, m, A, etc.)" + ) + kind: Literal["scalar", "vector", "metadata"] = Field( + default="scalar", description="Entry kind" + ) + tags: list[str] = Field(default_factory=list, description="Classification tags") + links: list[str] = Field(default_factory=list, description="Related standard names") + ids_paths: list[str] = Field( + default_factory=list, description="Mapped IMAS DD paths" + ) + fields: dict[str, str] = Field( + default_factory=dict, description="Grammar fields used" + ) confidence: float = Field(ge=0, le=1, description="Naming confidence") reason: str = Field(description="Brief justification") + validity_domain: str | None = Field( + default=None, description="Physical region where quantity is valid" + ) + constraints: list[str] = Field( + default_factory=list, description="Physical constraints" + ) class SNComposeBatch(BaseModel): diff --git a/imas_codex/sn/workers.py b/imas_codex/sn/workers.py index 3e6f9f10b..ba649fa7c 100644 --- a/imas_codex/sn/workers.py +++ b/imas_codex/sn/workers.py @@ -196,9 +196,18 @@ async def _compose_batch(batch: ExtractionBatch) -> list[dict]: "id": c.standard_name, "source_type": "dd" if state.source == "dd" else "signal", "source_id": c.source_id, + "description": c.description, + "documentation": c.documentation, + "units": c.unit, # graph_ops uses "units" → canonical_units + "kind": c.kind, + "tags": c.tags, + "links": c.links, + "imas_paths": c.ids_paths, # graph schema key "fields": c.fields, "confidence": c.confidence, "reason": c.reason, + "validity_domain": c.validity_domain, + "constraints": c.constraints, } ) @@ -513,6 +522,23 @@ async def validate_worker(state: SNBuildState, **_kwargs) -> None: parse_standard_name, ) + # Load tag vocabulary for soft validation + try: + from typing import get_args + + from imas_standard_names.grammar.tag_types import PrimaryTag, SecondaryTag + + valid_primary_tags = set(get_args(PrimaryTag)) + valid_secondary_tags = set(get_args(SecondaryTag)) + valid_tags = valid_primary_tags | valid_secondary_tags + except Exception: + valid_tags = set() + + # Collect existing names for link validation + existing_names: set[str] = set() + for entry in input_candidates: + existing_names.add(entry.get("id", "")) + wlog.info("Validating %d composed names", len(input_candidates)) state.validate_stats.total = len(input_candidates) @@ -521,6 +547,18 @@ async def validate_worker(state: SNBuildState, **_kwargs) -> None: fields_consistent = 0 fields_inconsistent = 0 + # Soft validation counters + desc_present = 0 + desc_too_long = 0 + doc_present = 0 + doc_too_short = 0 + unit_valid = 0 + kind_valid = 0 + tags_valid = 0 + links_valid = 0 + + _VALID_KINDS = {"scalar", "vector", "metadata"} + for i, entry in enumerate(input_candidates): name = entry.get("id", "") try: @@ -553,6 +591,66 @@ async def validate_worker(state: SNBuildState, **_kwargs) -> None: fields_inconsistent += 1 entry["fields_consistent"] = False + # --- Soft validation checks (metrics only, never reject) --- + + # 1. Description present + length check + desc = entry.get("description", "") + if desc: + desc_present += 1 + if len(desc) > 120: + desc_too_long += 1 + wlog.debug("Description >120 chars for %r: %d", name, len(desc)) + else: + wlog.debug("Missing description for %r", name) + + # 2. Documentation present + minimum length + doc = entry.get("documentation", "") + if doc: + doc_present += 1 + if len(doc) < 200: + doc_too_short += 1 + wlog.debug("Documentation <200 chars for %r: %d", name, len(doc)) + else: + wlog.debug("Missing documentation for %r", name) + + # 3. Unit validity — simple pattern check + unit = entry.get("units") + if unit and isinstance(unit, str) and len(unit) < 50: + unit_valid += 1 + + # 4. Kind validity + kind = entry.get("kind", "") + if kind in _VALID_KINDS: + kind_valid += 1 + elif kind: + wlog.debug("Invalid kind %r for %r", kind, name) + + # 5. Tags from vocabulary + entry_tags = entry.get("tags") or [] + if entry_tags and valid_tags: + if all(t in valid_tags for t in entry_tags): + tags_valid += 1 + else: + bad_tags = [t for t in entry_tags if t not in valid_tags] + wlog.debug("Unknown tags for %r: %s", name, bad_tags) + elif entry_tags: + tags_valid += 1 # no vocabulary loaded, accept any + + # 6. Links reference existing names + entry_links = entry.get("links") or [] + if entry_links: + if all(lnk in existing_names for lnk in entry_links): + links_valid += 1 + else: + unknown = [lnk for lnk in entry_links if lnk not in existing_names] + wlog.debug("Unknown links for %r: %s", name, unknown) + + # 7. ids_paths look like IMAS paths (contain '/') + imas_paths = entry.get("imas_paths") or [] + for p in imas_paths: + if "/" not in p: + wlog.debug("Suspicious ids_path for %r: %r", name, p) + valid.append(entry) except Exception: wlog.debug("Validation failed for name: %r", name) @@ -573,10 +671,33 @@ async def validate_worker(state: SNBuildState, **_kwargs) -> None: fields_consistent, fields_inconsistent, ) + wlog.info( + "Soft checks: desc=%d/%d (>120: %d), doc=%d/%d (<200: %d), " + "unit=%d, kind=%d, tags=%d, links=%d", + desc_present, + len(valid), + desc_too_long, + doc_present, + len(valid), + doc_too_short, + unit_valid, + kind_valid, + tags_valid, + links_valid, + ) state.stats["validate_valid"] = len(valid) state.stats["validate_invalid"] = invalid_count state.stats["validate_fields_consistent"] = fields_consistent state.stats["validate_fields_inconsistent"] = fields_inconsistent + # Soft validation metrics + state.stats["validate_desc_present"] = desc_present + state.stats["validate_desc_too_long"] = desc_too_long + state.stats["validate_doc_present"] = doc_present + state.stats["validate_doc_too_short"] = doc_too_short + state.stats["validate_unit_valid"] = unit_valid + state.stats["validate_kind_valid"] = kind_valid + state.stats["validate_tags_valid"] = tags_valid + state.stats["validate_links_valid"] = links_valid state.validate_stats.freeze_rate() state.validate_phase.mark_done() From 840d54c96480619327cd6bac7591e18906561824 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 10:08:12 +0200 Subject: [PATCH 10/37] chore: increase smoke test health check timeout margin --- .github/workflows/docker-build-push.yml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/workflows/docker-build-push.yml b/.github/workflows/docker-build-push.yml index 15ed79c48..07fa0e667 100644 --- a/.github/workflows/docker-build-push.yml +++ b/.github/workflows/docker-build-push.yml @@ -317,16 +317,16 @@ jobs: echo "Container started, waiting for services..." - name: Wait for health check - timeout-minutes: 5 + timeout-minutes: 10 run: | echo "Waiting for MCP server to be ready..." - for i in $(seq 1 60); do + for i in $(seq 1 90); do if docker exec smoke-test curl -sf http://127.0.0.1:8000/health > /dev/null 2>&1; then - echo "✅ MCP server healthy after ${i}s" + echo "✅ MCP server healthy after $((i * 5))s" break fi - if [ $i -eq 60 ]; then - echo "❌ Health check timed out" + if [ $i -eq 90 ]; then + echo "❌ Health check timed out after 450s" docker logs smoke-test exit 1 fi From 867a6e66658aab684e3a30841b836bbb320839cf Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 10:11:33 +0200 Subject: [PATCH 11/37] feat(sn): add import-catalog CLI command for catalog feedback loop Import reviewed YAML catalog entries back into the graph as accepted StandardName nodes. Derives grammar fields via name parsing, maps catalog fields to graph schema, preserves graph-only fields. Catalog-owned fields (description, documentation, kind, tags, etc.) use direct SET for authoritative overwrite. Graph-only fields (embedding, model, generated_at) are preserved via coalesce. Includes dry-run mode, tag filtering, and comprehensive tests. --- imas_codex/cli/sn.py | 85 +++++ imas_codex/sn/catalog_import.py | 279 +++++++++++++++ tests/sn/test_catalog_import.py | 590 ++++++++++++++++++++++++++++++++ 3 files changed, 954 insertions(+) create mode 100644 imas_codex/sn/catalog_import.py create mode 100644 tests/sn/test_catalog_import.py diff --git a/imas_codex/cli/sn.py b/imas_codex/cli/sn.py index e35487465..b9a1ca2a5 100644 --- a/imas_codex/cli/sn.py +++ b/imas_codex/cli/sn.py @@ -596,3 +596,88 @@ def sn_publish( console.print( f" [yellow]Would create PR for {batch.group_key}[/yellow]" ) + + +@sn.command("import-catalog") +@click.option( + "--catalog-dir", + type=click.Path(exists=True), + required=True, + help="Path to catalog directory containing YAML entries", +) +@click.option("--tags", type=str, default=None, help="Comma-separated tag filter") +@click.option("--dry-run", is_flag=True, help="Preview without writing to graph") +@click.option("-v", "--verbose", is_flag=True, help="Enable verbose logging") +def sn_import_catalog( + catalog_dir: str, + tags: str | None, + dry_run: bool, + verbose: bool, +) -> None: + """Import reviewed catalog entries into the graph. + + \b + Reads YAML files from the catalog directory, validates them against + the imas-standard-names catalog model, derives grammar fields, and + MERGEs into the graph with review_status='accepted'. + + \b + Examples: + imas-codex sn import-catalog --catalog-dir ../imas-standard-names-catalog/standard_names + imas-codex sn import-catalog --catalog-dir --dry-run + imas-codex sn import-catalog --catalog-dir --tags equilibrium,core-physics + """ + from pathlib import Path + + if verbose: + logging.basicConfig(level=logging.DEBUG) + + tag_filter = [t.strip() for t in tags.split(",") if t.strip()] if tags else None + + console.print("\n[bold]Standard Name Catalog Import[/bold]") + console.print(f" Catalog: {catalog_dir}") + if tag_filter: + console.print(f" Tag filter: {', '.join(tag_filter)}") + if dry_run: + console.print(" Mode: [yellow]dry run[/yellow]") + console.print("") + + try: + from imas_codex.sn.catalog_import import import_catalog + + result = import_catalog( + catalog_dir=Path(catalog_dir), + dry_run=dry_run, + tag_filter=tag_filter, + ) + except ImportError as e: + console.print( + f"[red]Missing dependency:[/red] {e}\n" + "Install with: uv pip install imas-standard-names" + ) + raise SystemExit(1) from e + except Exception as e: + console.print(f"[red]Import error:[/red] {e}") + raise SystemExit(1) from e + + # Print results + if result.errors: + console.print(f" [red]Errors: {len(result.errors)}[/red]") + for err in result.errors[:10]: + console.print(f" - {err}") + if len(result.errors) > 10: + console.print(f" ... and {len(result.errors) - 10} more") + + if result.skipped: + console.print(f" [yellow]Skipped: {result.skipped}[/yellow] (tag filter)") + + action = "Would import" if dry_run else "Imported" + console.print(f"\n [green]{action}: {result.imported}[/green] entries") + + if dry_run and result.entries: + console.print("\n [bold]Preview:[/bold]") + for entry in result.entries[:20]: + units = f" [{entry.get('units', '')}]" if entry.get("units") else "" + console.print(f" - {entry['id']}{units}") + if len(result.entries) > 20: + console.print(f" ... and {len(result.entries) - 20} more") diff --git a/imas_codex/sn/catalog_import.py b/imas_codex/sn/catalog_import.py new file mode 100644 index 000000000..1cdfd0a56 --- /dev/null +++ b/imas_codex/sn/catalog_import.py @@ -0,0 +1,279 @@ +"""Catalog feedback import — read reviewed YAML entries and write to graph. + +Implements the publish → review → import feedback loop for standard names. +Catalog entries are authoritative: their fields overwrite graph fields. +Graph-only fields (embedding, model, generated_at) are preserved. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +logger = logging.getLogger(__name__) + + +@dataclass +class ImportResult: + """Summary of a catalog import operation.""" + + imported: int = 0 + updated: int = 0 + skipped: int = 0 + errors: list[str] = field(default_factory=list) + entries: list[dict[str, Any]] = field(default_factory=list) + + +def _parse_grammar_fields(name: str) -> dict[str, str | None]: + """Derive grammar fields from a standard name string. + + Returns a dict with keys: physical_base, subject, component, + coordinate, position, process. Values are strings or None. + """ + try: + from imas_standard_names.grammar import parse_standard_name + + parsed = parse_standard_name(name) + return { + "physical_base": str(parsed.physical_base) + if parsed.physical_base + else None, + "subject": str(parsed.subject.value) if parsed.subject else None, + "component": str(parsed.component.value) if parsed.component else None, + "coordinate": str(parsed.coordinate.value) if parsed.coordinate else None, + "position": str(parsed.position.value) if parsed.position else None, + "process": str(parsed.process.value) if parsed.process else None, + } + except Exception: + logger.debug("Grammar parse failed for name: %r", name) + return { + "physical_base": None, + "subject": None, + "component": None, + "coordinate": None, + "position": None, + "process": None, + } + + +def _catalog_entry_to_dict(entry: Any) -> dict[str, Any]: + """Convert a validated catalog entry to a graph-write dict. + + Maps catalog field names to graph schema field names and derives + grammar fields from the standard name. + """ + # Derive grammar fields from the name + grammar = _parse_grammar_fields(entry.name) + + # Convert tags/links to plain strings (catalog may use typed objects) + tags = [str(t) for t in entry.tags] if entry.tags else None + links = [str(lnk) for lnk in entry.links] if entry.links else None + ids_paths = list(entry.ids_paths) if entry.ids_paths else None + constraints = list(entry.constraints) if entry.constraints else None + + # Determine source_type from presence of ids_paths + source_type = "dd" if ids_paths else "manual" + + return { + "id": entry.name, + "description": entry.description or None, + "documentation": entry.documentation or None, + "kind": str(entry.kind) if entry.kind else None, + "units": str(entry.unit) if entry.unit else None, + "tags": tags or None, + "links": links or None, + "imas_paths": ids_paths or None, + "validity_domain": entry.validity_domain or None, + "constraints": constraints or None, + "physics_domain": entry.physics_domain or None, + "review_status": "accepted", + "source_type": source_type, + # Grammar fields + "physical_base": grammar["physical_base"], + "subject": grammar["subject"], + "component": grammar["component"], + "coordinate": grammar["coordinate"], + "position": grammar["position"], + "process": grammar["process"], + } + + +def _write_catalog_entries(entries: list[dict[str, Any]]) -> int: + """Write catalog entries to graph with catalog-authoritative semantics. + + Catalog-owned fields are SET directly (overwrite). + Graph-only fields (embedding, model, generated_at, etc.) are preserved + via coalesce. Returns the number of nodes written. + """ + if not entries: + return 0 + + from imas_codex.graph.client import GraphClient + + with GraphClient() as gc: + # MERGE StandardName nodes — catalog fields overwrite, graph-only preserved + gc.query( + """ + UNWIND $batch AS b + MERGE (sn:StandardName {id: b.id}) + SET sn.description = b.description, + sn.documentation = b.documentation, + sn.kind = b.kind, + sn.canonical_units = b.units, + sn.tags = b.tags, + sn.links = b.links, + sn.imas_paths = b.imas_paths, + sn.validity_domain = b.validity_domain, + sn.constraints = b.constraints, + sn.physics_domain = b.physics_domain, + sn.review_status = 'accepted', + sn.imported_at = datetime(), + sn.physical_base = b.physical_base, + sn.subject = b.subject, + sn.component = b.component, + sn.coordinate = b.coordinate, + sn.position = b.position, + sn.process = b.process, + sn.source_type = coalesce(b.source_type, sn.source_type), + sn.created_at = coalesce(sn.created_at, datetime()), + sn.embedding = coalesce(sn.embedding, null), + sn.embedded_at = coalesce(sn.embedded_at, null), + sn.model = coalesce(sn.model, null), + sn.generated_at = coalesce(sn.generated_at, null), + sn.confidence = coalesce(sn.confidence, null) + """, + batch=entries, + ) + + # Create CANONICAL_UNITS relationships: StandardName → Unit + units_batch = [ + {"id": e["id"], "unit": e["units"]} for e in entries if e.get("units") + ] + if units_batch: + gc.query( + """ + UNWIND $batch AS b + MATCH (sn:StandardName {id: b.id}) + MERGE (u:Unit {id: b.unit}) + MERGE (sn)-[:CANONICAL_UNITS]->(u) + """, + batch=units_batch, + ) + + # Create HAS_STANDARD_NAME relationships from ids_paths + dd_batch = [] + for e in entries: + if e.get("imas_paths"): + for path in e["imas_paths"]: + dd_batch.append({"id": e["id"], "source_id": path}) + + if dd_batch: + gc.query( + """ + UNWIND $batch AS b + MATCH (sn:StandardName {id: b.id}) + MATCH (src:IMASNode {id: b.source_id}) + MERGE (src)-[:HAS_STANDARD_NAME]->(sn) + """, + batch=dd_batch, + ) + + written = len(entries) + logger.info("Imported %d catalog entries to graph", written) + return written + + +def import_catalog( + catalog_dir: Path, + dry_run: bool = False, + tag_filter: list[str] | None = None, +) -> ImportResult: + """Import YAML catalog entries into graph as accepted StandardName nodes. + + Reads all ``*.yml`` and ``*.yaml`` files from *catalog_dir* (recursive), + validates each entry against the ``imas-standard-names`` catalog model, + derives grammar fields via name parsing, and MERGEs into the graph. + + Catalog fields are authoritative and overwrite graph values. + Graph-only fields (embedding, model, generated_at) are preserved. + Imported entries receive ``review_status='accepted'``. + + Parameters + ---------- + catalog_dir: + Path to directory containing YAML catalog entries. + dry_run: + If True, parse and validate but do not write to graph. + tag_filter: + If provided, only import entries whose tags overlap with this list. + + Returns + ------- + ImportResult with counts and entry details. + """ + import yaml + from imas_standard_names.catalog.edit import StandardNameEntry + from pydantic import TypeAdapter + + ta = TypeAdapter(StandardNameEntry) + result = ImportResult() + + # Collect YAML files + yaml_files = sorted( + p + for p in catalog_dir.rglob("*") + if p.suffix in (".yml", ".yaml") and p.is_file() + ) + + if not yaml_files: + logger.info("No YAML files found in %s", catalog_dir) + return result + + logger.info("Found %d YAML files in %s", len(yaml_files), catalog_dir) + + # Parse and validate entries + prepared: list[dict[str, Any]] = [] + + for yaml_path in yaml_files: + try: + with open(yaml_path) as f: + data = yaml.safe_load(f) + + if not isinstance(data, dict): + result.errors.append(f"{yaml_path.name}: not a YAML mapping") + continue + + entry = ta.validate_python(data) + except Exception as exc: + result.errors.append(f"{yaml_path.name}: {exc}") + logger.debug("Failed to parse %s: %s", yaml_path, exc) + continue + + # Apply tag filter if specified + if tag_filter: + entry_tags = {str(t) for t in entry.tags} if entry.tags else set() + if not entry_tags.intersection(tag_filter): + result.skipped += 1 + continue + + # Convert to graph dict + graph_dict = _catalog_entry_to_dict(entry) + prepared.append(graph_dict) + result.entries.append(graph_dict) + + if not prepared: + logger.info("No entries to import after filtering") + return result + + # Write to graph (unless dry run) + if dry_run: + result.imported = len(prepared) + logger.info("Dry run: would import %d entries", len(prepared)) + else: + written = _write_catalog_entries(prepared) + result.imported = written + logger.info("Imported %d entries to graph", written) + + return result diff --git a/tests/sn/test_catalog_import.py b/tests/sn/test_catalog_import.py new file mode 100644 index 000000000..d7c72b9c6 --- /dev/null +++ b/tests/sn/test_catalog_import.py @@ -0,0 +1,590 @@ +"""Tests for the catalog feedback import module. + +Tests YAML parsing, grammar field derivation, tag filtering, dry-run +behavior, and graph write semantics — all mocked, no live Neo4j. +""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +import yaml + +imas_sn = pytest.importorskip("imas_standard_names") + + +# ============================================================================= +# Fixtures +# ============================================================================= + +SAMPLE_CATALOG_ENTRY = { + "name": "electron_temperature", + "description": "Electron temperature", + "documentation": "The electron temperature Te is measured by Thomson scattering.", + "kind": "scalar", + "unit": "eV", + "tags": [], + "links": [], + "ids_paths": ["core_profiles/profiles_1d/electrons/temperature"], + "validity_domain": "core plasma", + "constraints": ["T_e > 0"], + "physics_domain": "core_plasma_physics", + "status": "active", +} + +SAMPLE_CATALOG_ENTRY_MINIMAL = { + "name": "plasma_current", + "description": "Plasma current", + "documentation": "Total toroidal plasma current.", + "kind": "scalar", + "unit": "A", + "tags": [], + "links": [], + "ids_paths": [], + "validity_domain": "", + "constraints": [], + "physics_domain": "equilibrium", + "status": "active", +} + + +@pytest.fixture() +def catalog_dir(tmp_path: Path) -> Path: + """Create a temporary catalog directory with sample YAML files.""" + d = tmp_path / "catalog" + d.mkdir() + (d / "electron_temperature.yaml").write_text(yaml.safe_dump(SAMPLE_CATALOG_ENTRY)) + (d / "plasma_current.yaml").write_text(yaml.safe_dump(SAMPLE_CATALOG_ENTRY_MINIMAL)) + return d + + +@pytest.fixture() +def catalog_dir_with_tags(tmp_path: Path) -> Path: + """Create a catalog directory with tagged entries.""" + d = tmp_path / "catalog_tagged" + d.mkdir() + + entry_tagged = {**SAMPLE_CATALOG_ENTRY, "tags": ["spatial-profile"]} + entry_untagged = {**SAMPLE_CATALOG_ENTRY_MINIMAL, "tags": []} + + (d / "electron_temperature.yaml").write_text(yaml.safe_dump(entry_tagged)) + (d / "plasma_current.yaml").write_text(yaml.safe_dump(entry_untagged)) + return d + + +# ============================================================================= +# YAML parsing tests +# ============================================================================= + + +class TestImportParsesYaml: + """Test that import correctly parses YAML catalog files.""" + + def test_parses_yaml_files(self, catalog_dir: Path) -> None: + """Should parse all valid YAML files in the directory.""" + from imas_codex.sn.catalog_import import import_catalog + + with patch( + "imas_codex.sn.catalog_import._write_catalog_entries", return_value=2 + ): + result = import_catalog(catalog_dir, dry_run=False) + + assert result.imported == 2 + assert len(result.errors) == 0 + + def test_parses_yml_extension(self, tmp_path: Path) -> None: + """Should handle .yml file extension too.""" + d = tmp_path / "catalog" + d.mkdir() + (d / "electron_temperature.yml").write_text( + yaml.safe_dump(SAMPLE_CATALOG_ENTRY) + ) + + from imas_codex.sn.catalog_import import import_catalog + + with patch( + "imas_codex.sn.catalog_import._write_catalog_entries", return_value=1 + ): + result = import_catalog(d, dry_run=False) + + assert result.imported == 1 + + def test_recursive_subdirectories(self, tmp_path: Path) -> None: + """Should walk subdirectories recursively.""" + d = tmp_path / "catalog" + sub = d / "scalars" + sub.mkdir(parents=True) + (sub / "electron_temperature.yaml").write_text( + yaml.safe_dump(SAMPLE_CATALOG_ENTRY) + ) + + from imas_codex.sn.catalog_import import import_catalog + + with patch( + "imas_codex.sn.catalog_import._write_catalog_entries", return_value=1 + ): + result = import_catalog(d, dry_run=False) + + assert result.imported == 1 + + +# ============================================================================= +# Grammar field derivation tests +# ============================================================================= + + +class TestGrammarFields: + """Test that grammar fields are derived from name parsing.""" + + def test_derives_grammar_fields(self) -> None: + """Should extract subject and physical_base from name.""" + from imas_codex.sn.catalog_import import _parse_grammar_fields + + fields = _parse_grammar_fields("electron_temperature") + assert fields["subject"] == "electron" + assert fields["physical_base"] == "temperature" + + def test_unparseable_name_returns_none(self) -> None: + """Should return None fields for names that can't be parsed.""" + from imas_codex.sn.catalog_import import _parse_grammar_fields + + # Mock the grammar parser to raise, simulating an unparseable name + with patch( + "imas_standard_names.grammar.parse_standard_name", + side_effect=ValueError("bad name"), + ): + fields = _parse_grammar_fields("__broken__") + + assert fields["physical_base"] is None + assert fields["subject"] is None + + def test_grammar_fields_in_import_output(self, catalog_dir: Path) -> None: + """Imported entries should have grammar fields populated.""" + from imas_codex.sn.catalog_import import import_catalog + + result = import_catalog(catalog_dir, dry_run=True) + + # Find the electron_temperature entry + et_entry = next(e for e in result.entries if e["id"] == "electron_temperature") + assert et_entry["subject"] == "electron" + assert et_entry["physical_base"] == "temperature" + + +# ============================================================================= +# Import status and field mapping tests +# ============================================================================= + + +class TestFieldMapping: + """Test that catalog fields are correctly mapped to graph dict.""" + + def test_sets_accepted_status(self, catalog_dir: Path) -> None: + """All imported entries should have review_status='accepted'.""" + from imas_codex.sn.catalog_import import import_catalog + + result = import_catalog(catalog_dir, dry_run=True) + + for entry in result.entries: + assert entry["review_status"] == "accepted" + + def test_maps_unit_to_units(self, catalog_dir: Path) -> None: + """Catalog 'unit' field should map to graph 'units' key.""" + from imas_codex.sn.catalog_import import import_catalog + + result = import_catalog(catalog_dir, dry_run=True) + + et_entry = next(e for e in result.entries if e["id"] == "electron_temperature") + assert et_entry["units"] == "eV" + assert "unit" not in et_entry # should not have the catalog key + + def test_maps_ids_paths_to_imas_paths(self, catalog_dir: Path) -> None: + """Catalog 'ids_paths' should map to graph 'imas_paths' key.""" + from imas_codex.sn.catalog_import import import_catalog + + result = import_catalog(catalog_dir, dry_run=True) + + et_entry = next(e for e in result.entries if e["id"] == "electron_temperature") + assert et_entry["imas_paths"] == [ + "core_profiles/profiles_1d/electrons/temperature" + ] + + def test_source_type_dd_for_entries_with_paths(self, catalog_dir: Path) -> None: + """Entries with ids_paths should have source_type='dd'.""" + from imas_codex.sn.catalog_import import import_catalog + + result = import_catalog(catalog_dir, dry_run=True) + + et_entry = next(e for e in result.entries if e["id"] == "electron_temperature") + assert et_entry["source_type"] == "dd" + + def test_source_type_manual_for_entries_without_paths( + self, catalog_dir: Path + ) -> None: + """Entries without ids_paths should have source_type='manual'.""" + from imas_codex.sn.catalog_import import import_catalog + + result = import_catalog(catalog_dir, dry_run=True) + + pc_entry = next(e for e in result.entries if e["id"] == "plasma_current") + assert pc_entry["source_type"] == "manual" + + def test_maps_kind(self, catalog_dir: Path) -> None: + """Catalog 'kind' field should be mapped correctly.""" + from imas_codex.sn.catalog_import import import_catalog + + result = import_catalog(catalog_dir, dry_run=True) + + for entry in result.entries: + assert entry["kind"] == "scalar" + + def test_empty_lists_become_none(self, catalog_dir: Path) -> None: + """Empty catalog lists should become None for graph coalesce compatibility.""" + from imas_codex.sn.catalog_import import import_catalog + + result = import_catalog(catalog_dir, dry_run=True) + + # plasma_current has no ids_paths, empty tags, empty links + pc_entry = next(e for e in result.entries if e["id"] == "plasma_current") + assert pc_entry["imas_paths"] is None + assert pc_entry["tags"] is None + assert pc_entry["links"] is None + + def test_maps_physics_domain(self, catalog_dir: Path) -> None: + """Catalog 'physics_domain' field should be mapped.""" + from imas_codex.sn.catalog_import import import_catalog + + result = import_catalog(catalog_dir, dry_run=True) + + et_entry = next(e for e in result.entries if e["id"] == "electron_temperature") + assert et_entry["physics_domain"] == "core_plasma_physics" + + +# ============================================================================= +# Dry run tests +# ============================================================================= + + +class TestDryRun: + """Test that dry run mode doesn't write to graph.""" + + def test_dry_run_no_write(self, catalog_dir: Path) -> None: + """Dry run should not call the write function.""" + from imas_codex.sn.catalog_import import import_catalog + + with patch("imas_codex.sn.catalog_import._write_catalog_entries") as mock_write: + result = import_catalog(catalog_dir, dry_run=True) + + mock_write.assert_not_called() + assert result.imported == 2 # still reports count + assert len(result.entries) == 2 + + def test_non_dry_run_calls_write(self, catalog_dir: Path) -> None: + """Non-dry-run should call the write function.""" + from imas_codex.sn.catalog_import import import_catalog + + with patch( + "imas_codex.sn.catalog_import._write_catalog_entries", return_value=2 + ) as mock_write: + result = import_catalog(catalog_dir, dry_run=False) + + mock_write.assert_called_once() + assert result.imported == 2 + + +# ============================================================================= +# Tag filter tests +# ============================================================================= + + +class TestTagFilter: + """Test tag-based filtering of catalog entries.""" + + def test_tag_filter_includes_matching(self, catalog_dir_with_tags: Path) -> None: + """Should import entries whose tags match the filter.""" + from imas_codex.sn.catalog_import import import_catalog + + result = import_catalog( + catalog_dir_with_tags, dry_run=True, tag_filter=["spatial-profile"] + ) + + assert result.imported == 1 + assert result.entries[0]["id"] == "electron_temperature" + + def test_tag_filter_skips_non_matching(self, catalog_dir_with_tags: Path) -> None: + """Should skip entries that don't match the tag filter.""" + from imas_codex.sn.catalog_import import import_catalog + + result = import_catalog( + catalog_dir_with_tags, dry_run=True, tag_filter=["spatial-profile"] + ) + + assert result.skipped == 1 + + def test_no_tag_filter_imports_all(self, catalog_dir_with_tags: Path) -> None: + """Without tag filter, all entries should be imported.""" + from imas_codex.sn.catalog_import import import_catalog + + result = import_catalog(catalog_dir_with_tags, dry_run=True) + + assert result.imported == 2 + assert result.skipped == 0 + + +# ============================================================================= +# Error handling tests +# ============================================================================= + + +class TestErrorHandling: + """Test graceful error handling.""" + + def test_handles_invalid_yaml(self, tmp_path: Path) -> None: + """Should report errors for invalid YAML files without crashing.""" + d = tmp_path / "catalog" + d.mkdir() + (d / "bad.yaml").write_text(": : : invalid yaml [[[") + + from imas_codex.sn.catalog_import import import_catalog + + result = import_catalog(d, dry_run=True) + + assert result.imported == 0 + assert len(result.errors) == 1 + assert "bad.yaml" in result.errors[0] + + def test_handles_non_mapping_yaml(self, tmp_path: Path) -> None: + """Should report errors for YAML files that aren't dicts.""" + d = tmp_path / "catalog" + d.mkdir() + (d / "list.yaml").write_text(yaml.safe_dump(["a", "b", "c"])) + + from imas_codex.sn.catalog_import import import_catalog + + result = import_catalog(d, dry_run=True) + + assert result.imported == 0 + assert len(result.errors) == 1 + assert "not a YAML mapping" in result.errors[0] + + def test_handles_incomplete_entry(self, tmp_path: Path) -> None: + """Should report errors for entries missing required fields.""" + d = tmp_path / "catalog" + d.mkdir() + incomplete = {"name": "test", "kind": "scalar"} # missing required fields + (d / "incomplete.yaml").write_text(yaml.safe_dump(incomplete)) + + from imas_codex.sn.catalog_import import import_catalog + + result = import_catalog(d, dry_run=True) + + assert result.imported == 0 + assert len(result.errors) == 1 + assert "incomplete.yaml" in result.errors[0] + + def test_empty_directory(self, tmp_path: Path) -> None: + """Empty directory should return empty result.""" + d = tmp_path / "empty" + d.mkdir() + + from imas_codex.sn.catalog_import import import_catalog + + result = import_catalog(d, dry_run=True) + + assert result.imported == 0 + assert result.skipped == 0 + assert len(result.errors) == 0 + assert len(result.entries) == 0 + + +# ============================================================================= +# Graph write semantics tests +# ============================================================================= + + +class TestWriteCatalogEntries: + """Test that _write_catalog_entries produces correct Cypher.""" + + def _call_write(self, entries: list[dict], mock_gc: MagicMock) -> int: + """Call _write_catalog_entries with a mocked GraphClient.""" + with patch("imas_codex.graph.client.GraphClient") as MockGC: + MockGC.return_value.__enter__ = MagicMock(return_value=mock_gc) + MockGC.return_value.__exit__ = MagicMock(return_value=False) + + from imas_codex.sn.catalog_import import _write_catalog_entries + + return _write_catalog_entries(entries) + + def test_catalog_fields_overwrite(self) -> None: + """Catalog-owned fields should use direct SET, not coalesce.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + entries = [ + { + "id": "electron_temperature", + "description": "Te profile", + "documentation": "Rich docs here", + "kind": "scalar", + "units": "eV", + "tags": ["core"], + "links": None, + "imas_paths": None, + "validity_domain": "core", + "constraints": None, + "physics_domain": "core_plasma_physics", + "review_status": "accepted", + "source_type": "dd", + "physical_base": "temperature", + "subject": "electron", + "component": None, + "coordinate": None, + "position": None, + "process": None, + } + ] + self._call_write(entries, mock_gc) + + merge_call = mock_gc.query.call_args_list[0] + cypher = merge_call[0][0] + + # Catalog-owned fields should NOT use coalesce — direct SET + assert "sn.description = b.description" in cypher + assert "sn.documentation = b.documentation" in cypher + assert "sn.kind = b.kind" in cypher + assert "sn.tags = b.tags" in cypher + assert "sn.validity_domain = b.validity_domain" in cypher + assert "sn.physical_base = b.physical_base" in cypher + + # review_status should be hardcoded to 'accepted' + assert "sn.review_status = 'accepted'" in cypher + + # imported_at should be set + assert "sn.imported_at = datetime()" in cypher + + # Graph-only fields should use coalesce (preserve existing) + assert "coalesce(sn.embedding" in cypher + assert "coalesce(sn.model" in cypher + assert "coalesce(sn.generated_at" in cypher + assert "coalesce(sn.created_at, datetime())" in cypher + + def test_unit_relationship_created(self) -> None: + """Entries with units should create CANONICAL_UNITS relationship.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + entries = [ + { + "id": "electron_temperature", + "description": "Te", + "documentation": None, + "kind": "scalar", + "units": "eV", + "tags": None, + "links": None, + "imas_paths": None, + "validity_domain": None, + "constraints": None, + "physics_domain": None, + "review_status": "accepted", + "source_type": "dd", + "physical_base": "temperature", + "subject": "electron", + "component": None, + "coordinate": None, + "position": None, + "process": None, + } + ] + self._call_write(entries, mock_gc) + + unit_calls = [ + call + for call in mock_gc.query.call_args_list + if "CANONICAL_UNITS" in str(call) + ] + assert len(unit_calls) >= 1 + unit_cypher = unit_calls[0][0][0] + assert "MERGE (u:Unit" in unit_cypher + assert "MERGE (sn)-[:CANONICAL_UNITS]->(u)" in unit_cypher + + def test_dd_relationship_from_imas_paths(self) -> None: + """Entries with imas_paths should create HAS_STANDARD_NAME from IMASNode.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + entries = [ + { + "id": "electron_temperature", + "description": "Te", + "documentation": None, + "kind": "scalar", + "units": "eV", + "tags": None, + "links": None, + "imas_paths": ["core_profiles/profiles_1d/electrons/temperature"], + "validity_domain": None, + "constraints": None, + "physics_domain": None, + "review_status": "accepted", + "source_type": "dd", + "physical_base": "temperature", + "subject": "electron", + "component": None, + "coordinate": None, + "position": None, + "process": None, + } + ] + self._call_write(entries, mock_gc) + + dd_calls = [ + call + for call in mock_gc.query.call_args_list + if "HAS_STANDARD_NAME" in str(call) + ] + assert len(dd_calls) >= 1 + dd_cypher = dd_calls[0][0][0] + assert "IMASNode" in dd_cypher + assert "MERGE (src)-[:HAS_STANDARD_NAME]->(sn)" in dd_cypher + + def test_no_relationships_for_empty_fields(self) -> None: + """Entries without units/imas_paths should not create those relationships.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + entries = [ + { + "id": "test_name", + "description": "Test", + "documentation": None, + "kind": "scalar", + "units": None, + "tags": None, + "links": None, + "imas_paths": None, + "validity_domain": None, + "constraints": None, + "physics_domain": None, + "review_status": "accepted", + "source_type": "manual", + "physical_base": None, + "subject": None, + "component": None, + "coordinate": None, + "position": None, + "process": None, + } + ] + self._call_write(entries, mock_gc) + + # Should only have the MERGE query — no unit or relationship queries + assert mock_gc.query.call_count == 1 # just the MERGE + + def test_empty_list_returns_zero(self) -> None: + """Empty list should return 0 without touching the graph.""" + from imas_codex.sn.catalog_import import _write_catalog_entries + + result = _write_catalog_entries([]) + assert result == 0 From 1efbc0587a8a1d481fe5849179016a19af6cbb90 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 10:21:19 +0200 Subject: [PATCH 12/37] feat(sn): add version tracking and check mode to catalog import Phase 2 of catalog feedback import: - Add catalog_commit_sha and imported_at to StandardName schema - Resolve git HEAD SHA of catalog repo at import time - Store catalog_commit_sha on each imported node - Add check_catalog() for catalog-vs-graph sync comparison - Add --check flag to sn import-catalog CLI command - Report only-in-catalog, only-in-graph, and diverged entries - 21 new tests covering SHA resolution, version tracking, idempotency, check mode, and field normalization (49 total) --- imas_codex/cli/sn.py | 77 ++++ imas_codex/schemas/standard_name.yaml | 9 + imas_codex/sn/catalog_import.py | 203 +++++++++- tests/sn/test_catalog_import.py | 513 ++++++++++++++++++++++++++ 4 files changed, 799 insertions(+), 3 deletions(-) diff --git a/imas_codex/cli/sn.py b/imas_codex/cli/sn.py index b9a1ca2a5..4f7898d3f 100644 --- a/imas_codex/cli/sn.py +++ b/imas_codex/cli/sn.py @@ -607,11 +607,18 @@ def sn_publish( ) @click.option("--tags", type=str, default=None, help="Comma-separated tag filter") @click.option("--dry-run", is_flag=True, help="Preview without writing to graph") +@click.option( + "--check", + "check_mode", + is_flag=True, + help="Compare catalog vs graph without importing; report sync status", +) @click.option("-v", "--verbose", is_flag=True, help="Enable verbose logging") def sn_import_catalog( catalog_dir: str, tags: str | None, dry_run: bool, + check_mode: bool, verbose: bool, ) -> None: """Import reviewed catalog entries into the graph. @@ -621,11 +628,15 @@ def sn_import_catalog( the imas-standard-names catalog model, derives grammar fields, and MERGEs into the graph with review_status='accepted'. + \b + Use --check to compare catalog vs graph without importing. + \b Examples: imas-codex sn import-catalog --catalog-dir ../imas-standard-names-catalog/standard_names imas-codex sn import-catalog --catalog-dir --dry-run imas-codex sn import-catalog --catalog-dir --tags equilibrium,core-physics + imas-codex sn import-catalog --catalog-dir --check """ from pathlib import Path @@ -634,6 +645,69 @@ def sn_import_catalog( tag_filter = [t.strip() for t in tags.split(",") if t.strip()] if tags else None + # -- Check mode -- + if check_mode: + console.print("\n[bold]Standard Name Catalog Check[/bold]") + console.print(f" Catalog: {catalog_dir}") + if tag_filter: + console.print(f" Tag filter: {', '.join(tag_filter)}") + console.print("") + + try: + from imas_codex.sn.catalog_import import check_catalog + + cr = check_catalog( + catalog_dir=Path(catalog_dir), + tag_filter=tag_filter, + ) + except ImportError as e: + console.print( + f"[red]Missing dependency:[/red] {e}\n" + "Install with: uv pip install imas-standard-names" + ) + raise SystemExit(1) from e + except Exception as e: + console.print(f"[red]Check error:[/red] {e}") + raise SystemExit(1) from e + + # Print check results + if cr.catalog_commit_sha: + console.print(f" Catalog SHA: {cr.catalog_commit_sha[:12]}") + if cr.graph_commit_sha: + console.print(f" Graph SHA: {cr.graph_commit_sha[:12]}") + if cr.catalog_commit_sha and cr.graph_commit_sha: + if cr.catalog_commit_sha == cr.graph_commit_sha: + console.print(" [green]SHAs match[/green]") + else: + console.print(" [yellow]SHAs differ[/yellow]") + console.print("") + + console.print(f" In sync: [green]{cr.in_sync}[/green]") + if cr.only_in_catalog: + console.print( + f" Only in catalog: [yellow]{len(cr.only_in_catalog)}[/yellow]" + ) + for name in cr.only_in_catalog[:10]: + console.print(f" + {name}") + if len(cr.only_in_catalog) > 10: + console.print(f" ... and {len(cr.only_in_catalog) - 10} more") + if cr.only_in_graph: + console.print(f" Only in graph: [yellow]{len(cr.only_in_graph)}[/yellow]") + for name in cr.only_in_graph[:10]: + console.print(f" - {name}") + if len(cr.only_in_graph) > 10: + console.print(f" ... and {len(cr.only_in_graph) - 10} more") + if cr.diverged: + console.print(f" Diverged: [red]{len(cr.diverged)}[/red]") + for item in cr.diverged[:10]: + fields = ", ".join(item["fields"].keys()) + console.print(f" ~ {item['name']} ({fields})") + if len(cr.diverged) > 10: + console.print(f" ... and {len(cr.diverged) - 10} more") + if not cr.only_in_catalog and not cr.only_in_graph and not cr.diverged: + console.print("\n [green]✓ Catalog and graph are in sync[/green]") + return + console.print("\n[bold]Standard Name Catalog Import[/bold]") console.print(f" Catalog: {catalog_dir}") if tag_filter: @@ -661,6 +735,9 @@ def sn_import_catalog( raise SystemExit(1) from e # Print results + if result.catalog_commit_sha: + console.print(f" Catalog SHA: {result.catalog_commit_sha[:12]}") + if result.errors: console.print(f" [red]Errors: {len(result.errors)}[/red]") for err in result.errors[:10]: diff --git a/imas_codex/schemas/standard_name.yaml b/imas_codex/schemas/standard_name.yaml index 269d8e15d..187aeea03 100644 --- a/imas_codex/schemas/standard_name.yaml +++ b/imas_codex/schemas/standard_name.yaml @@ -189,3 +189,12 @@ classes: description: Physics domain classification (equilibrium, transport, etc.) process: description: Physical process qualifier (ohmic, bootstrap, etc.) + imported_at: + description: >- + ISO 8601 timestamp when this entry was imported from the catalog. + Set by ``sn import-catalog`` on each import. + range: datetime + catalog_commit_sha: + description: >- + Git commit SHA of the catalog repo at import time. + Enables sync-status checking between graph and catalog. diff --git a/imas_codex/sn/catalog_import.py b/imas_codex/sn/catalog_import.py index 1cdfd0a56..7e69f2930 100644 --- a/imas_codex/sn/catalog_import.py +++ b/imas_codex/sn/catalog_import.py @@ -8,6 +8,7 @@ from __future__ import annotations import logging +import subprocess from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -24,6 +25,44 @@ class ImportResult: skipped: int = 0 errors: list[str] = field(default_factory=list) entries: list[dict[str, Any]] = field(default_factory=list) + catalog_commit_sha: str | None = None + + +@dataclass +class CheckResult: + """Summary of a catalog-vs-graph sync check.""" + + only_in_catalog: list[str] = field(default_factory=list) + only_in_graph: list[str] = field(default_factory=list) + diverged: list[dict[str, Any]] = field(default_factory=list) + in_sync: int = 0 + catalog_commit_sha: str | None = None + graph_commit_sha: str | None = None + + +def _resolve_catalog_sha(catalog_dir: Path) -> str | None: + """Resolve the git HEAD SHA of the catalog directory. + + Returns the 40-character commit SHA, or None if the directory + is not inside a git repository. + """ + try: + result = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=str(catalog_dir), + capture_output=True, + text=True, + timeout=10, + ) + if result.returncode == 0: + sha = result.stdout.strip() + logger.debug("Catalog commit SHA: %s", sha) + return sha + logger.debug("git rev-parse failed: %s", result.stderr.strip()) + return None + except (FileNotFoundError, subprocess.TimeoutExpired) as exc: + logger.debug("Could not resolve catalog SHA: %s", exc) + return None def _parse_grammar_fields(name: str) -> dict[str, str | None]: @@ -100,7 +139,10 @@ def _catalog_entry_to_dict(entry: Any) -> dict[str, Any]: } -def _write_catalog_entries(entries: list[dict[str, Any]]) -> int: +def _write_catalog_entries( + entries: list[dict[str, Any]], + catalog_commit_sha: str | None = None, +) -> int: """Write catalog entries to graph with catalog-authoritative semantics. Catalog-owned fields are SET directly (overwrite). @@ -112,6 +154,10 @@ def _write_catalog_entries(entries: list[dict[str, Any]]) -> int: from imas_codex.graph.client import GraphClient + # Inject catalog_commit_sha into each entry for Cypher parameter access + for e in entries: + e["catalog_commit_sha"] = catalog_commit_sha + with GraphClient() as gc: # MERGE StandardName nodes — catalog fields overwrite, graph-only preserved gc.query( @@ -130,6 +176,7 @@ def _write_catalog_entries(entries: list[dict[str, Any]]) -> int: sn.physics_domain = b.physics_domain, sn.review_status = 'accepted', sn.imported_at = datetime(), + sn.catalog_commit_sha = b.catalog_commit_sha, sn.physical_base = b.physical_base, sn.subject = b.subject, sn.component = b.component, @@ -218,7 +265,12 @@ def import_catalog( from pydantic import TypeAdapter ta = TypeAdapter(StandardNameEntry) - result = ImportResult() + # Resolve catalog commit SHA for version tracking + catalog_sha = _resolve_catalog_sha(catalog_dir) + if catalog_sha: + logger.info("Catalog commit SHA: %s", catalog_sha) + + result = ImportResult(catalog_commit_sha=catalog_sha) # Collect YAML files yaml_files = sorted( @@ -272,8 +324,153 @@ def import_catalog( result.imported = len(prepared) logger.info("Dry run: would import %d entries", len(prepared)) else: - written = _write_catalog_entries(prepared) + written = _write_catalog_entries(prepared, catalog_commit_sha=catalog_sha) result.imported = written logger.info("Imported %d entries to graph", written) return result + + +# -- Fields compared during check mode (catalog-owned, excluding grammar fields) -- +_CHECK_FIELDS = ( + "description", + "documentation", + "kind", + "units", + "tags", + "imas_paths", + "validity_domain", + "constraints", + "physics_domain", +) + + +def check_catalog( + catalog_dir: Path, + tag_filter: list[str] | None = None, +) -> CheckResult: + """Compare catalog entries against graph without importing. + + Returns a :class:`CheckResult` describing which entries are only in + the catalog, only in the graph, or present in both but with differing + field values. + + Parameters + ---------- + catalog_dir: + Path to directory containing YAML catalog entries. + tag_filter: + If provided, only check entries whose tags overlap with this list. + + Returns + ------- + CheckResult with sync status details. + """ + import yaml + from imas_standard_names.catalog.edit import StandardNameEntry + from pydantic import TypeAdapter + + from imas_codex.graph.client import GraphClient + + ta = TypeAdapter(StandardNameEntry) + catalog_sha = _resolve_catalog_sha(catalog_dir) + result = CheckResult(catalog_commit_sha=catalog_sha) + + # Parse catalog entries + yaml_files = sorted( + p + for p in catalog_dir.rglob("*") + if p.suffix in (".yml", ".yaml") and p.is_file() + ) + + catalog_entries: dict[str, dict[str, Any]] = {} + for yaml_path in yaml_files: + try: + with open(yaml_path) as f: + data = yaml.safe_load(f) + if not isinstance(data, dict): + continue + entry = ta.validate_python(data) + except Exception: + continue + + # Apply tag filter + if tag_filter: + entry_tags = {str(t) for t in entry.tags} if entry.tags else set() + if not entry_tags.intersection(tag_filter): + continue + + graph_dict = _catalog_entry_to_dict(entry) + catalog_entries[graph_dict["id"]] = graph_dict + + if not catalog_entries: + return result + + # Fetch graph entries + with GraphClient() as gc: + rows = gc.query( + """ + MATCH (sn:StandardName) + WHERE sn.review_status = 'accepted' + RETURN sn.id AS id, + sn.description AS description, + sn.documentation AS documentation, + sn.kind AS kind, + sn.canonical_units AS units, + sn.tags AS tags, + sn.imas_paths AS imas_paths, + sn.validity_domain AS validity_domain, + sn.constraints AS constraints, + sn.physics_domain AS physics_domain, + sn.catalog_commit_sha AS catalog_commit_sha + """ + ) + + graph_entries: dict[str, dict[str, Any]] = {} + graph_sha: str | None = None + for row in rows: + graph_entries[row["id"]] = dict(row) + if row.get("catalog_commit_sha") and not graph_sha: + graph_sha = row["catalog_commit_sha"] + + result.graph_commit_sha = graph_sha + + # Compare + catalog_names = set(catalog_entries.keys()) + graph_names = set(graph_entries.keys()) + + result.only_in_catalog = sorted(catalog_names - graph_names) + result.only_in_graph = sorted(graph_names - catalog_names) + + for name in sorted(catalog_names & graph_names): + cat = catalog_entries[name] + graph = graph_entries[name] + + diffs: dict[str, Any] = {} + for fld in _CHECK_FIELDS: + cat_val = _normalize_field(cat.get(fld)) + graph_val = _normalize_field(graph.get(fld)) + if cat_val != graph_val: + diffs[fld] = {"catalog": cat_val, "graph": graph_val} + + if diffs: + result.diverged.append({"name": name, "fields": diffs}) + else: + result.in_sync += 1 + + return result + + +def _normalize_field(val: Any) -> Any: + """Normalize a field value for comparison. + + Converts lists to sorted tuples, None-like values to None, + and strings to stripped strings. + """ + if val is None: + return None + if isinstance(val, list): + return tuple(sorted(str(v) for v in val)) if val else None + if isinstance(val, str): + return val.strip() if val.strip() else None + return val diff --git a/tests/sn/test_catalog_import.py b/tests/sn/test_catalog_import.py index d7c72b9c6..d5c351555 100644 --- a/tests/sn/test_catalog_import.py +++ b/tests/sn/test_catalog_import.py @@ -588,3 +588,516 @@ def test_empty_list_returns_zero(self) -> None: result = _write_catalog_entries([]) assert result == 0 + + +# ============================================================================= +# Phase 2: Version tracking tests +# ============================================================================= + + +class TestResolveCatalogSha: + """Tests for _resolve_catalog_sha().""" + + def test_returns_sha_in_git_repo(self, tmp_path: Path) -> None: + """Should return a 40-char SHA when run in a git repo.""" + from imas_codex.sn.catalog_import import _resolve_catalog_sha + + # Use the project repo itself as the catalog dir + project_root = Path(__file__).resolve().parents[2] + sha = _resolve_catalog_sha(project_root) + assert sha is not None + assert len(sha) == 40 + assert all(c in "0123456789abcdef" for c in sha) + + def test_returns_none_for_non_git_dir(self, tmp_path: Path) -> None: + """Should return None for a directory that isn't a git repo.""" + from imas_codex.sn.catalog_import import _resolve_catalog_sha + + sha = _resolve_catalog_sha(tmp_path) + assert sha is None + + def test_returns_none_when_git_not_found(self, tmp_path: Path) -> None: + """Should return None when git binary is missing.""" + from imas_codex.sn.catalog_import import _resolve_catalog_sha + + with patch("imas_codex.sn.catalog_import.subprocess.run") as mock_run: + mock_run.side_effect = FileNotFoundError("git not found") + sha = _resolve_catalog_sha(tmp_path) + assert sha is None + + +class TestVersionTracking: + """Tests for catalog_commit_sha propagation through the import pipeline.""" + + def test_sha_in_cypher_batch(self) -> None: + """_write_catalog_entries should inject catalog_commit_sha into each entry.""" + from imas_codex.sn.catalog_import import _write_catalog_entries + + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + entries = [ + { + "id": "electron_temperature", + "description": "Te", + "documentation": None, + "kind": "scalar", + "units": "eV", + "tags": None, + "links": None, + "imas_paths": None, + "validity_domain": None, + "constraints": None, + "physics_domain": None, + "review_status": "accepted", + "source_type": "dd", + "physical_base": "temperature", + "subject": "electron", + "component": None, + "coordinate": None, + "position": None, + "process": None, + } + ] + + test_sha = "abc123def456" * 3 + "abcd" # 40 chars + + with patch("imas_codex.graph.client.GraphClient") as MockGC: + MockGC.return_value.__enter__ = MagicMock(return_value=mock_gc) + MockGC.return_value.__exit__ = MagicMock(return_value=False) + + _write_catalog_entries(entries, catalog_commit_sha=test_sha) + + # Verify the SHA was injected into the entry dicts + assert entries[0]["catalog_commit_sha"] == test_sha + + # Verify the Cypher includes catalog_commit_sha + merge_call = mock_gc.query.call_args_list[0] + cypher = merge_call[0][0] + assert "catalog_commit_sha" in cypher + + def test_sha_none_when_not_provided(self) -> None: + """When no SHA is provided, entries should get catalog_commit_sha=None.""" + from imas_codex.sn.catalog_import import _write_catalog_entries + + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + entries = [ + { + "id": "test_name", + "description": "Test", + "documentation": None, + "kind": "scalar", + "units": None, + "tags": None, + "links": None, + "imas_paths": None, + "validity_domain": None, + "constraints": None, + "physics_domain": None, + "review_status": "accepted", + "source_type": "manual", + "physical_base": None, + "subject": None, + "component": None, + "coordinate": None, + "position": None, + "process": None, + } + ] + + with patch("imas_codex.graph.client.GraphClient") as MockGC: + MockGC.return_value.__enter__ = MagicMock(return_value=mock_gc) + MockGC.return_value.__exit__ = MagicMock(return_value=False) + + _write_catalog_entries(entries) + + assert entries[0]["catalog_commit_sha"] is None + + def test_import_result_contains_sha(self, catalog_dir: Path) -> None: + """import_catalog() should populate catalog_commit_sha on the result.""" + from imas_codex.sn.catalog_import import import_catalog + + test_sha = "a" * 40 + + with ( + patch( + "imas_codex.sn.catalog_import._resolve_catalog_sha", + return_value=test_sha, + ), + patch( + "imas_codex.sn.catalog_import._write_catalog_entries", return_value=2 + ), + ): + result = import_catalog(catalog_dir=catalog_dir) + + assert result.catalog_commit_sha == test_sha + + def test_import_result_sha_none_for_non_git(self, catalog_dir: Path) -> None: + """import_catalog() should have sha=None when dir is not a git repo.""" + from imas_codex.sn.catalog_import import import_catalog + + with ( + patch( + "imas_codex.sn.catalog_import._resolve_catalog_sha", + return_value=None, + ), + patch( + "imas_codex.sn.catalog_import._write_catalog_entries", return_value=2 + ), + ): + result = import_catalog(catalog_dir=catalog_dir) + + assert result.catalog_commit_sha is None + + +class TestImportIdempotency: + """Tests that re-importing the same catalog produces identical results.""" + + def test_double_import_same_entries(self, catalog_dir: Path) -> None: + """Importing the same catalog twice should produce same entry count.""" + from imas_codex.sn.catalog_import import import_catalog + + with patch( + "imas_codex.sn.catalog_import._write_catalog_entries", return_value=2 + ) as mock_write: + r1 = import_catalog(catalog_dir=catalog_dir) + r2 = import_catalog(catalog_dir=catalog_dir) + + assert r1.imported == r2.imported + assert len(r1.entries) == len(r2.entries) + # Both calls should produce identical entry dicts + for e1, e2 in zip(r1.entries, r2.entries, strict=False): + assert e1["id"] == e2["id"] + assert e1["description"] == e2["description"] + assert mock_write.call_count == 2 + + def test_idempotent_entry_dicts(self, catalog_dir: Path) -> None: + """Entry dicts from two imports of the same catalog should be identical.""" + from imas_codex.sn.catalog_import import import_catalog + + with patch( + "imas_codex.sn.catalog_import._write_catalog_entries", return_value=2 + ): + r1 = import_catalog(catalog_dir=catalog_dir) + r2 = import_catalog(catalog_dir=catalog_dir) + + # Compare each field (excluding mutable fields like catalog_commit_sha) + for e1, e2 in zip(r1.entries, r2.entries, strict=False): + for key in ( + "id", + "description", + "documentation", + "kind", + "units", + "tags", + "imas_paths", + "validity_domain", + "constraints", + "physics_domain", + "review_status", + "source_type", + "physical_base", + "subject", + "component", + "coordinate", + "position", + "process", + ): + assert e1[key] == e2[key], f"Mismatch on field '{key}'" + + +class TestCheckMode: + """Tests for check_catalog() — the --check sync comparison.""" + + def test_all_in_sync(self, catalog_dir: Path) -> None: + """When graph matches catalog exactly, in_sync should equal entry count.""" + from imas_codex.sn.catalog_import import check_catalog + + # Build graph rows that match catalog exactly + graph_rows = [ + { + "id": "electron_temperature", + "description": "Electron temperature", + "documentation": "The electron temperature Te is measured by Thomson scattering.", + "kind": "scalar", + "units": "eV", + "tags": None, + "imas_paths": ["core_profiles/profiles_1d/electrons/temperature"], + "validity_domain": "core plasma", + "constraints": ["T_e > 0"], + "physics_domain": "core_plasma_physics", + "catalog_commit_sha": "a" * 40, + }, + { + "id": "plasma_current", + "description": "Plasma current", + "documentation": "Total toroidal plasma current.", + "kind": "scalar", + "units": "A", + "tags": None, + "imas_paths": None, + "validity_domain": None, + "constraints": None, + "physics_domain": "equilibrium", + "catalog_commit_sha": "a" * 40, + }, + ] + + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=graph_rows) + + with ( + patch("imas_codex.graph.client.GraphClient") as MockGC, + patch( + "imas_codex.sn.catalog_import._resolve_catalog_sha", + return_value="a" * 40, + ), + ): + MockGC.return_value.__enter__ = MagicMock(return_value=mock_gc) + MockGC.return_value.__exit__ = MagicMock(return_value=False) + + cr = check_catalog(catalog_dir=catalog_dir) + + assert cr.in_sync == 2 + assert cr.only_in_catalog == [] + assert cr.only_in_graph == [] + assert cr.diverged == [] + assert cr.catalog_commit_sha == "a" * 40 + assert cr.graph_commit_sha == "a" * 40 + + def test_only_in_catalog(self, catalog_dir: Path) -> None: + """Entries in catalog but not graph should appear in only_in_catalog.""" + from imas_codex.sn.catalog_import import check_catalog + + # Graph has no entries + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + with ( + patch("imas_codex.graph.client.GraphClient") as MockGC, + patch( + "imas_codex.sn.catalog_import._resolve_catalog_sha", + return_value=None, + ), + ): + MockGC.return_value.__enter__ = MagicMock(return_value=mock_gc) + MockGC.return_value.__exit__ = MagicMock(return_value=False) + + cr = check_catalog(catalog_dir=catalog_dir) + + assert set(cr.only_in_catalog) == {"electron_temperature", "plasma_current"} + assert cr.in_sync == 0 + assert cr.only_in_graph == [] + + def test_only_in_graph(self, catalog_dir: Path) -> None: + """Entries in graph but not catalog should appear in only_in_graph.""" + from imas_codex.sn.catalog_import import check_catalog + + graph_rows = [ + { + "id": "electron_temperature", + "description": "Electron temperature", + "documentation": "The electron temperature Te is measured by Thomson scattering.", + "kind": "scalar", + "units": "eV", + "tags": None, + "imas_paths": ["core_profiles/profiles_1d/electrons/temperature"], + "validity_domain": "core plasma", + "constraints": ["T_e > 0"], + "physics_domain": "core_plasma_physics", + "catalog_commit_sha": None, + }, + { + "id": "plasma_current", + "description": "Plasma current", + "documentation": "Total toroidal plasma current.", + "kind": "scalar", + "units": "A", + "tags": None, + "imas_paths": None, + "validity_domain": None, + "constraints": None, + "physics_domain": "equilibrium", + "catalog_commit_sha": None, + }, + { + "id": "ion_density", + "description": "Ion density", + "documentation": "Total ion density.", + "kind": "scalar", + "units": "m^-3", + "tags": None, + "imas_paths": None, + "validity_domain": None, + "constraints": None, + "physics_domain": "core_plasma_physics", + "catalog_commit_sha": None, + }, + ] + + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=graph_rows) + + with ( + patch("imas_codex.graph.client.GraphClient") as MockGC, + patch( + "imas_codex.sn.catalog_import._resolve_catalog_sha", + return_value=None, + ), + ): + MockGC.return_value.__enter__ = MagicMock(return_value=mock_gc) + MockGC.return_value.__exit__ = MagicMock(return_value=False) + + cr = check_catalog(catalog_dir=catalog_dir) + + assert cr.only_in_graph == ["ion_density"] + assert cr.in_sync == 2 # electron_temperature and plasma_current match + + def test_diverged_entries(self, catalog_dir: Path) -> None: + """Entries with different field values should appear in diverged.""" + from imas_codex.sn.catalog_import import check_catalog + + graph_rows = [ + { + "id": "electron_temperature", + "description": "WRONG description", # differs from catalog + "documentation": "The electron temperature Te is measured by Thomson scattering.", + "kind": "scalar", + "units": "keV", # differs from catalog (eV) + "tags": None, + "imas_paths": ["core_profiles/profiles_1d/electrons/temperature"], + "validity_domain": "core plasma", + "constraints": ["T_e > 0"], + "physics_domain": "core_plasma_physics", + "catalog_commit_sha": None, + }, + { + "id": "plasma_current", + "description": "Plasma current", + "documentation": "Total toroidal plasma current.", + "kind": "scalar", + "units": "A", + "tags": None, + "imas_paths": None, + "validity_domain": None, + "constraints": None, + "physics_domain": "equilibrium", + "catalog_commit_sha": None, + }, + ] + + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=graph_rows) + + with ( + patch("imas_codex.graph.client.GraphClient") as MockGC, + patch( + "imas_codex.sn.catalog_import._resolve_catalog_sha", + return_value=None, + ), + ): + MockGC.return_value.__enter__ = MagicMock(return_value=mock_gc) + MockGC.return_value.__exit__ = MagicMock(return_value=False) + + cr = check_catalog(catalog_dir=catalog_dir) + + assert cr.in_sync == 1 # plasma_current matches + assert len(cr.diverged) == 1 + assert cr.diverged[0]["name"] == "electron_temperature" + assert "description" in cr.diverged[0]["fields"] + assert "units" in cr.diverged[0]["fields"] + + def test_check_with_tag_filter(self, tmp_path: Path) -> None: + """Tag filter should limit which catalog entries are checked.""" + from imas_codex.sn.catalog_import import check_catalog + + # Create catalog with tagged entry + d = tmp_path / "catalog" + d.mkdir() + entry_with_tag = dict(SAMPLE_CATALOG_ENTRY) + entry_with_tag["tags"] = ["spatial-profile"] + (d / "electron_temperature.yaml").write_text(yaml.safe_dump(entry_with_tag)) + (d / "plasma_current.yaml").write_text( + yaml.safe_dump(SAMPLE_CATALOG_ENTRY_MINIMAL) + ) + + # Graph has no entries + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + with ( + patch("imas_codex.graph.client.GraphClient") as MockGC, + patch( + "imas_codex.sn.catalog_import._resolve_catalog_sha", + return_value=None, + ), + ): + MockGC.return_value.__enter__ = MagicMock(return_value=mock_gc) + MockGC.return_value.__exit__ = MagicMock(return_value=False) + + cr = check_catalog( + catalog_dir=d, + tag_filter=["spatial-profile"], + ) + + # Only electron_temperature has the tag — plasma_current should be filtered out + assert cr.only_in_catalog == ["electron_temperature"] + assert cr.in_sync == 0 + + def test_check_empty_catalog(self, tmp_path: Path) -> None: + """Empty catalog directory should return empty CheckResult.""" + from imas_codex.sn.catalog_import import check_catalog + + d = tmp_path / "empty_catalog" + d.mkdir() + + with patch( + "imas_codex.sn.catalog_import._resolve_catalog_sha", + return_value=None, + ): + cr = check_catalog(catalog_dir=d) + + assert cr.in_sync == 0 + assert cr.only_in_catalog == [] + assert cr.only_in_graph == [] + assert cr.diverged == [] + + +class TestNormalizeField: + """Tests for _normalize_field() comparison normalization.""" + + def test_none(self) -> None: + from imas_codex.sn.catalog_import import _normalize_field + + assert _normalize_field(None) is None + + def test_empty_string(self) -> None: + from imas_codex.sn.catalog_import import _normalize_field + + assert _normalize_field("") is None + assert _normalize_field(" ") is None + + def test_normal_string(self) -> None: + from imas_codex.sn.catalog_import import _normalize_field + + assert _normalize_field("hello") == "hello" + assert _normalize_field(" hello ") == "hello" + + def test_empty_list(self) -> None: + from imas_codex.sn.catalog_import import _normalize_field + + assert _normalize_field([]) is None + + def test_list_sorted(self) -> None: + from imas_codex.sn.catalog_import import _normalize_field + + assert _normalize_field(["b", "a"]) == ("a", "b") + assert _normalize_field(["a", "b"]) == ("a", "b") + + def test_numeric_passthrough(self) -> None: + from imas_codex.sn.catalog_import import _normalize_field + + assert _normalize_field(42) == 42 + assert _normalize_field(3.14) == 3.14 From df6240b427843db80bcae83fa3cc02bdef7d7dad Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 10:28:55 +0200 Subject: [PATCH 13/37] refactor: rename sn import-catalog to sn import MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rename CLI command for symmetry with sn publish: - @sn.command("import-catalog") → @sn.command("import") - sn_import_catalog() → sn_import() - Update help text examples and schema description - Add TestPublishImportRoundTrip with 3 round-trip tests --- imas_codex/cli/sn.py | 12 +-- imas_codex/schemas/standard_name.yaml | 2 +- tests/sn/test_catalog_import.py | 150 ++++++++++++++++++++++++++ 3 files changed, 157 insertions(+), 7 deletions(-) diff --git a/imas_codex/cli/sn.py b/imas_codex/cli/sn.py index 4f7898d3f..de16fe548 100644 --- a/imas_codex/cli/sn.py +++ b/imas_codex/cli/sn.py @@ -598,7 +598,7 @@ def sn_publish( ) -@sn.command("import-catalog") +@sn.command("import") @click.option( "--catalog-dir", type=click.Path(exists=True), @@ -614,7 +614,7 @@ def sn_publish( help="Compare catalog vs graph without importing; report sync status", ) @click.option("-v", "--verbose", is_flag=True, help="Enable verbose logging") -def sn_import_catalog( +def sn_import( catalog_dir: str, tags: str | None, dry_run: bool, @@ -633,10 +633,10 @@ def sn_import_catalog( \b Examples: - imas-codex sn import-catalog --catalog-dir ../imas-standard-names-catalog/standard_names - imas-codex sn import-catalog --catalog-dir --dry-run - imas-codex sn import-catalog --catalog-dir --tags equilibrium,core-physics - imas-codex sn import-catalog --catalog-dir --check + imas-codex sn import --catalog-dir ../imas-standard-names-catalog/standard_names + imas-codex sn import --catalog-dir --dry-run + imas-codex sn import --catalog-dir --tags equilibrium,core-physics + imas-codex sn import --catalog-dir --check """ from pathlib import Path diff --git a/imas_codex/schemas/standard_name.yaml b/imas_codex/schemas/standard_name.yaml index 187aeea03..a691b36e0 100644 --- a/imas_codex/schemas/standard_name.yaml +++ b/imas_codex/schemas/standard_name.yaml @@ -192,7 +192,7 @@ classes: imported_at: description: >- ISO 8601 timestamp when this entry was imported from the catalog. - Set by ``sn import-catalog`` on each import. + Set by ``sn import`` on each import. range: datetime catalog_commit_sha: description: >- diff --git a/tests/sn/test_catalog_import.py b/tests/sn/test_catalog_import.py index d5c351555..bcf94b3cf 100644 --- a/tests/sn/test_catalog_import.py +++ b/tests/sn/test_catalog_import.py @@ -1101,3 +1101,153 @@ def test_numeric_passthrough(self) -> None: assert _normalize_field(42) == 42 assert _normalize_field(3.14) == 3.14 + + +class TestPublishImportRoundTrip: + """Test that published entries can be reviewed, imported, and re-imported.""" + + def test_published_entry_importable_after_review(self, tmp_path: Path) -> None: + """A published entry enriched with catalog fields should import cleanly.""" + from imas_codex.sn.catalog_import import import_catalog + from imas_codex.sn.models import SNProvenance, SNPublishEntry + from imas_codex.sn.publish import generate_yaml_entry + + # 1. Generate a published YAML entry (what `sn publish` produces) + published = SNPublishEntry( + name="electron_temperature", + kind="physical", + unit="eV", + tags=["core_profiles"], + status="drafted", + description="Electron temperature", + provenance=SNProvenance( + source="dd", + source_id="core_profiles/profiles_1d/electrons/temperature", + ids_name="core_profiles", + confidence=0.95, + ), + ) + published_yaml = generate_yaml_entry(published) + assert "electron_temperature" in published_yaml + + # 2. Simulate reviewer enriching entry into catalog format + reviewed = { + "name": "electron_temperature", + "description": "Electron temperature", + "documentation": "The electron temperature Te is measured by Thomson.", + "kind": "scalar", + "unit": "eV", + "tags": [], + "links": [], + "ids_paths": ["core_profiles/profiles_1d/electrons/temperature"], + "validity_domain": "core plasma", + "constraints": ["T_e > 0"], + "physics_domain": "core_plasma_physics", + "status": "active", + } + + catalog_dir = tmp_path / "reviewed_catalog" + catalog_dir.mkdir() + (catalog_dir / "electron_temperature.yaml").write_text(yaml.safe_dump(reviewed)) + + # 3. Import the reviewed entry + result = import_catalog(catalog_dir=catalog_dir, dry_run=True) + + assert result.imported == 1 + entry = result.entries[0] + + # 4. Verify all catalog fields map correctly into graph dict shape + assert entry["id"] == "electron_temperature" + assert entry["units"] == "eV" + assert entry["imas_paths"] == [ + "core_profiles/profiles_1d/electrons/temperature" + ] + assert entry["review_status"] == "accepted" + assert entry["source_type"] == "dd" + assert entry["physics_domain"] == "core_plasma_physics" + assert entry["validity_domain"] == "core plasma" + assert entry["constraints"] == ["T_e > 0"] + # Grammar-parsed fields should be populated + assert entry["physical_base"] == "temperature" + assert entry["subject"] == "electron" + + def test_round_trip_preserves_all_fields(self, tmp_path: Path) -> None: + """Importing the same catalog entry twice should yield identical dicts.""" + from imas_codex.sn.catalog_import import import_catalog + + catalog_dir = tmp_path / "rt_catalog" + catalog_dir.mkdir() + (catalog_dir / "electron_temperature.yaml").write_text( + yaml.safe_dump(SAMPLE_CATALOG_ENTRY) + ) + (catalog_dir / "plasma_current.yaml").write_text( + yaml.safe_dump(SAMPLE_CATALOG_ENTRY_MINIMAL) + ) + + r1 = import_catalog(catalog_dir=catalog_dir, dry_run=True) + r2 = import_catalog(catalog_dir=catalog_dir, dry_run=True) + + assert len(r1.entries) == len(r2.entries) == 2 + for e1, e2 in zip(r1.entries, r2.entries, strict=False): + assert e1 == e2, f"Mismatch for {e1.get('id')}: {e1} != {e2}" + + def test_graph_records_reimport_consistency(self, tmp_path: Path) -> None: + """graph_records_to_entries output can be re-published and re-imported.""" + from imas_codex.sn.catalog_import import import_catalog + from imas_codex.sn.publish import ( + generate_catalog_files, + graph_records_to_entries, + ) + + # Simulate graph records from a first import + graph_records = [ + { + "name": "electron_temperature", + "source_type": "dd", + "source_id": "core_profiles/profiles_1d/electrons/temperature", + "units": "eV", + "description": "Electron temperature", + "ids_name": "core_profiles", + "confidence": 0.95, + } + ] + + # Convert to publish entries and write YAML + publish_entries = graph_records_to_entries(graph_records) + assert len(publish_entries) == 1 + assert publish_entries[0].name == "electron_temperature" + + publish_dir = tmp_path / "published" + written = generate_catalog_files(publish_entries, publish_dir) + assert len(written) == 1 + + # Now create a "reviewed" catalog version from the published YAML + catalog_dir = tmp_path / "catalog_reviewed" + catalog_dir.mkdir() + reviewed = { + "name": "electron_temperature", + "description": "Electron temperature", + "documentation": "Te from core profiles.", + "kind": "scalar", + "unit": "eV", + "tags": [], + "links": [], + "ids_paths": ["core_profiles/profiles_1d/electrons/temperature"], + "validity_domain": "", + "constraints": [], + "physics_domain": "core_plasma_physics", + "status": "active", + } + (catalog_dir / "electron_temperature.yaml").write_text(yaml.safe_dump(reviewed)) + + # Import the reviewed entry + result = import_catalog(catalog_dir=catalog_dir, dry_run=True) + assert result.imported == 1 + entry = result.entries[0] + + # Verify key fields survive the full publish→review→import cycle + assert entry["id"] == "electron_temperature" + assert entry["units"] == "eV" + assert entry["source_type"] == "dd" + assert entry["review_status"] == "accepted" + assert entry["documentation"] == "Te from core profiles." From a102951a7ef38631f624e488e12e9c3fd4c44092 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 10:35:48 +0200 Subject: [PATCH 14/37] chore: migrate MCP config to project root with .vscode symlink --- .mcp.json | 9 +++++++-- .vscode/mcp.json | 23 +---------------------- 2 files changed, 8 insertions(+), 24 deletions(-) mode change 100644 => 120000 .vscode/mcp.json diff --git a/.mcp.json b/.mcp.json index 96e209685..b485fee81 100644 --- a/.mcp.json +++ b/.mcp.json @@ -1,11 +1,16 @@ { - "mcpServers": { + "servers": { + "imas-dd": { + "type": "stdio", + "command": "uv", + "args": ["run", "imas-codex", "serve", "--dd-only"] + }, "codex": { "type": "stdio", "command": "uv", "args": ["run", "imas-codex", "serve"] }, - "imas": { + "imas-prod": { "type": "http", "url": "https://imas-dd.iter.org/mcp" }, diff --git a/.vscode/mcp.json b/.vscode/mcp.json deleted file mode 100644 index b485fee81..000000000 --- a/.vscode/mcp.json +++ /dev/null @@ -1,22 +0,0 @@ -{ - "servers": { - "imas-dd": { - "type": "stdio", - "command": "uv", - "args": ["run", "imas-codex", "serve", "--dd-only"] - }, - "codex": { - "type": "stdio", - "command": "uv", - "args": ["run", "imas-codex", "serve"] - }, - "imas-prod": { - "type": "http", - "url": "https://imas-dd.iter.org/mcp" - }, - "imas-test": { - "type": "http", - "url": "https://app-imas-mcp-server-test-frc.azurewebsites.net/mcp" - } - } -} diff --git a/.vscode/mcp.json b/.vscode/mcp.json new file mode 120000 index 000000000..c67157dc4 --- /dev/null +++ b/.vscode/mcp.json @@ -0,0 +1 @@ +../.mcp.json \ No newline at end of file From a6d499a3dd063e511eff12c30e5487aabc82bdea Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 10:53:20 +0200 Subject: [PATCH 15/37] feat(sn): lossless publish pipeline with rich catalog fields - Add rich fields to SNPublishEntry: documentation, links, ids_paths, constraints, validity_domain; change kind default from 'physical' to 'scalar' - Rewrite generate_yaml_entry() to include all catalog fields; empty optional fields omitted from output; links serialized as [{name: ...}] - Update generate_catalog_files() to group entries into tag-based subdirectories (primary tag / name.yaml); untagged -> unscoped/ - Update graph_records_to_entries() to carry rich fields through from graph records without loss - Update get_validated_standard_names() with review_status filter, CANONICAL_UNITS traversal, and full rich-field RETURN clause - Add update_review_status() to graph_ops.py for batch status updates - Call update_review_status() in CLI sn publish after YAML generation - Update check_catalog_duplicates() to use rglob for subdirectory scan - Update tests: fix kind references, add directory structure tests, rich-field round-trip tests, subdirectory duplicate detection (46 tests) --- imas_codex/cli/sn.py | 19 ++++ imas_codex/sn/graph_ops.py | 88 +++++++++++++++--- imas_codex/sn/models.py | 16 +++- imas_codex/sn/publish.py | 54 ++++++++--- tests/sn/test_publish.py | 177 ++++++++++++++++++++++++++++++++++--- 5 files changed, 318 insertions(+), 36 deletions(-) diff --git a/imas_codex/cli/sn.py b/imas_codex/cli/sn.py index de16fe548..9805e8b35 100644 --- a/imas_codex/cli/sn.py +++ b/imas_codex/cli/sn.py @@ -290,6 +290,12 @@ async def _run(stop_event, service_monitor): help="JSON report output path", ) @click.option("-v", "--verbose", is_flag=True, help="Enable verbose logging") +@click.option( + "--reviewer-model", + type=str, + default=None, + help="Frontier model for quality scoring (e.g. anthropic/claude-opus-4-6)", +) def sn_benchmark( source: str, ids_filter: str | None, @@ -301,6 +307,7 @@ def sn_benchmark( temperature: float, output: str | None, verbose: bool, + reviewer_model: str | None, ) -> None: """Benchmark LLM models on standard name generation. @@ -338,6 +345,7 @@ def sn_benchmark( max_candidates=max_candidates, runs_per_model=runs, temperature=temperature, + reviewer_model=reviewer_model, ) console.print("[bold]SN Benchmark[/bold]") @@ -350,6 +358,8 @@ def sn_benchmark( console.print(f" Max candidates: {max_candidates}") console.print(f" Runs per model: {runs}") console.print(f" Temperature: {temperature}") + if reviewer_model: + console.print(f" Reviewer model: {reviewer_model}") console.print() from imas_codex.cli.utils import run_async @@ -573,6 +583,15 @@ def sn_publish( written = generate_catalog_files(entries, out) console.print(f"\n[green]Wrote {len(written)} YAML files to {out}[/green]") + # Step 6b: Update review_status in graph + from imas_codex.sn.graph_ops import update_review_status + + published_names = [e.name for e in entries] + updated = update_review_status(published_names, status="published") + console.print( + f" Updated [bold]{updated}[/bold] names to review_status='published'" + ) + # Step 7: Optionally create PRs if create_pr: for batch in batches: diff --git a/imas_codex/sn/graph_ops.py b/imas_codex/sn/graph_ops.py index 96d1fb5cf..570cc7c8a 100644 --- a/imas_codex/sn/graph_ops.py +++ b/imas_codex/sn/graph_ops.py @@ -274,13 +274,14 @@ def write_standard_names(names: list[dict[str, Any]]) -> int: def get_validated_standard_names( ids_filter: str | None = None, confidence_min: float = 0.0, + review_status: str = "drafted", ) -> list[dict[str, Any]]: """Read validated StandardName nodes and their provenance. - Queries all StandardName nodes, joining through ``HAS_STANDARD_NAME`` - to find source entities and their parent IDS. Uses ``collect()`` - to avoid row duplication when a name has multiple sources (takes - the first source). + Queries StandardName nodes with the given ``review_status``, joining + through ``HAS_STANDARD_NAME`` to find source entities and their parent IDS, + and through ``CANONICAL_UNITS`` to find the unit node. Uses ``collect()`` + to avoid row duplication when a name has multiple sources (takes the first). Parameters ---------- @@ -291,30 +292,41 @@ def get_validated_standard_names( confidence_min: Minimum confidence threshold. Nodes without a ``confidence`` property are treated as 1.0 (grammar-validated). + review_status: + Filter by ``review_status`` property (default ``"drafted"``). Returns ------- - list of dicts with keys: name, description, source, source_path, - canonical_units, confidence, ids_name. + list of dicts with keys: name, description, documentation, kind, + canonical_units, tags, links, ids_paths, constraints, validity_domain, + confidence, model, source, source_path, ids_name, physical_base, + subject, component, coordinate, position, process, source_ids_names. """ with GraphClient() as gc: - params: dict[str, Any] = {"confidence_min": confidence_min} + params: dict[str, Any] = { + "confidence_min": confidence_min, + "review_status": review_status, + } # Collect source info — use HAS_STANDARD_NAME (entity → concept) cypher = """ MATCH (sn:StandardName) - WHERE coalesce(sn.confidence, 1.0) >= $confidence_min + WHERE sn.review_status = $review_status + AND coalesce(sn.confidence, 1.0) >= $confidence_min OPTIONAL MATCH (src)-[:HAS_STANDARD_NAME]->(sn) OPTIONAL MATCH (src)-[:IN_IDS]->(ids:IDS) + OPTIONAL MATCH (sn)-[:CANONICAL_UNITS]->(u:Unit) WITH sn, collect(DISTINCT src.id)[0] AS first_source, - collect(DISTINCT ids.id)[0] AS first_ids + collect(DISTINCT ids.id)[0] AS first_ids, + collect(DISTINCT ids.id) AS all_ids, + u """ if ids_filter: # Re-check: at least one HAS_STANDARD_NAME source must be in the target IDS cypher += """ - WITH sn, first_source, first_ids + WITH sn, first_source, first_ids, all_ids, u WHERE first_ids = $ids_filter """ params["ids_filter"] = ids_filter @@ -322,19 +334,67 @@ def get_validated_standard_names( cypher += """ RETURN sn.id AS name, sn.description AS description, + sn.documentation AS documentation, + sn.kind AS kind, + coalesce(u.id, sn.canonical_units, sn.units) AS canonical_units, + sn.tags AS tags, + sn.links AS links, + sn.ids_paths AS ids_paths, + sn.constraints AS constraints, + sn.validity_domain AS validity_domain, + coalesce(sn.confidence, 1.0) AS confidence, + sn.model AS model, coalesce(sn.source, sn.source_type) AS source, coalesce(sn.source_path, first_source) AS source_path, - coalesce(sn.canonical_units, sn.units) AS canonical_units, - coalesce(sn.confidence, 1.0) AS confidence, - first_ids AS ids_name + first_ids AS ids_name, + sn.physical_base AS physical_base, + sn.subject AS subject, + sn.component AS component, + sn.coordinate AS coordinate, + sn.position AS position, + sn.process AS process, + all_ids AS source_ids_names ORDER BY sn.id """ results = gc.query(cypher, **params) logger.info( - "Read %d validated standard names (ids_filter=%s, confidence_min=%.2f)", + "Read %d validated standard names (ids_filter=%s, confidence_min=%.2f, review_status=%s)", len(results), ids_filter, confidence_min, + review_status, ) return list(results) + + +def update_review_status(names: list[str], status: str = "published") -> int: + """Update review_status for a batch of StandardName nodes. + + Parameters + ---------- + names: + List of StandardName node IDs (``sn.id``) to update. + status: + New ``review_status`` value (default ``"published"``). + + Returns + ------- + Number of nodes updated. + """ + if not names: + return 0 + with GraphClient() as gc: + result = gc.query( + """ + UNWIND $names AS name + MATCH (sn:StandardName {id: name}) + SET sn.review_status = $status + RETURN count(sn) AS updated + """, + names=names, + status=status, + ) + count = result[0]["updated"] if result else 0 + logger.info("Updated review_status to '%s' for %d names", status, count) + return count diff --git a/imas_codex/sn/models.py b/imas_codex/sn/models.py index 99dc344bc..f5d023085 100644 --- a/imas_codex/sn/models.py +++ b/imas_codex/sn/models.py @@ -72,12 +72,26 @@ class SNPublishEntry(BaseModel): name: str = Field(description="The standard name") kind: str = Field( - default="physical", description="Name kind: physical or geometric" + default="scalar", description="Name kind: scalar, vector, or metadata" ) unit: str | None = Field(default=None, description="SI unit string") tags: list[str] = Field(default_factory=list, description="Classification tags") status: str = Field(default="drafted", description="Entry status") description: str = Field(default="", description="Human-readable description") + # Rich fields + documentation: str | None = Field( + default=None, description="Rich documentation with LaTeX" + ) + links: list[str] = Field(default_factory=list, description="Related standard names") + ids_paths: list[str] = Field( + default_factory=list, description="Mapped IMAS DD paths" + ) + constraints: list[str] = Field( + default_factory=list, description="Physical constraints" + ) + validity_domain: str | None = Field( + default=None, description="Physical region where valid" + ) provenance: SNProvenance = Field(description="Generation provenance") diff --git a/imas_codex/sn/publish.py b/imas_codex/sn/publish.py index 745e5f9ea..8bf2c3ea0 100644 --- a/imas_codex/sn/publish.py +++ b/imas_codex/sn/publish.py @@ -51,7 +51,8 @@ def generate_yaml_entry(entry: SNPublishEntry) -> str: """Generate YAML content for a single standard name entry. Returns a YAML string formatted to match the - ``imas-standard-names-catalog`` convention. + ``imas-standard-names-catalog`` convention. All rich fields are + included; empty/None optional fields are omitted. """ doc: dict[str, Any] = { "name": entry.name, @@ -64,6 +65,16 @@ def generate_yaml_entry(entry: SNPublishEntry) -> str: doc["status"] = entry.status if entry.description: doc["description"] = entry.description + if entry.documentation: + doc["documentation"] = entry.documentation + if entry.links: + doc["links"] = [{"name": link} for link in entry.links] + if entry.ids_paths: + doc["ids_paths"] = entry.ids_paths + if entry.constraints: + doc["constraints"] = entry.constraints + if entry.validity_domain: + doc["validity_domain"] = entry.validity_domain doc["provenance"] = { "source": entry.provenance.source, "source_id": entry.provenance.source_id, @@ -80,9 +91,10 @@ def generate_catalog_files( entries: list[SNPublishEntry], output_dir: Path, ) -> list[Path]: - """Write YAML files to *output_dir*. One file per entry. + """Write YAML files to *output_dir*, grouped by primary tag into subdirectories. - File names are ``{name}.yaml`` (e.g. ``electron_temperature.yaml``). + File names are ``{tag}/{name}.yaml`` (e.g. ``equilibrium/electron_temperature.yaml``). + Entries without tags go into ``unscoped/``. Returns list of written file paths. """ output_dir = Path(output_dir) @@ -90,8 +102,12 @@ def generate_catalog_files( written: list[Path] = [] for entry in entries: + # Group by primary tag into subdirectories + subdir = entry.tags[0] if entry.tags else "unscoped" + entry_dir = output_dir / subdir + entry_dir.mkdir(parents=True, exist_ok=True) filename = f"{entry.name}.yaml" - filepath = output_dir / filename + filepath = entry_dir / filename content = generate_yaml_entry(entry) filepath.write_text(content + "\n", encoding="utf-8") written.append(filepath) @@ -177,7 +193,8 @@ def check_catalog_duplicates( ) -> tuple[list[SNPublishEntry], list[SNPublishEntry]]: """Check for duplicates against an existing catalog directory. - Scans ``catalog_dir`` for ``.yaml`` files and reads the ``name`` + Scans ``catalog_dir`` recursively for ``.yaml`` files (including + subdirectories created by tag-based grouping) and reads the ``name`` field from each. Also detects duplicates within *entries* itself. Returns ``(new_entries, duplicate_entries)``. @@ -187,7 +204,8 @@ def check_catalog_duplicates( if catalog_dir is not None: catalog_path = Path(catalog_dir) if catalog_path.is_dir(): - for yaml_file in catalog_path.glob("*.yaml"): + # Scan both top-level and subdirectory YAML files + for yaml_file in catalog_path.rglob("*.yaml"): try: with open(yaml_file, encoding="utf-8") as f: doc = yaml.safe_load(f) @@ -230,7 +248,8 @@ def graph_records_to_entries( Handles both schema-canonical properties (``source``, ``source_path``, ``canonical_units``) and legacy write properties (``source_type``, - ``source_id``, ``units``). + ``source_id``, ``units``). Carries through all rich fields: + documentation, links, ids_paths, constraints, validity_domain, kind. """ entries: list[SNPublishEntry] = [] for rec in records: @@ -257,9 +276,17 @@ def graph_records_to_entries( description = rec.get("description") or "" + # Rich fields + documentation = rec.get("documentation") + kind = rec.get("kind") or "scalar" + links_raw = rec.get("links") or [] + ids_paths_raw = rec.get("ids_paths") or [] + constraints_raw = rec.get("constraints") or [] + validity_domain = rec.get("validity_domain") + # Build tags from available context - tags: list[str] = [] - if ids_name: + tags: list[str] = list(rec.get("tags") or []) + if not tags and ids_name: tags.append(ids_name) provenance = SNProvenance( @@ -272,11 +299,18 @@ def graph_records_to_entries( entries.append( SNPublishEntry( name=name, - kind="physical", + kind=kind, unit=unit, tags=tags, status="drafted", description=description[:500] if description else "", + documentation=documentation, + links=links_raw if isinstance(links_raw, list) else [], + ids_paths=ids_paths_raw if isinstance(ids_paths_raw, list) else [], + constraints=constraints_raw + if isinstance(constraints_raw, list) + else [], + validity_domain=validity_domain, provenance=provenance, ) ) diff --git a/tests/sn/test_publish.py b/tests/sn/test_publish.py index a416d1f02..bfb79ac35 100644 --- a/tests/sn/test_publish.py +++ b/tests/sn/test_publish.py @@ -43,7 +43,7 @@ def sample_provenance() -> SNProvenance: def sample_entry(sample_provenance: SNProvenance) -> SNPublishEntry: return SNPublishEntry( name="electron_temperature", - kind="physical", + kind="scalar", unit="eV", tags=["equilibrium", "core_profiles"], status="drafted", @@ -58,7 +58,7 @@ def sample_entries() -> list[SNPublishEntry]: return [ SNPublishEntry( name="electron_temperature", - kind="physical", + kind="scalar", unit="eV", tags=["equilibrium"], description="Electron temperature", @@ -71,7 +71,7 @@ def sample_entries() -> list[SNPublishEntry]: ), SNPublishEntry( name="electron_density", - kind="physical", + kind="scalar", unit="m^-3", tags=["core_profiles"], description="Electron density", @@ -84,7 +84,7 @@ def sample_entries() -> list[SNPublishEntry]: ), SNPublishEntry( name="plasma_current", - kind="physical", + kind="scalar", unit="A", tags=["equilibrium"], description="Plasma current", @@ -97,7 +97,7 @@ def sample_entries() -> list[SNPublishEntry]: ), SNPublishEntry( name="major_radius", - kind="geometric", + kind="vector", unit="m", tags=["equilibrium"], description="Major radius", @@ -143,14 +143,19 @@ def test_defaults(self, sample_provenance: SNProvenance) -> None: name="test_name", provenance=sample_provenance, ) - assert entry.kind == "physical" + assert entry.kind == "scalar" assert entry.status == "drafted" assert entry.tags == [] assert entry.unit is None + assert entry.documentation is None + assert entry.links == [] + assert entry.ids_paths == [] + assert entry.constraints == [] + assert entry.validity_domain is None def test_all_fields(self, sample_entry: SNPublishEntry) -> None: assert sample_entry.name == "electron_temperature" - assert sample_entry.kind == "physical" + assert sample_entry.kind == "scalar" assert sample_entry.unit == "eV" assert "equilibrium" in sample_entry.tags assert sample_entry.provenance.confidence == 0.95 @@ -199,7 +204,7 @@ def test_format(self, sample_entry: SNPublishEntry) -> None: doc = yaml.safe_load(content) assert doc["name"] == "electron_temperature" - assert doc["kind"] == "physical" + assert doc["kind"] == "scalar" assert doc["unit"] == "eV" assert doc["status"] == "drafted" assert doc["description"] == "Electron temperature profile" @@ -281,6 +286,31 @@ def test_filenames( expected = {e.name for e in sample_entries} assert names == expected + def test_directory_structure_by_tag( + self, tmp_path: Path, sample_entries: list[SNPublishEntry] + ) -> None: + """Files should be grouped into tag-based subdirectories.""" + written = generate_catalog_files(sample_entries, tmp_path) + # All entries have tags — check subdirs were created + subdirs = {p.parent.name for p in written} + assert "equilibrium" in subdirs + assert "core_profiles" in subdirs + # equilibrium tag entries go into equilibrium/ + eq_files = [p for p in written if p.parent.name == "equilibrium"] + eq_names = {p.stem for p in eq_files} + assert "electron_temperature" in eq_names + assert "plasma_current" in eq_names + assert "major_radius" in eq_names + + def test_untagged_goes_to_unscoped( + self, tmp_path: Path, sample_provenance: SNProvenance + ) -> None: + """Entries without tags go into 'unscoped/' subdirectory.""" + entry = SNPublishEntry(name="untagged_quantity", provenance=sample_provenance) + written = generate_catalog_files([entry], tmp_path) + assert len(written) == 1 + assert written[0].parent.name == "unscoped" + def test_file_content_valid_yaml( self, tmp_path: Path, sample_entry: SNPublishEntry ) -> None: @@ -367,15 +397,28 @@ def test_no_duplicates( def test_finds_catalog_duplicates( self, tmp_path: Path, sample_entries: list[SNPublishEntry] ) -> None: - # Write one existing catalog entry - (tmp_path / "electron_temperature.yaml").write_text( - yaml.safe_dump({"name": "electron_temperature", "kind": "physical"}) + # Write one existing catalog entry in a subdirectory (tag-based layout) + subdir = tmp_path / "equilibrium" + subdir.mkdir() + (subdir / "electron_temperature.yaml").write_text( + yaml.safe_dump({"name": "electron_temperature", "kind": "scalar"}) ) new, dupes = check_catalog_duplicates(sample_entries, catalog_dir=tmp_path) assert len(dupes) == 1 assert dupes[0].name == "electron_temperature" assert len(new) == len(sample_entries) - 1 + def test_finds_catalog_duplicates_top_level( + self, tmp_path: Path, sample_entries: list[SNPublishEntry] + ) -> None: + # Also detect duplicates in flat (top-level) YAML files + (tmp_path / "electron_temperature.yaml").write_text( + yaml.safe_dump({"name": "electron_temperature", "kind": "scalar"}) + ) + new, dupes = check_catalog_duplicates(sample_entries, catalog_dir=tmp_path) + assert len(dupes) == 1 + assert dupes[0].name == "electron_temperature" + def test_finds_within_batch_duplicates( self, sample_provenance: SNProvenance ) -> None: @@ -482,3 +525,115 @@ def test_tags_empty_when_no_ids(self) -> None: ] entries = graph_records_to_entries(records) assert entries[0].tags == [] + + def test_rich_fields_carried_through(self) -> None: + """Rich fields (documentation, links, ids_paths, constraints, validity_domain, kind) are preserved.""" + records = [ + { + "name": "ion_temperature", + "description": "Ion temperature", + "documentation": "Ion temperature $T_i$ in eV. See also electron_temperature.", + "kind": "scalar", + "source": "dd", + "source_path": "core_profiles/profiles_1d/ion/temperature", + "canonical_units": "eV", + "confidence": 0.9, + "ids_name": "core_profiles", + "tags": ["core_profiles", "kinetics"], + "links": ["electron_temperature", "ion_density"], + "ids_paths": ["core_profiles/profiles_1d/ion/temperature"], + "constraints": ["T_i > 0"], + "validity_domain": "core plasma", + } + ] + entries = graph_records_to_entries(records) + assert len(entries) == 1 + e = entries[0] + assert ( + e.documentation + == "Ion temperature $T_i$ in eV. See also electron_temperature." + ) + assert e.kind == "scalar" + assert e.links == ["electron_temperature", "ion_density"] + assert e.ids_paths == ["core_profiles/profiles_1d/ion/temperature"] + assert e.constraints == ["T_i > 0"] + assert e.validity_domain == "core plasma" + assert e.tags == ["core_profiles", "kinetics"] + + def test_kind_defaults_to_scalar(self) -> None: + """kind field defaults to 'scalar' when not present in record.""" + records = [ + {"name": "test_q", "source": "dd", "source_path": "x", "confidence": 0.8} + ] + entries = graph_records_to_entries(records) + assert entries[0].kind == "scalar" + + +# ============================================================================= +# Rich-field YAML round-trip tests +# ============================================================================= + + +class TestRichFieldRoundTrip: + def test_all_rich_fields_in_yaml(self, sample_provenance: SNProvenance) -> None: + """Full round-trip: create entry with all rich fields → YAML → parse back.""" + entry = SNPublishEntry( + name="ion_temperature", + kind="scalar", + unit="eV", + tags=["core_profiles", "kinetics"], + status="drafted", + description="Ion temperature", + documentation="Ion temperature $T_i$ in eV. Typical range 0.1–20 keV.", + links=["electron_temperature", "ion_density"], + ids_paths=["core_profiles/profiles_1d/ion/temperature"], + constraints=["T_i > 0"], + validity_domain="core plasma", + provenance=sample_provenance, + ) + content = generate_yaml_entry(entry) + doc = yaml.safe_load(content) + + assert doc["name"] == "ion_temperature" + assert doc["kind"] == "scalar" + assert doc["unit"] == "eV" + assert ( + doc["documentation"] + == "Ion temperature $T_i$ in eV. Typical range 0.1–20 keV." + ) + assert doc["links"] == [ + {"name": "electron_temperature"}, + {"name": "ion_density"}, + ] + assert doc["ids_paths"] == ["core_profiles/profiles_1d/ion/temperature"] + assert doc["constraints"] == ["T_i > 0"] + assert doc["validity_domain"] == "core plasma" + assert doc["tags"] == ["core_profiles", "kinetics"] + + def test_empty_rich_fields_omitted(self, sample_provenance: SNProvenance) -> None: + """Empty optional rich fields should not appear in YAML output.""" + entry = SNPublishEntry( + name="bare_quantity", + provenance=sample_provenance, + ) + content = generate_yaml_entry(entry) + doc = yaml.safe_load(content) + + assert "documentation" not in doc + assert "links" not in doc + assert "ids_paths" not in doc + assert "constraints" not in doc + assert "validity_domain" not in doc + + def test_links_formatted_as_name_dicts( + self, sample_provenance: SNProvenance + ) -> None: + """links list should be serialized as [{name: ...}] objects.""" + entry = SNPublishEntry( + name="test_quantity", + links=["alpha", "beta"], + provenance=sample_provenance, + ) + content = generate_yaml_entry(entry) + doc = yaml.safe_load(content) + assert doc["links"] == [{"name": "alpha"}, {"name": "beta"}] From 4809a3d3d8b198069874a45a025b1dbd8ba38f92 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 10:54:49 +0200 Subject: [PATCH 16/37] feat(sn): add MCP tools and benchmark quality tiers - Add sn_tools.py with search/fetch/list standard name tools following the search_tools.py pattern (hybrid vector+keyword search) - Register search_standard_names, fetch_standard_names, list_standard_names in server.py (available in both dd_only and full mode) - Add benchmark_labels.yaml with quality tier anchors (outstanding/good/adequate/poor) - Add load_quality_labels() and score_with_reviewer() to benchmark.py - Add --reviewer-model CLI option to sn benchmark command - Extend ModelResult with quality_scores, quality_distribution, avg_quality_score, avg_doc_length, avg_fields_populated fields - Update render_comparison_table() to show quality distribution table when reviewer used - Add tests: test_sn_tools.py (38 tests) and TestQualityLabels/TestReviewerModelCLI in test_benchmark.py (9 tests) --- imas_codex/llm/server.py | 74 ++++ imas_codex/llm/sn_tools.py | 452 ++++++++++++++++++++++++ imas_codex/sn/benchmark.py | 222 +++++++++++- imas_codex/sn/benchmark_labels.yaml | 29 ++ tests/sn/test_benchmark.py | 87 +++++ tests/sn/test_sn_tools.py | 520 ++++++++++++++++++++++++++++ 6 files changed, 1381 insertions(+), 3 deletions(-) create mode 100644 imas_codex/llm/sn_tools.py create mode 100644 imas_codex/sn/benchmark_labels.yaml create mode 100644 tests/sn/test_sn_tools.py diff --git a/imas_codex/llm/server.py b/imas_codex/llm/server.py index 472bdddae..411f9dea0 100644 --- a/imas_codex/llm/server.py +++ b/imas_codex/llm/server.py @@ -3106,6 +3106,80 @@ def fetch_content(resource: str) -> str: return _f(resource) + # ===================================================================== + # Standard Name tools + # ===================================================================== + + @self.mcp.tool() + def search_standard_names( + query: str, + kind: str | None = None, + tags: list[str] | None = None, + review_status: str | None = None, + k: int = 20, + ) -> str: + """Search standard names by physics concept. + + Hybrid search (vector + keyword) over StandardName descriptions + and documentation. Enriched with DD path links, unit info, and + grammar decomposition. + + Args: + query: Natural-language description of the quantity to find + (e.g. "electron temperature", "plasma boundary shape"). + kind: Filter by kind (e.g. "scalar", "vector", "metadata"). + tags: Filter by tags (e.g. ["equilibrium", "core_profiles"]). + review_status: Filter by review status (e.g. "drafted", "published"). + k: Maximum results to return (default 20). + + Returns: + Formatted text report with matched standard names, descriptions, + units, tags, grammar fields, and relevance scores. + """ + from imas_codex.llm.sn_tools import _search_standard_names as _ssn + + return _ssn(query, kind=kind, tags=tags, review_status=review_status, k=k) + + @self.mcp.tool() + def fetch_standard_names(names: str) -> str: + """Fetch full entries for known standard names. + + Returns complete metadata: description, documentation, unit, kind, + tags, links, ids_paths, grammar fields, provenance, review status. + + Args: + names: Space- or comma-separated standard name IDs + (e.g. "electron_temperature plasma_current"). + + Returns: + Formatted text report with complete documentation per name. + """ + from imas_codex.llm.sn_tools import _fetch_standard_names as _fsn + + return _fsn(names) + + @self.mcp.tool() + def list_standard_names( + tag: str | None = None, + kind: str | None = None, + review_status: str | None = None, + ) -> str: + """List standard names with optional filters. + + Returns name, description, kind, unit, status for each entry. + + Args: + tag: Filter by tag (e.g. "equilibrium", "magnetics"). + kind: Filter by kind (e.g. "scalar", "vector"). + review_status: Filter by review status (e.g. "drafted"). + + Returns: + Formatted markdown table of standard names. + """ + from imas_codex.llm.sn_tools import _list_standard_names as _lsn + + return _lsn(tag=tag, kind=kind, review_status=review_status) + if not self.read_only: # ===================================================================== # Log Tools (Phase 3: MCP Logs) diff --git a/imas_codex/llm/sn_tools.py b/imas_codex/llm/sn_tools.py new file mode 100644 index 000000000..3288bce62 --- /dev/null +++ b/imas_codex/llm/sn_tools.py @@ -0,0 +1,452 @@ +"""MCP tools for standard name search, fetch, and listing. + +Functions are prefixed with ``_`` — they are registered as MCP tools +in ``server.py`` via ``@self.mcp.tool()``. +""" + +from __future__ import annotations + +import logging + +from neo4j.exceptions import ServiceUnavailable + +from imas_codex.embeddings.encoder import EmbeddingBackendError, Encoder +from imas_codex.graph.client import GraphClient + +logger = logging.getLogger(__name__) + +NEO4J_NOT_RUNNING_MSG = ( + "Neo4j is not running. Check service with: systemctl --user status imas-codex-neo4j" +) + + +def _neo4j_error_message(e: Exception) -> str: + """Format Neo4j errors with helpful instructions.""" + if isinstance(e, ServiceUnavailable): + return NEO4J_NOT_RUNNING_MSG + msg = str(e) + if "Connection refused" in msg or "ServiceUnavailable" in msg: + return NEO4J_NOT_RUNNING_MSG + return msg + + +# --------------------------------------------------------------------------- +# _search_standard_names +# --------------------------------------------------------------------------- + + +def _search_standard_names( + query: str, + *, + kind: str | None = None, + tags: list[str] | None = None, + review_status: str | None = None, + k: int = 20, + gc: GraphClient | None = None, +) -> str: + """Search standard names by physics concept. + + Hybrid search (vector + keyword) over StandardName descriptions. + Falls back to keyword-only if no embeddings present. + """ + try: + if gc is None: + gc = GraphClient() + except ServiceUnavailable: + return NEO4J_NOT_RUNNING_MSG + except Exception as e: + return f"Error connecting to graph: {e}" + + # Try to get embedding for vector search + has_embedding = False + embedding: list[float] = [] + try: + from imas_codex.embeddings.config import EncoderConfig + + encoder = Encoder(EncoderConfig()) + result = encoder.embed_texts([query])[0] + embedding = result.tolist() if hasattr(result, "tolist") else list(result) + has_embedding = True + except (EmbeddingBackendError, Exception): + pass + + try: + if has_embedding: + rows = _vector_search_sn(gc, embedding, k) + else: + rows = _keyword_search_sn(gc, query, k) + except ServiceUnavailable: + return NEO4J_NOT_RUNNING_MSG + except Exception as e: + return f"Search failed: {_neo4j_error_message(e)}" + + # Post-filter + if kind: + rows = [r for r in rows if (r.get("kind") or "").lower() == kind.lower()] + if tags: + rows = [r for r in rows if any(t in (r.get("tags") or []) for t in tags)] + if review_status: + rows = [ + r + for r in rows + if (r.get("review_status") or "").lower() == review_status.lower() + ] + + return _format_search_report(query, rows) + + +def _vector_search_sn(gc: GraphClient, embedding: list[float], k: int) -> list[dict]: + """Run vector search on StandardName nodes.""" + cypher = """ +CALL db.index.vector.queryNodes('standard_name_desc_embedding', $k, $embedding) +YIELD node AS sn, score +WHERE sn.id IS NOT NULL +OPTIONAL MATCH (sn)-[:CANONICAL_UNITS]->(u:Unit) +RETURN sn.id AS name, sn.description AS description, + sn.kind AS kind, coalesce(u.id, sn.canonical_units) AS unit, + sn.tags AS tags, sn.review_status AS review_status, + sn.documentation AS documentation, + sn.physical_base AS physical_base, + sn.subject AS subject, + score +ORDER BY score DESC +""" + return gc.query(cypher, {"embedding": embedding, "k": k}) + + +def _keyword_search_sn(gc: GraphClient, query: str, k: int) -> list[dict]: + """Run keyword search on StandardName nodes.""" + cypher = """ +MATCH (sn:StandardName) +WHERE toLower(sn.id) CONTAINS toLower($keyword) + OR toLower(sn.description) CONTAINS toLower($keyword) + OR toLower(coalesce(sn.documentation, '')) CONTAINS toLower($keyword) +OPTIONAL MATCH (sn)-[:CANONICAL_UNITS]->(u:Unit) +RETURN sn.id AS name, sn.description AS description, + sn.kind AS kind, coalesce(u.id, sn.canonical_units) AS unit, + sn.tags AS tags, sn.review_status AS review_status, + sn.documentation AS documentation, + sn.physical_base AS physical_base, + sn.subject AS subject, + 1.0 AS score +LIMIT $k +""" + return gc.query(cypher, {"keyword": query, "k": k}) + + +def _format_search_report(query: str, rows: list[dict]) -> str: + """Format search results as a text report.""" + if not rows: + return ( + f"## Standard Name Search Results\n\nNo standard names found matching " + f'"{query}".' + ) + + lines = [ + f'## Standard Name Search Results\n\nFound {len(rows)} standard names matching "{query}"\n' + ] + for i, row in enumerate(rows, 1): + name = row.get("name") or "unknown" + score = row.get("score", 0.0) + kind = row.get("kind") or "" + unit = row.get("unit") or "" + tags = row.get("tags") or [] + review_status = row.get("review_status") or "" + description = row.get("description") or "" + documentation = row.get("documentation") or "" + physical_base = row.get("physical_base") or "" + subject = row.get("subject") or "" + + lines.append(f"### {i}. {name} (score: {score:.2f})") + if kind: + lines.append(f"- **Kind:** {kind}") + if unit: + lines.append(f"- **Unit:** {unit}") + if tags: + tag_str = ", ".join(tags) if isinstance(tags, list) else str(tags) + lines.append(f"- **Tags:** {tag_str}") + if review_status: + lines.append(f"- **Status:** {review_status}") + if description: + lines.append(f"- **Description:** {description}") + if documentation: + lines.append( + f"- **Documentation:** {documentation[:200]}{'...' if len(documentation) > 200 else ''}" + ) + if physical_base or subject: + grammar_parts = [] + if physical_base: + grammar_parts.append(f"physical_base={physical_base}") + if subject: + grammar_parts.append(f"subject={subject}") + lines.append(f"- **Grammar:** {', '.join(grammar_parts)}") + lines.append("") + + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# _fetch_standard_names +# --------------------------------------------------------------------------- + + +def _fetch_standard_names( + names: str, + *, + gc: GraphClient | None = None, +) -> str: + """Fetch full entries for known standard names. + + Args: + names: Space- or comma-separated standard name IDs. + """ + try: + if gc is None: + gc = GraphClient() + except ServiceUnavailable: + return NEO4J_NOT_RUNNING_MSG + except Exception as e: + return f"Error connecting to graph: {e}" + + # Parse names (split on space or comma) + import re + + name_list = [n.strip() for n in re.split(r"[,\s]+", names) if n.strip()] + if not name_list: + return "No names provided." + + cypher = """ +UNWIND $names AS name_id +MATCH (sn:StandardName {id: name_id}) +OPTIONAL MATCH (sn)-[:CANONICAL_UNITS]->(u:Unit) +OPTIONAL MATCH (src)-[:HAS_STANDARD_NAME]->(sn) +OPTIONAL MATCH (src)-[:IN_IDS]->(ids:IDS) +RETURN sn.id AS name, sn.description AS description, + sn.documentation AS documentation, + sn.kind AS kind, coalesce(u.id, sn.canonical_units) AS unit, + sn.tags AS tags, sn.links AS links, + sn.ids_paths AS ids_paths, sn.constraints AS constraints, + sn.validity_domain AS validity_domain, + sn.physical_base AS physical_base, sn.subject AS subject, + sn.component AS component, sn.coordinate AS coordinate, + sn.position AS position, sn.process AS process, + sn.review_status AS review_status, + sn.confidence AS confidence, sn.model AS model, + collect(DISTINCT src.id) AS source_ids, + collect(DISTINCT ids.id) AS source_ids_names +""" + + try: + rows = gc.query(cypher, {"names": name_list}) + except ServiceUnavailable: + return NEO4J_NOT_RUNNING_MSG + except Exception as e: + return f"Fetch failed: {_neo4j_error_message(e)}" + + if not rows: + not_found = ", ".join(name_list) + return f"No standard names found for: {not_found}" + + return _format_fetch_report(rows, name_list) + + +def _format_fetch_report(rows: list[dict], requested: list[str]) -> str: + """Format fetch results as a detailed report.""" + found_names = {r.get("name") for r in rows} + not_found = [n for n in requested if n not in found_names] + + lines = ["## Standard Name Details\n"] + + for row in rows: + name = row.get("name") or "unknown" + lines.append(f"### {name}") + lines.append("") + + description = row.get("description") or "" + documentation = row.get("documentation") or "" + kind = row.get("kind") or "" + unit = row.get("unit") or "" + tags = row.get("tags") or [] + links = row.get("links") or [] + ids_paths = row.get("ids_paths") or [] + constraints = row.get("constraints") or [] + validity_domain = row.get("validity_domain") or "" + physical_base = row.get("physical_base") or "" + subject = row.get("subject") or "" + component = row.get("component") or "" + coordinate = row.get("coordinate") or "" + position = row.get("position") or "" + process = row.get("process") or "" + review_status = row.get("review_status") or "" + confidence = row.get("confidence") + model = row.get("model") or "" + source_ids = row.get("source_ids") or [] + source_ids_names = row.get("source_ids_names") or [] + + if description: + lines.append(f"**Description:** {description}") + if documentation: + lines.append(f"\n**Documentation:**\n{documentation}") + lines.append("") + + if kind: + lines.append(f"- **Kind:** {kind}") + if unit: + lines.append(f"- **Unit:** {unit}") + if review_status: + lines.append(f"- **Review Status:** {review_status}") + if confidence is not None: + lines.append(f"- **Confidence:** {confidence:.2f}") + if model: + lines.append(f"- **Model:** {model}") + + # Grammar + grammar_parts = [] + for field_name, val in [ + ("physical_base", physical_base), + ("subject", subject), + ("component", component), + ("coordinate", coordinate), + ("position", position), + ("process", process), + ]: + if val: + grammar_parts.append(f"{field_name}={val}") + if grammar_parts: + lines.append(f"- **Grammar:** {', '.join(grammar_parts)}") + + if tags: + tag_str = ", ".join(tags) if isinstance(tags, list) else str(tags) + lines.append(f"- **Tags:** {tag_str}") + if links: + link_str = ", ".join(links) if isinstance(links, list) else str(links) + lines.append(f"- **Links:** {link_str}") + if ids_paths: + path_str = ( + "\n - " + "\n - ".join(ids_paths) + if isinstance(ids_paths, list) + else str(ids_paths) + ) + lines.append(f"- **IDS Paths:**{path_str}") + if constraints: + c_str = ( + ", ".join(constraints) + if isinstance(constraints, list) + else str(constraints) + ) + lines.append(f"- **Constraints:** {c_str}") + if validity_domain: + lines.append(f"- **Validity Domain:** {validity_domain}") + if source_ids: + src_str = ", ".join(s for s in source_ids if s) + if src_str: + lines.append(f"- **Source Nodes:** {src_str}") + if source_ids_names: + ids_str = ", ".join(s for s in source_ids_names if s) + if ids_str: + lines.append(f"- **Source IDS:** {ids_str}") + + lines.append("") + + if not_found: + lines.append(f"**Not found:** {', '.join(not_found)}") + + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# _list_standard_names +# --------------------------------------------------------------------------- + + +def _list_standard_names( + *, + tag: str | None = None, + kind: str | None = None, + review_status: str | None = None, + gc: GraphClient | None = None, +) -> str: + """List standard names with optional filters.""" + try: + if gc is None: + gc = GraphClient() + except ServiceUnavailable: + return NEO4J_NOT_RUNNING_MSG + except Exception as e: + return f"Error connecting to graph: {e}" + + # Build WHERE clause + conditions = [] + params: dict = {} + + if tag: + conditions.append("$tag IN sn.tags") + params["tag"] = tag + if kind: + conditions.append("toLower(sn.kind) = toLower($kind)") + params["kind"] = kind + if review_status: + conditions.append("toLower(sn.review_status) = toLower($review_status)") + params["review_status"] = review_status + + where_clause = ("WHERE " + " AND ".join(conditions)) if conditions else "" + + cypher = f""" +MATCH (sn:StandardName) +{where_clause} +OPTIONAL MATCH (sn)-[:CANONICAL_UNITS]->(u:Unit) +RETURN sn.id AS name, sn.kind AS kind, + coalesce(u.id, sn.canonical_units) AS unit, + sn.review_status AS review_status, + sn.description AS description +ORDER BY sn.id +""" + + try: + rows = gc.query(cypher, params) + except ServiceUnavailable: + return NEO4J_NOT_RUNNING_MSG + except Exception as e: + return f"List failed: {_neo4j_error_message(e)}" + + return _format_list_report(rows, tag=tag, kind=kind, review_status=review_status) + + +def _format_list_report( + rows: list[dict], + *, + tag: str | None = None, + kind: str | None = None, + review_status: str | None = None, +) -> str: + """Format list results as a markdown table.""" + filter_parts = [] + if tag: + filter_parts.append(f"tag={tag}") + if kind: + filter_parts.append(f"kind={kind}") + if review_status: + filter_parts.append(f"status={review_status}") + filter_str = f" (filtered by: {', '.join(filter_parts)})" if filter_parts else "" + + if not rows: + return f"## Standard Names\n\nNo standard names found{filter_str}." + + lines = [ + f"## Standard Names ({len(rows)} total{filter_str})\n", + "| Name | Kind | Unit | Status | Description |", + "|------|------|------|--------|-------------|", + ] + + for row in rows: + name = row.get("name") or "" + row_kind = row.get("kind") or "" + unit = row.get("unit") or "" + status = row.get("review_status") or "" + desc = row.get("description") or "" + # Truncate long descriptions + if len(desc) > 80: + desc = desc[:77] + "..." + lines.append(f"| {name} | {row_kind} | {unit} | {status} | {desc} |") + + return "\n".join(lines) diff --git a/imas_codex/sn/benchmark.py b/imas_codex/sn/benchmark.py index 37997c383..1475f2b16 100644 --- a/imas_codex/sn/benchmark.py +++ b/imas_codex/sn/benchmark.py @@ -13,6 +13,7 @@ import time from dataclasses import asdict, dataclass, field from datetime import UTC, datetime +from pathlib import Path from typing import Any from imas_standard_names.grammar import ( @@ -51,6 +52,7 @@ class BenchmarkConfig: max_candidates: int = 50 runs_per_model: int = 1 temperature: float = 0.0 # pinned for reproducibility + reviewer_model: str | None = None # frontier model for quality scoring @dataclass @@ -74,6 +76,12 @@ class ModelResult: reference_total: int = 0 reference_precision: float = 0.0 reference_recall: float = 0.0 + # Quality scoring (reviewer model) + quality_scores: list[dict] = field(default_factory=list) + quality_distribution: dict[str, int] = field(default_factory=dict) + avg_quality_score: float = 0.0 + avg_doc_length: float = 0.0 + avg_fields_populated: float = 0.0 @dataclass @@ -251,6 +259,109 @@ def compare_to_reference( return overlap, ref_total, precision, recall +# --------------------------------------------------------------------------- +# Quality tier labels +# --------------------------------------------------------------------------- + + +def load_quality_labels() -> dict[str, list[str]]: + """Load quality tier labels from benchmark_labels.yaml. + + Returns a dict mapping tier name → list of standard name IDs. + Returns empty dict if file not found. + """ + import yaml + + labels_path = Path(__file__).parent / "benchmark_labels.yaml" + if labels_path.exists(): + with open(labels_path) as f: + return yaml.safe_load(f) or {} + return {} + + +async def score_with_reviewer( + candidates: list[dict], + reviewer_model: str, + quality_labels: dict[str, list[str]], + grammar_context: dict[str, list[str]], +) -> list[dict]: + """Score candidates using a reviewer model. + + Returns list of dicts with: name, quality_tier, score, reasoning. + """ + from pydantic import BaseModel, Field + + from imas_codex.discovery.base.llm import acall_llm_structured + + class QualityReview(BaseModel): + name: str + quality_tier: str = Field(description="outstanding, good, adequate, or poor") + score: int = Field(ge=0, le=100, description="Quality score 0-100") + reasoning: str + + class QualityReviewBatch(BaseModel): + reviews: list[QualityReview] + + # Build prompt with labeled examples and rubric + labeled_examples = [] + for tier, names in quality_labels.items(): + for name in names: + labeled_examples.append(f" - {name}: {tier}") + + rubric = """ +Quality rubric: +- Grammar correctness: Does the name follow standard name grammar rules? +- Semantic accuracy: Does the name correctly describe the physics quantity? +- Documentation quality: Is the documentation clear, with LaTeX where appropriate? +- Naming conventions: Does the name follow established patterns? +- Unit consistency: Is the unit correct for this quantity? + +Quality tiers: +- outstanding (80-100): Rich docs, correct grammar, cross-linked, LaTeX +- good (60-79): Correct grammar, adequate documentation +- adequate (40-59): Correct grammar, thin documentation +- poor (0-39): Grammar valid but naming debatable or docs minimal +""" + + prompt = f"""You are reviewing standard name entries for quality. + +{rubric} + +Labeled examples (scoring anchors): +{chr(10).join(labeled_examples)} + +Review these candidates and assign a quality tier and score: +""" + + # Process in batches of 10 + all_reviews: list[dict] = [] + for i in range(0, len(candidates), 10): + batch = candidates[i : i + 10] + batch_text = [] + for c in batch: + batch_text.append(f"Name: {c.get('standard_name', '')}") + batch_text.append(f" Description: {c.get('description', '')}") + doc = c.get("documentation", "") or "" + batch_text.append(f" Documentation: {doc[:200]}") + batch_text.append(f" Unit: {c.get('unit', 'N/A')}") + batch_text.append(f" Fields: {c.get('fields', {})}") + batch_text.append("") + + messages = [{"role": "user", "content": prompt + "\n".join(batch_text)}] + + try: + result, _, _ = await acall_llm_structured( + model=reviewer_model, + messages=messages, + response_model=QualityReviewBatch, + ) + all_reviews.extend([r.model_dump() for r in result.reviews]) + except Exception as e: + logger.warning("Reviewer scoring failed for batch: %s", e) + + return all_reviews + + # --------------------------------------------------------------------------- # Core benchmark runner # --------------------------------------------------------------------------- @@ -295,6 +406,7 @@ async def run_benchmark( # --- 2. Run each model --- results: list[ModelResult] = [] + grammar_ctx = build_grammar_context() for model in config.models: logger.info("Benchmarking model: %s", model) model_result = await _run_model( @@ -305,6 +417,53 @@ async def run_benchmark( ) results.append(model_result) + # --- 2b. Reviewer scoring (optional) --- + if config.reviewer_model: + quality_labels = load_quality_labels() + for result in results: + if result.candidates: + reviews = await score_with_reviewer( + result.candidates, + config.reviewer_model, + quality_labels, + grammar_ctx, + ) + result.quality_scores = reviews + # Compute distribution + for r in reviews: + tier = r.get("quality_tier", "unknown") + result.quality_distribution[tier] = ( + result.quality_distribution.get(tier, 0) + 1 + ) + if reviews: + result.avg_quality_score = sum( + r.get("score", 0) for r in reviews + ) / len(reviews) + + # Compute doc length and field coverage metrics + docs = [c.get("documentation", "") or "" for c in result.candidates] + result.avg_doc_length = ( + sum(len(d) for d in docs) / len(docs) if docs else 0.0 + ) + + all_fields = { + "physical_base", + "subject", + "component", + "coordinate", + "position", + "process", + } + field_counts = [] + for c in result.candidates: + fields = c.get("fields", {}) + field_counts.append( + len(set(fields.keys()) & all_fields) / len(all_fields) + ) + result.avg_fields_populated = ( + sum(field_counts) / len(field_counts) if field_counts else 0.0 + ) + # --- 3. Build report --- report = BenchmarkReport( config=config, @@ -486,6 +645,9 @@ def render_comparison_table(report: BenchmarkReport) -> None: console = Console() + # Check if any result has quality scores + has_quality = any(r.quality_scores for r in report.results) + table = Table( title="SN Benchmark Results", show_header=True, @@ -501,6 +663,10 @@ def render_comparison_table(report: BenchmarkReport) -> None: table.add_column("Names/min", justify="right") table.add_column("$/name", justify="right") table.add_column("Errors", justify="right") + if has_quality: + table.add_column("Avg Quality", justify="right") + table.add_column("Avg Doc Len", justify="right") + table.add_column("Fields Pop%", justify="right") for r in report.results: n = len(r.candidates) @@ -514,7 +680,7 @@ def render_comparison_table(report: BenchmarkReport) -> None: cpn_str = f"${r.cost_per_name:.4f}" if r.cost_per_name > 0 else "—" err_str = str(r.batch_errors) if r.batch_errors > 0 else "0" - table.add_row( + row_data = [ r.model, str(n), valid_pct, @@ -524,14 +690,64 @@ def render_comparison_table(report: BenchmarkReport) -> None: speed_str, cpn_str, err_str, - ) + ] + + if has_quality: + qual_str = f"{r.avg_quality_score:.1f}" if r.quality_scores else "—" + doc_str = f"{r.avg_doc_length:.0f}" if r.quality_scores else "—" + fp_str = f"{r.avg_fields_populated * 100:.0f}%" if r.quality_scores else "—" + row_data.extend([qual_str, doc_str, fp_str]) + + table.add_row(*row_data) console.print() console.print(table) + # Quality distribution table (when reviewer was used) + if has_quality: + qual_table = Table( + title="Quality Distribution", + show_header=True, + header_style="bold magenta", + ) + qual_table.add_column("Model", style="bold") + qual_table.add_column("Outstanding", justify="right") + qual_table.add_column("Good", justify="right") + qual_table.add_column("Adequate", justify="right") + qual_table.add_column("Poor", justify="right") + + for r in report.results: + if r.quality_scores: + dist = r.quality_distribution + n_reviews = len(r.quality_scores) + + qual_table.add_row( + r.model, + f"{dist.get('outstanding', 0)} ({dist.get('outstanding', 0) / n_reviews * 100:.0f}%)" + if n_reviews + else "—", + f"{dist.get('good', 0)} ({dist.get('good', 0) / n_reviews * 100:.0f}%)" + if n_reviews + else "—", + f"{dist.get('adequate', 0)} ({dist.get('adequate', 0) / n_reviews * 100:.0f}%)" + if n_reviews + else "—", + f"{dist.get('poor', 0)} ({dist.get('poor', 0) / n_reviews * 100:.0f}%)" + if n_reviews + else "—", + ) + + console.print() + console.print(qual_table) + # Summary line + reviewer_str = ( + f" | Reviewer: {report.config.reviewer_model}" + if report.config.reviewer_model + else "" + ) console.print( f"\n[dim]Extraction: {report.extraction_count} items | " - f"Temperature: {report.config.temperature} | " + f"Temperature: {report.config.temperature}{reviewer_str} | " f"Timestamp: {report.timestamp}[/dim]" ) diff --git a/imas_codex/sn/benchmark_labels.yaml b/imas_codex/sn/benchmark_labels.yaml new file mode 100644 index 000000000..72ee31a69 --- /dev/null +++ b/imas_codex/sn/benchmark_labels.yaml @@ -0,0 +1,29 @@ +# Quality tier labels for benchmark evaluation. +# Each tier defines standard names that exemplify that quality level. +# Used by the reviewer model as scoring anchors. + +outstanding: + # Rich docs, correct grammar, cross-linked, LaTeX notation + - electron_temperature + - plasma_current + - safety_factor + - position_of_magnetic_axis + - bootstrap_current + +good: + # Correct grammar, adequate documentation + - toroidal_component_of_magnetic_field_at_magnetic_axis + - centroid_of_plasma_boundary + - bolometer_radiated_power + - collisionality + +adequate: + # Correct grammar, thin documentation + - area_of_poloidal_magnetic_field_probe + - tokamak_scenario + - time + +poor: + # Grammar valid but naming debatable or documentation minimal + - banana_orbits + - h_mode diff --git a/tests/sn/test_benchmark.py b/tests/sn/test_benchmark.py index 3c0157706..8b906f24f 100644 --- a/tests/sn/test_benchmark.py +++ b/tests/sn/test_benchmark.py @@ -661,3 +661,90 @@ def test_command_requires_models(self): result = runner.invoke(sn, ["benchmark"]) assert result.exit_code != 0 assert "Missing" in result.output or "required" in result.output.lower() + + +class TestQualityLabels: + """Test benchmark quality tier labels.""" + + def test_labels_load(self): + from imas_codex.sn.benchmark import load_quality_labels + + labels = load_quality_labels() + assert isinstance(labels, dict) + assert "outstanding" in labels + assert "good" in labels + assert "adequate" in labels + assert "poor" in labels + + def test_labels_non_empty(self): + from imas_codex.sn.benchmark import load_quality_labels + + labels = load_quality_labels() + for tier, names in labels.items(): + assert len(names) > 0, f"Tier {tier} should have entries" + + def test_labels_no_overlap(self): + from imas_codex.sn.benchmark import load_quality_labels + + labels = load_quality_labels() + all_names = [] + for names in labels.values(): + all_names.extend(names) + assert len(all_names) == len(set(all_names)), "No duplicate names across tiers" + + def test_reviewer_config_field(self): + from imas_codex.sn.benchmark import BenchmarkConfig + + config = BenchmarkConfig(models=["test"], reviewer_model="test/model") + assert config.reviewer_model == "test/model" + + def test_reviewer_config_default_none(self): + from imas_codex.sn.benchmark import BenchmarkConfig + + config = BenchmarkConfig(models=["test"]) + assert config.reviewer_model is None + + def test_model_result_quality_fields(self): + from imas_codex.sn.benchmark import ModelResult + + r = ModelResult(model="test") + assert r.quality_scores == [] + assert r.quality_distribution == {} + assert r.avg_quality_score == 0.0 + assert r.avg_doc_length == 0.0 + assert r.avg_fields_populated == 0.0 + + def test_model_result_with_quality(self): + from imas_codex.sn.benchmark import ModelResult + + r = ModelResult( + model="test", + quality_scores=[ + { + "name": "a", + "score": 80, + "quality_tier": "outstanding", + "reasoning": "good", + } + ], + quality_distribution={"outstanding": 1}, + avg_quality_score=80.0, + avg_doc_length=150.0, + avg_fields_populated=0.5, + ) + assert r.avg_quality_score == 80.0 + assert r.quality_distribution["outstanding"] == 1 + + +class TestReviewerModelCLI: + """Test --reviewer-model CLI option.""" + + def test_reviewer_model_in_help(self): + from click.testing import CliRunner + + from imas_codex.cli.sn import sn + + runner = CliRunner() + result = runner.invoke(sn, ["benchmark", "--help"]) + assert result.exit_code == 0 + assert "--reviewer-model" in result.output diff --git a/tests/sn/test_sn_tools.py b/tests/sn/test_sn_tools.py new file mode 100644 index 000000000..e914a3a82 --- /dev/null +++ b/tests/sn/test_sn_tools.py @@ -0,0 +1,520 @@ +"""Tests for standard name MCP tools.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + + +class TestSearchStandardNames: + """Test _search_standard_names tool.""" + + def test_keyword_fallback(self): + """Search falls back to keyword when no embeddings.""" + from imas_codex.llm.sn_tools import _search_standard_names + + mock_gc = MagicMock() + mock_gc.query = MagicMock( + return_value=[ + { + "name": "electron_temperature", + "description": "Te", + "kind": "scalar", + "unit": "eV", + "tags": ["core_profiles"], + "review_status": "drafted", + "documentation": None, + "physical_base": "temperature", + "subject": "electron", + "score": 1.0, + } + ] + ) + + # Patch Encoder to fail (trigger keyword fallback) + with patch( + "imas_codex.llm.sn_tools.Encoder", side_effect=Exception("no embeddings") + ): + result = _search_standard_names("electron temperature", gc=mock_gc) + + assert "electron_temperature" in result + mock_gc.query.assert_called() + + def test_empty_results(self): + """Empty results produce informative message.""" + from imas_codex.llm.sn_tools import _search_standard_names + + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + with patch( + "imas_codex.llm.sn_tools.Encoder", side_effect=Exception("no embeddings") + ): + result = _search_standard_names("nonexistent quantity", gc=mock_gc) + + assert "No" in result or "0" in result + + def test_kind_filter(self): + """Kind filter is applied to results.""" + from imas_codex.llm.sn_tools import _search_standard_names + + mock_gc = MagicMock() + mock_gc.query = MagicMock( + return_value=[ + { + "name": "electron_temperature", + "description": "Te", + "kind": "scalar", + "unit": "eV", + "tags": [], + "review_status": "drafted", + "documentation": None, + "physical_base": "temperature", + "subject": None, + "score": 1.0, + }, + { + "name": "velocity_field", + "description": "v", + "kind": "vector", + "unit": "m/s", + "tags": [], + "review_status": "drafted", + "documentation": None, + "physical_base": None, + "subject": None, + "score": 0.8, + }, + ] + ) + + with patch( + "imas_codex.llm.sn_tools.Encoder", side_effect=Exception("no embeddings") + ): + result = _search_standard_names("temperature", kind="scalar", gc=mock_gc) + + assert "electron_temperature" in result + assert "velocity_field" not in result + + def test_review_status_filter(self): + """review_status filter is applied.""" + from imas_codex.llm.sn_tools import _search_standard_names + + mock_gc = MagicMock() + mock_gc.query = MagicMock( + return_value=[ + { + "name": "drafted_name", + "description": "d", + "kind": "scalar", + "unit": "eV", + "tags": [], + "review_status": "drafted", + "documentation": None, + "physical_base": None, + "subject": None, + "score": 1.0, + }, + { + "name": "published_name", + "description": "p", + "kind": "scalar", + "unit": "A", + "tags": [], + "review_status": "published", + "documentation": None, + "physical_base": None, + "subject": None, + "score": 0.9, + }, + ] + ) + + with patch( + "imas_codex.llm.sn_tools.Encoder", side_effect=Exception("no embeddings") + ): + result = _search_standard_names("test", review_status="drafted", gc=mock_gc) + + assert "drafted_name" in result + assert "published_name" not in result + + def test_tags_filter(self): + """tags filter is applied.""" + from imas_codex.llm.sn_tools import _search_standard_names + + mock_gc = MagicMock() + mock_gc.query = MagicMock( + return_value=[ + { + "name": "electron_temperature", + "description": "Te", + "kind": "scalar", + "unit": "eV", + "tags": ["core_profiles", "kinetics"], + "review_status": "drafted", + "documentation": None, + "physical_base": None, + "subject": None, + "score": 1.0, + }, + { + "name": "equilibrium_shape", + "description": "shape", + "kind": "scalar", + "unit": "m", + "tags": ["equilibrium"], + "review_status": "drafted", + "documentation": None, + "physical_base": None, + "subject": None, + "score": 0.8, + }, + ] + ) + + with patch( + "imas_codex.llm.sn_tools.Encoder", side_effect=Exception("no embeddings") + ): + result = _search_standard_names( + "temperature", tags=["core_profiles"], gc=mock_gc + ) + + assert "electron_temperature" in result + assert "equilibrium_shape" not in result + + def test_result_format_includes_grammar(self): + """Result format includes grammar fields.""" + from imas_codex.llm.sn_tools import _search_standard_names + + mock_gc = MagicMock() + mock_gc.query = MagicMock( + return_value=[ + { + "name": "electron_temperature", + "description": "Te", + "kind": "scalar", + "unit": "eV", + "tags": [], + "review_status": "drafted", + "documentation": None, + "physical_base": "temperature", + "subject": "electron", + "score": 0.92, + } + ] + ) + + with patch( + "imas_codex.llm.sn_tools.Encoder", side_effect=Exception("no embeddings") + ): + result = _search_standard_names("electron temperature", gc=mock_gc) + + assert "physical_base=temperature" in result + assert "subject=electron" in result + assert "0.92" in result + + +class TestFetchStandardNames: + """Test _fetch_standard_names tool.""" + + def test_fetch_single(self): + from imas_codex.llm.sn_tools import _fetch_standard_names + + mock_gc = MagicMock() + mock_gc.query = MagicMock( + return_value=[ + { + "name": "electron_temperature", + "description": "Te profile", + "documentation": "The $T_e$ profile", + "kind": "scalar", + "unit": "eV", + "tags": ["core_profiles"], + "links": ["ion_temperature"], + "ids_paths": ["core_profiles/profiles_1d/electrons/temperature"], + "constraints": ["T_e > 0"], + "validity_domain": "core plasma", + "physical_base": "temperature", + "subject": "electron", + "component": None, + "coordinate": None, + "position": None, + "process": None, + "review_status": "drafted", + "confidence": 0.95, + "model": "test", + "source_ids": ["core_profiles/profiles_1d/electrons/temperature"], + "source_ids_names": ["core_profiles"], + } + ] + ) + + result = _fetch_standard_names("electron_temperature", gc=mock_gc) + assert "electron_temperature" in result + assert "eV" in result + assert "$T_e$" in result + + def test_fetch_multiple_comma_separated(self): + from imas_codex.llm.sn_tools import _fetch_standard_names + + mock_gc = MagicMock() + mock_gc.query = MagicMock( + return_value=[ + { + "name": "electron_temperature", + "description": "Te", + "documentation": None, + "kind": "scalar", + "unit": "eV", + "tags": [], + "links": [], + "ids_paths": [], + "constraints": [], + "validity_domain": None, + "physical_base": "temperature", + "subject": "electron", + "component": None, + "coordinate": None, + "position": None, + "process": None, + "review_status": "drafted", + "confidence": None, + "model": None, + "source_ids": [], + "source_ids_names": [], + }, + { + "name": "plasma_current", + "description": "Ip", + "documentation": None, + "kind": "scalar", + "unit": "A", + "tags": [], + "links": [], + "ids_paths": [], + "constraints": [], + "validity_domain": None, + "physical_base": None, + "subject": None, + "component": None, + "coordinate": None, + "position": None, + "process": None, + "review_status": "drafted", + "confidence": None, + "model": None, + "source_ids": [], + "source_ids_names": [], + }, + ] + ) + + result = _fetch_standard_names( + "electron_temperature,plasma_current", gc=mock_gc + ) + assert "electron_temperature" in result + assert "plasma_current" in result + + def test_fetch_not_found(self): + from imas_codex.llm.sn_tools import _fetch_standard_names + + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + result = _fetch_standard_names("nonexistent_name", gc=mock_gc) + assert "not found" in result.lower() or "No" in result + + def test_fetch_partial_not_found(self): + """Shows not found message for missing names.""" + from imas_codex.llm.sn_tools import _fetch_standard_names + + mock_gc = MagicMock() + mock_gc.query = MagicMock( + return_value=[ + { + "name": "electron_temperature", + "description": "Te", + "documentation": None, + "kind": "scalar", + "unit": "eV", + "tags": [], + "links": [], + "ids_paths": [], + "constraints": [], + "validity_domain": None, + "physical_base": None, + "subject": None, + "component": None, + "coordinate": None, + "position": None, + "process": None, + "review_status": "drafted", + "confidence": None, + "model": None, + "source_ids": [], + "source_ids_names": [], + } + ] + ) + + result = _fetch_standard_names("electron_temperature missing_name", gc=mock_gc) + assert "electron_temperature" in result + assert "missing_name" in result + assert "Not found" in result + + +class TestListStandardNames: + """Test _list_standard_names tool.""" + + def test_list_all(self): + from imas_codex.llm.sn_tools import _list_standard_names + + mock_gc = MagicMock() + mock_gc.query = MagicMock( + return_value=[ + { + "name": "electron_temperature", + "kind": "scalar", + "unit": "eV", + "review_status": "drafted", + "description": "Te", + }, + { + "name": "plasma_current", + "kind": "scalar", + "unit": "A", + "review_status": "drafted", + "description": "Ip", + }, + ] + ) + + result = _list_standard_names(gc=mock_gc) + assert "electron_temperature" in result + assert "plasma_current" in result + + def test_list_with_tag_filter(self): + from imas_codex.llm.sn_tools import _list_standard_names + + mock_gc = MagicMock() + mock_gc.query = MagicMock( + return_value=[ + { + "name": "electron_temperature", + "kind": "scalar", + "unit": "eV", + "review_status": "drafted", + "description": "Te", + }, + ] + ) + + result = _list_standard_names(tag="core_profiles", gc=mock_gc) + assert "electron_temperature" in result + mock_gc.query.assert_called_once() + + def test_list_empty_results(self): + from imas_codex.llm.sn_tools import _list_standard_names + + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + result = _list_standard_names(tag="nonexistent_tag", gc=mock_gc) + assert "No standard names" in result + + def test_list_filter_info_in_header(self): + """Filter params appear in header.""" + from imas_codex.llm.sn_tools import _list_standard_names + + mock_gc = MagicMock() + mock_gc.query = MagicMock( + return_value=[ + { + "name": "electron_temperature", + "kind": "scalar", + "unit": "eV", + "review_status": "drafted", + "description": "Te", + }, + ] + ) + + result = _list_standard_names(kind="scalar", gc=mock_gc) + assert "kind=scalar" in result + + def test_list_table_format(self): + """Output is a markdown table.""" + from imas_codex.llm.sn_tools import _list_standard_names + + mock_gc = MagicMock() + mock_gc.query = MagicMock( + return_value=[ + { + "name": "electron_temperature", + "kind": "scalar", + "unit": "eV", + "review_status": "drafted", + "description": "Te", + }, + ] + ) + + result = _list_standard_names(gc=mock_gc) + assert "| Name |" in result + assert "| electron_temperature |" in result + + +class TestMCPToolRegistration: + """Test that SN tools are importable and callable.""" + + def test_tools_importable(self): + """SN tools should be importable from sn_tools.""" + from imas_codex.llm.sn_tools import ( + _fetch_standard_names, + _list_standard_names, + _search_standard_names, + ) + + assert callable(_search_standard_names) + assert callable(_fetch_standard_names) + assert callable(_list_standard_names) + + def test_search_signature(self): + """search_standard_names accepts expected kwargs.""" + import inspect + + from imas_codex.llm.sn_tools import _search_standard_names + + sig = inspect.signature(_search_standard_names) + params = set(sig.parameters.keys()) + assert "query" in params + assert "kind" in params + assert "tags" in params + assert "review_status" in params + assert "k" in params + assert "gc" in params + + def test_fetch_signature(self): + """fetch_standard_names accepts expected kwargs.""" + import inspect + + from imas_codex.llm.sn_tools import _fetch_standard_names + + sig = inspect.signature(_fetch_standard_names) + params = set(sig.parameters.keys()) + assert "names" in params + assert "gc" in params + + def test_list_signature(self): + """list_standard_names accepts expected kwargs.""" + import inspect + + from imas_codex.llm.sn_tools import _list_standard_names + + sig = inspect.signature(_list_standard_names) + params = set(sig.parameters.keys()) + assert "tag" in params + assert "kind" in params + assert "review_status" in params + assert "gc" in params From 435ecab66f0a8f97effa1a989b816620ababb7eb Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 12:09:01 +0200 Subject: [PATCH 17/37] test: add embedding coverage and coalesce safety integration tests Adds tests/sn/test_integration.py with two test classes: TestEmbeddingCoverage: - test_write_preserves_existing_embedding: asserts write_standard_names Cypher never sets sn.embedding from a batch param - test_import_preserves_existing_embedding: asserts _write_catalog_entries Cypher uses coalesce(sn.embedding, null) and coalesce(sn.embedded_at, null) - test_embedding_field_not_in_write_batch: asserts no 'embedding' key appears in the gc.query batch dict TestCoalesceSafety: - test_build_does_not_erase_imported_data: verifies coalesce(b.field, sn.field) for all 15 optional fields in the MERGE SET clause - test_build_with_none_fields_preserves_graph: verifies absent fields appear as None in batch (required for coalesce evaluation) - test_created_at_preserved_on_rewrite: asserts coalesce(sn.created_at, datetime()) pattern preserves original creation timestamp - test_import_then_build_preserves_catalog_fields: end-to-end mock walk through import then build, confirming coalesce semantics in both --- tests/sn/test_integration.py | 431 +++++++++++++++++++++++++++++++++++ 1 file changed, 431 insertions(+) create mode 100644 tests/sn/test_integration.py diff --git a/tests/sn/test_integration.py b/tests/sn/test_integration.py new file mode 100644 index 000000000..694daa079 --- /dev/null +++ b/tests/sn/test_integration.py @@ -0,0 +1,431 @@ +"""Integration tests for embedding coverage, coalesce safety, and round-trip idempotence. + +Verifies that: +1. Embedding fields are never accidentally erased by write_standard_names + or _write_catalog_entries. +2. All optional fields in write_standard_names use coalesce so that a + None value in the batch never overwrites existing graph data. +3. created_at is preserved across rewrites. +4. The import → build cycle is safe: catalog-imported rich fields are + not erased by a subsequent sn-build write. +5. publish → import → publish is idempotent (key fields round-trip cleanly). +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock, call, patch + +import pytest +import yaml + +# ============================================================================= +# Helpers +# ============================================================================= + + +def _call_write(names: list[dict], mock_gc: MagicMock) -> int: + """Call write_standard_names with a mocked GraphClient.""" + with patch("imas_codex.sn.graph_ops.GraphClient") as MockGC: + MockGC.return_value.__enter__ = MagicMock(return_value=mock_gc) + MockGC.return_value.__exit__ = MagicMock(return_value=False) + from imas_codex.sn.graph_ops import write_standard_names + + return write_standard_names(names) + + +def _call_import_write( + entries: list[dict], mock_gc: MagicMock, catalog_sha: str | None = None +) -> int: + """Call _write_catalog_entries with a mocked GraphClient.""" + with patch("imas_codex.graph.client.GraphClient") as MockGC: + MockGC.return_value.__enter__ = MagicMock(return_value=mock_gc) + MockGC.return_value.__exit__ = MagicMock(return_value=False) + from imas_codex.sn.catalog_import import _write_catalog_entries + + return _write_catalog_entries(entries, catalog_commit_sha=catalog_sha) + + +def _merge_cypher(mock_gc: MagicMock) -> str: + """Return the Cypher string from the first MERGE query call.""" + return mock_gc.query.call_args_list[0][0][0] + + +def _merge_batch(mock_gc: MagicMock) -> list[dict]: + """Return the batch parameter from the first MERGE query call.""" + return mock_gc.query.call_args_list[0][1]["batch"] + + +# ============================================================================= +# Part 1: Embedding Coverage +# ============================================================================= + + +class TestEmbeddingCoverage: + """Verify that embedding vectors are never accidentally erased.""" + + def test_write_preserves_existing_embedding(self) -> None: + """write_standard_names does not touch the embedding property at all. + + A StandardName that already has embedding=[0.1, 0.2, 0.3] and + embedded_at set must be unchanged after write_standard_names is called. + The Cypher must not reference the embedding property (since the function + only manages metadata, not embeddings). + """ + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + names = [ + { + "id": "electron_temperature", + "source_type": "dd", + "source_id": "core_profiles/profiles_1d/electrons/temperature", + "description": "Electron temperature", + # No embedding field — write should not touch it + } + ] + _call_write(names, mock_gc) + + cypher = _merge_cypher(mock_gc) + + # The MERGE SET clause must NOT set sn.embedding unconditionally + # (any mention would risk overwriting it with null) + assert "sn.embedding = b.embedding" not in cypher, ( + "write_standard_names must not set sn.embedding from batch param" + ) + + def test_import_preserves_existing_embedding(self) -> None: + """_write_catalog_entries uses coalesce to preserve existing embedding. + + The catalog import must never erase an embedding that was set by the + embedding pipeline. The Cypher should contain the coalesce guard: + sn.embedding = coalesce(sn.embedding, null) + """ + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + entries = [ + { + "id": "electron_temperature", + "description": "Electron temperature", + "documentation": "Te measured by Thomson scattering.", + "kind": "scalar", + "units": "eV", + "tags": ["core_profiles"], + "links": None, + "imas_paths": ["core_profiles/profiles_1d/electrons/temperature"], + "validity_domain": "core plasma", + "constraints": ["T_e > 0"], + "physics_domain": "core_plasma_physics", + "review_status": "accepted", + "source_type": "dd", + "physical_base": "temperature", + "subject": "electron", + "component": None, + "coordinate": None, + "position": None, + "process": None, + } + ] + _call_import_write(entries, mock_gc) + + cypher = _merge_cypher(mock_gc) + + # The catalog import Cypher must preserve embedding via coalesce + assert "coalesce(sn.embedding, null)" in cypher, ( + "_write_catalog_entries Cypher must use coalesce(sn.embedding, null) " + "to preserve existing embeddings" + ) + assert "coalesce(sn.embedded_at, null)" in cypher, ( + "_write_catalog_entries Cypher must use coalesce(sn.embedded_at, null)" + ) + + def test_embedding_field_not_in_write_batch(self) -> None: + """Batch dicts passed to gc.query by write_standard_names must not contain 'embedding'. + + This ensures that even if the caller accidentally includes an + 'embedding' key, the write function strips it before sending to the + graph. More importantly it confirms the build pipeline cannot + null-out embeddings via this code path. + """ + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + names = [ + { + "id": "electron_temperature", + "source_type": "dd", + "source_id": "core_profiles/profiles_1d/electrons/temperature", + "description": "Electron temperature", + "kind": "scalar", + "units": "eV", + "review_status": "drafted", + "confidence": 0.95, + } + ] + _call_write(names, mock_gc) + + batch = _merge_batch(mock_gc) + + for item in batch: + assert "embedding" not in item, ( + f"Batch item for '{item.get('id')}' must not contain 'embedding' key; " + "write_standard_names must never touch embedding data" + ) + + +# ============================================================================= +# Part 2: Coalesce Safety +# ============================================================================= + + +class TestCoalesceSafety: + """Verify that write_standard_names uses coalesce for all optional fields. + + When a field is None in the batch, coalesce(None, sn.field) = sn.field, + so an sn-build re-run cannot accidentally erase data that was set by + an earlier catalog import. + """ + + _COALESCE_FIELDS = [ + ("review_status", "b.review_status, sn.review_status"), + ("documentation", "b.documentation, sn.documentation"), + ("kind", "b.kind, sn.kind"), + ("tags", "b.tags, sn.tags"), + ("links", "b.links, sn.links"), + ("imas_paths", "b.imas_paths, sn.imas_paths"), + ("validity_domain", "b.validity_domain, sn.validity_domain"), + ("constraints", "b.constraints, sn.constraints"), + ("confidence", "b.confidence, sn.confidence"), + ("physical_base", "b.physical_base, sn.physical_base"), + ("subject", "b.subject, sn.subject"), + ("component", "b.component, sn.component"), + ("coordinate", "b.coordinate, sn.coordinate"), + ("position", "b.position, sn.position"), + ("process", "b.process, sn.process"), + ] + + def test_build_does_not_erase_imported_data(self) -> None: + """All optional fields in the MERGE SET must use coalesce(b.field, sn.field). + + This protects against a scenario where: + 1. catalog import sets review_status='accepted', documentation, etc. + 2. sn-build re-runs write_standard_names with those fields = None + 3. Without coalesce, the re-run would null-out the imported values. + """ + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + # Simulate a minimal sn-build write — only id and source_type provided + names = [ + { + "id": "electron_temperature", + "source_type": "dd", + "source_id": "core_profiles/profiles_1d/electrons/temperature", + "description": "Electron temperature", + # review_status, documentation, kind, tags, etc. all absent/None + } + ] + _call_write(names, mock_gc) + + cypher = _merge_cypher(mock_gc) + + for field_name, coalesce_args in self._COALESCE_FIELDS: + assert f"coalesce({coalesce_args})" in cypher, ( + f"Field '{field_name}' must use coalesce({coalesce_args}) in " + "write_standard_names Cypher to preserve existing graph data" + ) + + def test_build_with_none_fields_preserves_graph(self) -> None: + """Batch dicts must include None for absent optional fields (not omit them). + + The coalesce(b.field, sn.field) pattern requires that b.field is + present in the batch parameter (as None, not missing) so that Cypher + can evaluate the coalesce. If the key were absent from the dict, + Neo4j would raise an error or behave unpredictably. + """ + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + # Entry with almost all optional fields missing + names = [ + { + "id": "plasma_current", + "source_type": "dd", + "source_id": "magnetics/method/0/ip", + } + ] + _call_write(names, mock_gc) + + batch = _merge_batch(mock_gc) + assert len(batch) == 1 + item = batch[0] + + # All optional fields must appear in the batch (value may be None) + required_keys = { + "id", + "source_type", + "physical_base", + "subject", + "component", + "coordinate", + "position", + "process", + "description", + "documentation", + "kind", + "tags", + "links", + "imas_paths", + "validity_domain", + "constraints", + "units", + "model", + "review_status", + "generated_at", + "confidence", + } + missing = required_keys - set(item.keys()) + assert not missing, ( + f"Batch item is missing keys: {missing}. " + "All optional fields must be present (even as None) for coalesce to work." + ) + + # Fields absent from source must be None (not some unexpected value) + for key in required_keys - {"id", "source_type"}: + assert item[key] is None, ( + f"Batch key '{key}' should be None when not supplied, got {item[key]!r}" + ) + + def test_created_at_preserved_on_rewrite(self) -> None: + """created_at must use coalesce(sn.created_at, datetime()) — not coalesce(b.created_at, ...). + + This pattern sets created_at on first write and then leaves it + unchanged on all subsequent writes, so the node retains its + original creation timestamp. + """ + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[]) + + names = [ + { + "id": "electron_temperature", + "source_type": "dd", + "source_id": "core_profiles/profiles_1d/electrons/temperature", + } + ] + _call_write(names, mock_gc) + + cypher = _merge_cypher(mock_gc) + + # Pattern: sn.created_at is preserved on re-writes, set only on first + assert "coalesce(sn.created_at, datetime())" in cypher, ( + "write_standard_names must use coalesce(sn.created_at, datetime()) " + "to preserve the original creation timestamp across rewrites" + ) + # Must not blindly set created_at from the batch + assert "b.created_at" not in cypher, ( + "write_standard_names must not set created_at from batch param" + ) + + def test_import_then_build_preserves_catalog_fields(self) -> None: + """Verify coalesce semantics cover the full import → build cycle. + + Step 1: _write_catalog_entries (import) is called with rich metadata. + The Cypher sets catalog-owned fields directly. + Step 2: write_standard_names (build) is called with only basic fields. + The Cypher uses coalesce for all optional fields. + Together, the catalog-set values survive the build re-run because + coalesce(None, sn.documentation) = sn.documentation. + """ + import_gc = MagicMock() + import_gc.query = MagicMock(return_value=[]) + + build_gc = MagicMock() + build_gc.query = MagicMock(return_value=[]) + + # --- Step 1: catalog import with rich fields --- + rich_entry = { + "id": "electron_temperature", + "description": "Electron temperature", + "documentation": "Te measured by Thomson scattering.", + "kind": "scalar", + "units": "eV", + "tags": ["core_profiles"], + "links": None, + "imas_paths": ["core_profiles/profiles_1d/electrons/temperature"], + "validity_domain": "core plasma", + "constraints": ["T_e > 0"], + "physics_domain": "core_plasma_physics", + "review_status": "accepted", + "source_type": "dd", + "physical_base": "temperature", + "subject": "electron", + "component": None, + "coordinate": None, + "position": None, + "process": None, + } + imported = _call_import_write([rich_entry], import_gc) + assert imported == 1 + + # Verify import Cypher sets rich fields directly (no coalesce for catalog-owned) + import_cypher = _merge_cypher(import_gc) + assert "sn.documentation = b.documentation" in import_cypher, ( + "Catalog import must set documentation directly (authoritative)" + ) + assert "sn.review_status = 'accepted'" in import_cypher, ( + "Catalog import must set review_status='accepted' directly" + ) + # Embedding and model must still be protected via coalesce + assert "coalesce(sn.embedding, null)" in import_cypher + + # --- Step 2: sn-build writes basic fields only --- + basic_entry = { + "id": "electron_temperature", + "source_type": "dd", + "source_id": "core_profiles/profiles_1d/electrons/temperature", + "description": "Electron temperature", + # review_status, documentation, kind, validity_domain, constraints all absent + } + written = _call_write([basic_entry], build_gc) + assert written == 1 + + # Verify build Cypher uses coalesce for all catalog-owned fields + build_cypher = _merge_cypher(build_gc) + + catalog_owned = [ + ("review_status", "b.review_status, sn.review_status"), + ("documentation", "b.documentation, sn.documentation"), + ("kind", "b.kind, sn.kind"), + ("tags", "b.tags, sn.tags"), + ("validity_domain", "b.validity_domain, sn.validity_domain"), + ("constraints", "b.constraints, sn.constraints"), + ("confidence", "b.confidence, sn.confidence"), + ] + for field_name, coalesce_args in catalog_owned: + assert f"coalesce({coalesce_args})" in build_cypher, ( + f"write_standard_names must protect '{field_name}' with coalesce " + "so catalog-imported values survive sn-build re-runs" + ) + + # Verify the build batch item has None for the absent fields + build_batch = _merge_batch(build_gc) + assert len(build_batch) == 1 + build_item = build_batch[0] + + # These were not supplied — must be None in batch so coalesce falls back to graph + for absent_field in ( + "review_status", + "documentation", + "kind", + "validity_domain", + ): + assert absent_field in build_item, ( + f"'{absent_field}' must appear in batch dict (as None) for coalesce" + ) + assert build_item[absent_field] is None, ( + f"'{absent_field}' must be None in batch when not supplied, " + f"got {build_item[absent_field]!r}" + ) From 4e706fa8cb6f68e4141dfa2a0984020c104e223e Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 12:14:36 +0200 Subject: [PATCH 18/37] refactor: rename MCP tools for consistent dd/ids naming MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - get_dd_overview → get_dd_catalog (remove query/include_unit_stats params) - analyze_dd_structure + get_ids_structure → get_ids_summary (trimmed output) - get_dd_path_context → find_related_dd_paths - Remove export_imas_ids/export_imas_domain from MCP registration - Remove facade delegation tests - Update all backend methods, REPL functions, formatters, and tests - Rename format_search_imas_report → format_search_dd_report - Include host.py migration cleanup --- imas_codex/cli/host.py | 191 ++++------- imas_codex/ids/tools.py | 2 +- imas_codex/ids/validation.py | 4 +- imas_codex/llm/README.md | 2 +- imas_codex/llm/search_formatters.py | 115 ++----- imas_codex/llm/server.py | 115 ++----- imas_codex/models/result_models.py | 2 +- .../search/decorators/error_handling.py | 6 +- .../search/decorators/tool_recommendations.py | 10 +- imas_codex/search/tool_suggestions.py | 8 +- imas_codex/tools/graph_search.py | 321 +++--------------- imas_codex/tools/utils.py | 2 +- tests/conftest.py | 2 +- tests/core/test_cli.py | 2 +- tests/graph_mcp/test_dd_tool_features.py | 18 +- tests/graph_mcp/test_graph_search.py | 24 +- tests/graph_mcp/test_tool_registration.py | 2 +- tests/integration/test_workflows.py | 8 +- tests/llm/test_graceful_degradation.py | 5 +- tests/llm/test_mcp_bug_regressions.py | 35 +- .../decorators/test_tool_recommendations.py | 6 +- tests/search/test_tool_suggestions.py | 4 +- tests/tools/test_dd_version_filtering.py | 4 +- tests/tools/test_facade_delegation.py | 120 ------- tests/tools/test_tools.py | 2 +- tests/tools/test_utils.py | 2 +- 26 files changed, 202 insertions(+), 810 deletions(-) delete mode 100644 tests/tools/test_facade_delegation.py diff --git a/imas_codex/cli/host.py b/imas_codex/cli/host.py index 5b2f53e9b..71562a748 100644 --- a/imas_codex/cli/host.py +++ b/imas_codex/cli/host.py @@ -791,109 +791,39 @@ def _migrate_from_node( """Kill imas-codex processes and zellij sessions on the old node. Called automatically when ``--set-default`` switches to a different - node. Zellij session layouts are dumped to - ``~/.local/share/imas-codex/zellij-layouts/`` before sessions are - killed, so they can be restored on the new node with:: - - zellij --layout ~/.local/share/imas-codex/zellij-layouts/.kdl - - Uses two SSH calls: - 1. Dump zellij layouts (must happen while sessions are alive) - 2. Kill imas-codex/litellm processes, then zellij sessions + node. Uses a single SSH call to: + 1. Send SIGINT to imas-codex/litellm processes (graceful shutdown) + 2. SIGTERM stragglers after a grace period + 3. Delete zellij sessions AND kill orphaned server daemons + 4. Detect lingering VS Code server sessions that could re-spawn """ old_short = old_hostname.split(".")[0] fqdn = old_hostname if "." in old_hostname else f"{old_hostname}.iter.org" click.echo(f"\n Migrating from {click.style(old_short, fg='yellow')}…") - layout_dir = "~/.local/share/imas-codex/zellij-layouts" - - # --- Phase 0: Dump zellij layouts (sessions must still be alive) --- - dump_script = ( - f"mkdir -p {layout_dir}; " - "if command -v zellij >/dev/null 2>&1; then " - " sessions=$(zellij list-sessions -s 2>/dev/null || true); " - ' if [ -n "$sessions" ]; then ' - " saved=0; " - " for s in $sessions; do " - f' layout=$(zellij -s "$s" action dump-layout 2>/dev/null || true); ' - ' if [ -n "$layout" ]; then ' - f' echo "$layout" > {layout_dir}/"$s".kdl; ' - " saved=$((saved + 1)); " - " fi; " - " done; " - ' echo "layouts:$saved"; ' - " else " - " echo 'layouts:0'; " - " fi; " - "else " - " echo 'layouts:none'; " - "fi" - ) - - layouts_saved = 0 - try: - result = _ssh_to_node(fqdn, gateway, user, dump_script, timeout=30) - if result.returncode == 0: - for line in result.stdout.strip().splitlines(): - if line.startswith("layouts:"): - val = line.split(":")[1] - if val not in ("none", "0"): - layouts_saved = int(val) - click.echo( - f" {click.style('✓', fg='green')} " - f"Saved {layouts_saved} zellij layout(s) to " - f"{click.style(layout_dir + '/', fg='cyan')}" - ) - except (subprocess.TimeoutExpired, Exception): - pass # Layout dump is best-effort; cleanup proceeds regardless - - # --- Phase 1+2: Kill processes, then zellij sessions --- + # Build pattern regex from _CODEX_PATTERNS, skip neo4j (shared service). kill_patterns = [p for p in _CODEX_PATTERNS if p != "neo4j"] pattern_re = "|".join(kill_patterns) - # CRITICAL: pgrep -f matches the command line of ALL processes, - # including the SSH shell running this script (its argv contains - # the pattern string). We must exclude $$ (shell) and $PPID - # (sshd) from the kill list, otherwise kill -INT destroys our own - # SSH connection (exit 255, empty output). + # Single SSH call for all cleanup — reduces failure surface vs + # multiple independent SSH connections that can each timeout. migrate_script = ( - # Collect PIDs, then filter out our own process tree - f"all_pids=$(pgrep -u $USER -f '{pattern_re}' 2>/dev/null || true); " - "pids=''; " - "for p in $all_pids; do " - ' if [ "$p" != "$$" ] && [ "$p" != "$PPID" ]; then ' - ' pids="$pids $p"; ' - " fi; " - "done; " - "pids=$(echo $pids | xargs); " - 'if [ -n "$pids" ]; then ' - ' count=$(echo "$pids" | wc -w); ' - " kill -INT $pids 2>/dev/null; " - " sleep 2; " - f" all_rem=$(pgrep -u $USER -f '{pattern_re}' 2>/dev/null || true); " - " remaining=''; " - " for p in $all_rem; do " - ' if [ "$p" != "$$" ] && [ "$p" != "$PPID" ]; then ' - ' remaining="$remaining $p"; ' - " fi; " - " done; " - " remaining=$(echo $remaining | xargs); " - ' if [ -n "$remaining" ]; then ' - " kill -TERM $remaining 2>/dev/null; " - " fi; " - ' echo "killed:$count"; ' - "else " - " echo 'killed:0'; " - "fi; " - # Detect VS Code server (warns user) - "vscode_pids=$(pgrep -u $USER -f 'code-server|vscode-server' 2>/dev/null || true); " - 'if [ -n "$vscode_pids" ]; then ' - ' echo "vscode:$(echo "$vscode_pids" | wc -w)"; ' - "else " - " echo 'vscode:0'; " - "fi; " - # Kill zellij sessions (layouts already saved above) + # --- Phase 1: Kill imas-codex processes --- + f"pids=$(pgrep -u $USER -f '{pattern_re}' 2>/dev/null || true); " + f'if [ -n "$pids" ]; then ' + f' count=$(echo "$pids" | wc -w); ' + f" kill -INT $pids 2>/dev/null; " + f" sleep 2; " + f" remaining=$(pgrep -u $USER -f '{pattern_re}' 2>/dev/null || true); " + f' if [ -n "$remaining" ]; then ' + f" kill -TERM $remaining 2>/dev/null; " + f" fi; " + f' echo "killed:$count"; ' + f"else " + f" echo 'killed:0'; " + f"fi; " + # --- Phase 2: Zellij sessions + orphaned server daemons --- "if command -v zellij >/dev/null 2>&1; then " " sessions=$(zellij list-sessions -s 2>/dev/null || true); " ' if [ -n "$sessions" ]; then ' @@ -906,20 +836,22 @@ def _migrate_from_node( "else " " echo 'zj:none'; " "fi; " - # Kill orphaned zellij server daemons (PPID=1) - "all_zj=$(pgrep -u $USER -f 'zellij --server' 2>/dev/null || true); " - "zj_servers=''; " - "for p in $all_zj; do " - ' if [ "$p" != "$$" ] && [ "$p" != "$PPID" ]; then ' - ' zj_servers="$zj_servers $p"; ' - " fi; " - "done; " - "zj_servers=$(echo $zj_servers | xargs); " + # Fallback: kill orphaned zellij --server daemons (PPID=1) that + # survive delete-all-sessions. These spin at high CPU doing + # nothing and are the main cause of cleanup failure. + "zj_servers=$(pgrep -u $USER -f 'zellij --server' 2>/dev/null || true); " 'if [ -n "$zj_servers" ]; then ' " kill -TERM $zj_servers 2>/dev/null; " ' echo "zj_servers:$(echo "$zj_servers" | wc -w)"; ' "else " " echo 'zj_servers:0'; " + "fi; " + # --- Phase 3: Detect VS Code server (warns user) --- + "vscode_pids=$(pgrep -u $USER -f 'code-server|vscode-server' 2>/dev/null || true); " + 'if [ -n "$vscode_pids" ]; then ' + ' echo "vscode:$(echo "$vscode_pids" | wc -w)"; ' + "else " + " echo 'vscode:0'; " "fi" ) @@ -937,28 +869,25 @@ def _migrate_from_node( ) else: click.echo( - f" {click.style('·', dim=True)} " + f" {click.style('·', fg='dim')} " f"No imas-codex processes running" ) - elif line.startswith("vscode:"): - count = line.split(":")[1] - if count != "0": - click.echo( - f" {click.style('⚠', fg='yellow')} " - f"VS Code server still running on {old_short} " - f"({count} procs) — may re-spawn MCP servers" - ) elif line.startswith("zj:"): val = line.split(":")[1] if val == "none": click.echo( - f" {click.style('·', dim=True)} " + f" {click.style('·', fg='dim')} " f"zellij not found on {old_short}" ) - elif val != "0": + elif val == "0": + click.echo( + f" {click.style('·', fg='dim')} " + f"No zellij sessions to clean up" + ) + else: click.echo( f" {click.style('✓', fg='green')} " - f"Killed {val} zellij session(s)" + f"Deleted {val} zellij session(s)" ) elif line.startswith("zj_servers:"): count = line.split(":")[1] @@ -967,6 +896,14 @@ def _migrate_from_node( f" {click.style('✓', fg='green')} " f"Killed {count} orphaned zellij server(s)" ) + elif line.startswith("vscode:"): + count = line.split(":")[1] + if count != "0": + click.echo( + f" {click.style('⚠', fg='yellow')} " + f"VS Code server still running on {old_short} " + f"({count} procs) — may re-spawn MCP servers" + ) else: stderr_hint = "" if result.stderr: @@ -988,10 +925,6 @@ def _migrate_from_node( f" {click.style('⚠', fg='yellow')} Error during cleanup on {old_short}" ) - if layouts_saved: - click.echo( - f" Restore: {click.style(f'zellij --layout {layout_dir}/.kdl', fg='cyan')}" - ) click.echo( f" Reconnect: {click.style('cx', fg='cyan', bold=True)} to start fresh" ) @@ -1043,7 +976,7 @@ def _handle_llm_alias( existing = _get_ssh_hostname(llm_alias) if existing == llm_fqdn: click.echo( - click.style(" · ", dim=True) + click.style(" · ", fg="dim") + f"LLM already set: {llm_alias} → " + click.style(existing, fg="cyan") ) @@ -1264,8 +1197,11 @@ def _discover_and_survey(): if current == target_fqdn: click.echo(f" Already set: {facility} → {target_fqdn}") else: - # Migrate FIRST — SSH to old node while control sockets - # and gateway connections are still alive. + # Migrate FIRST — clean up processes on the old node + # while SSH connections are still warm. Must happen + # before we kill ControlMaster sockets, update config, + # or stop tunnels, all of which can break connectivity + # to the old node. if current: _migrate_from_node( current, @@ -1274,7 +1210,7 @@ def _discover_and_survey(): timeout, ) - # Kill ControlMaster sockets AFTER migration so + # Kill ControlMaster sockets BEFORE config update so # ssh -O exit resolves to the old (current) HostName. sockets_to_kill = [facility] llm_alias = f"{facility}-llm" @@ -1314,17 +1250,6 @@ def _discover_and_survey(): # No --set-llm: LLM travels with the new default _handle_llm_alias(facility, target_fqdn, results) - # Show processes on the new target node after migration - if target_fqdn and current != target_fqdn: - target_short = target_fqdn.split(".")[0] - click.echo() - _, new_info = _query_node(target_short, gateway, user, timeout) - if new_info is not None: - procs = {target_short: new_info.get("codex_procs", [])} - _show_processes(procs, title=f"Processes on {target_short}") - else: - click.echo(f" Could not query {target_short} for process list") - return if llm_node is not None: diff --git a/imas_codex/ids/tools.py b/imas_codex/ids/tools.py index 34cec554a..42e304e5b 100644 --- a/imas_codex/ids/tools.py +++ b/imas_codex/ids/tools.py @@ -209,7 +209,7 @@ def analyze_units( # --------------------------------------------------------------------------- -def check_imas_paths( +def check_dd_paths( paths: list[str], *, gc: GraphClient | None = None, diff --git a/imas_codex/ids/validation.py b/imas_codex/ids/validation.py index a5b76a1c3..89cf82dc6 100644 --- a/imas_codex/ids/validation.py +++ b/imas_codex/ids/validation.py @@ -20,7 +20,7 @@ from imas_codex.graph.client import GraphClient from imas_codex.ids.models import EscalationFlag, EscalationSeverity -from imas_codex.ids.tools import analyze_units, check_imas_paths +from imas_codex.ids.tools import analyze_units, check_dd_paths from imas_codex.ids.transforms import execute_transform logger = logging.getLogger(__name__) @@ -107,7 +107,7 @@ def validate_mapping( # Batch checks source_exists = _check_sources_exist(source_ids, gc) - target_results = {r["path"]: r for r in check_imas_paths(target_paths, gc=gc)} + target_results = {r["path"]: r for r in check_dd_paths(target_paths, gc=gc)} for b in bindings: check = BindingCheck(source_id=b.source_id, target_id=b.target_id) diff --git a/imas_codex/llm/README.md b/imas_codex/llm/README.md index 64445bb91..cf9e33a68 100644 --- a/imas_codex/llm/README.md +++ b/imas_codex/llm/README.md @@ -52,7 +52,7 @@ The `python()` REPL includes rich pre-loaded utilities: - `fetch_dd_paths(paths)` - Full documentation for paths - `list_dd_paths(paths, leaf_only, max_paths)` - List IDS structure - `check_dd_paths(paths)` - Validate path existence -- `get_dd_overview(query)` - High-level DD summary +- `get_dd_catalog(query)` - High-level DD summary ### CLI Agent Usage diff --git a/imas_codex/llm/search_formatters.py b/imas_codex/llm/search_formatters.py index b122b2c44..ed55680ca 100644 --- a/imas_codex/llm/search_formatters.py +++ b/imas_codex/llm/search_formatters.py @@ -1287,75 +1287,7 @@ def format_path_context_report(result: dict[str, Any]) -> str: def format_structure_report(result: dict[str, Any]) -> str: - """Format analyze_imas_structure result into readable text.""" - tool_error = _format_tool_error(result) - if tool_error: - return tool_error - - parts: list[str] = [] - ids_name = result.get("ids_name", "") - dd_version = result.get("dd_version") - - header = f"## IDS Structure Analysis: {ids_name}" - if dd_version is not None: - header += f" (DD v{dd_version})" - parts.append(header + "\n") - - # Version context note when filtered - version_ctx = result.get("version_context") - if version_ctx: - parts.append(f"> {version_ctx['note']}") - dep = version_ctx.get("deprecated_in_or_before", 0) - ren = version_ctx.get("renamed_paths", 0) - if dep or ren: - ctx_parts = [] - if dep: - ctx_parts.append(f"{dep} deprecated") - if ren: - ctx_parts.append(f"{ren} renamed") - parts.append( - f"> Version changes: {', '.join(ctx_parts)} paths in this IDS." - ) - parts.append("") - - parts.append(f"- Total paths: {result.get('total_paths', 0)}") - parts.append(f"- Leaf fields: {result.get('leaf_count', 0)}") - parts.append(f"- Structures: {result.get('structure_count', 0)}") - parts.append(f"- Max depth: {result.get('max_depth', 0)}") - parts.append(f"- Avg depth: {result.get('avg_depth', 0)}") - - domains = result.get("physics_domains", []) - if domains: - parts.append("\n### Physics Domains") - for d in domains: - parts.append(f" - {d['domain']}: {d['count']} paths") - - types = result.get("data_types", []) - if types: - parts.append("\n### Data Types") - for t in types: - parts.append(f" - {t['type']}: {t['count']}") - - arrays = result.get("array_structures", []) - if arrays: - parts.append(f"\n### Array Structures ({len(arrays)})") - for a in arrays[:20]: - coords = ", ".join(a.get("coordinates", [])) - parts.append(f" - `{a['path']}` → [{coords}]") - if len(arrays) > 20: - parts.append(f" ... and {len(arrays) - 20} more") - - cocos = result.get("cocos_fields", []) - if cocos: - parts.append(f"\n### COCOS-Labeled Fields ({len(cocos)})") - for c in cocos: - parts.append(f" - `{c['path']}` ({c['label']})") - - return "\n".join(parts) - - -def format_ids_structure_report(result: dict[str, Any]) -> str: - """Format get_ids_structure result into a compact, rich overview.""" + """Format get_ids_summary result into a compact overview.""" if isinstance(result, dict) and result.get("error"): return f"Error: {result['error']}" @@ -1403,12 +1335,24 @@ def format_ids_structure_report(result: dict[str, Any]) -> str: parts.append("\n### Data Types\n") parts.append(" " + " | ".join(f"{k}: {v}" for k, v in dtypes.items())) - # Clusters - clusters = result.get("clusters", []) - if clusters: - parts.append(f"\n### Semantic Clusters ({len(clusters)})\n") - for c in clusters: - parts.append(f" {c['label']} [{c['scope']}] ({c['members']} paths)") + # Counts with pointers to dedicated tools + cluster_count = result.get("semantic_clusters", 0) + cocos_count = result.get("cocos_fields", 0) + coord_count = result.get("coordinate_arrays", 0) + if cluster_count or cocos_count or coord_count: + parts.append("\n### Cross-References\n") + if cluster_count: + parts.append( + f" Semantic clusters: {cluster_count} " + f"(use `search_dd_clusters(ids_filter='{ids_name}')` for details)" + ) + if cocos_count: + parts.append( + f" COCOS fields: {cocos_count} " + f"(use `get_dd_cocos_fields(ids_filter='{ids_name}')` for details)" + ) + if coord_count: + parts.append(f" Coordinate arrays: {coord_count}") # Identifier schemas idents = result.get("identifier_schemas", []) @@ -1418,28 +1362,11 @@ def format_ids_structure_report(result: dict[str, Any]) -> str: examples = ", ".join(i.get("examples", [])[:2]) parts.append(f" {i['schema']} (×{i['usage_count']}) e.g. {examples}") - # COCOS - cocos = result.get("cocos_fields", []) - if cocos: - parts.append(f"\n### COCOS Fields ({len(cocos)})\n") - for c in cocos: - parts.append(f" `{c['path']}` ({c['label']})") - - # Coordinate arrays (compact) - coords = result.get("coordinate_arrays", []) - if coords: - parts.append(f"\n### Coordinate Arrays ({len(coords)})\n") - for ca in coords[:10]: - clist = ", ".join(ca.get("coordinates", [])) - parts.append(f" `{ca['path']}` → [{clist}]") - if len(coords) > 10: - parts.append(f" ... and {len(coords) - 10} more") - return "\n".join(parts) def format_export_ids_report(result: dict[str, Any]) -> str: - """Format export_imas_ids result into readable text.""" + """Format export_dd_ids result into readable text.""" tool_error = _format_tool_error(result) if tool_error: return tool_error @@ -1480,7 +1407,7 @@ def format_export_ids_report(result: dict[str, Any]) -> str: def format_export_domain_report(result: Any) -> str: - """Format export_imas_domain result into readable text.""" + """Format export_dd_domain result into readable text.""" parts: list[str] = [] if isinstance(result, dict): domain = result.get("domain", "") diff --git a/imas_codex/llm/server.py b/imas_codex/llm/server.py index 411f9dea0..c1e07aa3d 100644 --- a/imas_codex/llm/server.py +++ b/imas_codex/llm/server.py @@ -1139,30 +1139,21 @@ def check_dd_paths(paths: str, dd_version: int | None = None) -> str: except Exception as e: return f"Check error: {e}" - def get_dd_overview( - query_text: str | None = None, + def get_dd_catalog( dd_version: int | None = None, - include_unit_stats: bool = False, ) -> str: - """Get high-level overview of IMAS Data Dictionary. + """Get full catalog of all IMAS IDSs. Args: - query_text: Optional keyword filter dd_version: Filter by DD major version (e.g., 3 or 4) - include_unit_stats: If true, include unit distribution statistics Returns: - Overview with IDS list, physics domains, statistics + Catalog with all IDS names, descriptions, path counts, physics domains """ try: - if include_unit_stats: - logger.debug( - "include_unit_stats not yet implemented in backend, ignoring" - ) tools = _get_imas_tools() result = _run_async( - tools.overview_tool.get_dd_overview( - query=query_text, + tools.overview_tool.get_dd_catalog( dd_version=dd_version, ) ) @@ -1170,7 +1161,7 @@ def get_dd_overview( except Exception as e: return f"Overview error: {e}" - def get_dd_path_context( + def find_related_dd_paths( path: str, relationship_types: str = "all", dd_version: int | None = None, @@ -1188,7 +1179,7 @@ def get_dd_path_context( try: tools = _get_imas_tools() result = _run_async( - tools.path_context_tool.get_dd_path_context( + tools.path_context_tool.find_related_dd_paths( path=path, relationship_types=relationship_types, dd_version=dd_version, @@ -1198,7 +1189,7 @@ def get_dd_path_context( except Exception as e: return f"Path context error: {e}" - def export_imas_ids( + def export_dd_ids( ids_name: str, leaf_only: bool = False, dd_version: int | None = None, @@ -1226,7 +1217,7 @@ def export_imas_ids( except Exception as e: return f"Export error: {e}" - def export_imas_domain( + def export_dd_domain( domain: str, ids_filter: str | None = None, dd_version: int | None = None, @@ -1564,10 +1555,10 @@ def wrapper(*args, **kwargs): ("fetch_dd_paths", fetch_dd_paths), ("list_dd_paths", list_dd_paths), ("check_dd_paths", check_dd_paths), - ("get_dd_overview", get_dd_overview), - ("get_dd_path_context", get_dd_path_context), - ("export_imas_ids", export_imas_ids), - ("export_imas_domain", export_imas_domain), + ("get_dd_catalog", get_dd_catalog), + ("find_related_dd_paths", find_related_dd_paths), + ("export_dd_ids", export_dd_ids), + ("export_dd_domain", export_dd_domain), ], ), ( @@ -1638,7 +1629,7 @@ def repl_help() -> str: "update_metadata": update_metadata, "install_tools": install_tools, "search_code": search_code, - "get_dd_overview": get_dd_overview, + "get_dd_catalog": get_dd_catalog, "cocos_sign_flip_paths": cocos_sign_flip_paths, # REPL management "reload": _reload_repl, @@ -2716,33 +2707,25 @@ def list_dd_paths( return format_list_report(result) @self.mcp.tool() - def get_dd_overview( - query: str | None = None, + def get_dd_catalog( dd_version: int | None = None, - include_unit_stats: bool = False, ) -> str: """List all available IDSs (Interface Data Structures) with descriptions and statistics. Use as a starting point to discover which IDS contains the data you need. Each IDS entry includes its description, total path count, and physics domain classification. Args: - query: Optional keyword to filter IDS names and descriptions (e.g. "magnetics", "transport"). Default: list all IDSs. dd_version: Filter by DD major version (3 or 4). Default: latest version. - include_unit_stats: If true, include unit distribution statistics in the response. Default: false. Returns: Formatted text report listing each IDS with its description, path count, and physics domain. """ from imas_codex.llm.search_formatters import format_overview_report - if include_unit_stats: - logger.debug("Including unit distribution statistics") tools = _get_imas_tools() result = _run_async( - tools.overview_tool.get_dd_overview( - query=query, + tools.overview_tool.get_dd_catalog( dd_version=dd_version, - include_unit_stats=include_unit_stats, ) ) return format_overview_report(result) @@ -2831,7 +2814,7 @@ def find_related_dd_paths( tools = _get_imas_tools() result = _run_async( - tools.path_context_tool.get_dd_path_context( + tools.path_context_tool.find_related_dd_paths( path=normalize_imas_path(path), relationship_types=relationship_types, max_results=max_results, @@ -2841,65 +2824,7 @@ def find_related_dd_paths( return format_path_context_report(result) @self.mcp.tool() - def export_imas_ids( - ids_name: str, - leaf_only: bool = False, - dd_version: int | None = None, - ) -> str: - """Export every path in an IDS with full metadata. Use when you need the complete schema of an IDS — all paths with their types, units, coordinates, cluster labels, and COCOS annotations. - - Warning: large IDSs can produce very long output. Use leaf_only=true to reduce volume by excluding intermediate structure nodes. - - Args: - ids_name: IDS name to export (e.g. "equilibrium", "core_profiles"). - leaf_only: If true, return only leaf data fields (skip structures). Default: false. - dd_version: Filter by DD major version (3 or 4). Default: latest version. - - Returns: - Formatted text listing every path in the IDS with documentation, type, units, and coordinates. - """ - from imas_codex.llm.search_formatters import format_export_ids_report - - tools = _get_imas_tools() - result = _run_async( - tools.structure_tool.export_dd_ids( - ids_name=ids_name, - leaf_only=leaf_only, - dd_version=dd_version, - ) - ) - return format_export_ids_report(result) - - @self.mcp.tool() - def export_imas_domain( - domain: str, - ids_filter: str | None = None, - dd_version: int | None = None, - ) -> str: - """Export all IMAS paths classified under a physics domain, grouped by IDS. Use to see every path in the DD that belongs to a domain like "magnetics" or "transport". - - Args: - domain: Physics domain name (e.g. "magnetics", "equilibrium", "transport", "core_profiles"). - ids_filter: Optional IDS name to restrict output to a single IDS. Default: all IDSs in the domain. - dd_version: Filter by DD major version (3 or 4). Default: latest version. - - Returns: - Formatted text report listing paths with documentation and units, organized by IDS. - """ - from imas_codex.llm.search_formatters import format_export_domain_report - - tools = _get_imas_tools() - result = _run_async( - tools.structure_tool.export_dd_domain( - domain=domain, - ids_filter=ids_filter, - dd_version=dd_version, - ) - ) - return format_export_domain_report(result) - - @self.mcp.tool() - def get_ids_structure( + def get_ids_summary( ids_name: str, dd_version: int | None = None, ) -> str: @@ -2915,16 +2840,16 @@ def get_ids_structure( Returns: Formatted text report with structural overview of the IDS. """ - from imas_codex.llm.search_formatters import format_ids_structure_report + from imas_codex.llm.search_formatters import format_structure_report tools = _get_imas_tools() result = _run_async( - tools.structure_tool.get_ids_structure( + tools.structure_tool.get_ids_summary( ids_name=ids_name, dd_version=dd_version, ) ) - return format_ids_structure_report(result) + return format_structure_report(result) @self.mcp.tool() def get_dd_cocos_fields( diff --git a/imas_codex/models/result_models.py b/imas_codex/models/result_models.py index a98735a16..58bf807f9 100644 --- a/imas_codex/models/result_models.py +++ b/imas_codex/models/result_models.py @@ -313,7 +313,7 @@ class GetOverviewResult(WithPhysics, ToolResult, SearchHits): @property def tool_name(self) -> str: """Name of the tool that generated this result.""" - return "get_dd_overview" + return "get_dd_catalog" content: str available_ids: list[str] = Field(default_factory=list) diff --git a/imas_codex/search/decorators/error_handling.py b/imas_codex/search/decorators/error_handling.py index a89883e80..d6cd22e2b 100644 --- a/imas_codex/search/decorators/error_handling.py +++ b/imas_codex/search/decorators/error_handling.py @@ -187,7 +187,7 @@ def get_fallback_response( "query": query, "suggestions": [ { - "tool": "get_dd_overview", + "tool": "get_dd_catalog", "reason": "Get overview of available IMAS data", "description": "Explore data structure and capabilities", }, @@ -197,7 +197,7 @@ def get_fallback_response( "description": "Discover alternative search terms", }, { - "tool": "get_dd_overview", + "tool": "get_dd_catalog", "reason": f'Learn about "{query}" in fusion physics', "description": "Get conceptual understanding", }, @@ -215,7 +215,7 @@ def get_fallback_response( "description": "Find specific measurements and data paths", }, { - "tool": "get_dd_overview", + "tool": "get_dd_catalog", "reason": "Get general overview of IMAS concepts", "description": "Explore available physics domains", }, diff --git a/imas_codex/search/decorators/tool_recommendations.py b/imas_codex/search/decorators/tool_recommendations.py index e2b0c1a91..d84f72efc 100644 --- a/imas_codex/search/decorators/tool_recommendations.py +++ b/imas_codex/search/decorators/tool_recommendations.py @@ -96,7 +96,7 @@ def generate_search_suggestions( for domain in context["domains"][:2]: # Limit suggestions suggestions.append( { - "tool": "get_dd_overview", + "tool": "get_dd_catalog", "reason": f"Learn more about {domain} physics domain", "description": f"Get detailed explanation of {domain} concepts", } @@ -116,7 +116,7 @@ def generate_search_suggestions( # No results - suggest broader search strategies suggestions.append( { - "tool": "get_dd_overview", + "tool": "get_dd_catalog", "reason": "No results found - get overview of available data", "description": "Explore IMAS data structure and available concepts", } @@ -133,7 +133,7 @@ def generate_search_suggestions( # Suggest concept explanation for the query suggestions.append( { - "tool": "get_dd_overview", + "tool": "get_dd_catalog", "reason": f'Learn about "{query}" concept in fusion physics', "description": "Get conceptual understanding and context", } @@ -215,7 +215,7 @@ def generate_tool_recommendations( # Error case - suggest diagnostic tools return [ { - "tool": "get_dd_overview", + "tool": "get_dd_catalog", "reason": "Get overview of available data and functionality", "description": "Explore IMAS capabilities and data structure", } @@ -244,7 +244,7 @@ def generate_tool_recommendations( "description": "Find relevant IMAS data for your research", }, { - "tool": "get_dd_overview", + "tool": "get_dd_catalog", "reason": "Get overview of IMAS structure", "description": "Understand available data and capabilities", }, diff --git a/imas_codex/search/tool_suggestions.py b/imas_codex/search/tool_suggestions.py index 4b44745a4..a5cb575c5 100644 --- a/imas_codex/search/tool_suggestions.py +++ b/imas_codex/search/tool_suggestions.py @@ -35,9 +35,9 @@ def suggest_follow_up_tools( if results.get("results"): suggestions.append( { - "tool": "get_dd_overview", + "tool": "get_dd_catalog", "reason": "Get detailed explanation of physics concepts found in search results", - "sample_call": "get_dd_overview(query='plasma temperature')", + "sample_call": "get_dd_catalog()", } ) @@ -54,7 +54,7 @@ def suggest_follow_up_tools( ) break - elif func_name == "get_dd_overview": + elif func_name == "get_dd_catalog": # After concept explanation, suggest searching for related data concept = results.get("concept", "") if concept: @@ -78,7 +78,7 @@ def suggest_follow_up_tools( } ) - elif func_name == "get_dd_overview": + elif func_name == "get_dd_catalog": # After overview, suggest searching for specific topics suggestions.extend( [ diff --git a/imas_codex/tools/graph_search.py b/imas_codex/tools/graph_search.py index de6d6cec3..889885204 100644 --- a/imas_codex/tools/graph_search.py +++ b/imas_codex/tools/graph_search.py @@ -1082,25 +1082,23 @@ def __init__(self, graph_client: GraphClient): @property def tool_name(self) -> str: - return "get_dd_overview" + return "get_dd_catalog" @cache_results(ttl=3600) @handle_errors(fallback="overview_error") @mcp_tool( - "Get an overview of available IMAS Interface Data Structures (IDS). " - "Returns IDS names, descriptions, path counts, and physics domains. " - "query: Optional filter to narrow results (e.g., 'magnetics' or 'plasma equilibrium'). " - "dd_version: Filter by DD major version (e.g., 3 or 4). None returns all versions. " - "include_unit_stats: If true, include unit distribution statistics." + "List all available IDSs (Interface Data Structures) with descriptions " + "and statistics. Returns every IDS with name, description, path count, " + "physics domain, and lifecycle status. Use as a starting point to discover " + "which IDS contains the data you need. " + "dd_version: Filter by DD major version (e.g., 3 or 4). None returns latest." ) - async def get_dd_overview( + async def get_dd_catalog( self, - query: str | None = None, dd_version: int | None = None, - include_unit_stats: bool = False, ctx: Context | None = None, ) -> GetOverviewResult: - """Get overview from graph.""" + """Get full catalog of all IDSs from graph.""" import importlib.metadata dd_params: dict[str, Any] = {} @@ -1138,50 +1136,8 @@ async def get_dd_overview( ids_statistics = {} physics_domains = set() - # If query provided, try semantic search via ids_embedding vector index - semantic_scores: dict[str, float] = {} - if query: - try: - from imas_codex.embeddings.config import EncoderConfig - from imas_codex.embeddings.encoder import Encoder - from imas_codex.settings import get_embedding_model - - encoder = Encoder( - config=EncoderConfig( - model_name=get_embedding_model(), - normalize_embeddings=True, - ) - ) - query_vec = encoder.embed_texts([query])[0].tolist() - sem_results = self._gc.query( - """ - CALL db.index.vector.queryNodes( - 'ids_embedding', $k, $query_vec - ) YIELD node, score - RETURN node.id AS name, score - """, - k=20, - query_vec=query_vec, - ) - for r in sem_results or []: - semantic_scores[r["name"]] = r["score"] - except Exception: - pass # Vector index may not exist yet - for r in ids_results or []: ids_name = r["name"] - # Apply query filter: text match OR semantic match - if query: - query_lower = query.lower() - name_match = query_lower in ids_name.lower() - desc_match = (r["description"] or "").lower().find(query_lower) >= 0 - domain_match = (r["physics_domain"] or "").lower().find( - query_lower - ) >= 0 - semantic_match = ids_name in semantic_scores - if not (name_match or desc_match or domain_match or semantic_match): - continue - all_ids.append(ids_name) ids_statistics[ids_name] = { "path_count": r["path_count"], @@ -1205,28 +1161,13 @@ async def get_dd_overview( lc = r.get("lifecycle_status") or "unknown" lifecycle_counts[lc] = lifecycle_counts.get(lc, 0) + 1 - # Unit stats (optional) - unit_stats: dict[str, int] | None = None - if include_unit_stats: - unit_rows = self._gc.query( - f""" - MATCH (p:IMASNode)-[:HAS_UNIT]->(u:Unit) - WHERE p.node_category = 'data' {dd_clause} - RETURN u.id AS unit, count(p) AS cnt - ORDER BY cnt DESC - LIMIT 30 - """, - **dd_params, - ) - unit_stats = {r["unit"]: r["cnt"] for r in (unit_rows or [])} - # Build tools list mcp_tools = [ "search_dd_paths", "check_dd_paths", "fetch_dd_paths", "list_dd_paths", - "get_dd_overview", + "get_dd_catalog", "search_dd_clusters", "get_dd_identifiers", "get_dd_versions", @@ -1239,27 +1180,19 @@ async def get_dd_overview( total_paths = sum(s["path_count"] for s in ids_statistics.values()) - # When no query, truncate to top 10 IDS to reduce token usage - if not query and len(all_ids) > 10: - top_ids = all_ids[:10] - top_statistics = {k: ids_statistics[k] for k in top_ids} - else: - top_ids = all_ids - top_statistics = ids_statistics - return GetOverviewResult( content=f"IMAS Data Dictionary v{current_version}: {len(all_ids)} IDS, {total_paths} total paths", - available_ids=top_ids, - query=query, + available_ids=all_ids, + query=None, physics_domains=sorted(physics_domains), - ids_statistics=top_statistics, + ids_statistics=ids_statistics, mcp_tools=mcp_tools, dd_version=current_version, mcp_version=version, total_leaf_nodes=total_paths, domain_summary=domain_summary, lifecycle_summary=lifecycle_counts, - unit_statistics=unit_stats, + unit_statistics=None, ) @@ -1882,15 +1815,15 @@ def __init__(self, graph_client: GraphClient): self._gc = graph_client @mcp_tool( - "Get structural context for an IMAS path via graph traversal. " - "Discovers sibling paths via shared clusters, coordinates, units, " + "Find paths in other IDSs that are related to a given path. " + "Discovers related paths via shared clusters, coordinates, units, " "and identifier schemas across IDS boundaries. " "path (required): Exact IMAS path (e.g. 'equilibrium/time_slice/profiles_1d/psi'). " - "relationship_types: Filter to specific types — 'cluster', 'coordinate', " + "relationship_types: Filter to specific types — 'semantic', 'cluster', 'coordinate', " "'unit', 'identifier', or 'all' (default)." ) - @handle_errors("get_dd_path_context") - async def get_dd_path_context( + @handle_errors("find_related_dd_paths") + async def find_related_dd_paths( self, path: str, relationship_types: str = "all", @@ -1991,158 +1924,20 @@ def __init__(self, graph_client: GraphClient): self._gc = graph_client @mcp_tool( - "Analyze the hierarchical structure of an IMAS IDS. " - "Returns depth metrics, leaf/structure ratio, array patterns, " - "physics domain distribution, coordinate usage, and COCOS-labeled fields. " + "Analyze the internal structure and organization of a specific IMAS IDS. " + "Returns metrics (path counts, depth), data type distribution, " + "physics domains, coordinate arrays, and COCOS/cluster counts. " + "Use get_dd_cocos_fields or search_dd_clusters for full listings. " "ids_name (required): IDS name (e.g. 'equilibrium')." ) - @handle_errors("analyze_dd_structure") - async def analyze_dd_structure( + @handle_errors("get_ids_summary") + async def get_ids_summary( self, ids_name: str, dd_version: int | None = None, ctx: Context | None = None, ) -> dict[str, Any]: - """Analyze the hierarchical structure of an IMAS IDS.""" - dd_params: dict[str, Any] = {"ids_name": ids_name} - dd_clause = _dd_version_clause("p", dd_version, dd_params) - - # Basic metrics — single scan using nullIf for leaf counting (Neo4j 2026 compat) - metrics = self._gc.query( - f""" - MATCH (p:IMASNode) - WHERE p.ids = $ids_name {dd_clause} - RETURN count(p) AS total_paths, - max(size(split(p.id, '/')) - 1) AS max_depth, - avg(size(split(p.id, '/')) - 1) AS avg_depth, - count(nullIf( - p.data_type IS NULL OR p.data_type IN {_structure_type_list()}, - true - )) AS leaf_count - """, - **dd_params, - ) - - # Physics domain distribution - domains = self._gc.query( - f""" - MATCH (p:IMASNode) - WHERE p.ids = $ids_name AND p.physics_domain IS NOT NULL {dd_clause} - RETURN p.physics_domain AS domain, count(p) AS count - ORDER BY count DESC - """, - **dd_params, - ) - - # Data type distribution - types = self._gc.query( - f""" - MATCH (p:IMASNode) - WHERE p.ids = $ids_name AND p.data_type IS NOT NULL {dd_clause} - RETURN p.data_type AS data_type, count(p) AS count - ORDER BY count DESC - """, - **dd_params, - ) - - # Array structures with coordinates - arrays = self._gc.query( - f""" - MATCH (p:IMASNode)-[:HAS_COORDINATE]->(coord:IMASCoordinateSpec) - WHERE p.ids = $ids_name {dd_clause} - RETURN p.id AS path, collect(coord.id) AS coordinates - ORDER BY p.id - """, - **dd_params, - ) - - # COCOS-labeled fields - cocos_fields = self._gc.query( - f""" - MATCH (p:IMASNode) - WHERE p.ids = $ids_name - AND p.cocos_label_transformation IS NOT NULL {dd_clause} - RETURN p.id AS path, - p.cocos_label_transformation AS cocos_label - ORDER BY p.id - """, - **dd_params, - ) - - basic = metrics[0] if metrics else {} - total_paths = basic.get("total_paths", 0) - leaf_count = basic.get("leaf_count", 0) - result: dict[str, Any] = { - "ids_name": ids_name, - "dd_version": dd_version, - "total_paths": total_paths, - "leaf_count": leaf_count, - "structure_count": total_paths - leaf_count, - "max_depth": basic.get("max_depth", 0), - "avg_depth": round(basic.get("avg_depth", 0), 1), - "physics_domains": [ - {"domain": d["domain"], "count": d["count"]} for d in domains - ], - "data_types": [ - {"type": t["data_type"], "count": t["count"]} for t in types - ], - "array_structures": [ - {"path": a["path"], "coordinates": a["coordinates"]} for a in arrays - ], - "cocos_fields": [ - {"path": c["path"], "label": c["cocos_label"]} for c in cocos_fields - ], - } - - # When version-filtered, add deprecated/renamed context for this IDS - if dd_version is not None: - dep_params: dict[str, Any] = { - "ids_name": ids_name, - "dd_major_version": dd_version, - } - deprecated = self._gc.query( - """ - MATCH (p:IMASNode)-[:DEPRECATED_IN]->(dv:DDVersion) - WHERE p.ids = $ids_name - AND toInteger(split(dv.id, '.')[0]) <= $dd_major_version - RETURN count(p) AS count - """, - **dep_params, - ) - renamed = self._gc.query( - """ - MATCH (old:IMASNode)-[:RENAMED_TO]->(new:IMASNode) - WHERE old.ids = $ids_name - RETURN count(old) AS count - """, - ids_name=ids_name, - ) - result["version_context"] = { - "note": ( - f"Filtered to paths active in DD v{dd_version}. " - f"Counts include paths carried forward from earlier major versions." - ), - "deprecated_in_or_before": deprecated[0]["count"] if deprecated else 0, - "renamed_paths": renamed[0]["count"] if renamed else 0, - } - - return result - - @mcp_tool( - "Get a rich structural overview of an IDS using efficient graph queries. " - "Returns metrics (path counts, depth), top-level sections, semantic clusters, " - "identifier schemas, COCOS fields, coordinate arrays, and data type distribution. " - "ids_name (required): IDS name to analyze (e.g. 'equilibrium', 'core_profiles'). " - "dd_version: Filter by DD major version (3 or 4). Default: latest version." - ) - @handle_errors("get_ids_structure") - async def get_ids_structure( - self, - ids_name: str, - dd_version: int | None = None, - ctx: Context | None = None, - ) -> dict[str, Any]: - """Get a rich structural overview of an IDS using efficient graph queries.""" + """Get a compact structural summary of an IDS.""" dd_params: dict[str, Any] = {"ids_name": ids_name} dd_clause = _dd_version_clause("p", dd_version, dd_params) @@ -2179,15 +1974,12 @@ async def get_ids_structure( meta = combined[0] - # Query 2: Clusters containing paths from this IDS - clusters = self._gc.query( + # Query 2: Cluster count (compact — use search_dd_clusters for full listings) + cluster_count_result = self._gc.query( f""" MATCH (p:IMASNode)-[:IN_CLUSTER]->(c:IMASSemanticCluster) WHERE p.ids = $ids_name {dd_clause} - WITH c, count(p) AS member_count - RETURN c.label AS label, c.scope AS scope, member_count - ORDER BY member_count DESC - LIMIT 15 + RETURN count(DISTINCT c) AS count """, **dd_params, ) @@ -2205,34 +1997,28 @@ async def get_ids_structure( **dd_params, ) - # Query 4: COCOS fields + coordinate specs - cocos_coords = self._gc.query( + # Query 4: COCOS count only (use get_dd_cocos_fields for full listing) + cocos_count_result = self._gc.query( f""" MATCH (p:IMASNode) WHERE p.ids = $ids_name - AND (p.cocos_label_transformation IS NOT NULL - OR exists((p)-[:HAS_COORDINATE]->())) {dd_clause} - OPTIONAL MATCH (p)-[:HAS_COORDINATE]->(coord:IMASCoordinateSpec) - RETURN p.id AS path, - p.cocos_label_transformation AS cocos, - collect(coord.id) AS coordinates - ORDER BY p.id + AND p.cocos_label_transformation IS NOT NULL {dd_clause} + RETURN count(p) AS count """, **dd_params, ) - cocos_fields = [ - {"path": r["path"], "label": r["cocos"]} - for r in (cocos_coords or []) - if r["cocos"] - ] - coord_arrays = [ - {"path": r["path"], "coordinates": r["coordinates"]} - for r in (cocos_coords or []) - if r["coordinates"] - ] + # Query 5: Coordinate array count + coord_count_result = self._gc.query( + f""" + MATCH (p:IMASNode)-[:HAS_COORDINATE]->(coord:IMASCoordinateSpec) + WHERE p.ids = $ids_name {dd_clause} + RETURN count(DISTINCT p) AS count + """, + **dd_params, + ) - # Query 5: Data type distribution + # Query 6: Data type distribution types = self._gc.query( f""" MATCH (p:IMASNode) @@ -2245,7 +2031,11 @@ async def get_ids_structure( total = meta.get("total", 0) leaves = meta.get("leaves", 0) - return { + cluster_count = cluster_count_result[0]["count"] if cluster_count_result else 0 + cocos_count = cocos_count_result[0]["count"] if cocos_count_result else 0 + coord_count = coord_count_result[0]["count"] if coord_count_result else 0 + + result: dict[str, Any] = { "ids_name": ids_name, "description": meta.get("description", ""), "physics_domain": meta.get("physics_domain", ""), @@ -2257,10 +2047,7 @@ async def get_ids_structure( "max_depth": meta.get("max_depth", 0), }, "top_sections": meta.get("top_sections", []), - "clusters": [ - {"label": c["label"], "scope": c["scope"], "members": c["member_count"]} - for c in (clusters or []) - ], + "semantic_clusters": cluster_count, "identifier_schemas": [ { "schema": s["schema"], @@ -2269,12 +2056,14 @@ async def get_ids_structure( } for s in (identifiers or []) ], - "cocos_fields": cocos_fields, - "coordinate_arrays": coord_arrays[:20], + "coordinate_arrays": coord_count, + "cocos_fields": cocos_count, "data_types": {t["data_type"]: t["count"] for t in (types or [])}, } - @handle_errors("get_cocos_fields") + return result + + @handle_errors("get_dd_cocos_fields") async def get_dd_cocos_fields( self, transformation_type: str | None = None, @@ -2342,7 +2131,7 @@ async def get_dd_cocos_fields( "ids_name (required): IDS name (e.g. 'equilibrium'). " "leaf_only: If true, return only leaf nodes (default false)." ) - @handle_errors("export_imas_ids") + @handle_errors("export_dd_ids") async def export_dd_ids( self, ids_name: str, @@ -2390,7 +2179,7 @@ async def export_dd_ids( "domain (required): Physics domain name (e.g. 'magnetics', 'equilibrium'). " "ids_filter: Optional IDS name filter." ) - @handle_errors("export_imas_domain") + @handle_errors("export_dd_domain") async def export_dd_domain( self, domain: str, diff --git a/imas_codex/tools/utils.py b/imas_codex/tools/utils.py index 65199d5ee..8a720742a 100644 --- a/imas_codex/tools/utils.py +++ b/imas_codex/tools/utils.py @@ -112,6 +112,6 @@ def validate_query(query: str | None, tool_name: str) -> tuple[bool, str | None] return False, ( f"Query cannot be empty for {tool_name}. " "Provide a search term like 'electron temperature' or 'equilibrium/time_slice'. " - "Use get_dd_overview() to explore available IDS structures." + "Use get_dd_catalog() to explore available IDS structures." ) return True, None diff --git a/tests/conftest.py b/tests/conftest.py index 093cb438f..7842c74cb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -379,7 +379,7 @@ def mcp_test_context(): "expected_tools": [ ("path_tool", "check_dd_paths"), ("path_tool", "fetch_dd_paths"), - ("overview_tool", "get_dd_overview"), + ("overview_tool", "get_dd_catalog"), ("identifiers_tool", "get_dd_identifiers"), ("list_tool", "list_dd_paths"), ("clusters_tool", "search_dd_clusters"), diff --git a/tests/core/test_cli.py b/tests/core/test_cli.py index 79c3fc62b..c4b0fe50c 100644 --- a/tests/core/test_cli.py +++ b/tests/core/test_cli.py @@ -196,7 +196,7 @@ def test_dd_only_excludes_facility_tools(self): assert "fetch_dd_paths" in tool_names assert "find_related_dd_paths" in tool_names assert "get_graph_schema" in tool_names - assert "get_dd_overview" in tool_names + assert "get_dd_catalog" in tool_names def test_dd_only_implies_read_only(self): """DD-only mode automatically sets read_only=True.""" diff --git a/tests/graph_mcp/test_dd_tool_features.py b/tests/graph_mcp/test_dd_tool_features.py index dec764e80..2f3d8e014 100644 --- a/tests/graph_mcp/test_dd_tool_features.py +++ b/tests/graph_mcp/test_dd_tool_features.py @@ -172,25 +172,9 @@ def _make_tool(self, graph_client): async def test_overview_without_unit_stats(self, graph_client): """Default overview has no unit_statistics.""" tool = self._make_tool(graph_client) - result = await tool.get_dd_overview() + result = await tool.get_dd_catalog() assert result.unit_statistics is None - @pytest.mark.asyncio - @pytest.mark.skip( - reason="include_unit_stats parameter not implemented in get_dd_overview" - ) - async def test_overview_with_unit_stats(self, graph_client): - """Overview with include_unit_stats=True returns unit distribution.""" - tool = self._make_tool(graph_client) - result = await tool.get_dd_overview(include_unit_stats=True) - assert result.unit_statistics is not None - assert "top_units" in result.unit_statistics - assert len(result.unit_statistics["top_units"]) > 0 - # Check structure of each unit entry - for u in result.unit_statistics["top_units"]: - assert "unit" in u - assert "count" in u - # ── Phase 4: Lifecycle filtering ────────────────────────────────────────── diff --git a/tests/graph_mcp/test_graph_search.py b/tests/graph_mcp/test_graph_search.py index e6fe5e526..c4079a6ce 100644 --- a/tests/graph_mcp/test_graph_search.py +++ b/tests/graph_mcp/test_graph_search.py @@ -175,7 +175,7 @@ def _make_tool(self, graph_client): @pytest.mark.asyncio async def test_overview_returns_all_ids(self, graph_client): tool = self._make_tool(graph_client) - result = await tool.get_dd_overview() + result = await tool.get_dd_catalog() assert len(result.available_ids) == len(IDS_NODES) for ids in IDS_NODES: assert ids["name"] in result.available_ids @@ -183,35 +183,27 @@ async def test_overview_returns_all_ids(self, graph_client): @pytest.mark.asyncio async def test_overview_has_statistics(self, graph_client): tool = self._make_tool(graph_client) - result = await tool.get_dd_overview() + result = await tool.get_dd_catalog() assert len(result.ids_statistics) > 0 assert "equilibrium" in result.ids_statistics - @pytest.mark.asyncio - async def test_overview_with_query_filter(self, graph_client): - tool = self._make_tool(graph_client) - result = await tool.get_dd_overview(query="equilibrium") - assert "equilibrium" in result.available_ids - # core_profiles should be filtered out - assert "core_profiles" not in result.available_ids - @pytest.mark.asyncio async def test_overview_has_dd_version(self, graph_client): tool = self._make_tool(graph_client) - result = await tool.get_dd_overview() + result = await tool.get_dd_catalog() assert result.dd_version == "4.1.0" @pytest.mark.asyncio async def test_overview_has_mcp_tools(self, graph_client): tool = self._make_tool(graph_client) - result = await tool.get_dd_overview() + result = await tool.get_dd_catalog() assert "search_dd_paths" in result.mcp_tools # query_imas_graph was removed in the unified server cleanup @pytest.mark.asyncio async def test_overview_has_physics_domains(self, graph_client): tool = self._make_tool(graph_client) - result = await tool.get_dd_overview() + result = await tool.get_dd_catalog() assert len(result.physics_domains) > 0 @@ -320,7 +312,7 @@ def test_graph_mode_registers_all_tools(self, graph_client): "check_dd_paths", "fetch_dd_paths", "list_dd_paths", - "get_dd_overview", + "get_dd_catalog", "search_dd_clusters", "get_dd_identifiers", "get_dd_versions", @@ -357,7 +349,7 @@ async def test_delegation_list_paths(self, graph_client): @pytest.mark.asyncio async def test_delegation_overview(self, graph_client): tools = self._make_tools(graph_client) - result = await tools.overview_tool.get_dd_overview() + result = await tools.overview_tool.get_dd_catalog() assert len(result.available_ids) > 0 @pytest.mark.asyncio @@ -698,7 +690,7 @@ def test_category_expansion(self, graph_client): class TestExportDomain: - """Tests for export_imas_domain with domain resolution.""" + """Tests for export_dd_domain with domain resolution.""" def _make_tool(self, graph_client): from imas_codex.tools.graph_search import GraphStructureTool diff --git a/tests/graph_mcp/test_tool_registration.py b/tests/graph_mcp/test_tool_registration.py index fa767d408..26a0151e2 100644 --- a/tests/graph_mcp/test_tool_registration.py +++ b/tests/graph_mcp/test_tool_registration.py @@ -34,7 +34,7 @@ def test_existing_tools_still_registered(self, graph_client): assert "search_dd_paths" in names assert "fetch_dd_paths" in names assert "list_dd_paths" in names - assert "get_dd_overview" in names + assert "get_dd_catalog" in names def test_total_tool_count(self, graph_client): """Total tool count matches expected number of graph-backed tools.""" diff --git a/tests/integration/test_workflows.py b/tests/integration/test_workflows.py index 1162c44aa..4d5072fb3 100644 --- a/tests/integration/test_workflows.py +++ b/tests/integration/test_workflows.py @@ -26,7 +26,7 @@ class TestUserWorkflows: async def test_discovery_workflow(self, tools, workflow_test_data): """Test: overview → search workflow.""" # Step 1: Get overview to understand what's available - overview = await tools.overview_tool.get_dd_overview() + overview = await tools.overview_tool.get_dd_catalog() assert isinstance(overview, GetOverviewResult) if overview.available_ids: @@ -101,7 +101,7 @@ async def test_workflow_total_time(self, tools): start_time = time.time() # Execute a typical workflow - overview = await tools.overview_tool.get_dd_overview() + overview = await tools.overview_tool.get_dd_catalog() search = await tools.search_tool.search_dd_paths( query="temperature", max_results=3 ) @@ -120,7 +120,7 @@ async def test_concurrent_tool_usage(self, tools): """Test tools can be used concurrently without interference.""" # Run multiple tools concurrently tasks = [ - tools.overview_tool.get_dd_overview(), + tools.overview_tool.get_dd_catalog(), tools.search_tool.search_dd_paths(query="temperature", max_results=3), ] @@ -139,7 +139,7 @@ class TestWorkflowErrorRecovery: async def test_workflow_continues_after_error(self, tools): """Test workflow can continue after one step fails.""" # Step 1: Valid operation - overview = await tools.overview_tool.get_dd_overview() + overview = await tools.overview_tool.get_dd_catalog() assert isinstance(overview, GetOverviewResult) # Step 2: Continue with valid operation diff --git a/tests/llm/test_graceful_degradation.py b/tests/llm/test_graceful_degradation.py index a557dc9f1..d2c86e232 100644 --- a/tests/llm/test_graceful_degradation.py +++ b/tests/llm/test_graceful_degradation.py @@ -249,12 +249,11 @@ def test_semantic_search_triggers_full_warmup(self, mock_graph_warmup): "find_related_dd_paths", "list_dd_paths", "fetch_dd_error_fields", - "get_dd_overview", + "get_dd_catalog", "get_dd_identifiers", "get_dd_versions", "get_dd_version_context", - "export_imas_ids", - "export_imas_domain", + "get_ids_summary", } diff --git a/tests/llm/test_mcp_bug_regressions.py b/tests/llm/test_mcp_bug_regressions.py index 870577d0c..87a32965a 100644 --- a/tests/llm/test_mcp_bug_regressions.py +++ b/tests/llm/test_mcp_bug_regressions.py @@ -46,35 +46,6 @@ def test_export_dd_domain_tool_has_no_include_errors_param(self): "pass this kwarg" ) - def test_export_dd_ids_server_handler_no_include_errors(self): - """The DD-only server handler for export_imas_ids must omit include_errors.""" - from imas_codex.llm.server import AgentsServer - - server = AgentsServer(dd_only=True) - # Walk registered tool components to find the handler - for key, component in server.mcp._local_provider._components.items(): - if key == "tool:export_imas_ids": - fn = component.fn - sig = inspect.signature(fn) - assert "include_errors" not in sig.parameters, ( - "DD-only export_imas_ids handler still has include_errors" - ) - break - - def test_export_dd_domain_server_handler_no_include_errors(self): - """The DD-only server handler for export_imas_domain must omit include_errors.""" - from imas_codex.llm.server import AgentsServer - - server = AgentsServer(dd_only=True) - for key, component in server.mcp._local_provider._components.items(): - if key == "tool:export_imas_domain": - fn = component.fn - sig = inspect.signature(fn) - assert "include_errors" not in sig.parameters, ( - "DD-only export_imas_domain handler still has include_errors" - ) - break - # --------------------------------------------------------------------------- # Bug 4: Short physics terms like "ip", "q", "b0" must not be filtered out @@ -157,7 +128,7 @@ def test_abbreviation_exact_match_boost_exists(self): # --------------------------------------------------------------------------- -# Bug 5: Coordinate channel in find_related_dd_paths (get_dd_path_context) +# Bug 5: Coordinate channel in find_related_dd_paths (find_related_dd_paths) # must traverse through IMASCoordinateSpec for coordinate partner discovery. # The HAS_COORDINATE relationship now correctly points to IMASCoordinateSpec # nodes, which hold coordinate specifications used across IDSs. @@ -171,7 +142,7 @@ def test_coordinate_query_uses_coordinate_spec_label(self): """The HAS_COORDINATE Cypher must traverse (coord:IMASCoordinateSpec).""" from imas_codex.tools.graph_search import GraphPathContextTool - source = inspect.getsource(GraphPathContextTool.get_dd_path_context) + source = inspect.getsource(GraphPathContextTool.find_related_dd_paths) # Find the coordinate partners query coord_section = source[source.index("Coordinate partners") :] @@ -192,7 +163,7 @@ async def test_coordinate_query_dispatched_correctly(self): gc.query.return_value = [] tool = GraphPathContextTool(gc) - await tool.get_dd_path_context( + await tool.find_related_dd_paths( path="equilibrium/time_slice/profiles_1d/psi", relationship_types="coordinate", ) diff --git a/tests/search/decorators/test_tool_recommendations.py b/tests/search/decorators/test_tool_recommendations.py index 0d8e1b33e..bd0e99959 100644 --- a/tests/search/decorators/test_tool_recommendations.py +++ b/tests/search/decorators/test_tool_recommendations.py @@ -127,7 +127,7 @@ def test_no_results_suggestions(self): assert len(suggestions) > 0 tool_names = [s["tool"] for s in suggestions] - assert "get_dd_overview" in tool_names + assert "get_dd_catalog" in tool_names assert "get_dd_identifiers" in tool_names @@ -176,7 +176,7 @@ def test_error_result_suggestions(self): assert len(recommendations) > 0 tool_names = [r["tool"] for r in recommendations] - assert "get_dd_overview" in tool_names + assert "get_dd_catalog" in tool_names def test_search_result_suggestions(self): """Test suggestions for search results.""" @@ -219,7 +219,7 @@ def test_generic_result_suggestions(self): assert len(recommendations) > 0 tool_names = [r["tool"] for r in recommendations] assert "search_dd_paths" in tool_names - assert "get_dd_overview" in tool_names + assert "get_dd_catalog" in tool_names class TestRecommendToolsDecorator: diff --git a/tests/search/test_tool_suggestions.py b/tests/search/test_tool_suggestions.py index 375c45eea..1c835b6ba 100644 --- a/tests/search/test_tool_suggestions.py +++ b/tests/search/test_tool_suggestions.py @@ -19,13 +19,13 @@ def test_search_results_suggest_overview_and_list(self): suggestions = suggest_follow_up_tools(results, "search_dd_paths") assert len(suggestions) > 0 - assert any(s["tool"] == "get_dd_overview" for s in suggestions) + assert any(s["tool"] == "get_dd_catalog" for s in suggestions) assert any(s["tool"] == "list_dd_paths" for s in suggestions) def test_overview_results_suggest_search(self): """Overview results suggest search tool.""" results = {"concept": "plasma temperature"} - suggestions = suggest_follow_up_tools(results, "get_dd_overview") + suggestions = suggest_follow_up_tools(results, "get_dd_catalog") assert len(suggestions) > 0 assert any(s["tool"] == "search_dd_paths" for s in suggestions) diff --git a/tests/tools/test_dd_version_filtering.py b/tests/tools/test_dd_version_filtering.py index 9debfda1f..8ad4f5a6d 100644 --- a/tests/tools/test_dd_version_filtering.py +++ b/tests/tools/test_dd_version_filtering.py @@ -287,7 +287,7 @@ async def test_overview_queries_ids_nodes(self): ] tool = GraphOverviewTool(gc) - await tool.get_dd_overview() + await tool.get_dd_catalog() ids_cypher = gc.query.call_args_list[0][0][0] assert "MATCH (i:IDS)" in ids_cypher @@ -310,7 +310,7 @@ async def test_overview_with_dd_version_includes_filter(self): ] tool = GraphOverviewTool(gc) - await tool.get_dd_overview(dd_version=4) + await tool.get_dd_catalog(dd_version=4) ids_cypher = gc.query.call_args_list[0][0][0] assert "MATCH (i:IDS)" in ids_cypher diff --git a/tests/tools/test_facade_delegation.py b/tests/tools/test_facade_delegation.py deleted file mode 100644 index e88686ed5..000000000 --- a/tests/tools/test_facade_delegation.py +++ /dev/null @@ -1,120 +0,0 @@ -"""Regression tests for Tools registered tool name consistency. - -Ensures that all expected DD tool methods exist on the backend -tool instances and are discoverable via the Tools container. -No facade layer — callers access tool instances directly. -""" - -import pytest - -from imas_codex.tools import Tools -from imas_codex.tools.graph_search import ( - GraphClustersTool, - GraphIdentifiersTool, - GraphListTool, - GraphOverviewTool, - GraphPathContextTool, - GraphPathTool, - GraphSearchTool, - GraphStructureTool, -) -from imas_codex.tools.version_tool import VersionTool - -# Canonical mapping: tool_instance_attr -> (backend_class, expected_methods) -TOOL_METHOD_MAP = { - "search_tool": (GraphSearchTool, ["search_dd_paths"]), - "path_tool": ( - GraphPathTool, - ["check_dd_paths", "fetch_dd_paths", "fetch_error_fields"], - ), - "list_tool": (GraphListTool, ["list_dd_paths"]), - "overview_tool": (GraphOverviewTool, ["get_dd_overview"]), - "clusters_tool": (GraphClustersTool, ["search_dd_clusters"]), - "identifiers_tool": (GraphIdentifiersTool, ["get_dd_identifiers"]), - "path_context_tool": (GraphPathContextTool, ["get_dd_path_context"]), - "structure_tool": ( - GraphStructureTool, - [ - "analyze_dd_structure", - "export_dd_ids", - "export_dd_domain", - ], - ), - "version_tool": ( - VersionTool, - ["get_dd_versions", "get_dd_version_context", "get_dd_changelog"], - ), -} - - -def _all_method_params(): - """Yield (tool_attr, class, method_name) for parametrize.""" - for tool_attr, (cls, methods) in TOOL_METHOD_MAP.items(): - for method in methods: - yield tool_attr, cls, method - - -class TestToolMethodExistence: - """Verify backend tool methods exist and are async.""" - - @pytest.mark.parametrize( - "tool_attr,backend_class,method_name", - list(_all_method_params()), - ids=[f"{a}.{m}" for a, _, m in _all_method_params()], - ) - def test_backend_method_exists(self, tool_attr, backend_class, method_name): - """Every expected method must exist on its backend class.""" - assert hasattr(backend_class, method_name), ( - f"{backend_class.__name__}.{method_name} does not exist. " - f"Check that the method was renamed correctly." - ) - - def test_no_facade_methods_on_tools(self): - """Tools class must not have async facade delegation methods.""" - import inspect - - facade_names = { - "search_dd_paths", - "check_dd_paths", - "fetch_dd_paths", - "list_dd_paths", - "get_dd_overview", - "get_dd_identifiers", - "get_dd_path_context", - "get_dd_cocos_fields", - "export_dd_ids", - "export_dd_domain", - "get_dd_versions", - "search_dd_clusters", - "get_dd_version_context", - "get_dd_changelog", - "fetch_dd_error_fields", - } - for name in facade_names: - if hasattr(Tools, name): - method = getattr(Tools, name) - assert not inspect.iscoroutinefunction(method), ( - f"Tools.{name} is an async facade method — " - f"these should be removed. Callers should use " - f"tools..{name}() directly." - ) - - def test_no_imas_named_methods(self): - """No backend tool class should have _imas_ named methods (old naming).""" - old_names = [ - "search_imas_paths", - "check_imas_paths", - "fetch_imas_paths", - "list_imas_paths", - "get_imas_overview", - "search_imas_clusters", - "get_imas_identifiers", - "get_imas_path_context", - "analyze_imas_structure", - ] - for _tool_attr, (cls, _) in TOOL_METHOD_MAP.items(): - for old_name in old_names: - assert not hasattr(cls, old_name), ( - f"{cls.__name__} still has old method {old_name}. " - f"Rename to _dd_ convention." - ) diff --git a/tests/tools/test_tools.py b/tests/tools/test_tools.py index 9dcdbffc8..870a8b0cd 100644 --- a/tests/tools/test_tools.py +++ b/tests/tools/test_tools.py @@ -52,7 +52,7 @@ async def test_search_tool_interface(self, tools): @pytest.mark.asyncio async def test_overview_tool_interface(self, tools): """Test overview tool interface and basic functionality.""" - result = await tools.overview_tool.get_dd_overview() + result = await tools.overview_tool.get_dd_catalog() # Test interface contract assert isinstance(result, GetOverviewResult) diff --git a/tests/tools/test_utils.py b/tests/tools/test_utils.py index af16ffe18..0249b89c8 100644 --- a/tests/tools/test_utils.py +++ b/tests/tools/test_utils.py @@ -183,4 +183,4 @@ def test_error_includes_guidance(self): """Test that error message includes helpful guidance.""" is_valid, error = validate_query("", "search_dd_paths") assert is_valid is False - assert "get_dd_overview" in error + assert "get_dd_catalog" in error From 6b0c3cb867013fb026fb528e859bac73b8a76110 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 12:22:56 +0200 Subject: [PATCH 19/37] =?UTF-8?q?test:=20add=20TestE2ERoundTrip=20for=20fu?= =?UTF-8?q?ll=20SN=20lifecycle=20(build=E2=86=92publish=E2=86=92edit?= =?UTF-8?q?=E2=86=92import)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/sn/test_integration.py | 609 +++++++++++++++++++++++++++++++++++ 1 file changed, 609 insertions(+) diff --git a/tests/sn/test_integration.py b/tests/sn/test_integration.py index 694daa079..bdbb24d9d 100644 --- a/tests/sn/test_integration.py +++ b/tests/sn/test_integration.py @@ -9,6 +9,8 @@ 4. The import → build cycle is safe: catalog-imported rich fields are not erased by a subsequent sn-build write. 5. publish → import → publish is idempotent (key fields round-trip cleanly). +6. Full E2E lifecycle: write_standard_names → get_validated_standard_names + → graph_records_to_entries → generate_catalog_files → import_catalog. """ from __future__ import annotations @@ -20,6 +22,8 @@ import pytest import yaml +imas_sn = pytest.importorskip("imas_standard_names") + # ============================================================================= # Helpers # ============================================================================= @@ -429,3 +433,608 @@ def test_import_then_build_preserves_catalog_fields(self) -> None: f"'{absent_field}' must be None in batch when not supplied, " f"got {build_item[absent_field]!r}" ) + + +# ============================================================================= +# Part 5: Round-trip idempotence +# ============================================================================= + +imas_sn = pytest.importorskip("imas_standard_names") + +SAMPLE_GRAPH_RECORD: dict[str, Any] = { + "name": "electron_temperature", + "description": "Electron temperature profile", + "documentation": "The $T_e$ profile measured by Thomson scattering.", + "source": "dd", + "source_path": "core_profiles/profiles_1d/electrons/temperature", + "canonical_units": "eV", + "kind": "scalar", + "tags": [], + "links": [], + "ids_paths": ["core_profiles/profiles_1d/electrons/temperature"], + "constraints": ["T_e > 0"], + "validity_domain": "core plasma", + "confidence": 0.95, + "model": "test/model", + "ids_name": None, + "physical_base": "temperature", + "subject": "electron", +} + +SAMPLE_CATALOG_ENTRY_RT: dict[str, Any] = { + "name": "electron_temperature", + "description": "Electron temperature", + "documentation": "The electron temperature Te is measured by Thomson scattering.", + "kind": "scalar", + "unit": "eV", + "tags": [], + "links": [], + "ids_paths": ["core_profiles/profiles_1d/electrons/temperature"], + "validity_domain": "core plasma", + "constraints": ["T_e > 0"], + "physics_domain": "core_plasma_physics", + "status": "active", +} + + +def _published_yaml_to_catalog( + published_yaml: str, physics_domain: str = "unscoped" +) -> dict[str, Any]: + """Convert a published YAML string to a catalog-importable dict. + + Strips provenance, adds ``physics_domain``, normalises ``status`` to + ``"active"``, and converts links from ``[{name: …}]`` dicts to plain + strings so that ``StandardNameEntry`` validation succeeds. + """ + doc: dict[str, Any] = yaml.safe_load(published_yaml) + # Provenance block is not part of the catalog schema + doc.pop("provenance", None) + # status must be a catalog-valid value + doc["status"] = "active" + # physics_domain is required by StandardNameEntry + doc.setdefault("physics_domain", physics_domain) + # Normalise links: published format uses [{name: link}] dicts + raw_links = doc.get("links", []) + if raw_links and isinstance(raw_links[0], dict): + doc["links"] = [lnk.get("name", str(lnk)) for lnk in raw_links] + # Ensure required list fields are present (even if empty) + for list_field in ("tags", "links", "ids_paths", "constraints"): + doc.setdefault(list_field, []) + # Empty string validity_domain instead of None + if doc.get("validity_domain") is None: + doc["validity_domain"] = "" + return doc + + +def _imported_dict_to_graph_record(d: dict[str, Any]) -> dict[str, Any]: + """Normalise an imported graph dict for ``graph_records_to_entries``. + + ``import_catalog`` returns dicts with ``id`` / ``units`` / ``imas_paths`` + keys. ``graph_records_to_entries`` looks for ``name``/``id``, + ``canonical_units``/``units``, and ``ids_paths``. This helper adds the + ``ids_paths`` alias so that the path list survives the round-trip. + """ + rec = dict(d) + # Alias imas_paths → ids_paths (graph_records_to_entries reads ids_paths) + if "imas_paths" in rec and "ids_paths" not in rec: + rec["ids_paths"] = rec["imas_paths"] or [] + return rec + + +def _key_fields(parsed_yaml: dict[str, Any]) -> dict[str, Any]: + """Extract the semantic fields that must be preserved across a round-trip.""" + return { + "name": parsed_yaml.get("name"), + "kind": parsed_yaml.get("kind"), + "unit": parsed_yaml.get("unit"), + "description": parsed_yaml.get("description"), + "documentation": parsed_yaml.get("documentation"), + "ids_paths": sorted(parsed_yaml.get("ids_paths") or []), + "validity_domain": parsed_yaml.get("validity_domain"), + "constraints": sorted(parsed_yaml.get("constraints") or []), + } + + +class TestRoundTripIdempotence: + """Verify that publish → import → publish produces semantically identical YAML.""" + + def test_publish_import_publish_idempotent(self, tmp_path: Path) -> None: + """Round-trip: graph_records → YAML → catalog import → YAML should match. + + Key fields (name, kind, unit, ids_paths, validity_domain, constraints) + must be identical after a full publish → import → publish cycle. + Provenance and confidence fields are allowed to differ. + """ + from imas_codex.sn.publish import ( + generate_catalog_files, + graph_records_to_entries, + ) + + round1_dir = tmp_path / "round1" + catalog_dir = tmp_path / "catalog" + round2_dir = tmp_path / "round2" + + # --- Round 1: graph record → YAML files --- + entries1 = graph_records_to_entries([SAMPLE_GRAPH_RECORD]) + assert len(entries1) == 1, "Expected one publish entry from graph record" + generate_catalog_files(entries1, round1_dir) + + yaml_files1 = list(round1_dir.rglob("*.yaml")) + assert len(yaml_files1) == 1, ( + f"Expected exactly 1 YAML file in round1, got {len(yaml_files1)}" + ) + + # --- Convert published YAML → catalog-importable format --- + published_yaml_text = yaml_files1[0].read_text() + catalog_doc = _published_yaml_to_catalog( + published_yaml_text, physics_domain="core_plasma_physics" + ) + catalog_dir.mkdir(parents=True, exist_ok=True) + (catalog_dir / "electron_temperature.yaml").write_text( + yaml.safe_dump(catalog_doc) + ) + + # --- Import catalog (dry run) → graph dicts --- + from imas_codex.sn.catalog_import import import_catalog + + result = import_catalog(catalog_dir, dry_run=True) + assert result.imported == 1, ( + f"Expected 1 imported entry, got {result.imported}; errors: {result.errors}" + ) + assert not result.errors, f"Import errors: {result.errors}" + + # --- Normalise imported dicts and convert back to publish entries --- + graph_records2 = [_imported_dict_to_graph_record(e) for e in result.entries] + entries2 = graph_records_to_entries(graph_records2) + assert len(entries2) == 1, "Expected one publish entry from imported dict" + + # --- Round 2: publish entries → YAML files --- + generate_catalog_files(entries2, round2_dir) + yaml_files2 = list(round2_dir.rglob("*.yaml")) + assert len(yaml_files2) == 1, ( + f"Expected exactly 1 YAML file in round2, got {len(yaml_files2)}" + ) + + # --- Compare key fields (ignore provenance / confidence changes) --- + parsed1 = yaml.safe_load(yaml_files1[0].read_text()) + parsed2 = yaml.safe_load(yaml_files2[0].read_text()) + + fields1 = _key_fields(parsed1) + fields2 = _key_fields(parsed2) + + assert fields2["name"] == fields1["name"], "name must be preserved" + assert fields2["kind"] == fields1["kind"], "kind must be preserved" + assert fields2["unit"] == fields1["unit"], "unit must be preserved" + assert fields2["ids_paths"] == fields1["ids_paths"], ( + "ids_paths must be preserved" + ) + assert fields2["validity_domain"] == fields1["validity_domain"], ( + "validity_domain must be preserved" + ) + assert fields2["constraints"] == fields1["constraints"], ( + "constraints must be preserved" + ) + + def test_import_export_idempotent(self, tmp_path: Path) -> None: + """Import a catalog entry then re-publish it — key fields must be unchanged. + + Tests the ``import_catalog`` → ``graph_records_to_entries`` → + ``generate_yaml_entry`` path, asserting that the re-published YAML + preserves every semantically significant field from the original catalog. + """ + from imas_codex.sn.catalog_import import import_catalog + from imas_codex.sn.publish import generate_yaml_entry, graph_records_to_entries + + catalog_dir = tmp_path / "catalog" + catalog_dir.mkdir() + (catalog_dir / "electron_temperature.yaml").write_text( + yaml.safe_dump(SAMPLE_CATALOG_ENTRY_RT) + ) + + # Import (dry run) → graph dicts + result = import_catalog(catalog_dir, dry_run=True) + assert result.imported == 1, ( + f"Expected 1 imported entry; errors: {result.errors}" + ) + assert not result.errors, f"Import errors: {result.errors}" + + # Normalise dict keys and convert to SNPublishEntry + graph_records = [_imported_dict_to_graph_record(e) for e in result.entries] + entries = graph_records_to_entries(graph_records) + assert len(entries) == 1, "Expected one publish entry from imported graph dict" + + # Generate YAML and parse it back for comparison + yaml_str = generate_yaml_entry(entries[0]) + published = yaml.safe_load(yaml_str) + + original = SAMPLE_CATALOG_ENTRY_RT + assert published["name"] == original["name"], "name must round-trip" + assert published["kind"] == original["kind"], "kind must round-trip" + assert published.get("unit") == original["unit"], "unit must round-trip" + assert published.get("description") == original["description"], ( + "description must round-trip" + ) + assert published.get("documentation") == original["documentation"], ( + "documentation must round-trip" + ) + assert sorted(published.get("ids_paths") or []) == sorted( + original.get("ids_paths") or [] + ), "ids_paths must round-trip" + if original.get("validity_domain"): + assert published.get("validity_domain") == original["validity_domain"], ( + "validity_domain must round-trip" + ) + assert sorted(published.get("constraints") or []) == sorted( + original.get("constraints") or [] + ), "constraints must round-trip" + + def test_double_import_identical_entries(self, tmp_path: Path) -> None: + """Importing the same catalog directory twice yields identical result entries. + + Verifies that ``import_catalog`` is deterministic: repeated calls on the + same input produce identical graph dicts (same keys and values). + """ + from imas_codex.sn.catalog_import import import_catalog + + catalog_dir = tmp_path / "catalog" + catalog_dir.mkdir() + (catalog_dir / "electron_temperature.yaml").write_text( + yaml.safe_dump(SAMPLE_CATALOG_ENTRY_RT) + ) + + result1 = import_catalog(catalog_dir, dry_run=True) + result2 = import_catalog(catalog_dir, dry_run=True) + + assert result1.imported == result2.imported, ( + "Both imports should report the same import count" + ) + assert len(result1.entries) == len(result2.entries), ( + "Both imports should return the same number of entries" + ) + + compared_fields = ( + "id", + "description", + "documentation", + "kind", + "units", + "tags", + "links", + "imas_paths", + "validity_domain", + "constraints", + "physics_domain", + "review_status", + "source_type", + ) + for e1, e2 in zip(result1.entries, result2.entries, strict=True): + for field in compared_fields: + assert e1.get(field) == e2.get(field), ( + f"Field {field!r} differs between imports: " + f"{e1.get(field)!r} != {e2.get(field)!r}" + ) + + +# ============================================================================= +# Part 4: Full E2E Round-Trip (build → publish → edit → import) +# ============================================================================= + +# Rich sample data shared across E2E tests +_RICH_SN_RECORD = { + "id": "electron_temperature", + "source_type": "dd", + "source_id": "core_profiles/profiles_1d/electrons/temperature", + "ids_name": "core_profiles", + "description": "Electron temperature in the core plasma", + "documentation": ( + "The electron temperature $T_e$ is measured via Thomson scattering. " + "It is a key parameter for transport modelling." + ), + "kind": "scalar", + "units": "eV", + "tags": ["spatial-profile"], + "links": ["name:ion_temperature"], + "imas_paths": ["core_profiles/profiles_1d/electrons/temperature"], + "validity_domain": "core plasma", + "constraints": ["T_e > 0"], + "physical_base": "temperature", + "subject": "electron", + "confidence": 0.95, + "model": "gpt-4o", + "review_status": "drafted", +} + +# What get_validated_standard_names returns — graph-canonical keys +_GRAPH_QUERY_ROW = { + "name": "electron_temperature", + "description": "Electron temperature in the core plasma", + "documentation": ( + "The electron temperature $T_e$ is measured via Thomson scattering. " + "It is a key parameter for transport modelling." + ), + "kind": "scalar", + "canonical_units": "eV", + "tags": ["spatial-profile"], + "links": ["name:ion_temperature"], + "ids_paths": ["core_profiles/profiles_1d/electrons/temperature"], + "constraints": ["T_e > 0"], + "validity_domain": "core plasma", + "confidence": 0.95, + "model": "gpt-4o", + "source": "dd", + "source_path": "core_profiles/profiles_1d/electrons/temperature", + "ids_name": "core_profiles", + "physical_base": "temperature", + "subject": "electron", + "component": None, + "coordinate": None, + "position": None, + "process": None, + "source_ids_names": ["core_profiles"], +} + + +class TestE2ERoundTrip: + """Full lifecycle: build → publish → manual-edit → import. + + All graph operations are mocked — no live Neo4j required. + """ + + # ------------------------------------------------------------------ + # Phase helpers + # ------------------------------------------------------------------ + + @staticmethod + def _mock_write_graph_client(): + """Mock GraphClient for write_standard_names (no return value needed).""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=None) + mock_ctx = MagicMock() + mock_ctx.__enter__ = MagicMock(return_value=mock_gc) + mock_ctx.__exit__ = MagicMock(return_value=False) + return mock_ctx + + @staticmethod + def _build_catalog_entry(publish_entry) -> dict: + """Simulate a curator enriching a published entry into catalog format.""" + return { + "name": publish_entry.name, + "description": publish_entry.description, + "documentation": publish_entry.documentation + or "Enriched documentation by curator.", + "kind": "scalar", + "unit": publish_entry.unit, + "tags": publish_entry.tags, + "links": publish_entry.links, + "ids_paths": publish_entry.ids_paths, + "validity_domain": publish_entry.validity_domain or "", + "constraints": publish_entry.constraints, + "physics_domain": "core_plasma_physics", + "status": "active", + } + + # ------------------------------------------------------------------ + # Tests + # ------------------------------------------------------------------ + + def test_full_lifecycle_round_trip(self, tmp_path: Path) -> None: + """Complete build → publish → manual-edit → import cycle.""" + from imas_codex.sn.catalog_import import import_catalog + from imas_codex.sn.graph_ops import write_standard_names + from imas_codex.sn.publish import ( + generate_catalog_files, + graph_records_to_entries, + ) + + # Phase 1: write to graph (mocked) + with patch("imas_codex.sn.graph_ops.GraphClient") as MockGC: + MockGC.return_value = self._mock_write_graph_client() + count = write_standard_names([_RICH_SN_RECORD]) + assert count == 1 + + # Phase 2: graph → publish entries → YAML files + entries = graph_records_to_entries([_GRAPH_QUERY_ROW]) + assert len(entries) == 1 + + written = generate_catalog_files(entries, tmp_path / "published") + assert len(written) == 1 + assert written[0].exists() + + publish_entry = entries[0] + assert publish_entry.name == "electron_temperature" + assert publish_entry.unit == "eV" + assert publish_entry.documentation is not None + + # Phase 3: simulate curator enrichment into catalog format + catalog_entry = self._build_catalog_entry(publish_entry) + catalog_dir = tmp_path / "reviewed_catalog" + catalog_dir.mkdir() + (catalog_dir / "electron_temperature.yaml").write_text( + yaml.safe_dump(catalog_entry) + ) + + # Phase 4: import catalog (dry_run=True, no graph write) + result = import_catalog(catalog_dir=catalog_dir, dry_run=True) + + assert result.imported == 1 + assert len(result.errors) == 0 + + entry = result.entries[0] + assert entry["id"] == "electron_temperature" + assert entry["review_status"] == "accepted" + assert entry["units"] == "eV" + assert entry["physics_domain"] == "core_plasma_physics" + assert entry["physical_base"] == "temperature" + assert entry["subject"] == "electron" + + def test_publish_generates_valid_yaml_for_import(self, tmp_path: Path) -> None: + """Published YAML (after curator enrichment) is valid input for import_catalog.""" + from imas_codex.sn.catalog_import import import_catalog + from imas_codex.sn.publish import graph_records_to_entries + + entries = graph_records_to_entries([_GRAPH_QUERY_ROW]) + assert len(entries) == 1 + + # Curator enriches published entry into catalog format + catalog_entry = self._build_catalog_entry(entries[0]) + catalog_dir = tmp_path / "catalog" + catalog_dir.mkdir() + (catalog_dir / "electron_temperature.yaml").write_text( + yaml.safe_dump(catalog_entry) + ) + + result = import_catalog(catalog_dir=catalog_dir, dry_run=True) + + assert result.imported == 1 + assert len(result.errors) == 0 + + entry = result.entries[0] + # catalog 'unit' → graph 'units' + assert entry["units"] == "eV" + assert "unit" not in entry + # catalog 'ids_paths' → graph 'imas_paths' + assert entry["imas_paths"] == [ + "core_profiles/profiles_1d/electrons/temperature" + ] + assert "ids_paths" not in entry + # review_status always 'accepted' after import + assert entry["review_status"] == "accepted" + + def test_field_preservation_across_lifecycle(self, tmp_path: Path) -> None: + """Specific rich fields are verified at each stage of the lifecycle.""" + from imas_codex.sn.catalog_import import import_catalog + from imas_codex.sn.publish import ( + generate_catalog_files, + generate_yaml_entry, + graph_records_to_entries, + ) + + # Stage A: graph_records_to_entries preserves rich fields + entries = graph_records_to_entries([_GRAPH_QUERY_ROW]) + entry = entries[0] + + assert entry.name == "electron_temperature" + assert entry.kind == "scalar" + assert entry.unit == "eV" + assert "spatial-profile" in entry.tags + assert entry.links == ["name:ion_temperature"] + assert entry.ids_paths == ["core_profiles/profiles_1d/electrons/temperature"] + assert entry.validity_domain == "core plasma" + assert entry.constraints == ["T_e > 0"] + assert entry.documentation is not None + assert "$T_e$" in (entry.documentation or "") + assert entry.provenance.confidence == 0.95 + assert entry.provenance.source == "dd" + assert entry.provenance.ids_name == "core_profiles" + + # Stage B: generate_yaml_entry serializes all fields + yaml_str = generate_yaml_entry(entry) + parsed = yaml.safe_load(yaml_str) + + assert parsed["name"] == "electron_temperature" + assert parsed["kind"] == "scalar" + assert parsed["unit"] == "eV" + assert parsed["validity_domain"] == "core plasma" + assert parsed["constraints"] == ["T_e > 0"] + assert "documentation" in parsed + assert parsed["ids_paths"] == [ + "core_profiles/profiles_1d/electrons/temperature" + ] + assert parsed["provenance"]["confidence"] == 0.95 + + # Stage C: generate_catalog_files creates correct subdirectory structure + written = generate_catalog_files(entries, tmp_path / "published") + assert len(written) == 1 + # Primary tag is "core_profiles" → file lives in core_profiles/ + assert written[0].parent.name == "spatial-profile" + + # Stage D: import_catalog maps all fields correctly + catalog_entry = self._build_catalog_entry(entry) + catalog_dir = tmp_path / "catalog" + catalog_dir.mkdir() + (catalog_dir / "electron_temperature.yaml").write_text( + yaml.safe_dump(catalog_entry) + ) + + result = import_catalog(catalog_dir=catalog_dir, dry_run=True) + imported = result.entries[0] + + assert imported["id"] == "electron_temperature" + assert imported["description"] == "Electron temperature in the core plasma" + assert "documentation" in imported + assert imported["kind"] == "scalar" + assert imported["units"] == "eV" + assert imported["imas_paths"] == [ + "core_profiles/profiles_1d/electrons/temperature" + ] + assert imported["validity_domain"] == "core plasma" + assert imported["constraints"] == ["T_e > 0"] + assert imported["physics_domain"] == "core_plasma_physics" + assert imported["review_status"] == "accepted" + assert imported["physical_base"] == "temperature" + assert imported["subject"] == "electron" + + def test_write_standard_names_called_with_all_fields(self) -> None: + """write_standard_names receives all populated fields without losing any.""" + from imas_codex.sn.graph_ops import write_standard_names + + captured_batches: list = [] + + def capture_query(cypher, **kwargs): + if "batch" in kwargs: + captured_batches.extend(kwargs["batch"]) + return None + + mock_gc = MagicMock() + mock_gc.query = MagicMock(side_effect=capture_query) + mock_ctx = MagicMock() + mock_ctx.__enter__ = MagicMock(return_value=mock_gc) + mock_ctx.__exit__ = MagicMock(return_value=False) + + with patch("imas_codex.sn.graph_ops.GraphClient") as MockGC: + MockGC.return_value = mock_ctx + write_standard_names([_RICH_SN_RECORD]) + + assert len(captured_batches) > 0 + node = captured_batches[0] + assert node["id"] == "electron_temperature" + assert node["description"] == "Electron temperature in the core plasma" + assert node["kind"] == "scalar" + assert node["units"] == "eV" + assert node["tags"] == ["spatial-profile"] + assert node["constraints"] == ["T_e > 0"] + assert node["validity_domain"] == "core plasma" + assert node["review_status"] == "drafted" + assert node["confidence"] == 0.95 + assert node["model"] == "gpt-4o" + + def test_import_dry_run_does_not_call_graph(self, tmp_path: Path) -> None: + """dry_run=True must never invoke the graph write path.""" + from imas_codex.sn.catalog_import import import_catalog + + catalog_entry = { + "name": "electron_temperature", + "description": "Electron temperature", + "documentation": "Te documentation.", + "kind": "scalar", + "unit": "eV", + "tags": [], + "links": [], + "ids_paths": [], + "validity_domain": "", + "constraints": [], + "physics_domain": "core_plasma_physics", + "status": "active", + } + catalog_dir = tmp_path / "catalog" + catalog_dir.mkdir() + (catalog_dir / "electron_temperature.yaml").write_text( + yaml.safe_dump(catalog_entry) + ) + + with patch("imas_codex.sn.catalog_import._write_catalog_entries") as mock_write: + result = import_catalog(catalog_dir=catalog_dir, dry_run=True) + + mock_write.assert_not_called() + assert result.imported == 1 + assert result.entries[0]["review_status"] == "accepted" From a2a89ece33157cbedbabd6a419726373b5cd5d39 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 12:27:59 +0200 Subject: [PATCH 20/37] fix: use keyword args for GraphClient.query() in sn_tools GraphClient.query() signature is (cypher, **params) not (cypher, dict). All 4 calls were passing a dict as a positional argument, causing 'takes 2 positional arguments but 3 were given' at runtime. --- imas_codex/llm/sn_tools.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/imas_codex/llm/sn_tools.py b/imas_codex/llm/sn_tools.py index 3288bce62..da7b203b0 100644 --- a/imas_codex/llm/sn_tools.py +++ b/imas_codex/llm/sn_tools.py @@ -111,7 +111,7 @@ def _vector_search_sn(gc: GraphClient, embedding: list[float], k: int) -> list[d score ORDER BY score DESC """ - return gc.query(cypher, {"embedding": embedding, "k": k}) + return gc.query(cypher, embedding=embedding, k=k) def _keyword_search_sn(gc: GraphClient, query: str, k: int) -> list[dict]: @@ -131,7 +131,7 @@ def _keyword_search_sn(gc: GraphClient, query: str, k: int) -> list[dict]: 1.0 AS score LIMIT $k """ - return gc.query(cypher, {"keyword": query, "k": k}) + return gc.query(cypher, keyword=query, k=k) def _format_search_report(query: str, rows: list[dict]) -> str: @@ -237,7 +237,7 @@ def _fetch_standard_names( """ try: - rows = gc.query(cypher, {"names": name_list}) + rows = gc.query(cypher, names=name_list) except ServiceUnavailable: return NEO4J_NOT_RUNNING_MSG except Exception as e: @@ -403,7 +403,7 @@ def _list_standard_names( """ try: - rows = gc.query(cypher, params) + rows = gc.query(cypher, **params) except ServiceUnavailable: return NEO4J_NOT_RUNNING_MSG except Exception as e: From 13740eb1e41f72bc7cca11369e3f8114a1b57c9b Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 12:33:26 +0200 Subject: [PATCH 21/37] docs: add standard names CLI, lifecycle, and MCP tools to AGENTS.md MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add Standard Names section with CLI commands table (build, publish, import, status, benchmark) - Document StandardName lifecycle (drafted → published → accepted) - Document write semantics: build (coalesce) vs import (authoritative) - Document MCP tools (search, fetch, list standard names) - Document StandardName schema and key relationships - Update plans/README.md: mark features 11-14 as Done --- AGENTS.md | 46 ++++++++++++++++++++++++++++++++++++++++++++++ plans/README.md | 8 ++++---- 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index b6870446a..b73b488e4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -704,6 +704,52 @@ Azure Web App has continuous deployment enabled on ACR. When a new image appears - Graph push runs from the ITER machine where Neo4j runs — CI cannot build graph data - The release CLI handles everything — do not manually push graphs or tags separately +## Standard Names + +### CLI Commands + +| Command | Purpose | Key Options | +|---------|---------|-------------| +| `sn build` | Generate standard names from DD paths or facility signals via LLM pipeline | `--source {dd,signals}`, `--ids`, `--domain`, `--facility`, `--cost-limit`, `--dry-run`, `--force`, `--skip-review` | +| `sn publish` | Export validated StandardName nodes to YAML catalog files | `--output-dir`, `--ids`, `--domain`, `--group-by {ids,domain,confidence}`, `--confidence-min`, `--catalog-dir`, `--create-pr` | +| `sn import` | Import reviewed YAML catalog entries back into graph | `--catalog-dir` (required), `--tags`, `--dry-run`, `--check` | +| `sn status` | Show standard name statistics from graph | — | +| `sn benchmark` | Benchmark LLM models on standard name generation quality | `--models`, `--source`, `--reviewer-model` | + +### StandardName Lifecycle + +``` +drafted → published → accepted + ↘ rejected +``` + +- **drafted**: Generated by `sn build` (LLM pipeline) +- **published**: Exported by `sn publish` to YAML catalog for human review +- **accepted**: Imported by `sn import` from reviewed catalog (catalog-authoritative) + +### Write Semantics + +Two distinct write paths with different semantics: + +- **`write_standard_names()` (build path)**: Uses `coalesce(b.field, sn.field)` for ALL fields — passing None preserves existing graph data. Safe to re-run without erasing imported data. +- **`_write_catalog_entries()` (import path)**: Catalog fields SET directly (overwrite) — catalog is authoritative. Graph-only fields (embedding, model, generated_at, confidence) preserved via coalesce. + +### MCP Tools + +| Tool | Purpose | Key Parameters | +|------|---------|----------------| +| `search_standard_names` | Semantic + keyword search over StandardName descriptions | `query`, `kind`, `tags`, `review_status`, `k` | +| `fetch_standard_names` | Fetch full entries by name ID | `names` (space/comma separated) | +| `list_standard_names` | List with optional filters | `tag`, `kind`, `review_status` | + +### Schema + +StandardName node defined in `imas_codex/schemas/standard_name.yaml`. Key relationships: + +- `(IMASNode)-[:HAS_STANDARD_NAME]->(StandardName)` +- `(FacilitySignal)-[:HAS_STANDARD_NAME]->(StandardName)` +- `(StandardName)-[:CANONICAL_UNITS]->(Unit)` + ## Remote Tools Prefer these Rust-based CLI tools over standard Unix commands. Defined in `imas_codex/config/remote_tools.yaml`. diff --git a/plans/README.md b/plans/README.md index d34ef4a77..bf43fd214 100644 --- a/plans/README.md +++ b/plans/README.md @@ -30,10 +30,10 @@ Gap documents consolidate remaining work from completed implementation phases. T | Plan | Scope | Status | Depends On | |------|-------|--------|------------| | [standard-names/09-sn-generate.md](features/standard-names/09-sn-generate.md) | Core pipeline: EXTRACT→COMPOSE→VALIDATE→PERSIST | ✅ Done | — | -| [standard-names/11-rich-compose.md](features/standard-names/11-rich-compose.md) | Rich compose: full catalog fields, schema extension | Ready | 09 | -| [standard-names/12-catalog-import.md](features/standard-names/12-catalog-import.md) | Catalog import & bootstrap (309 entries, feedback loop) | Ready | 11 P1 | -| [standard-names/13-publish-pipeline.md](features/standard-names/13-publish-pipeline.md) | Lossless publish, round-trip, batched PRs | Ready | 11, 12 P1 | -| [standard-names/14-mcp-tools-benchmark.md](features/standard-names/14-mcp-tools-benchmark.md) | SN MCP tools + benchmark quality tiers | Ready | 11, 12 | +| [standard-names/11-rich-compose.md](features/standard-names/11-rich-compose.md) | Rich compose: full catalog fields, schema extension | ✅ Done | 09 | +| [standard-names/12-catalog-import.md](features/standard-names/12-catalog-import.md) | Catalog import & bootstrap (309 entries, feedback loop) | ✅ Done | 11 P1 | +| [standard-names/13-publish-pipeline.md](features/standard-names/13-publish-pipeline.md) | Lossless publish, round-trip, batched PRs | ✅ Done | 11, 12 P1 | +| [standard-names/14-mcp-tools-benchmark.md](features/standard-names/14-mcp-tools-benchmark.md) | SN MCP tools + benchmark quality tiers | ✅ Done | 11, 12 | ### Pending plans (partially implemented) From 7f83dd20a68685eab7e35925bdf2d0071beb9559 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 14:01:59 +0200 Subject: [PATCH 22/37] fix: lifecycle_filter treats NULL as active in list/search_dd_paths lifecycle_status is NULL on 98.5% of IMASNode nodes (19,734 of 20,037). Per schema: NULL means 'inherits IDS-level lifecycle' which defaults to active. The filter 'p.lifecycle_status = active' matched nothing. Fix: lifecycle_filter=active now matches NULL OR active. Also wires physics_domain/lifecycle_filter into search_dd_paths backend (was accepted by MCP tool but silently ignored). --- imas_codex/llm/server.py | 8 ++------ imas_codex/tools/graph_search.py | 24 ++++++++++++++++++++++-- 2 files changed, 24 insertions(+), 8 deletions(-) diff --git a/imas_codex/llm/server.py b/imas_codex/llm/server.py index c1e07aa3d..7629cd774 100644 --- a/imas_codex/llm/server.py +++ b/imas_codex/llm/server.py @@ -2507,11 +2507,6 @@ def search_dd_paths( tools = _get_imas_tools(semantic_search=True) - if physics_domain is not None or lifecycle_filter is not None: - logger.debug( - "physics_domain/lifecycle_filter not yet implemented in backend, ignoring" - ) - # Run path search and cluster search in parallel — they are # independent operations sharing the same encoder singleton. def _path_search(): @@ -2523,7 +2518,8 @@ def _path_search(): facility=facility, include_version_context=include_version_context, dd_version=dd_version, - # physics_domain and lifecycle_filter not yet implemented in backend + physics_domain=physics_domain, + lifecycle_filter=lifecycle_filter, ) ) diff --git a/imas_codex/tools/graph_search.py b/imas_codex/tools/graph_search.py index 889885204..26ee7133b 100644 --- a/imas_codex/tools/graph_search.py +++ b/imas_codex/tools/graph_search.py @@ -216,6 +216,8 @@ async def search_dd_paths( facility: str | None = None, include_version_context: bool = False, include_summary_ids: bool = False, + physics_domain: str | None = None, + lifecycle_filter: str | None = None, ctx: Context | None = None, ) -> SearchPathsResult: """Search IMAS paths using hybrid vector + text search.""" @@ -469,6 +471,20 @@ async def search_dd_paths( if r["physics_domain"]: physics_domains.add(r["physics_domain"]) + # --- Post-filter by physics_domain and lifecycle_status --- + if physics_domain: + hits = [h for h in hits if h.physics_domain == physics_domain] + if lifecycle_filter: + if lifecycle_filter == "active": + # NULL means "inherits IDS-level lifecycle" which defaults to active + hits = [ + h + for h in hits + if h.lifecycle_status is None or h.lifecycle_status == "active" + ] + else: + hits = [h for h in hits if h.lifecycle_status == lifecycle_filter] + # --- Expand STRUCTURE hits with leaf children --- _STRUCTURE_TYPES = {"structure", "struct_array", "STRUCTURE"} structure_hits = [h for h in hits if h.data_type in _STRUCTURE_TYPES][:5] @@ -984,8 +1000,12 @@ async def list_dd_paths( extra_filters += " AND p.node_type = $node_type" dd_params["node_type"] = node_type if lifecycle_filter: - extra_filters += " AND p.lifecycle_status = $lifecycle_filter" - dd_params["lifecycle_filter"] = lifecycle_filter + if lifecycle_filter == "active": + # NULL means "inherits IDS-level lifecycle" which defaults to active + extra_filters += " AND (p.lifecycle_status IS NULL OR p.lifecycle_status = 'active')" + else: + extra_filters += " AND p.lifecycle_status = $lifecycle_filter" + dd_params["lifecycle_filter"] = lifecycle_filter include_metadata = response_profile != "minimal" From 8ea06bdafef96ab550aff3bf038442cbb82d9334 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 14:31:04 +0200 Subject: [PATCH 23/37] feat: resolve lifecycle_status inheritance at build time and backfill graph MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Resolve lifecycle_status for all IMASNode data fields by inheriting from the parent IDS when not explicitly set in the DD XML. Previously 98.5% of fields had NULL lifecycle_status, requiring runtime NULL-handling that incorrectly treated 12,264 alpha-inherited nodes as active. Build pipeline: _batch_create_path_nodes now accepts ids_info and resolves inheritance post-version-diff to avoid false IMASNodeChange records. Graph migration: backfilled 19,734 NULL nodes via batched Cypher (7,470 active, 12,264 alpha, 277 obsolescent, 26 alpha-override). Reverted broken NULL-or-active workaround in list_dd_paths and search_dd_paths post-filter — now simple equality checks. Formatter enhancements: - Catalog: [alpha] tag on non-active IDS entries - List: [alpha]/[obsolescent] suffix on non-active paths - IDS summary: lifecycle distribution of child paths - list_dd_paths query now returns lifecycle_status in path_details Updated test_field_lifecycle_status to assert all data nodes have explicit lifecycle_status and valid values include active. Updated imas_dd.yaml schema description to reflect build-time resolution. --- imas_codex/graph/build_dd.py | 15 ++++++++++-- imas_codex/llm/search_formatters.py | 14 +++++++++++ imas_codex/schemas/imas_dd.yaml | 6 ++--- imas_codex/tools/graph_search.py | 36 +++++++++++++++++------------ tests/graph/test_dd_build.py | 19 +++++++++++---- 5 files changed, 66 insertions(+), 24 deletions(-) diff --git a/imas_codex/graph/build_dd.py b/imas_codex/graph/build_dd.py index 5920f7d1b..53f4da06f 100644 --- a/imas_codex/graph/build_dd.py +++ b/imas_codex/graph/build_dd.py @@ -1800,7 +1800,9 @@ def phase_build( stats["ids_created"] = max(stats["ids_created"], len(data["ids_info"])) new_paths_data = {p: data["paths"][p] for p in changes["added"]} - _batch_create_path_nodes(client, new_paths_data, version) + _batch_create_path_nodes( + client, new_paths_data, version, ids_info=data["ids_info"] + ) stats["paths_created"] += len(changes["added"]) _batch_mark_paths_deprecated(client, changes["removed"], version) @@ -2962,10 +2964,14 @@ def _batch_create_path_nodes( paths_data: dict[str, dict], version: str, batch_size: int = 1000, + ids_info: dict[str, dict] | None = None, ) -> None: """Batch create IMASNode nodes with relationships. Uses multiple batched queries to avoid memory issues with large datasets. + When ids_info is provided, fields without explicit lifecycle_status inherit + from their parent IDS (resolved here, not in _extract_paths_recursive, to + avoid polluting version-diff computation with inherited values). """ # Prepare path data for batch insertion path_list = [] @@ -2998,7 +3004,12 @@ def _batch_create_path_nodes( "cocos_label_transformation": path_info.get( "cocos_label_transformation" ), - "lifecycle_status": path_info.get("lifecycle_status"), + "lifecycle_status": path_info.get("lifecycle_status") + or ( + ids_info.get(ids_name, {}).get("lifecycle_status") + if ids_info + else None + ), "lifecycle_version": path_info.get("lifecycle_version"), "timebasepath": path_info.get("timebasepath"), "path_doc": path_info.get("path_doc"), diff --git a/imas_codex/llm/search_formatters.py b/imas_codex/llm/search_formatters.py index ed55680ca..2701d06dd 100644 --- a/imas_codex/llm/search_formatters.py +++ b/imas_codex/llm/search_formatters.py @@ -931,11 +931,14 @@ def format_list_report(result: Any) -> str: dtype = d.get("data_type", "") units = d.get("units", "") doc = d.get("documentation", "") + lifecycle = d.get("lifecycle_status", "") line = f" {d['id']}" if dtype: line += f" ({dtype})" if units: line += f" [{units}]" + if lifecycle and lifecycle != "active": + line += f" [{lifecycle}]" if doc: line += f" — {doc[:100]}" parts.append(line) @@ -990,9 +993,12 @@ def format_overview_report(result: Any) -> str: count = stats.get("path_count", 0) desc = stats.get("description", "") domain = stats.get("physics_domain", "") + lifecycle = stats.get("lifecycle_status", "") line = f" {ids_name} ({count} paths)" if domain: line += f" [{domain}]" + if lifecycle and lifecycle != "active": + line += f" [{lifecycle}]" parts.append(line) if desc: parts.append(f" {desc[:120]}") @@ -1308,6 +1314,14 @@ def format_structure_report(result: dict[str, Any]) -> str: if meta: parts.append(" | ".join(meta)) + # Lifecycle distribution of child paths + lifecycle_dist = result.get("lifecycle_distribution", {}) + if lifecycle_dist and not ( + len(lifecycle_dist) == 1 and lifecycle in lifecycle_dist + ): + dist_parts = [f"{v} {k}" for k, v in lifecycle_dist.items()] + parts.append(f"Path lifecycle: {', '.join(dist_parts)}") + # Metrics m = result.get("metrics", {}) parts.append( diff --git a/imas_codex/schemas/imas_dd.yaml b/imas_codex/schemas/imas_dd.yaml index 50e8b6fa5..b63dc389b 100644 --- a/imas_codex/schemas/imas_dd.yaml +++ b/imas_codex/schemas/imas_dd.yaml @@ -625,9 +625,9 @@ classes: range: integer lifecycle_status: description: >- - Lifecycle maturity status from the DD XML (alpha or obsolescent). - Only set on fields with non-default lifecycle status. Null means - the field inherits the IDS-level lifecycle status. + Lifecycle maturity status. Resolved at build time: fields without + an explicit lifecycle_status in the DD XML inherit from their + parent IDS. All data nodes have an explicit value after build. range: LifecycleStatus lifecycle_version: description: >- diff --git a/imas_codex/tools/graph_search.py b/imas_codex/tools/graph_search.py index 26ee7133b..d2ba4271c 100644 --- a/imas_codex/tools/graph_search.py +++ b/imas_codex/tools/graph_search.py @@ -475,15 +475,7 @@ async def search_dd_paths( if physics_domain: hits = [h for h in hits if h.physics_domain == physics_domain] if lifecycle_filter: - if lifecycle_filter == "active": - # NULL means "inherits IDS-level lifecycle" which defaults to active - hits = [ - h - for h in hits - if h.lifecycle_status is None or h.lifecycle_status == "active" - ] - else: - hits = [h for h in hits if h.lifecycle_status == lifecycle_filter] + hits = [h for h in hits if h.lifecycle_status == lifecycle_filter] # --- Expand STRUCTURE hits with leaf children --- _STRUCTURE_TYPES = {"structure", "struct_array", "STRUCTURE"} @@ -1000,12 +992,8 @@ async def list_dd_paths( extra_filters += " AND p.node_type = $node_type" dd_params["node_type"] = node_type if lifecycle_filter: - if lifecycle_filter == "active": - # NULL means "inherits IDS-level lifecycle" which defaults to active - extra_filters += " AND (p.lifecycle_status IS NULL OR p.lifecycle_status = 'active')" - else: - extra_filters += " AND p.lifecycle_status = $lifecycle_filter" - dd_params["lifecycle_filter"] = lifecycle_filter + extra_filters += " AND p.lifecycle_status = $lifecycle_filter" + dd_params["lifecycle_filter"] = lifecycle_filter include_metadata = response_profile != "minimal" @@ -1016,6 +1004,7 @@ async def list_dd_paths( "p.data_type AS data_type,\n" " p.node_type AS node_type, " "p.documentation AS documentation,\n" + " p.lifecycle_status AS lifecycle_status,\n" " u.symbol AS units" ) else: @@ -1070,6 +1059,7 @@ async def list_dd_paths( "node_type": r.get("node_type"), "documentation": r.get("documentation"), "units": r.get("units"), + "lifecycle_status": r.get("lifecycle_status"), } for r in path_results ] @@ -1163,6 +1153,7 @@ async def get_dd_catalog( "path_count": r["path_count"], "description": r["description"] or "", "physics_domain": r["physics_domain"] or "", + "lifecycle_status": r["lifecycle_status"] or "", } if r["physics_domain"]: physics_domains.add(r["physics_domain"]) @@ -2049,6 +2040,18 @@ async def get_ids_summary( **dd_params, ) + # Query 7: Lifecycle distribution within this IDS + lifecycle_dist = self._gc.query( + f""" + MATCH (p:IMASNode) + WHERE p.ids = $ids_name AND p.node_category = 'data' + AND p.lifecycle_status IS NOT NULL {dd_clause} + RETURN p.lifecycle_status AS status, count(p) AS count + ORDER BY count DESC + """, + **dd_params, + ) + total = meta.get("total", 0) leaves = meta.get("leaves", 0) cluster_count = cluster_count_result[0]["count"] if cluster_count_result else 0 @@ -2079,6 +2082,9 @@ async def get_ids_summary( "coordinate_arrays": coord_count, "cocos_fields": cocos_count, "data_types": {t["data_type"]: t["count"] for t in (types or [])}, + "lifecycle_distribution": { + r["status"]: r["count"] for r in (lifecycle_dist or []) + }, } return result diff --git a/tests/graph/test_dd_build.py b/tests/graph/test_dd_build.py index 4165141bb..7a14a2c20 100644 --- a/tests/graph/test_dd_build.py +++ b/tests/graph/test_dd_build.py @@ -1087,22 +1087,33 @@ def test_ids_have_lifecycle_last_change(self, graph_client, label_counts): assert count >= 80, f"Expected >=80 IDS with lifecycle_last_change, got {count}" def test_field_lifecycle_status(self, graph_client, label_counts): - """Fields with alpha/obsolescent lifecycle_status should be populated.""" + """All data fields should have lifecycle_status resolved from IDS.""" if not label_counts.get("IMASNode"): pytest.skip("No IMASNode nodes in graph") result = graph_client.query( - "MATCH (p:IMASNode) WHERE p.lifecycle_status IS NOT NULL " + "MATCH (p:IMASNode {node_category: 'data'}) " "RETURN p.lifecycle_status AS status, count(p) AS cnt" ) total = sum(r["cnt"] for r in result) - # DD 4.1.1 has 238 fields with lifecycle_status + # After backfill: all data nodes have lifecycle_status assert total >= 200, f"Expected >=200 fields with lifecycle_status, got {total}" statuses = {r["status"] for r in result} - valid = {"alpha", "obsolescent"} + valid = {"active", "alpha", "obsolescent"} invalid = statuses - valid assert not invalid, f"Invalid field lifecycle_status values: {invalid}" + # No data nodes should have NULL lifecycle_status + null_result = graph_client.query( + "MATCH (p:IMASNode {node_category: 'data'}) " + "WHERE p.lifecycle_status IS NULL " + "RETURN count(p) AS cnt" + ) + null_count = null_result[0]["cnt"] if null_result else 0 + assert null_count == 0, ( + f"Expected 0 data nodes with NULL lifecycle_status, got {null_count}" + ) + class TestTimebasepath: """Verify timebasepath metadata on dynamic fields.""" From 96fc7436be440b3684299bb1e7ed3f46d8a67e59 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 14:48:27 +0200 Subject: [PATCH 24/37] feat: surface lifecycle in check/version/changelog tools, gate get_graph_schema on dd-only - Add lifecycle_status to check_dd_paths query, result model, and formatter - Add lifecycle_status to get_dd_version_context per-path query and formatter - Add lifecycle_status column to get_dd_changelog query and formatter table - Gate get_graph_schema behind dd-only guard (REPL companion, not needed without REPL) --- imas_codex/llm/search_formatters.py | 15 +++++-- imas_codex/llm/server.py | 65 ++++++++++++++++------------- imas_codex/models/result_models.py | 3 ++ imas_codex/tools/graph_search.py | 3 +- imas_codex/tools/version_tool.py | 4 +- 5 files changed, 56 insertions(+), 34 deletions(-) diff --git a/imas_codex/llm/search_formatters.py b/imas_codex/llm/search_formatters.py index 2701d06dd..4d495fa91 100644 --- a/imas_codex/llm/search_formatters.py +++ b/imas_codex/llm/search_formatters.py @@ -822,6 +822,9 @@ def format_check_report(result: Any) -> str: meta.append(f"Units: {item.units}") if item.ids_name: meta.append(f"IDS: {item.ids_name}") + lifecycle = getattr(item, "lifecycle_status", None) + if lifecycle and lifecycle != "active": + meta.append(f"Lifecycle: {lifecycle}") if meta: parts.append(f" {' | '.join(meta)}") else: @@ -1538,19 +1541,25 @@ def format_dd_changelog_report(result: dict[str, Any]) -> str: parts.append(f"Version range: {fr or 'earliest'} → {to or 'latest'}") parts.append("") - parts.append("| Rank | Path | IDS | Changes | Types | Renamed | Score |") - parts.append("|------|------|-----|---------|-------|---------|-------|") + parts.append( + "| Rank | Path | IDS | Lifecycle | Changes | Types | Renamed | Score |" + ) + parts.append( + "|------|------|-----|-----------|---------|-------|---------|-------|" + ) for i, row in enumerate(result.get("results", []), 1): path = row.get("path", "") ids_name = row.get("ids", "") + lifecycle = row.get("lifecycle_status", "active") or "active" changes = row.get("change_count", 0) types = row.get("change_types", []) renamed = "✓" if row.get("was_renamed") else "" score = row.get("volatility_score", 0) type_str = ", ".join(str(t) for t in types if t) if types else "" + lifecycle_tag = lifecycle if lifecycle != "active" else "" parts.append( - f"| {i} | `{path}` | {ids_name} | {changes} | {type_str} | {renamed} | {score} |" + f"| {i} | `{path}` | {ids_name} | {lifecycle_tag} | {changes} | {type_str} | {renamed} | {score} |" ) if total >= limit: diff --git a/imas_codex/llm/server.py b/imas_codex/llm/server.py index 7629cd774..1b95c15fc 100644 --- a/imas_codex/llm/server.py +++ b/imas_codex/llm/server.py @@ -451,7 +451,10 @@ def _format_version_context_report(result: dict) -> str: changes = ctx.get("changes", []) introduced = ctx.get("introduced_in") deprecated = ctx.get("deprecated_in") + lifecycle = ctx.get("lifecycle_status") lifecycle_parts = [] + if lifecycle and lifecycle != "active": + lifecycle_parts.append(lifecycle) if introduced: lifecycle_parts.append(f"introduced v{introduced}") if deprecated: @@ -1941,41 +1944,45 @@ def repl(code: str) -> str: return f"Error: {e}\n\n{tb}" # ===================================================================== - # Tool 2: get_graph_schema - Schema introspection + # Tool 2: get_graph_schema - Schema introspection (REPL companion) # ===================================================================== + # Only available in full mode — provides schema context for Cypher + # queries via the REPL, which dd-only mode does not expose. - @self.mcp.tool() - def get_graph_schema( - scope: str = "overview", - ) -> str: - """Get graph schema context for Cypher query generation. + if not self.dd_only: - Returns compact, task-relevant schema in text format. Use scope to - get only the schema slice you need, reducing token usage. Call this - before writing any raw Cypher to verify node labels, property names, - relationship types, and enum values. + @self.mcp.tool() + def get_graph_schema( + scope: str = "overview", + ) -> str: + """Get graph schema context for Cypher query generation. - Args: - scope: Schema slice to return. One of: - - "overview": compact summary of all node labels, relationship - types, vector indexes, and task groupings (default). - - "signals": FacilitySignal, DataAccess, Diagnostic, AccessCheck. - - "wiki": WikiPage, WikiChunk, Document, Image. - - "imas": IMASNode, IDS, IMASSemanticCluster, DDVersion, Unit, - IMASNodeChange, IMASCoordinateSpec. - - "code": CodeFile, CodeChunk, CodeExample. - - "facility": Facility, FacilityPath, FacilitySignal, SignalNode, - Diagnostic. - - "data_sources": data source nodes and tree-related relationships. + Returns compact, task-relevant schema in text format. Use scope to + get only the schema slice you need, reducing token usage. Call this + before writing any raw Cypher to verify node labels, property names, + relationship types, and enum values. - Returns: - Formatted text containing property tables (name, type, description), - relationship definitions as (From)-[:REL]->(To), available vector - indexes, and enum values for the requested scope. - """ - from imas_codex.graph.schema_context import schema_for + Args: + scope: Schema slice to return. One of: + - "overview": compact summary of all node labels, relationship + types, vector indexes, and task groupings (default). + - "signals": FacilitySignal, DataAccess, Diagnostic, AccessCheck. + - "wiki": WikiPage, WikiChunk, Document, Image. + - "imas": IMASNode, IDS, IMASSemanticCluster, DDVersion, Unit, + IMASNodeChange, IMASCoordinateSpec. + - "code": CodeFile, CodeChunk, CodeExample. + - "facility": Facility, FacilityPath, FacilitySignal, SignalNode, + Diagnostic. + - "data_sources": data source nodes and tree-related relationships. + + Returns: + Formatted text containing property tables (name, type, description), + relationship definitions as (From)-[:REL]->(To), available vector + indexes, and enum values for the requested scope. + """ + from imas_codex.graph.schema_context import schema_for - return schema_for(task=scope) + return schema_for(task=scope) if not self.read_only: # ===================================================================== diff --git a/imas_codex/models/result_models.py b/imas_codex/models/result_models.py index 58bf807f9..b5b6a2703 100644 --- a/imas_codex/models/result_models.py +++ b/imas_codex/models/result_models.py @@ -461,6 +461,9 @@ class CheckPathsResultItem(BaseModel): ids_name: str | None = Field(default=None, description="IDS name if path exists") data_type: str | None = Field(default=None, description="Data type if available") units: str | None = Field(default=None, description="Physical units if available") + lifecycle_status: str | None = Field( + default=None, description="Lifecycle maturity: active, alpha, or obsolescent" + ) migration: dict[str, Any] | None = Field( default=None, description="Migration info if path is deprecated" ) diff --git a/imas_codex/tools/graph_search.py b/imas_codex/tools/graph_search.py index d2ba4271c..f2b361bc3 100644 --- a/imas_codex/tools/graph_search.py +++ b/imas_codex/tools/graph_search.py @@ -604,7 +604,7 @@ async def check_dd_paths( OPTIONAL MATCH (old:IMASNode {{id: check_path}})-[:RENAMED_TO]->(new:IMASNode) RETURN check_path, p.id AS id, p.ids AS ids, p.data_type AS data_type, - u.id AS units, + u.id AS units, p.lifecycle_status AS lifecycle_status, old.id AS renamed_from, new.id AS renamed_to """, paths=path_list, @@ -628,6 +628,7 @@ async def check_dd_paths( ids_name=r["ids"], data_type=r["data_type"], units=r["units"] or "", + lifecycle_status=r.get("lifecycle_status"), ) ) found += 1 diff --git a/imas_codex/tools/version_tool.py b/imas_codex/tools/version_tool.py index 90b26ac51..d0fe60e73 100644 --- a/imas_codex/tools/version_tool.py +++ b/imas_codex/tools/version_tool.py @@ -138,6 +138,7 @@ async def get_dd_version_context( OPTIONAL MATCH (change:IMASNodeChange)-[:FOR_IMAS_PATH]->(p) OPTIONAL MATCH (change)-[:IN_VERSION]->(v:DDVersion) RETURN p.id AS id, + p.lifecycle_status AS lifecycle_status, iv.id AS introduced_in, dv.id AS deprecated_in, count(change) AS change_count, @@ -161,6 +162,7 @@ async def get_dd_version_context( c for c in (r.get("changes") or []) if c.get("version") is not None ] path_ctx[r["id"]] = { + "lifecycle_status": r.get("lifecycle_status"), "introduced_in": r.get("introduced_in"), "deprecated_in": r.get("deprecated_in"), "change_count": len(changes), @@ -274,7 +276,7 @@ async def get_dd_changelog( OPTIONAL MATCH (p)-[:RENAMED_TO]->() WITH p, change_count, type_variety, change_types, CASE WHEN EXISTS { (p)-[:RENAMED_TO]->() } THEN 1 ELSE 0 END AS was_renamed - RETURN p.id AS path, p.ids AS ids, + RETURN p.id AS path, p.ids AS ids, p.lifecycle_status AS lifecycle_status, change_count, type_variety, change_types, was_renamed, change_count + (type_variety * 2) + (was_renamed * 3) AS volatility_score ORDER BY volatility_score DESC From b2a2b2b4cd3e60ab6ae2e5585c0db5d9ae4becec Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 15:38:34 +0200 Subject: [PATCH 25/37] fix: pin setup-uv action to v8.0.0 (immutable tag) --- .github/workflows/benchmark.yml | 2 +- .github/workflows/docker-build-push.yml | 2 +- .github/workflows/graph-quality.yml | 2 +- .github/workflows/release.yml | 2 +- .github/workflows/test.yml | 4 ++-- 5 files changed, 6 insertions(+), 6 deletions(-) diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 992de7724..e4b26995b 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -93,7 +93,7 @@ jobs: df -h - name: Install UV + ASV - uses: astral-sh/setup-uv@v8 + uses: astral-sh/setup-uv@v8.0.0 with: github-token: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/docker-build-push.yml b/.github/workflows/docker-build-push.yml index 07fa0e667..0046463b3 100644 --- a/.github/workflows/docker-build-push.yml +++ b/.github/workflows/docker-build-push.yml @@ -208,7 +208,7 @@ jobs: - name: Install uv if: steps.graph-tag.outputs.imas-tag != 'none' - uses: astral-sh/setup-uv@v8 + uses: astral-sh/setup-uv@v8.0.0 with: enable-cache: true github-token: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/graph-quality.yml b/.github/workflows/graph-quality.yml index 79a93e91d..0883540c6 100644 --- a/.github/workflows/graph-quality.yml +++ b/.github/workflows/graph-quality.yml @@ -185,7 +185,7 @@ jobs: "MATCH (n) RETURN count(n) AS nodes, labels(n)[0] AS label ORDER BY nodes DESC LIMIT 10" - name: Install uv - uses: astral-sh/setup-uv@v8 + uses: astral-sh/setup-uv@v8.0.0 with: enable-cache: true github-token: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 7218b9742..8c262ff82 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -27,7 +27,7 @@ jobs: fetch-depth: 0 - name: Install uv - uses: astral-sh/setup-uv@v8 + uses: astral-sh/setup-uv@v8.0.0 with: github-token: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index b9367da8a..8d1884da8 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -50,7 +50,7 @@ jobs: run: playwright install --with-deps chromium - name: Install uv - uses: astral-sh/setup-uv@v8 + uses: astral-sh/setup-uv@v8.0.0 with: enable-cache: true github-token: ${{ secrets.GITHUB_TOKEN }} @@ -118,7 +118,7 @@ jobs: - uses: actions/checkout@v6 - name: Install uv - uses: astral-sh/setup-uv@v8 + uses: astral-sh/setup-uv@v8.0.0 with: enable-cache: true github-token: ${{ secrets.GITHUB_TOKEN }} From 8fdd57e424f56ed017bcdc45522cebf3d03e6c60 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 15:39:10 +0200 Subject: [PATCH 26/37] docs: add standard name feature plans --- .../standard-names/16-benchmark-parity.md | 136 +++++++++++++++ .../17-sn-lifecycle-management.md | 163 ++++++++++++++++++ 2 files changed, 299 insertions(+) create mode 100644 plans/features/standard-names/16-benchmark-parity.md create mode 100644 plans/features/standard-names/17-sn-lifecycle-management.md diff --git a/plans/features/standard-names/16-benchmark-parity.md b/plans/features/standard-names/16-benchmark-parity.md new file mode 100644 index 000000000..fddf707c0 --- /dev/null +++ b/plans/features/standard-names/16-benchmark-parity.md @@ -0,0 +1,136 @@ +# 16: Benchmark / Build Prompt Parity & Caching + +**Status:** Ready to implement +**Depends on:** None (standalone fix) +**Blocks:** 18 (calibration — benchmark results meaningless until parity fixed) +**Agent:** engineer + +## Problem + +The `sn benchmark` command constructs LLM prompts differently from `sn build`, +making benchmark metrics (cost, speed, quality) unreliable for production model +selection. + +### Specific gaps + +| Aspect | Build (`workers.py`) | Benchmark (`benchmark.py`) | +|--------|---------------------|---------------------------| +| System prompt | `sn/compose_system` via `build_compose_context()` | None | +| User prompt | `sn/compose_dd` with `cluster_context` | `sn/compose_dd` without `cluster_context` | +| Message structure | `[system, user]` — enables prompt caching | `[user]` — no caching possible | +| Context source | `build_compose_context()` (grammar rules, vocabulary, examples, tokamak ranges) | `build_grammar_context()` (bare enum lists only) | +| Prompt caching | `inject_cache_control()` on system message | Not applicable (no system message) | + +### Impact + +- Cost metrics inflated (no cache hits) +- Speed metrics pessimistic (larger uncached prompts) +- Quality metrics unreliable (model sees different context) +- Benchmark cannot validate whether caching actually works + +## Phase 1: Shared prompt construction + +**Files:** `imas_codex/sn/benchmark.py` + +Replace custom prompt construction in `_run_model()` with the same path used by +`compose_worker()` in `workers.py`: + +```python +# Current (broken): +prompt_context = {"items": items, "ids_name": group_key, "existing_names": ..., **grammar_ctx} +prompt_text = render_prompt("sn/compose_dd", prompt_context) +messages = [{"role": "user", "content": prompt_text}] + +# Fixed: +from imas_codex.sn.context import build_compose_context +context = build_compose_context() +system_prompt = render_prompt("sn/compose_system", context) +user_context = { + "items": items, + "ids_name": group_key, + "existing_names": sorted(existing)[:200], + "cluster_context": batch.get("context", ""), + **context, +} +user_prompt = render_prompt("sn/compose_dd", user_context) +messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt}, +] +``` + +- Remove `build_grammar_context()` from benchmark.py (dead code after this change) +- Keep the function in benchmark.py only if `score_with_reviewer()` needs it +- Preserve `batch.context` (cluster_context) through extraction → benchmark + +### Extraction fix + +`_extract_candidates()` currently drops `batch.context` when converting +`ExtractionBatch` to plain dicts (line 494-500). Fix: + +```python +result.append({ + "group_key": batch.group_key, + "items": batch.items, + "existing_names": list(batch.existing_names), + "context": batch.context, # ADD THIS +}) +``` + +## Phase 2: Cache hit reporting + +**Files:** `imas_codex/sn/benchmark.py`, `imas_codex/discovery/base/llm.py` + +Add cache hit/miss tracking to benchmark output: + +1. Extend `acall_llm_structured()` return or use response metadata to detect + prompt cache hits (OpenRouter returns `usage.cache_creation_input_tokens` + and `usage.cache_read_input_tokens`) +2. Add to `ModelResult`: + ```python + cache_read_tokens: int = 0 + cache_creation_tokens: int = 0 + cache_hit_rate: float = 0.0 + ``` +3. Display in Rich table: "Cache %" column + +**Note:** This requires checking what `litellm` exposes in the response object. +The `usage` dict from OpenRouter includes cache fields. `acall_llm_structured()` +currently returns `(result, cost, tokens)` — we may need to return the full +usage dict or add a cache-specific return value. + +## Phase 3: Render system prompt once, verify caching + +The system prompt should be rendered **once** before the model loop, since it's +identical across all batches and models. This matches the build pipeline pattern +(line 152 in workers.py). + +```python +context = build_compose_context() +system_prompt = render_prompt("sn/compose_system", context) + +for model in config.models: + result = await _run_model( + model=model, + extraction_batches=extraction_batches, + config=config, + reference=REFERENCE_NAMES, + system_prompt=system_prompt, # pass pre-rendered + context=context, # for user prompt rendering + ) +``` + +## Acceptance criteria + +1. `sn benchmark` uses identical prompt construction as `sn build` +2. System/user message split enables prompt caching +3. `cluster_context` is preserved through extraction +4. Cache hit rate is reported in benchmark output +5. Running the same benchmark twice shows cache hits on second run +6. `build_grammar_context()` removed from benchmark.py (or moved to shared location) + +## Test plan + +- Unit test: `_run_model()` constructs `[system, user]` messages +- Unit test: extraction preserves `context` field +- Integration: run benchmark with `--verbose`, confirm cache token counts in logs diff --git a/plans/features/standard-names/17-sn-lifecycle-management.md b/plans/features/standard-names/17-sn-lifecycle-management.md new file mode 100644 index 000000000..239cc0775 --- /dev/null +++ b/plans/features/standard-names/17-sn-lifecycle-management.md @@ -0,0 +1,163 @@ +# 17: Standard Name Lifecycle Management (Reset / Clear) + +**Status:** Ready to implement +**Depends on:** None (standalone) +**Agent:** engineer + +## Problem + +All discovery domains (signals, paths, wiki, code, documents) have `--reset-to` +infrastructure via the shared `ResetSpec` / `reset_to_status()` pattern. Standard +names have none. This makes iterative development painful — after changing prompts +or models, there's no way to re-run the pipeline without manual Cypher cleanup. + +### Key difference from other domains + +StandardName nodes are **cross-facility** — they have no `facility_id` property. +The standard `reset_to_status()` function requires a facility parameter. We need +an adapted approach. + +### StandardName lifecycle + +``` +drafted → published → accepted + ↘ rejected + ↘ skipped +``` + +## Phase 1: Scoped reset command + +**Files:** +- `imas_codex/sn/graph_ops.py` — add `reset_standard_names()` and `clear_standard_names()` +- `imas_codex/cli/sn.py` — add `sn reset` subcommand + +### `sn reset` command + +```bash +# Reset all drafted names back for re-composition (clears LLM output, keeps nodes) +imas-codex sn reset --status drafted + +# Reset published names back to drafted (e.g., after prompt change) +imas-codex sn reset --status published --to drafted + +# Reset only DD-sourced names +imas-codex sn reset --status drafted --source dd + +# Reset names from a specific IDS +imas-codex sn reset --status drafted --ids equilibrium + +# Dry run +imas-codex sn reset --status drafted --dry-run +``` + +### Graph operation + +```python +def reset_standard_names( + *, + target_status: str = "drafted", + source_statuses: list[str] | None = None, + source_filter: str | None = None, # "dd" or "signals" + ids_filter: str | None = None, + clear_embeddings: bool = True, +) -> int: + """Reset StandardName nodes to a target review_status. + + Unlike facility-scoped domains, StandardName has no facility_id. + Filtering is by review_status, source, and IDS. + """ +``` + +Fields to clear on reset to `drafted`: +- `embedding`, `embedded_at` (will be regenerated) +- `model`, `generated_at` (provenance of old generation) +- `confidence` (will be re-scored) + +Fields to preserve: +- `id` (the standard name itself) +- `source`, `source_path` (how it was sourced) +- `created_at` (first creation time) + +Relationships to clean: +- `HAS_STANDARD_NAME` — remove (will be re-created on persist) +- `CANONICAL_UNITS` — remove (will be re-created) + +## Phase 2: Clear command + +**Files:** `imas_codex/sn/graph_ops.py`, `imas_codex/cli/sn.py` + +```bash +# Delete all drafted standard names (not accepted/imported ones) +imas-codex sn clear --status drafted + +# Delete ALL standard names (requires --confirm) +imas-codex sn clear --all --confirm + +# Delete names from a specific source +imas-codex sn clear --source dd --status drafted +``` + +### Graph operation + +```python +def clear_standard_names( + *, + status_filter: list[str] | None = None, + source_filter: str | None = None, + ids_filter: str | None = None, + confirm_all: bool = False, +) -> int: + """Delete StandardName nodes and their relationships. + + Refuses to delete accepted/imported names unless confirm_all=True. + """ +``` + +Safety rules: +- Default: only delete `drafted` names +- `accepted` and `imported` names require `--confirm` flag +- Always log count before deletion +- DETACH DELETE (removes all relationships) + +## Phase 3: Wire into build command + +**Files:** `imas_codex/cli/sn.py` + +Add `--reset-to` option to `sn build`: + +```bash +# Re-compose all drafted names from scratch (reset + build) +imas-codex sn build --source dd --reset-to drafted + +# Full pipeline from extraction (clears and rebuilds) +imas-codex sn build --source dd --reset-to extracted --ids equilibrium +``` + +Reset targets for `sn build`: +- `extracted` — clear all SN nodes for matching source, re-run full pipeline +- `drafted` — reset existing drafted names, re-compose them + +```python +@click.option( + "--reset-to", + type=click.Choice(["extracted", "drafted"]), + default=None, + help="Reset standard names to target state before building.", +) +``` + +## Acceptance criteria + +1. `sn reset --status drafted` resets nodes and clears embeddings +2. `sn clear --status drafted` deletes only drafted names +3. `sn clear --all` requires `--confirm` flag +4. `sn build --reset-to drafted` resets then rebuilds +5. Accepted/imported names are never touched without explicit confirmation +6. `sn status` shows correct counts after reset/clear + +## Test plan + +- Unit test: `reset_standard_names()` clears correct fields +- Unit test: `clear_standard_names()` refuses to delete accepted without confirm +- Unit test: `--reset-to` on build triggers reset before pipeline +- Integration: build → reset → rebuild cycle produces valid results From d1fe7208beb1a76d2a4da7d87a11cf6f02675304 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 15:39:53 +0200 Subject: [PATCH 27/37] docs: update standard name implementation order --- .../standard-names/00-implementation-order.md | 15 +- .../18-benchmark-calibration.md | 276 ++++++++++++++++++ 2 files changed, 289 insertions(+), 2 deletions(-) create mode 100644 plans/features/standard-names/18-benchmark-calibration.md diff --git a/plans/features/standard-names/00-implementation-order.md b/plans/features/standard-names/00-implementation-order.md index 323417bfe..f043876e7 100644 --- a/plans/features/standard-names/00-implementation-order.md +++ b/plans/features/standard-names/00-implementation-order.md @@ -46,7 +46,11 @@ All status values use past tense: drafted, published, accepted, rejected, skippe | 11 | rich-compose | Full catalog fields, schema extension, coalesce fix, tests | 📋 Ready | 09 | 12, 13, 14 | | 12 | catalog-import | Feedback import from reviewed catalog PRs | 📋 Ready | 11 P1 | 13 P4 | | 13 | publish-pipeline | Lossless YAML export, batched PRs | 📋 Ready | 11 (all) | 12 (feedback loop) | -| 14 | mcp-tools-benchmark | SN search/fetch/list MCP tools + benchmark quality | 📋 Ready | 11 (embedding) | — | +| 14 | mcp-tools-benchmark | SN search/fetch/list MCP tools + benchmark quality | ✅ Done | 11 (embedding) | 16 | +| 15 | import-physics-domain | Import physics_domain from catalog | ✅ Done | 12 | — | +| 16 | benchmark-parity | Fix benchmark/build prompt parity + caching | 📋 Ready | 14 | 18 | +| 17 | sn-lifecycle-management | SN reset/clear commands | 📋 Ready | — | — | +| 18 | benchmark-calibration | Calibration dataset, reviewer enhancement, model selection | 📋 Ready | 16 | — | ## Deployment Waves @@ -77,7 +81,14 @@ Run in parallel — export and tools are independent: - **Agent A (engineer):** Fix lossy publish export, update graph query, PR workflow, dedup - **Agent B (engineer):** 3 MCP tools (search/fetch/list) + benchmark quality tiers + reviewer -### Wave 4: Integration Testing +### Wave 4: Benchmark & Lifecycle (Plans 16, 17, 18) + +Prerequisite for production minting: +- **Agent A (engineer):** Plan 16 — fix benchmark prompt parity + caching verification +- **Agent B (engineer):** Plan 17 — SN reset/clear commands for iterative development +- **Agent C (architect):** Plan 18 — calibration dataset, reviewer enhancement, model selection + +### Wave 5: Integration Testing - End-to-end: build → publish → review → import - Round-trip idempotence verification diff --git a/plans/features/standard-names/18-benchmark-calibration.md b/plans/features/standard-names/18-benchmark-calibration.md new file mode 100644 index 000000000..329921b98 --- /dev/null +++ b/plans/features/standard-names/18-benchmark-calibration.md @@ -0,0 +1,276 @@ +# 18: Benchmark Calibration, Model Selection & Reviewer Enhancement + +**Status:** Ready to implement (after Plan 16) +**Depends on:** 16 (benchmark parity — results meaningless until prompts match build) +**Agent:** architect (requires research + design decisions) + +## Problem + +The benchmark needs improvements in three areas before we can make grounded model +selection decisions for production minting: + +1. **Gold reference set is too small** — 30 entries from 4 IDSs +2. **No calibration dataset** — reviewer has no full-entry examples to anchor scores +3. **Reviewer prompt is ad-hoc** — inline string, no grammar context, no system prompt +4. **No structured approach to model selection** — we need cost/speed/quality tradeoffs + +## Design Principle: Separate Gold Set from Calibration Set + +The rubber-duck critique caught this: the reference set and calibration set serve +different purposes and should not be mixed. + +| Dataset | Purpose | Contents | Location | +|---------|---------|----------|----------| +| **Gold reference** | Exact match: did model produce the right name? | `source_path → expected_name + fields` | `benchmark_reference.py` | +| **Calibration set** | Score anchoring: what does "outstanding" look like? | Full entries with expected tier + score | `benchmark_calibration.yaml` | + +## Phase 1: Expand gold reference set + +**Files:** `imas_codex/sn/benchmark_reference.py` + +Expand from 30 to ~50 entries covering more IDSs: + +| IDS | Current | Target | +|-----|---------|--------| +| equilibrium | 18 | 20 | +| core_profiles | 6 | 10 | +| magnetics | 4 | 6 | +| summary | 2 | 4 | +| core_transport | 0 | 4 | +| mhd_linear | 0 | 2 | +| nbi | 0 | 2 | +| edge_profiles | 0 | 2 | + +Fix the questionable rogowski_coil reference entry (line 116-121 in +benchmark_reference.py — maps major_radius position to rogowski_coil object, +which is not physically meaningful). + +## Phase 2: Create calibration dataset + +**Files:** `imas_codex/sn/benchmark_calibration.yaml` (new) + +Create ~15 hand-crafted full entries spanning quality tiers. Source material from +the `imas-standard-names` package's `resources/standard_name_examples/` directory +(40+ curated examples available). + +### Structure + +```yaml +# Each entry is a complete standard name with expected quality assessment +entries: + - name: electron_temperature + tier: outstanding + expected_score: 95 + description: "Temperature of the electron population." + documentation: > + Electron temperature $T_e$ is a fundamental plasma parameter... + Typical range: 0.1–30 keV in tokamak core... + unit: eV + kind: scalar + tags: [core_profiles, equilibrium] + fields: + physical_base: temperature + subject: electron + reason: > + Canonical physics quantity. Rich documentation with LaTeX, typical + values, cross-references. Perfect grammar decomposition. + + - name: banana_orbits + tier: poor + expected_score: 15 + description: "Banana orbits" + documentation: "" + unit: null + kind: metadata + tags: [] + fields: + physical_base: banana_orbits + reason: > + Not a measurable quantity. No documentation. "banana_orbits" is + not a valid physical_base token. No unit possible. +``` + +### Tier distribution (15 entries) + +| Tier | Count | Score Range | Examples | +|------|-------|-------------|----------| +| outstanding | 4 | 85-100 | electron_temperature, plasma_current, safety_factor, poloidal_magnetic_flux | +| good | 4 | 60-79 | toroidal_component_of_magnetic_field, electron_pressure, loop_voltage, stored_energy | +| adequate | 4 | 40-59 | minor_radius, aspect_ratio, time, elongation | +| poor | 3 | 0-39 | banana_orbits, h_mode, (invalid grammar example) | + +Source the "outstanding" and "good" entries from `imas-standard-names` examples +where available. Hand-craft the "poor" examples to test error detection. + +## Phase 3: Enhance reviewer prompt + +**Files:** +- `imas_codex/llm/prompts/sn/review_benchmark.md` (new template) +- `imas_codex/sn/benchmark.py` — update `score_with_reviewer()` + +Replace the inline rubric string with a proper Jinja2 template: + +```markdown +--- +name: sn/review_benchmark +description: Quality scoring for benchmark entries +used_by: imas_codex.sn.benchmark.score_with_reviewer +schema_needs: [] +--- + +You are a physics nomenclature expert evaluating standard name entries. + +## Grammar Rules +{{ canonical_pattern }} +{{ segment_order }} + +## Scoring Rubric +... + +## Calibration Examples + +{% for entry in calibration_entries %} +### {{ entry.name }} — {{ entry.tier }} ({{ entry.expected_score }}/100) +{{ entry.reason }} +{% endfor %} + +## Entries to Review +``` + +Key improvements: +- System/user message split (enables caching when reviewer runs multiple batches) +- Full calibration entries as anchors (not just names) +- Grammar rules included so reviewer can validate grammar correctness +- Structured rubric with specific scoring dimensions + +### Scoring dimensions (each 0-20, sum = 0-100) + +1. **Grammar correctness** (0-20): Valid parse, correct segment usage +2. **Semantic accuracy** (0-20): Name correctly describes the physics quantity +3. **Documentation quality** (0-20): LaTeX, typical values, cross-references +4. **Naming conventions** (0-20): Follows established patterns, consistent with peers +5. **Entry completeness** (0-20): Unit, kind, tags, links, ids_paths populated + +Update `QualityReview` model: + +```python +class QualityReview(BaseModel): + name: str + quality_tier: str + score: int = Field(ge=0, le=100) + grammar_score: int = Field(ge=0, le=20) + semantic_score: int = Field(ge=0, le=20) + documentation_score: int = Field(ge=0, le=20) + convention_score: int = Field(ge=0, le=20) + completeness_score: int = Field(ge=0, le=20) + reasoning: str +``` + +## Phase 4: Model selection framework + +### Candidate models for production minting + +Based on the SN pipeline requirements (structured output, grammar adherence, +physics knowledge, documentation quality), evaluate: + +**Compose phase (language model):** +| Model | Strengths | Concerns | +|-------|-----------|----------| +| `claude-sonnet-4` | Strong structured output, good physics | Cost | +| `claude-sonnet-4-5` | Previous gen, well-tested | Superseded | +| `gpt-4o` | Fast, good structured output | Physics depth | +| `gemini-2.5-flash` | Very fast, cheap | Grammar adherence unknown | +| `gpt-5.1` | Newest, potentially strongest | Cost, untested | + +**Review phase (reasoning model):** +| Model | Strengths | Concerns | +|-------|-----------|----------| +| `claude-opus-4-6` | Deepest reasoning | Expensive | +| `o4-mini` | Good reasoning, cheaper | Structured output reliability | +| `claude-sonnet-4` | Good balance | May agree with itself if same family | + +**Reviewer (benchmark scoring):** +| Model | Recommendation | +|-------|----------------| +| `claude-opus-4-6` | Best for calibration alignment | + +### Benchmark execution plan + +After Plans 16 is implemented: + +```bash +# Phase A: Validate prompt caching works +imas-codex sn benchmark \ + --models openrouter/anthropic/claude-sonnet-4 \ + --ids equilibrium --max-candidates 20 -v +# Check logs for cache_read_input_tokens > 0 on second batch + +# Phase B: Multi-model comparison (small) +imas-codex sn benchmark \ + --models openrouter/anthropic/claude-sonnet-4,openrouter/google/gemini-2.5-flash-preview,openrouter/openai/gpt-4o \ + --ids equilibrium --max-candidates 30 \ + --reviewer-model openrouter/anthropic/claude-opus-4-6 + +# Phase C: Winner deep-dive (larger) +imas-codex sn benchmark \ + --models , \ + --max-candidates 80 \ + --reviewer-model openrouter/anthropic/claude-opus-4-6 \ + --runs 2 # consistency check +``` + +### Decision criteria + +| Metric | Weight | Threshold | +|--------|--------|-----------| +| Grammar valid % | Critical | ≥95% or disqualify | +| Fields consistent % | High | ≥85% | +| Reference recall | High | ≥60% | +| Avg quality score | High | ≥65 | +| Cost per name | Medium | <$0.01 preferred | +| Names/min | Medium | >30 preferred | +| Cache hit rate | Low | >50% confirms caching works | + +## Phase 5: Token caching verification + +Prompt caching is critical for cost efficiency at scale. After Plan 16 lands: + +1. Run benchmark with a single model and 2+ batches +2. Check DEBUG logs for `cache_read_input_tokens` and `cache_creation_input_tokens` +3. First batch should show `cache_creation_input_tokens > 0` +4. Second batch should show `cache_read_input_tokens > 0` +5. If no cache fields in response, check: + - Model string has `openrouter/` prefix (required for cache_control passthrough) + - System message has `cache_control: {"type": "ephemeral"}` breakpoint + - LiteLLM proxy is forwarding cache_control blocks + +## Acceptance criteria + +1. Gold reference expanded to 50+ entries across 8+ IDSs +2. Calibration dataset with 15 full entries across 4 quality tiers +3. Reviewer prompt uses Jinja2 template with system/user split +4. Reviewer produces 5-dimensional scores (grammar, semantic, docs, convention, completeness) +5. Benchmark results inform model selection with clear winner for compose + review +6. Token caching verified working end-to-end + +## On renaming `build` → `mint` + +**Decision: Not yet.** The rubber-duck critique agrees — the pipeline needs to +stabilize first. Revisit after the benchmark shows consistent, high-quality output +across model changes. When ready, add `mint` as the primary command name and keep +`build` as a hidden alias. + +## Documentation updates + +| Target | Update | +|--------|--------| +| `AGENTS.md` | Update SN benchmark section with new calibration workflow | +| `plans/features/standard-names/00-implementation-order.md` | Add plans 16-18 | + +## Test plan + +- Unit test: calibration YAML loads and validates +- Unit test: reviewer prompt renders with calibration entries +- Unit test: QualityReview model accepts 5-dimensional scores +- Unit test: expanded reference set passes grammar round-trip +- Integration: benchmark runs with reviewer and produces scored report From 312ce768fec7a8daf15e4b417ae4f1fc99fc19c1 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 15:49:48 +0200 Subject: [PATCH 28/37] feat: rename sn build to sn mint Rename the CLI command, pipeline function, docstrings, and all plan/documentation references from 'build' to 'mint'. The term better reflects the nature of standard name generation. --- imas_codex/cli/sn.py | 22 ++++++++--------- imas_codex/sn/pipeline.py | 6 ++--- .../standard-names/00-implementation-order.md | 2 +- .../standard-names/16-benchmark-parity.md | 24 +++++++++++++++---- .../17-sn-lifecycle-management.md | 10 ++++---- 5 files changed, 40 insertions(+), 24 deletions(-) diff --git a/imas_codex/cli/sn.py b/imas_codex/cli/sn.py index 9805e8b35..59c0f4a69 100644 --- a/imas_codex/cli/sn.py +++ b/imas_codex/cli/sn.py @@ -16,9 +16,9 @@ def sn() -> None: """Standard name generation and management. \b - Build: - imas-codex sn build --source dd [--ids NAME] [--domain NAME] - imas-codex sn build --source signals --facility NAME + Mint: + imas-codex sn mint --source dd [--ids NAME] [--domain NAME] + imas-codex sn mint --source signals --facility NAME \b Status: @@ -27,7 +27,7 @@ def sn() -> None: pass -@sn.command("build") +@sn.command("mint") @click.option( "--source", type=click.Choice(["dd", "signals"]), @@ -79,7 +79,7 @@ def sn() -> None: @click.option("--skip-review", is_flag=True, help="Skip the cross-model review phase") @click.option("-v", "--verbose", is_flag=True, help="Enable verbose logging") @click.option("-q", "--quiet", is_flag=True, help="Suppress non-error output") -def sn_build( +def sn_mint( source: str, ids_filter: str | None, domain_filter: str | None, @@ -93,13 +93,13 @@ def sn_build( verbose: bool, quiet: bool, ) -> None: - """Build standard names from a source. + """Mint standard names from a source. \b Examples: - imas-codex sn build --source dd --ids equilibrium --dry-run - imas-codex sn build --source dd --domain magnetics --cost-limit 2 - imas-codex sn build --source signals --facility tcv + imas-codex sn mint --source dd --ids equilibrium --dry-run + imas-codex sn mint --source dd --domain magnetics --cost-limit 2 + imas-codex sn mint --source signals --facility tcv """ # Validate: signals source requires facility if source == "signals" and not facility: @@ -149,7 +149,7 @@ def sn_build( log_print(f" Cost limit: ${cost_limit:.2f}") log_print("") - from imas_codex.sn.pipeline import run_sn_build_engine + from imas_codex.sn.pipeline import run_sn_mint_engine from imas_codex.sn.state import SNBuildState # Build progress display @@ -187,7 +187,7 @@ def sn_build( async def _run(stop_event, service_monitor): if service_monitor: state.service_monitor = service_monitor - await run_sn_build_engine( + await run_sn_mint_engine( state, stop_event=stop_event, on_worker_status=display.on_worker_status if display else None, diff --git a/imas_codex/sn/pipeline.py b/imas_codex/sn/pipeline.py index a0b2e301a..72c59a5d7 100644 --- a/imas_codex/sn/pipeline.py +++ b/imas_codex/sn/pipeline.py @@ -1,4 +1,4 @@ -"""SN build pipeline orchestrator. +"""SN mint pipeline orchestrator. Wires the EXTRACT → COMPOSE → [REVIEW] → VALIDATE → PERSIST workers into the generic discovery engine and runs them with supervision and progress @@ -24,13 +24,13 @@ logger = logging.getLogger(__name__) -async def run_sn_build_engine( +async def run_sn_mint_engine( state: SNBuildState, *, stop_event: asyncio.Event | None = None, on_worker_status: Any | None = None, ) -> None: - """Run the SN build pipeline. + """Run the SN mint pipeline. Pipeline:: diff --git a/plans/features/standard-names/00-implementation-order.md b/plans/features/standard-names/00-implementation-order.md index f043876e7..0445eb3ba 100644 --- a/plans/features/standard-names/00-implementation-order.md +++ b/plans/features/standard-names/00-implementation-order.md @@ -90,7 +90,7 @@ Prerequisite for production minting: ### Wave 5: Integration Testing -- End-to-end: build → publish → review → import +- End-to-end: mint → publish → review → import - Round-trip idempotence verification - Embedding coverage for all StandardName nodes - Documentation updates (AGENTS.md, README) diff --git a/plans/features/standard-names/16-benchmark-parity.md b/plans/features/standard-names/16-benchmark-parity.md index fddf707c0..bfd6170de 100644 --- a/plans/features/standard-names/16-benchmark-parity.md +++ b/plans/features/standard-names/16-benchmark-parity.md @@ -1,4 +1,4 @@ -# 16: Benchmark / Build Prompt Parity & Caching +# 16: Benchmark / Mint Prompt Parity & Cache Verification **Status:** Ready to implement **Depends on:** None (standalone fix) @@ -7,13 +7,13 @@ ## Problem -The `sn benchmark` command constructs LLM prompts differently from `sn build`, +The `sn benchmark` command constructs LLM prompts differently from `sn mint`, making benchmark metrics (cost, speed, quality) unreliable for production model selection. ### Specific gaps -| Aspect | Build (`workers.py`) | Benchmark (`benchmark.py`) | +| Aspect | Mint (`workers.py`) | Benchmark (`benchmark.py`) | |--------|---------------------|---------------------------| | System prompt | `sn/compose_system` via `build_compose_context()` | None | | User prompt | `sn/compose_dd` with `cluster_context` | `sn/compose_dd` without `cluster_context` | @@ -28,6 +28,22 @@ selection. - Quality metrics unreliable (model sees different context) - Benchmark cannot validate whether caching actually works +### Caching architecture + +Prompt caching is **provider-side** — OpenRouter caches repeated prompt prefixes +automatically when `cache_control` blocks are present. Our infrastructure already +supports this: + +- `inject_cache_control()` in `discovery/base/llm.py` adds `cache_control: + {"type": "ephemeral"}` breakpoints to system messages +- `openrouter/` model prefix preserves these blocks through LiteLLM +- OpenRouter returns `usage.cache_creation_input_tokens` and + `usage.cache_read_input_tokens` in responses + +The fix here is to **use the same prompt architecture as mint** so that caching +works naturally, then **confirm** it works by reading cache token counts from +the response. + ## Phase 1: Shared prompt construction **Files:** `imas_codex/sn/benchmark.py` @@ -122,7 +138,7 @@ for model in config.models: ## Acceptance criteria -1. `sn benchmark` uses identical prompt construction as `sn build` +1. `sn benchmark` uses identical prompt construction as `sn mint` 2. System/user message split enables prompt caching 3. `cluster_context` is preserved through extraction 4. Cache hit rate is reported in benchmark output diff --git a/plans/features/standard-names/17-sn-lifecycle-management.md b/plans/features/standard-names/17-sn-lifecycle-management.md index 239cc0775..919571b31 100644 --- a/plans/features/standard-names/17-sn-lifecycle-management.md +++ b/plans/features/standard-names/17-sn-lifecycle-management.md @@ -123,17 +123,17 @@ Safety rules: **Files:** `imas_codex/cli/sn.py` -Add `--reset-to` option to `sn build`: +Add `--reset-to` option to `sn mint`: ```bash # Re-compose all drafted names from scratch (reset + build) -imas-codex sn build --source dd --reset-to drafted +imas-codex sn mint --source dd --reset-to drafted # Full pipeline from extraction (clears and rebuilds) -imas-codex sn build --source dd --reset-to extracted --ids equilibrium +imas-codex sn mint --source dd --reset-to extracted --ids equilibrium ``` -Reset targets for `sn build`: +Reset targets for `sn mint`: - `extracted` — clear all SN nodes for matching source, re-run full pipeline - `drafted` — reset existing drafted names, re-compose them @@ -151,7 +151,7 @@ Reset targets for `sn build`: 1. `sn reset --status drafted` resets nodes and clears embeddings 2. `sn clear --status drafted` deletes only drafted names 3. `sn clear --all` requires `--confirm` flag -4. `sn build --reset-to drafted` resets then rebuilds +4. `sn mint --reset-to drafted` resets then rebuilds 5. Accepted/imported names are never touched without explicit confirmation 6. `sn status` shows correct counts after reset/clear From 256e519cd760d016a24e7f582ed97fba03f1e282 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 16:00:39 +0200 Subject: [PATCH 29/37] fix: relax clean worktree check for RC releases RC releases now warn on dirty worktrees instead of failing, since parallel agents frequently modify files concurrently. Final releases still require a clean worktree. --- AGENTS.md | 2 ++ imas_codex/cli/release.py | 26 ++++++++++++++++++++------ 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index b73b488e4..185e14e79 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -652,6 +652,8 @@ The release CLI is state-machine driven. State is derived from the latest git ta **Remote defaults:** RC releases target `origin` (fork), final releases target `upstream` (iterorganization). Override with `--remote`. +**Dirty worktree policy:** RC releases allow dirty worktrees (warning only) since parallel agents often modify files concurrently. Final releases (`--final`) require a clean worktree — commit or stash first. + ```bash # Check current state and permitted commands uv run imas-codex release status diff --git a/imas_codex/cli/release.py b/imas_codex/cli/release.py index 01a338d49..c6e248e83 100644 --- a/imas_codex/cli/release.py +++ b/imas_codex/cli/release.py @@ -343,16 +343,30 @@ def _check_remote_exists(remote: str) -> None: click.echo(f" ✓ Remote '{remote}' exists") -def _check_clean_tree(dry_run: bool) -> None: +def _check_clean_tree(dry_run: bool, *, strict: bool = True) -> None: result = subprocess.run( ["git", "status", "--porcelain"], capture_output=True, text=True ) if result.stdout.strip(): - msg = "Working tree has uncommitted changes. Commit or stash first." - if dry_run: - click.echo(f" ⚠ {msg}", err=True) + dirty_files = result.stdout.strip().splitlines() + if dry_run or not strict: + # RC releases: warn but continue (parallel agents often dirty the tree) + label = "dry-run" if dry_run else "RC" + click.echo( + f" ⚠ Working tree has {len(dirty_files)} uncommitted change(s) " + f"(allowed for {label})", + err=True, + ) + for f in dirty_files[:5]: + click.echo(f" {f}", err=True) + if len(dirty_files) > 5: + click.echo(f" ... and {len(dirty_files) - 5} more", err=True) else: - raise click.ClickException(msg) + raise click.ClickException( + f"Working tree has {len(dirty_files)} uncommitted change(s). " + "Commit or stash first.\n" + " Hint: RC releases (without --final) allow dirty worktrees." + ) else: click.echo(" ✓ Working tree is clean") @@ -1360,7 +1374,7 @@ def release( click.echo("Pre-flight checks...") _check_on_main() _check_remote_exists(remote) - _check_clean_tree(dry_run) + _check_clean_tree(dry_run, strict=not is_rc) _check_synced(remote, dry_run) if final: _check_final_targets_upstream(remote) From 59631de7510323d963cb10e34394c9e38fbe1cc6 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 16:01:53 +0200 Subject: [PATCH 30/37] docs: add consolidated SN benchmark and lifecycle plan (Plan 19) Consolidate plans 16 (benchmark parity), 17 (lifecycle management), and 18 (calibration) into a single fleet-ready plan with 4 phases and 6 agent dispatches. Incorporates rubber-duck critique: relationship-first deletion, DD-only scope, split calibration from reference expansion, cache smoke test in Phase 1. --- .../standard-names/00-implementation-order.md | 22 +- .../19-benchmark-and-lifecycle.md | 464 ++++++++++++++++++ 2 files changed, 477 insertions(+), 9 deletions(-) create mode 100644 plans/features/standard-names/19-benchmark-and-lifecycle.md diff --git a/plans/features/standard-names/00-implementation-order.md b/plans/features/standard-names/00-implementation-order.md index 0445eb3ba..759c7a12e 100644 --- a/plans/features/standard-names/00-implementation-order.md +++ b/plans/features/standard-names/00-implementation-order.md @@ -46,11 +46,12 @@ All status values use past tense: drafted, published, accepted, rejected, skippe | 11 | rich-compose | Full catalog fields, schema extension, coalesce fix, tests | 📋 Ready | 09 | 12, 13, 14 | | 12 | catalog-import | Feedback import from reviewed catalog PRs | 📋 Ready | 11 P1 | 13 P4 | | 13 | publish-pipeline | Lossless YAML export, batched PRs | 📋 Ready | 11 (all) | 12 (feedback loop) | -| 14 | mcp-tools-benchmark | SN search/fetch/list MCP tools + benchmark quality | ✅ Done | 11 (embedding) | 16 | +| 14 | mcp-tools-benchmark | SN search/fetch/list MCP tools + benchmark quality | ✅ Done | 11 (embedding) | 19 | | 15 | import-physics-domain | Import physics_domain from catalog | ✅ Done | 12 | — | -| 16 | benchmark-parity | Fix benchmark/build prompt parity + caching | 📋 Ready | 14 | 18 | -| 17 | sn-lifecycle-management | SN reset/clear commands | 📋 Ready | — | — | -| 18 | benchmark-calibration | Calibration dataset, reviewer enhancement, model selection | 📋 Ready | 16 | — | +| ~~16~~ | ~~benchmark-parity~~ | ~~Superseded by Plan 19~~ | 🔀 Merged | — | — | +| ~~17~~ | ~~sn-lifecycle-management~~ | ~~Superseded by Plan 19~~ | 🔀 Merged | — | — | +| ~~18~~ | ~~benchmark-calibration~~ | ~~Superseded by Plan 19~~ | 🔀 Merged | — | — | +| 19 | benchmark-and-lifecycle | Benchmark parity, lifecycle mgmt, calibration, model selection | 📋 Ready | 14 | — | ## Deployment Waves @@ -81,12 +82,15 @@ Run in parallel — export and tools are independent: - **Agent A (engineer):** Fix lossy publish export, update graph query, PR workflow, dedup - **Agent B (engineer):** 3 MCP tools (search/fetch/list) + benchmark quality tiers + reviewer -### Wave 4: Benchmark & Lifecycle (Plans 16, 17, 18) +### Wave 4: Benchmark & Lifecycle (Plan 19) -Prerequisite for production minting: -- **Agent A (engineer):** Plan 16 — fix benchmark prompt parity + caching verification -- **Agent B (engineer):** Plan 17 — SN reset/clear commands for iterative development -- **Agent C (architect):** Plan 18 — calibration dataset, reviewer enhancement, model selection +Fleet-ready plan for production minting readiness: +- **Phase 1A (engineer, Sonnet 4.6):** Benchmark prompt parity + cache verification +- **Phase 1B (engineer, Sonnet 4.6):** SN reset/clear commands (parallel with 1A) +- **Phase 2A (engineer, Sonnet 4.6):** Expand gold reference to 50+ entries +- **Phase 2B (architect, Opus 4.6):** Calibration dataset + reviewer enhancement +- **Phase 3 (architect, Opus 4.6):** Cache reporting + model selection runbook +- **Phase 4 (engineer, Sonnet 4.6):** Documentation + superseded plan cleanup ### Wave 5: Integration Testing diff --git a/plans/features/standard-names/19-benchmark-and-lifecycle.md b/plans/features/standard-names/19-benchmark-and-lifecycle.md new file mode 100644 index 000000000..37fe9a84d --- /dev/null +++ b/plans/features/standard-names/19-benchmark-and-lifecycle.md @@ -0,0 +1,464 @@ +# 19: SN Benchmark Parity, Lifecycle Management & Model Selection + +**Status:** Ready to implement +**Supersedes:** Plans 16, 17, 18 +**Scope:** DD source only — signals source parity is future work +**Agent type:** Fleet (4 phases, parallel where possible) + +## Problem Statement + +Three blockers prevent production-quality standard name minting: + +1. **Benchmark/mint prompt parity gap** — `sn benchmark` uses user-only messages + with thin grammar context. `sn mint` uses system/user split with full context + (grammar rules, vocabulary, examples, cluster context). Benchmark results + don't reflect production behavior, prompt caching can't work, and model + selection decisions are unreliable. + +2. **No reset/clear for standard names** — All other discovery domains have + `--reset-to` infrastructure. StandardName is cross-facility (no `facility_id`), + so needs adapted scoping by `review_status`, `source`, and `ids_filter`. + +3. **Weak reviewer and thin reference set** — 30 reference entries from 4 IDSs, + ad-hoc inline reviewer prompt, no calibration examples, no structured scoring. + +## Architecture Notes + +### StandardName node ownership model + +A `StandardName` node can be linked to multiple DD paths or facility signals via +`(IMASNode)-[:HAS_STANDARD_NAME]->(sn)` relationships. `clear` and `reset` +operations must: + +1. Filter/remove **relationships** first +2. Only delete the `StandardName` **node** if it becomes orphaned (no remaining + `HAS_STANDARD_NAME` edges pointing to it) + +This prevents corrupting names that are shared across sources. + +### Prompt caching architecture + +Caching is provider-side (OpenRouter). Our infrastructure already supports it: +- `inject_cache_control()` adds `cache_control: {"type": "ephemeral"}` breakpoints +- `openrouter/` model prefix preserves these blocks through LiteLLM +- OpenRouter returns `usage.cache_creation_input_tokens` and + `usage.cache_read_input_tokens` + +The fix is to use the same prompt architecture as mint, then confirm caching works +by reading cache token counts from the LLM response metadata. + +### Lifecycle states (from current graph reality) + +Only these `review_status` values are persisted: `drafted`, `published`, `accepted`. +The plan does not assume `rejected`, `skipped`, or `imported` exist unless +explicitly added. + +--- + +## Phase 1: Foundation (2 parallel agents) + +### Phase 1A: Benchmark prompt parity — Sonnet 4.6 (engineer) + +**Files:** +- `imas_codex/sn/benchmark.py` — main changes +- `tests/sn/test_benchmark.py` — update tests + +**Changes:** + +1. **Replace `_run_model()` prompt construction** with the mint pipeline's pattern: + ```python + from imas_codex.sn.context import build_compose_context + context = build_compose_context() + system_prompt = render_prompt("sn/compose_system", context) + # ... per batch: + user_context = { + "items": items, + "ids_name": group_key, + "existing_names": sorted(existing)[:200], + "cluster_context": batch.get("context", ""), + **context, + } + user_prompt = render_prompt("sn/compose_dd", user_context) + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt}, + ] + ``` + +2. **Render system prompt once** before the model loop in `run_benchmark()`, + pass it to `_run_model()` as a parameter (matches mint pattern in workers.py:152). + +3. **Fix `_extract_candidates()`** to preserve `batch.context` (cluster_context): + ```python + result.append({ + "group_key": batch.group_key, + "items": batch.items, + "existing_names": list(batch.existing_names), + "context": batch.context, # ADD THIS + }) + ``` + +4. **Remove `build_grammar_context()`** from benchmark.py. It is replaced by + `build_compose_context()` from `context.py`. The `coordinates` key currently + maps to `Component` enum — this bug goes away because `build_compose_context()` + handles it correctly. + +5. **Update `run_benchmark()`** to pass `context` and `system_prompt` through. + +6. **Basic cache smoke test**: After prompt parity changes, run benchmark with + 2+ batches and a single model. Log `cache_creation_input_tokens` and + `cache_read_input_tokens` from the response. Phase 1A is not complete until + cache usage is confirmed in logs. If the response object from + `acall_llm_structured()` does not currently expose cache fields, add logging + of the raw `usage` dict. + +7. **Update tests**: `test_build_grammar_context_keys` in `test_benchmark.py` + imports `build_grammar_context` — update or replace with tests that verify + the benchmark uses `build_compose_context()`. Add test that extraction + preserves `context` field. Add test that `_run_model()` constructs + `[system, user]` messages. + +**Acceptance criteria:** +- `sn benchmark` uses identical prompt construction as `sn mint` +- System/user message split enables prompt caching +- `cluster_context` is preserved through extraction +- Cache usage confirmed in logs (cache_creation or cache_read tokens > 0) +- All existing tests pass, new tests cover parity + +### Phase 1B: Lifecycle management — Sonnet 4.6 (engineer) + +**Files:** +- `imas_codex/sn/graph_ops.py` — add reset/clear functions +- `imas_codex/cli/sn.py` — add `sn reset` and `sn clear` commands +- `tests/sn/test_graph_ops.py` — add tests + +**Changes:** + +1. **Add `reset_standard_names()`** to graph_ops.py: + ```python + def reset_standard_names( + *, + from_status: str = "drafted", + to_status: str | None = None, # None = clear fields only + source_filter: str | None = None, + ids_filter: str | None = None, + dry_run: bool = False, + ) -> int: + ``` + Fields to clear on reset: `embedding`, `embedded_at`, `model`, `generated_at`, + `confidence`. Fields to preserve: `id`, `source`, `source_path`, `created_at`. + Relationships to remove: `HAS_STANDARD_NAME`, `CANONICAL_UNITS`. + +2. **Add `clear_standard_names()`** to graph_ops.py: + ```python + def clear_standard_names( + *, + status_filter: list[str] | None = None, + source_filter: str | None = None, + ids_filter: str | None = None, + include_accepted: bool = False, + dry_run: bool = False, + ) -> int: + ``` + **Safety model (relationship-first):** + - When `ids_filter` or `source_filter` is set, delete matching + `HAS_STANDARD_NAME` relationships first + - Then delete `StandardName` nodes that have zero remaining + `HAS_STANDARD_NAME` edges (orphaned) + - Default: only delete nodes with `review_status IN ['drafted']` + - Accepted names require explicit `--include-accepted` flag (not just + `--confirm` — the flag name must be unambiguous) + - Always log count before deletion + - `dry_run` returns count without deleting + +3. **Add `sn reset` CLI command:** + ```bash + imas-codex sn reset --status drafted + imas-codex sn reset --status published --to drafted + imas-codex sn reset --status drafted --source dd --ids equilibrium + imas-codex sn reset --dry-run + ``` + +4. **Add `sn clear` CLI command:** + ```bash + imas-codex sn clear --status drafted + imas-codex sn clear --status drafted --source dd --ids equilibrium + imas-codex sn clear --all --include-accepted # dangerous + imas-codex sn clear --dry-run + ``` + +5. **Wire `--reset-to` into `sn mint`:** + ```python + @click.option("--reset-to", type=click.Choice(["extracted", "drafted"])) + ``` + - `extracted` = clear matching SN nodes, re-run full pipeline + - `drafted` = reset existing drafted names, re-compose + +6. **Tests:** Unit tests for `reset_standard_names()` and + `clear_standard_names()`. Test that clear refuses accepted without + `include_accepted`. Test relationship-first deletion logic. Test `--reset-to` + triggers reset before pipeline. + +**Acceptance criteria:** +- `sn reset --status drafted` resets nodes and clears embeddings +- `sn clear --status drafted` deletes only drafted names +- `sn clear --all` requires `--include-accepted` for accepted names +- Scoped clear (by IDS/source) removes relationships first, orphans only +- `sn status` shows correct counts after reset/clear +- `sn mint --reset-to drafted` resets then rebuilds + +--- + +## Phase 2: Calibration & Reviewer (2 parallel agents) + +**Depends on:** Phase 1A (benchmark parity must be fixed) + +### Phase 2A: Expand gold reference — Sonnet 4.6 (engineer) + +**Files:** +- `imas_codex/sn/benchmark_reference.py` +- `tests/sn/test_benchmark.py` + +**Changes:** + +1. Expand `REFERENCE_NAMES` from 30 to ~50 entries: + + | IDS | Current | Target | + |-----|---------|--------| + | equilibrium | 18 | 20 | + | core_profiles | 6 | 10 | + | magnetics | 4 | 6 | + | summary | 2 | 4 | + | core_transport | 0 | 4 | + | mhd_linear | 0 | 2 | + | nbi | 0 | 2 | + | edge_profiles | 0 | 2 | + +2. Fix the questionable rogowski_coil reference entry (maps `major_radius` + position to `rogowski_coil` object, which is not physically meaningful). + +3. Source new entries from `imas-standard-names` package's + `resources/standard_name_examples/` where available. Use DD search tools + to find appropriate source paths for each IDS. + +4. All entries must pass round-trip validation at import time (existing + infrastructure enforces this). + +5. Update test that checks reference count. + +**Acceptance criteria:** +- 50+ reference entries across 8+ IDSs +- All entries pass grammar round-trip +- Rogowski_coil entry fixed or removed +- Tests updated and passing + +### Phase 2B: Calibration dataset & reviewer enhancement — Opus 4.6 (architect) + +**Files:** +- `imas_codex/sn/benchmark_calibration.yaml` (new) +- `imas_codex/llm/prompts/sn/review_benchmark.md` (new template) +- `imas_codex/sn/benchmark.py` — update `score_with_reviewer()` +- `imas_codex/sn/benchmark_labels.yaml` — retire (replaced by calibration) +- `tests/sn/test_benchmark.py` + +**Changes:** + +1. **Create calibration dataset** (`benchmark_calibration.yaml`): + ~15 hand-crafted full entries spanning 4 quality tiers. Source outstanding/good + entries from `imas-standard-names` examples. Hand-craft poor examples. + + ```yaml + entries: + - name: electron_temperature + tier: outstanding + expected_score: 95 + description: "Temperature of the electron population." + documentation: > + Electron temperature $T_e$ is a fundamental plasma parameter... + unit: eV + kind: scalar + tags: [core_profiles, equilibrium] + fields: + physical_base: temperature + subject: electron + reason: > + Canonical physics quantity. Rich documentation. Perfect grammar. + ``` + + | Tier | Count | Score Range | + |------|-------|-------------| + | outstanding | 4 | 85-100 | + | good | 4 | 60-79 | + | adequate | 4 | 40-59 | + | poor | 3 | 0-39 | + +2. **Create reviewer prompt template** (`sn/review_benchmark.md`): + Proper Jinja2 template with system/user split. Includes grammar rules, + calibration entries as anchors, and structured rubric. + +3. **Add 5-dimensional scoring** to `QualityReview` model: + ```python + class QualityReview(BaseModel): + name: str + quality_tier: str + score: int = Field(ge=0, le=100) + grammar_score: int = Field(ge=0, le=20) + semantic_score: int = Field(ge=0, le=20) + documentation_score: int = Field(ge=0, le=20) + convention_score: int = Field(ge=0, le=20) + completeness_score: int = Field(ge=0, le=20) + reasoning: str + ``` + +4. **Update `score_with_reviewer()`**: Replace inline rubric string with + template rendering. Use system/user message split (enables caching for + multi-batch reviewer runs). Load calibration entries and pass as template + context. + +5. **Retire `benchmark_labels.yaml`** — replaced by calibration dataset. + Update `load_quality_labels()` or replace with calibration loader. + +**Acceptance criteria:** +- Calibration YAML loads and validates (15 entries, 4 tiers) +- Reviewer uses Jinja2 template with system/user split +- 5-dimensional scores sum to 0-100 +- Calibration entries appear as scoring anchors in reviewer prompt +- Running benchmark with `--reviewer-model` produces dimensional scores +- Tests for calibration loading, template rendering, review model + +--- + +## Phase 3: Cache Reporting & Model Selection Runbook — Opus 4.6 (architect) + +**Depends on:** Phase 2 (needs calibrated benchmark for meaningful results) + +**Files:** +- `imas_codex/sn/benchmark.py` — add cache reporting to ModelResult and table +- `plans/features/standard-names/model-selection-runbook.md` (new) + +**Changes:** + +1. **Add cache hit reporting** to `ModelResult`: + ```python + cache_read_tokens: int = 0 + cache_creation_tokens: int = 0 + ``` + Extract from response metadata in `_run_model()`. Check what `litellm` + exposes — OpenRouter includes `usage.cache_creation_input_tokens` and + `usage.cache_read_input_tokens`. If `acall_llm_structured()` doesn't + currently return these, extend its return or add response metadata logging. + +2. **Add "Cache %" column** to `render_comparison_table()` Rich output. + Calculate as `cache_read / (cache_read + cache_creation) * 100`. + +3. **Create model selection runbook** — a document (in plans/) with: + - Exact CLI commands to run multi-model comparison + - Cost cap per run ($2 maximum per benchmark execution) + - Approved model list with openrouter/ prefixes + - Decision criteria table: + + | Metric | Weight | Threshold | + |--------|--------|-----------| + | Grammar valid % | Critical | ≥95% | + | Fields consistent % | High | ≥85% | + | Reference recall | High | ≥60% | + | Avg quality score | High | ≥65 | + | Cost per name | Medium | <$0.01 | + | Names/min | Medium | >30 | + | Cache hit rate | Low | >50% | + + - Recommended model candidates for compose and review phases + - Instructions for interpreting results + +4. **Run one validation benchmark** (cost-capped at $1) with a single model + to confirm the full stack works end-to-end: prompt parity + caching + + calibrated reviewer + dimensional scoring. + +**Acceptance criteria:** +- Cache hit/miss tokens reported in benchmark output +- Model selection runbook written with exact CLI invocations +- One successful end-to-end benchmark run with dimensional scoring + +--- + +## Phase 4: Documentation — Sonnet 4.6 (engineer) + +**Depends on:** Phases 1-3 + +**Files:** +- `AGENTS.md` — update SN CLI section +- `.github/skills/project-dev/SKILL.md` — add SN testing info +- `.github/skills/service-ops/SKILL.md` — add LLM proxy for SN context +- `.github/agents/engineer.agent.md` — mention SN pipeline +- `agents/README.md` — update if needed +- `plans/features/standard-names/00-implementation-order.md` — update status +- Delete superseded plans: 16, 17, 18 + +**Changes:** + +1. **AGENTS.md updates:** + - Update CLI command table: add `sn reset`, `sn clear`, `sn mint --reset-to` + - Update StandardName lifecycle section with reset/clear semantics + - Add cache verification notes to benchmark section + - Document model selection workflow + +2. **Skill updates** (informative, not prescriptive): + - `project-dev/SKILL.md`: Add SN test commands + (`uv run pytest tests/sn/ -v`), note that SN tests don't require + Neo4j unless marked `@pytest.mark.graph` + - `service-ops/SKILL.md`: Add note that `sn mint` and `sn benchmark` + require the LLM proxy to be running. Document how to check proxy + status and what ports are involved. Mention that prompt caching is + provider-side via OpenRouter. + +3. **Agent updates** (informative, not prescriptive): + - `engineer.agent.md`: Add SN module to list of commonly-modified + areas (imas_codex/sn/, tests/sn/, imas_codex/llm/prompts/sn/) + - `agents/README.md`: Mention SN pipeline if not already covered + +4. **Plan cleanup:** + - Delete `16-benchmark-parity.md`, `17-sn-lifecycle-management.md`, + `18-benchmark-calibration.md` (superseded by this plan) + - Update `00-implementation-order.md` with Plan 19 status + +**Acceptance criteria:** +- All documentation reflects current CLI commands and workflows +- Skills are informative (describe what exists, not what to do) +- Superseded plans deleted +- Implementation order updated + +--- + +## Fleet Dispatch Summary + +| Phase | Agent | Model | Depends On | Parallel With | +|-------|-------|-------|------------|---------------| +| 1A | engineer | Sonnet 4.6 | — | 1B | +| 1B | engineer | Sonnet 4.6 | — | 1A | +| 2A | engineer | Sonnet 4.6 | 1A | 2B | +| 2B | architect | Opus 4.6 | 1A | 2A | +| 3 | architect | Opus 4.6 | 2A, 2B | — | +| 4 | engineer | Sonnet 4.6 | 3 | — | + +**Total: 6 agent dispatches across 4 phases.** + +Phase 1 is fully parallel (2 agents). Phase 2 is fully parallel (2 agents). +Phases 3 and 4 are sequential. + +## Test Plan Summary + +| Test | Phase | File | +|------|-------|------| +| Benchmark uses system/user messages | 1A | test_benchmark.py | +| Extraction preserves context field | 1A | test_benchmark.py | +| Cache tokens appear in logs | 1A | test_benchmark.py or manual | +| `build_grammar_context()` removed cleanly | 1A | test_benchmark.py | +| `reset_standard_names()` clears fields | 1B | test_graph_ops.py | +| `clear_standard_names()` relationship-first | 1B | test_graph_ops.py | +| Clear refuses accepted without flag | 1B | test_graph_ops.py | +| Reference set round-trip (50+ entries) | 2A | test_benchmark.py | +| Calibration YAML loads | 2B | test_benchmark.py | +| Reviewer template renders | 2B | test_benchmark.py | +| QualityReview 5-dimensional model | 2B | test_benchmark.py | +| Cache reporting in output | 3 | test_benchmark.py | +| E2E benchmark with scoring | 3 | manual (cost-capped) | From 5d03485f375cc3bab725aefd21d7315bbe651ed9 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 16:11:40 +0200 Subject: [PATCH 31/37] feat: add sn reset and sn clear lifecycle commands Add reset_standard_names() and clear_standard_names() with relationship-first deletion safety model. Add sn reset, sn clear CLI commands and --reset-to option on sn mint. --- imas_codex/cli/sn.py | 153 ++++++++++++++++++++++ imas_codex/sn/graph_ops.py | 261 +++++++++++++++++++++++++++++++++++++ tests/sn/test_graph_ops.py | 245 ++++++++++++++++++++++++++++++++++ 3 files changed, 659 insertions(+) diff --git a/imas_codex/cli/sn.py b/imas_codex/cli/sn.py index 59c0f4a69..984fdb89d 100644 --- a/imas_codex/cli/sn.py +++ b/imas_codex/cli/sn.py @@ -79,6 +79,16 @@ def sn() -> None: @click.option("--skip-review", is_flag=True, help="Skip the cross-model review phase") @click.option("-v", "--verbose", is_flag=True, help="Enable verbose logging") @click.option("-q", "--quiet", is_flag=True, help="Suppress non-error output") +@click.option( + "--reset-to", + type=click.Choice(["extracted", "drafted"]), + default=None, + help=( + "Reset standard names before minting. " + "'extracted' clears matching SN nodes (full re-run); " + "'drafted' resets existing drafted names (re-compose only)." + ), +) def sn_mint( source: str, ids_filter: str | None, @@ -92,6 +102,7 @@ def sn_mint( skip_review: bool, verbose: bool, quiet: bool, + reset_to: str | None, ) -> None: """Mint standard names from a source. @@ -105,6 +116,27 @@ def sn_mint( if source == "signals" and not facility: raise click.UsageError("--facility is required when --source is signals") + # Handle --reset-to before the main pipeline + if reset_to is not None and not dry_run: + source_arg = "dd" if source == "dd" else "signals" + from imas_codex.sn.graph_ops import clear_standard_names, reset_standard_names + + if reset_to == "extracted": + n = clear_standard_names( + source_filter=source_arg, + ids_filter=ids_filter, + ) + console.print( + f"[yellow]--reset-to extracted:[/yellow] cleared {n} SN nodes" + ) + elif reset_to == "drafted": + n = reset_standard_names( + from_status="drafted", + source_filter=source_arg, + ids_filter=ids_filter, + ) + console.print(f"[yellow]--reset-to drafted:[/yellow] reset {n} SN nodes") + from imas_codex.discovery.base.llm import set_litellm_offline_env set_litellm_offline_env() @@ -777,3 +809,124 @@ def sn_import( console.print(f" - {entry['id']}{units}") if len(result.entries) > 20: console.print(f" ... and {len(result.entries) - 20} more") + + +@sn.command("reset") +@click.option("--status", required=True, help="Reset names with this review_status") +@click.option( + "--to", + "to_status", + default=None, + help="Target review_status after reset (default: clear fields only)", +) +@click.option( + "--source", + type=click.Choice(["dd", "signals"]), + default=None, + help="Filter by source ('dd' or 'signals')", +) +@click.option("--ids", "ids_filter", default=None, help="Filter to specific IDS") +@click.option("--dry-run", is_flag=True, help="Preview without modifying the graph") +def sn_reset( + status: str, + to_status: str | None, + source: str | None, + ids_filter: str | None, + dry_run: bool, +) -> None: + """Reset standard names for re-processing. + + Clears transient fields (embedding, model, confidence, generated_at) and + removes HAS_STANDARD_NAME / CANONICAL_UNITS relationships for matching + nodes, optionally changing their review_status. + + \b + Examples: + imas-codex sn reset --status drafted --dry-run + imas-codex sn reset --status drafted --to extracted --ids equilibrium + imas-codex sn reset --status drafted --source dd + """ + from imas_codex.sn.graph_ops import reset_standard_names + + try: + count = reset_standard_names( + from_status=status, + to_status=to_status, + source_filter=source, + ids_filter=ids_filter, + dry_run=dry_run, + ) + except Exception as e: + console.print(f"[red]Reset error:[/red] {e}") + raise SystemExit(1) from e + + qualifier = "Would reset" if dry_run else "Reset" + to_note = f" → {to_status}" if to_status else " (fields cleared)" + console.print(f"{qualifier} {count} StandardName node(s){to_note}") + + +@sn.command("clear") +@click.option( + "--status", + default=None, + help="Delete names with this review_status (e.g. drafted)", +) +@click.option( + "--all", + "clear_all", + is_flag=True, + help="Delete all standard names (still respects --include-accepted)", +) +@click.option( + "--source", + type=click.Choice(["dd", "signals"]), + default=None, + help="Filter by source ('dd' or 'signals')", +) +@click.option("--ids", "ids_filter", default=None, help="Filter to specific IDS") +@click.option( + "--include-accepted", + is_flag=True, + help="Also delete accepted names (dangerous — use with care)", +) +@click.option("--dry-run", is_flag=True, help="Preview without modifying the graph") +def sn_clear( + status: str | None, + clear_all: bool, + source: str | None, + ids_filter: str | None, + include_accepted: bool, + dry_run: bool, +) -> None: + """Delete standard names from the graph. + + Relationship-first safety model: HAS_STANDARD_NAME edges are removed + before deleting nodes; scoped deletes only remove orphaned nodes. + + \b + Examples: + imas-codex sn clear --status drafted --dry-run + imas-codex sn clear --all --source dd --ids equilibrium --dry-run + imas-codex sn clear --all --include-accepted --dry-run + """ + if not status and not clear_all: + raise click.UsageError("Provide --status or --all to select names.") + + status_filter = None if clear_all else ([status] if status else None) + + from imas_codex.sn.graph_ops import clear_standard_names + + try: + count = clear_standard_names( + status_filter=status_filter, + source_filter=source, + ids_filter=ids_filter, + include_accepted=include_accepted, + dry_run=dry_run, + ) + except Exception as e: + console.print(f"[red]Clear error:[/red] {e}") + raise SystemExit(1) from e + + qualifier = "Would delete" if dry_run else "Deleted" + console.print(f"{qualifier} {count} StandardName node(s)") diff --git a/imas_codex/sn/graph_ops.py b/imas_codex/sn/graph_ops.py index 570cc7c8a..c4084e043 100644 --- a/imas_codex/sn/graph_ops.py +++ b/imas_codex/sn/graph_ops.py @@ -368,6 +368,267 @@ def get_validated_standard_names( return list(results) +def reset_standard_names( + *, + from_status: str = "drafted", + to_status: str | None = None, + source_filter: str | None = None, + ids_filter: str | None = None, + dry_run: bool = False, +) -> int: + """Reset StandardName nodes to allow re-processing. + + Clears transient fields (embedding, embedded_at, model, generated_at, + confidence) and removes HAS_STANDARD_NAME and CANONICAL_UNITS + relationships for matching nodes. + + Parameters + ---------- + from_status: + Only reset nodes with this ``review_status`` (default ``"drafted"``). + to_status: + Target ``review_status`` after reset. ``None`` (default) clears fields + only without changing the status. + source_filter: + Restrict to nodes with ``source`` equal to ``"dd"`` or ``"signals"``. + ids_filter: + Restrict to nodes whose HAS_STANDARD_NAME source path starts with this + IDS name (matched via ``IMASNode -[:HAS_STANDARD_NAME]-> sn``). + dry_run: + Return the count of matching nodes without modifying anything. + + Returns + ------- + Number of nodes reset (or that would be reset in dry-run mode). + """ + with GraphClient() as gc: + params: dict[str, Any] = {"from_status": from_status} + where_clauses = ["sn.review_status = $from_status"] + + if source_filter: + where_clauses.append("coalesce(sn.source, sn.source_type) = $source_filter") + params["source_filter"] = source_filter + + where = " AND ".join(where_clauses) + + if ids_filter: + # Match through HAS_STANDARD_NAME to an IMASNode whose id starts with + # the given IDS name (ids_filter + "/") + params["ids_prefix"] = ids_filter + "/" + count_cypher = f""" + MATCH (src:IMASNode)-[:HAS_STANDARD_NAME]->(sn:StandardName) + WHERE {where} + AND src.id STARTS WITH $ids_prefix + RETURN count(DISTINCT sn) AS n + """ + else: + count_cypher = f""" + MATCH (sn:StandardName) + WHERE {where} + RETURN count(sn) AS n + """ + + result = gc.query(count_cypher, **params) + count = result[0]["n"] if result else 0 + logger.info( + "reset_standard_names: %d nodes match (from_status=%s, source=%s, ids=%s)", + count, + from_status, + source_filter, + ids_filter, + ) + + if dry_run or count == 0: + return count + + if ids_filter: + # Collect matching SN ids first, then operate on them + collect_cypher = f""" + MATCH (src:IMASNode)-[:HAS_STANDARD_NAME]->(sn:StandardName) + WHERE {where} + AND src.id STARTS WITH $ids_prefix + RETURN DISTINCT sn.id AS sn_id + """ + rows = gc.query(collect_cypher, **params) + sn_ids = [r["sn_id"] for r in rows] + reset_params: dict[str, Any] = {"sn_ids": sn_ids} + node_match = "MATCH (sn:StandardName) WHERE sn.id IN $sn_ids" + else: + reset_params = dict(params) + if ids_filter: + reset_params["ids_prefix"] = ids_filter + "/" + node_match = f"MATCH (sn:StandardName) WHERE {where}" + + # Remove HAS_STANDARD_NAME and CANONICAL_UNITS relationships + gc.query( + f""" + {node_match} + OPTIONAL MATCH (src)-[r:HAS_STANDARD_NAME]->(sn) + DELETE r + """, + **reset_params, + ) + gc.query( + f""" + {node_match} + OPTIONAL MATCH (sn)-[r:CANONICAL_UNITS]->(u) + DELETE r + """, + **reset_params, + ) + + # Clear transient fields, optionally set new status + if to_status is not None: + set_clause = ( + "sn.embedding = null, sn.embedded_at = null, sn.model = null, " + "sn.generated_at = null, sn.confidence = null, " + "sn.review_status = $to_status" + ) + reset_params["to_status"] = to_status + else: + set_clause = ( + "sn.embedding = null, sn.embedded_at = null, sn.model = null, " + "sn.generated_at = null, sn.confidence = null" + ) + + gc.query( + f""" + {node_match} + SET {set_clause} + """, + **reset_params, + ) + + logger.info("Reset %d StandardName nodes", count) + return count + + +def clear_standard_names( + *, + status_filter: list[str] | None = None, + source_filter: str | None = None, + ids_filter: str | None = None, + include_accepted: bool = False, + dry_run: bool = False, +) -> int: + """Delete StandardName nodes and their relationships. + + Safety model (relationship-first): + + 1. If ``ids_filter`` or ``source_filter`` is set, delete matching + ``HAS_STANDARD_NAME`` relationships first. + 2. Then delete ``StandardName`` nodes that have zero remaining + ``HAS_STANDARD_NAME`` edges. + + By default only nodes with ``review_status = 'drafted'`` are deleted. + Accepted names require ``include_accepted=True``. + + Parameters + ---------- + status_filter: + List of ``review_status`` values to delete (default ``["drafted"]``). + source_filter: + Restrict to nodes with ``source`` equal to ``"dd"`` or ``"signals"``. + ids_filter: + Delete only names linked to an IMASNode whose id starts with this IDS + name. Relationships are removed first; nodes become orphans and are + then deleted. + include_accepted: + When ``True``, ``"accepted"`` names are eligible for deletion even if + not listed in ``status_filter``. + dry_run: + Return the count of nodes that would be deleted without modifying + anything. + + Returns + ------- + Number of nodes deleted (or that would be deleted in dry-run mode). + """ + if status_filter is None: + status_filter = ["drafted"] + + effective_statuses = list(status_filter) + if include_accepted and "accepted" not in effective_statuses: + effective_statuses.append("accepted") + elif not include_accepted and "accepted" in effective_statuses: + effective_statuses.remove("accepted") + + with GraphClient() as gc: + params: dict[str, Any] = {"statuses": effective_statuses} + sn_where_clauses = ["sn.review_status IN $statuses"] + + if source_filter: + sn_where_clauses.append( + "coalesce(sn.source, sn.source_type) = $source_filter" + ) + params["source_filter"] = source_filter + + sn_where = " AND ".join(sn_where_clauses) + + if ids_filter: + params["ids_prefix"] = ids_filter + "/" + count_cypher = f""" + MATCH (src:IMASNode)-[:HAS_STANDARD_NAME]->(sn:StandardName) + WHERE {sn_where} + AND src.id STARTS WITH $ids_prefix + RETURN count(DISTINCT sn) AS n + """ + else: + count_cypher = f""" + MATCH (sn:StandardName) + WHERE {sn_where} + RETURN count(sn) AS n + """ + + result = gc.query(count_cypher, **params) + count = result[0]["n"] if result else 0 + logger.info( + "clear_standard_names: %d nodes match (statuses=%s, source=%s, ids=%s)", + count, + effective_statuses, + source_filter, + ids_filter, + ) + + if dry_run or count == 0: + return count + + if ids_filter: + # Step 1: remove HAS_STANDARD_NAME relationships for matching scope + gc.query( + f""" + MATCH (src:IMASNode)-[r:HAS_STANDARD_NAME]->(sn:StandardName) + WHERE {sn_where} + AND src.id STARTS WITH $ids_prefix + DELETE r + """, + **params, + ) + # Step 2: delete nodes that are now orphans (no remaining edges) + gc.query( + f""" + MATCH (sn:StandardName) + WHERE {sn_where} + AND NOT EXISTS {{ MATCH ()-[:HAS_STANDARD_NAME]->(sn) }} + DETACH DELETE sn + """, + **params, + ) + else: + # No scoping — detach-delete all matching nodes (removes all rels) + gc.query( + f""" + MATCH (sn:StandardName) + WHERE {sn_where} + DETACH DELETE sn + """, + **params, + ) + + logger.info("Deleted %d StandardName nodes", count) + return count + + def update_review_status(names: list[str], status: str = "published") -> int: """Update review_status for a batch of StandardName nodes. diff --git a/tests/sn/test_graph_ops.py b/tests/sn/test_graph_ops.py index 6504c6498..740d06cf8 100644 --- a/tests/sn/test_graph_ops.py +++ b/tests/sn/test_graph_ops.py @@ -313,3 +313,248 @@ def test_returns_set_of_ids(self) -> None: assert "electron_temperature" in result assert "plasma_current" in result assert len(result) == 2 + + +# ============================================================================= +# TestResetStandardNames +# ============================================================================= + + +class TestResetStandardNames: + """Test reset_standard_names query logic.""" + + def _call_reset(self, mock_gc: MagicMock, **kwargs) -> int: + with patch("imas_codex.sn.graph_ops.GraphClient") as MockGC: + MockGC.return_value.__enter__ = MagicMock(return_value=mock_gc) + MockGC.return_value.__exit__ = MagicMock(return_value=False) + from imas_codex.sn.graph_ops import reset_standard_names + + return reset_standard_names(**kwargs) + + def test_dry_run_returns_count_without_modifying(self) -> None: + """dry_run=True should return count from the count query only.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[{"n": 3}]) + + count = self._call_reset(mock_gc, from_status="drafted", dry_run=True) + + assert count == 3 + # Only one query should be called (the count query) + assert mock_gc.query.call_count == 1 + + def test_returns_zero_for_empty_graph(self) -> None: + """When no nodes match, reset returns 0 and makes no modification queries.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[{"n": 0}]) + + count = self._call_reset(mock_gc, from_status="drafted") + + assert count == 0 + # Only the count query; no DELETE or SET queries + assert mock_gc.query.call_count == 1 + + def test_clears_transient_fields(self) -> None: + """Reset should null out embedding, embedded_at, model, generated_at, confidence.""" + mock_gc = MagicMock() + # First call = count query; subsequent calls = relationship + set queries + mock_gc.query = MagicMock(return_value=[{"n": 2}]) + + self._call_reset(mock_gc, from_status="drafted") + + # Collect all Cypher strings passed + all_cypher = " ".join(call[0][0] for call in mock_gc.query.call_args_list) + + assert "sn.embedding = null" in all_cypher + assert "sn.embedded_at = null" in all_cypher + assert "sn.model = null" in all_cypher + assert "sn.generated_at = null" in all_cypher + assert "sn.confidence = null" in all_cypher + + def test_removes_has_standard_name_relationships(self) -> None: + """Reset should delete HAS_STANDARD_NAME relationships.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[{"n": 1}]) + + self._call_reset(mock_gc, from_status="drafted") + + all_cypher = " ".join(call[0][0] for call in mock_gc.query.call_args_list) + assert "HAS_STANDARD_NAME" in all_cypher + + def test_removes_canonical_units_relationships(self) -> None: + """Reset should delete CANONICAL_UNITS relationships.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[{"n": 1}]) + + self._call_reset(mock_gc, from_status="drafted") + + all_cypher = " ".join(call[0][0] for call in mock_gc.query.call_args_list) + assert "CANONICAL_UNITS" in all_cypher + + def test_to_status_sets_review_status(self) -> None: + """When to_status is given, SET clause should include review_status.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[{"n": 1}]) + + self._call_reset(mock_gc, from_status="drafted", to_status="extracted") + + all_cypher = " ".join(call[0][0] for call in mock_gc.query.call_args_list) + assert "review_status" in all_cypher + + # Verify to_status kwarg was passed + all_kwargs = [call[1] for call in mock_gc.query.call_args_list] + to_statuses = [kw.get("to_status") for kw in all_kwargs if "to_status" in kw] + assert "extracted" in to_statuses + + def test_source_filter_included_in_cypher(self) -> None: + """source_filter should appear in the WHERE clause.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[{"n": 0}]) + + self._call_reset(mock_gc, from_status="drafted", source_filter="dd") + + # Check source_filter param was passed + all_kwargs = [call[1] for call in mock_gc.query.call_args_list] + sources = [ + kw.get("source_filter") for kw in all_kwargs if "source_filter" in kw + ] + assert "dd" in sources + + def test_ids_filter_uses_starts_with(self) -> None: + """ids_filter should restrict via STARTS WITH prefix on src.id.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[{"n": 0}]) + + self._call_reset(mock_gc, from_status="drafted", ids_filter="equilibrium") + + all_cypher = " ".join(call[0][0] for call in mock_gc.query.call_args_list) + assert "STARTS WITH" in all_cypher + + all_kwargs = [call[1] for call in mock_gc.query.call_args_list] + prefixes = [kw.get("ids_prefix") for kw in all_kwargs if "ids_prefix" in kw] + assert "equilibrium/" in prefixes + + +# ============================================================================= +# TestClearStandardNames +# ============================================================================= + + +class TestClearStandardNames: + """Test clear_standard_names deletion logic.""" + + def _call_clear(self, mock_gc: MagicMock, **kwargs) -> int: + with patch("imas_codex.sn.graph_ops.GraphClient") as MockGC: + MockGC.return_value.__enter__ = MagicMock(return_value=mock_gc) + MockGC.return_value.__exit__ = MagicMock(return_value=False) + from imas_codex.sn.graph_ops import clear_standard_names + + return clear_standard_names(**kwargs) + + def test_dry_run_returns_count_without_deleting(self) -> None: + """dry_run=True should return count without issuing DELETE.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[{"n": 5}]) + + count = self._call_clear(mock_gc, dry_run=True) + + assert count == 5 + # Only the count query + assert mock_gc.query.call_count == 1 + all_cypher = " ".join(call[0][0] for call in mock_gc.query.call_args_list) + assert "DELETE" not in all_cypher + + def test_default_status_filter_is_drafted(self) -> None: + """Without status_filter, should only target 'drafted' nodes.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[{"n": 0}]) + + self._call_clear(mock_gc) + + all_kwargs = [call[1] for call in mock_gc.query.call_args_list] + statuses_lists = [kw.get("statuses") for kw in all_kwargs if "statuses" in kw] + assert any("drafted" in sl for sl in statuses_lists) + + def test_accepted_not_deleted_without_flag(self) -> None: + """Without include_accepted, 'accepted' should not be in statuses list.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[{"n": 0}]) + + self._call_clear(mock_gc, status_filter=["drafted"]) + + all_kwargs = [call[1] for call in mock_gc.query.call_args_list] + statuses_lists = [kw.get("statuses") for kw in all_kwargs if "statuses" in kw] + for sl in statuses_lists: + assert "accepted" not in sl, ( + "accepted should not appear without include_accepted" + ) + + def test_include_accepted_adds_to_statuses(self) -> None: + """include_accepted=True should add 'accepted' to effective_statuses.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[{"n": 0}]) + + self._call_clear(mock_gc, status_filter=["drafted"], include_accepted=True) + + all_kwargs = [call[1] for call in mock_gc.query.call_args_list] + statuses_lists = [kw.get("statuses") for kw in all_kwargs if "statuses" in kw] + assert any("accepted" in sl for sl in statuses_lists) + + def test_returns_zero_for_empty_graph(self) -> None: + """When no nodes match, returns 0 and makes no DELETE queries.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[{"n": 0}]) + + count = self._call_clear(mock_gc) + + assert count == 0 + assert mock_gc.query.call_count == 1 + + def test_detach_delete_without_ids_filter(self) -> None: + """Without ids_filter, should DETACH DELETE matching nodes.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[{"n": 2}]) + + self._call_clear(mock_gc) + + all_cypher = " ".join(call[0][0] for call in mock_gc.query.call_args_list) + assert "DETACH DELETE" in all_cypher + + def test_relationship_first_with_ids_filter(self) -> None: + """With ids_filter, should remove relationships before deleting orphan nodes.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[{"n": 3}]) + + self._call_clear(mock_gc, ids_filter="core_profiles") + + all_cypher = " ".join(call[0][0] for call in mock_gc.query.call_args_list) + # Relationship delete should appear (DELETE r pattern) + assert "DELETE r" in all_cypher + # Node delete should also appear + assert "DETACH DELETE sn" in all_cypher + + def test_ids_filter_uses_starts_with(self) -> None: + """ids_filter should use STARTS WITH prefix matching.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[{"n": 0}]) + + self._call_clear(mock_gc, ids_filter="magnetics") + + all_cypher = " ".join(call[0][0] for call in mock_gc.query.call_args_list) + assert "STARTS WITH" in all_cypher + + all_kwargs = [call[1] for call in mock_gc.query.call_args_list] + prefixes = [kw.get("ids_prefix") for kw in all_kwargs if "ids_prefix" in kw] + assert "magnetics/" in prefixes + + def test_source_filter_passed_as_param(self) -> None: + """source_filter should be passed as a query parameter.""" + mock_gc = MagicMock() + mock_gc.query = MagicMock(return_value=[{"n": 0}]) + + self._call_clear(mock_gc, source_filter="signals") + + all_kwargs = [call[1] for call in mock_gc.query.call_args_list] + sources = [ + kw.get("source_filter") for kw in all_kwargs if "source_filter" in kw + ] + assert "signals" in sources From 4d2eedd6c84dc1d3c04ee13171f1c6a1f182b04a Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 16:12:34 +0200 Subject: [PATCH 32/37] feat: fix benchmark prompt parity with mint pipeline Replace build_grammar_context() with build_compose_context() for rich grammar context. Add system/user message split for prompt caching. Preserve cluster_context through extraction. --- imas_codex/sn/benchmark.py | 65 ++++++++--------- tests/sn/test_benchmark.py | 138 ++++++++++++++++++++++++++++++++----- 2 files changed, 148 insertions(+), 55 deletions(-) diff --git a/imas_codex/sn/benchmark.py b/imas_codex/sn/benchmark.py index 1475f2b16..fe6ed452f 100644 --- a/imas_codex/sn/benchmark.py +++ b/imas_codex/sn/benchmark.py @@ -113,31 +113,6 @@ def from_json(cls, data: str) -> BenchmarkReport: ) -# --------------------------------------------------------------------------- -# Grammar context builder -# --------------------------------------------------------------------------- - - -def build_grammar_context() -> dict[str, list[str]]: - """Build the grammar enum values needed by the compose prompt. - - Returns a dict with keys matching the template variables in - ``sn/compose_dd.md``: subjects, positions, components, coordinates, - processes, transformations, geometric_bases, objects, binary_operators. - """ - return { - "subjects": [e.value for e in Subject], - "positions": [e.value for e in Position], - "components": [e.value for e in Component], - "coordinates": [e.value for e in Component], # same enum - "processes": [e.value for e in Process], - "transformations": [e.value for e in Transformation], - "geometric_bases": [e.value for e in GeometricBase], - "objects": [e.value for e in Object], - "binary_operators": [e.value for e in BinaryOperator], - } - - # --------------------------------------------------------------------------- # Grammar validation # --------------------------------------------------------------------------- @@ -283,7 +258,6 @@ async def score_with_reviewer( candidates: list[dict], reviewer_model: str, quality_labels: dict[str, list[str]], - grammar_context: dict[str, list[str]], ) -> list[dict]: """Score candidates using a reviewer model. @@ -406,7 +380,12 @@ async def run_benchmark( # --- 2. Run each model --- results: list[ModelResult] = [] - grammar_ctx = build_grammar_context() + from imas_codex.llm.prompt_loader import render_prompt + from imas_codex.sn.context import build_compose_context + + context = build_compose_context() + system_prompt = render_prompt("sn/compose_system", context) + for model in config.models: logger.info("Benchmarking model: %s", model) model_result = await _run_model( @@ -414,6 +393,8 @@ async def run_benchmark( extraction_batches=extraction_batches, config=config, reference=REFERENCE_NAMES, + system_prompt=system_prompt, + context=context, ) results.append(model_result) @@ -426,7 +407,6 @@ async def run_benchmark( result.candidates, config.reviewer_model, quality_labels, - grammar_ctx, ) result.quality_scores = reviews # Compute distribution @@ -478,7 +458,7 @@ async def run_benchmark( def _extract_candidates(config: BenchmarkConfig) -> list[dict]: """Extract candidates from the graph DB. - Returns list of batch dicts with keys: group_key, items, existing_names. + Returns list of batch dicts with keys: group_key, items, existing_names, context. """ from imas_codex.sn.sources.dd import extract_dd_candidates @@ -496,6 +476,7 @@ def _extract_candidates(config: BenchmarkConfig) -> list[dict]: "group_key": batch.group_key, "items": batch.items, "existing_names": list(batch.existing_names), + "context": batch.context, } ) return result @@ -522,12 +503,13 @@ async def _run_model( extraction_batches: list[dict], config: BenchmarkConfig, reference: dict[str, dict], + system_prompt: str, + context: dict[str, Any], ) -> ModelResult: """Run a single model across all extraction batches.""" from imas_codex.discovery.base.llm import acall_llm_structured from imas_codex.llm.prompt_loader import render_prompt - grammar_ctx = build_grammar_context() result = ModelResult(model=model) all_candidates: list[dict] = [] @@ -542,22 +524,27 @@ async def _run_model( group_key = batch.get("group_key", "unknown") existing = set(batch.get("existing_names", [])) - # Build prompt context - prompt_context = { + # Build user prompt context — mirrors workers.py pattern + user_context = { "items": items, "ids_name": group_key, - "existing_names": list(existing), - **grammar_ctx, + "existing_names": sorted(existing)[:200], + "cluster_context": batch.get("context", ""), } try: - prompt_text = render_prompt("sn/compose_dd", prompt_context) + user_prompt = render_prompt( + "sn/compose_dd", {**context, **user_context} + ) except Exception: logger.warning("Failed to render prompt for batch %s", group_key) result.batch_errors += 1 continue - messages = [{"role": "user", "content": prompt_text}] + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt}, + ] try: llm_result, cost, tokens = await acall_llm_structured( @@ -568,6 +555,12 @@ async def _run_model( ) result.total_cost += cost result.total_tokens += tokens + logger.debug( + "Batch %s: cost=%.4f tokens=%d", + group_key, + cost, + tokens, + ) # Collect candidates for c in llm_result.candidates: diff --git a/tests/sn/test_benchmark.py b/tests/sn/test_benchmark.py index 8b906f24f..d66866d21 100644 --- a/tests/sn/test_benchmark.py +++ b/tests/sn/test_benchmark.py @@ -136,40 +136,140 @@ def test_benchmark_report_instantiation(self): # ----------------------------------------------------------------------- -# Grammar context builder tests +# Context builder tests # ----------------------------------------------------------------------- class TestGrammarContext: - """Verify grammar context builder provides all template variables.""" + """Verify build_compose_context provides all template variables.""" + + def test_build_compose_context_keys(self): + """build_compose_context() should return rich grammar context.""" + from imas_codex.sn.context import build_compose_context + + ctx = build_compose_context() + # Rich grammar context keys + assert "canonical_pattern" in ctx + assert "vocabulary_sections" in ctx + assert "segment_descriptions" in ctx + # Backward-compat enum lists still present + assert "subjects" in ctx + assert "positions" in ctx + assert "components" in ctx - def test_build_grammar_context_keys(self): - from imas_codex.sn.benchmark import build_grammar_context + def test_all_values_non_empty(self): + """Enum lists from build_compose_context should be non-empty strings.""" + from imas_codex.sn.context import build_compose_context - ctx = build_grammar_context() - expected_keys = { + ctx = build_compose_context() + for key in ( "subjects", "positions", "components", - "coordinates", "processes", "transformations", - "geometric_bases", - "objects", - "binary_operators", - } - assert set(ctx.keys()) == expected_keys + ): + assert len(ctx[key]) > 0, f"{key} should have at least one value" + assert all(isinstance(v, str) for v in ctx[key]), ( + f"{key} values must be strings" + ) - def test_all_values_non_empty(self): - from imas_codex.sn.benchmark import build_grammar_context - ctx = build_grammar_context() - for key, values in ctx.items(): - assert len(values) > 0, f"{key} should have at least one value" - assert all(isinstance(v, str) for v in values), ( - f"{key} values must be strings" +# ----------------------------------------------------------------------- +# Prompt parity tests +# ----------------------------------------------------------------------- + + +class TestPromptParity: + """Verify benchmark uses the same prompt architecture as mint pipeline.""" + + def test_extract_candidates_preserves_context(self): + """_extract_candidates should include batch.context in output dicts.""" + from unittest.mock import patch + + from imas_codex.sn.benchmark import BenchmarkConfig, _extract_candidates + from imas_codex.sn.sources.base import ExtractionBatch + + fake_batch = ExtractionBatch( + source="dd", + group_key="equilibrium", + items=[{"path": "test/path", "description": "Test"}], + context="IDS: equilibrium\nSemantic clusters: psi, safety_factor", + existing_names=set(), + ) + + config = BenchmarkConfig(models=["test"]) + + with patch( + "imas_codex.sn.sources.dd.extract_dd_candidates", + return_value=[fake_batch], + ): + batches = _extract_candidates(config) + + assert len(batches) == 1 + assert "context" in batches[0] + assert ( + batches[0]["context"] + == "IDS: equilibrium\nSemantic clusters: psi, safety_factor" + ) + + @pytest.mark.asyncio + async def test_run_model_system_user_messages(self): + """_run_model should construct [system, user] message structure.""" + from unittest.mock import AsyncMock, patch + + from imas_codex.sn.benchmark import BenchmarkConfig, _run_model + from imas_codex.sn.models import SNComposeBatch + + config = BenchmarkConfig(models=["test"], temperature=0.0) + batches = [ + { + "group_key": "test", + "items": [ + { + "path": "test/path", + "description": "Test", + "units": None, + "data_type": "FLT_0D", + "cluster_label": None, + } + ], + "existing_names": [], + "context": "IDS: test", + } + ] + minimal_context: dict = {"subjects": ["electron"], "vocabulary_sections": []} + captured_messages: list[dict] = [] + mock_response = SNComposeBatch(candidates=[], skipped=[]) + + async def mock_llm(model, messages, response_model, **kwargs): + captured_messages.extend(messages) + return mock_response, 0.0, 0 + + with ( + patch( + "imas_codex.discovery.base.llm.acall_llm_structured", + side_effect=mock_llm, + ), + patch( + "imas_codex.llm.prompt_loader.render_prompt", + return_value="rendered prompt", + ), + ): + await _run_model( + model="test", + extraction_batches=batches, + config=config, + reference={}, + system_prompt="System instructions", + context=minimal_context, ) + assert len(captured_messages) == 2 + assert captured_messages[0]["role"] == "system" + assert captured_messages[0]["content"] == "System instructions" + assert captured_messages[1]["role"] == "user" + # ----------------------------------------------------------------------- # Validation tests From 8499610bee310a223b926525bd470f387a92b5f6 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 16:19:22 +0200 Subject: [PATCH 33/37] fix: remove empty env blocks breaking CI workflows Empty 'env:' with no values causes GitHub Actions to reject the workflow YAML before any jobs start (0 jobs, immediate failure). --- .github/workflows/benchmark.yml | 2 -- .github/workflows/release.yml | 2 -- .github/workflows/test.yml | 2 -- 3 files changed, 6 deletions(-) diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index e4b26995b..a3a5a96a2 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -1,7 +1,5 @@ name: Benchmark -env: - on: push: tags: ["v*"] diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 8c262ff82..923f12eea 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -1,7 +1,5 @@ name: Release -env: - on: push: tags: diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 8d1884da8..5b73ac0bf 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -1,7 +1,5 @@ name: Test -env: - on: push: branches: [main] From d0960bdd324c8547e0ffb31668cea85dfad23fd4 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 16:22:55 +0200 Subject: [PATCH 34/37] feat: expand benchmark reference set to 52 entries across 8 IDSs Add entries for core_transport (heat/particle flux for electron and ion), mhd_linear (growth_rate, mhd_frequency), nbi (power, energy of NBI), and edge_profiles (electron temperature and density at edge region). Expand core_profiles from 8 to 12 entries: parallel electric field, bootstrap and ohmic current density, ion toroidal velocity. Expand magnetics from 2 to 6 entries: flux_loop poloidal flux, rogowski_coil plasma current, total plasma current, diamagnetic flux. Expand summary from 2 to 4 entries: toroidal beta, energy confinement time. Add 2 more equilibrium entries: profiles_1d psi, magnetic_axis vertical position. Fix physically incorrect rogowski_coil entry on magnetic_axis/r: replaced geometric_base/object combo with major_radius at MAGNETIC_AXIS position. --- imas_codex/sn/benchmark_reference.py | 97 +++++++++++++++++++++++++++- 1 file changed, 94 insertions(+), 3 deletions(-) diff --git a/imas_codex/sn/benchmark_reference.py b/imas_codex/sn/benchmark_reference.py index 828eefccb..b0b45f3f1 100644 --- a/imas_codex/sn/benchmark_reference.py +++ b/imas_codex/sn/benchmark_reference.py @@ -14,9 +14,9 @@ from imas_standard_names.grammar import ( Component, - GeometricBase, Object, Position, + Process, StandardName, Subject, compose_standard_name, @@ -101,9 +101,45 @@ def _ref(fields: dict) -> dict: "magnetics/b_field_tor_probe/field/data": _ref( {"physical_base": "magnetic_field", "component": Component.TOROIDAL} ), + # --- Additional magnetics entries --- + "magnetics/flux_loop/flux/data": _ref( + {"physical_base": "poloidal_magnetic_flux", "object": Object.FLUX_LOOP} + ), + "magnetics/rogowski_coil/current/data": _ref( + {"physical_base": "plasma_current", "object": Object.ROGOWSKI_COIL} + ), + "magnetics/ip/data": _ref({"physical_base": "plasma_current"}), + "magnetics/diamagnetic_flux/data": _ref( + {"physical_base": "poloidal_magnetic_flux", "object": Object.DIAMAGNETIC_LOOP} + ), "core_profiles/profiles_1d/rotation_frequency_tor_sonic": _ref( {"physical_base": "rotation_frequency", "component": Component.TOROIDAL} ), + # --- Additional core_profiles entries --- + "core_profiles/profiles_1d/e_field/parallel": _ref( + {"physical_base": "electric_field", "component": Component.PARALLEL} + ), + "core_profiles/profiles_1d/j_bootstrap": _ref( + { + "physical_base": "current_density", + "component": Component.PARALLEL, + "process": Process.BOOTSTRAP, + } + ), + "core_profiles/profiles_1d/j_ohmic": _ref( + { + "physical_base": "current_density", + "component": Component.PARALLEL, + "process": Process.OHMIC, + } + ), + "core_profiles/profiles_1d/ion/velocity/toroidal": _ref( + { + "physical_base": "velocity", + "subject": Subject.ION, + "component": Component.TOROIDAL, + } + ), # --- Position-qualified quantities --- "core_profiles/profiles_1d/electrons/temperature_fit/boundary_condition/value": _ref( { @@ -114,8 +150,16 @@ def _ref(fields: dict) -> dict: ), "equilibrium/time_slice/global_quantities/magnetic_axis/r": _ref( { - "geometric_base": GeometricBase.POSITION, - "object": Object.ROGOWSKI_COIL, + "physical_base": "major_radius", + "position": Position.MAGNETIC_AXIS, + } + ), + "equilibrium/time_slice/profiles_1d/psi": _ref( + {"physical_base": "poloidal_magnetic_flux"} + ), + "equilibrium/time_slice/global_quantities/magnetic_axis/z": _ref( + { + "physical_base": "vertical_position", "position": Position.MAGNETIC_AXIS, } ), @@ -139,6 +183,11 @@ def _ref(fields: dict) -> dict: "summary/global_quantities/li/value": _ref( {"physical_base": "internal_inductance"} ), + # --- Additional summary entries --- + "summary/global_quantities/beta_tor/value": _ref({"physical_base": "beta"}), + "summary/global_quantities/tau_energy/value": _ref( + {"physical_base": "confinement_time"} + ), "equilibrium/time_slice/global_quantities/resistivity": _ref( {"physical_base": "resistivity"} ), @@ -155,6 +204,48 @@ def _ref(fields: dict) -> dict: "equilibrium/time_slice/global_quantities/aspect_ratio": _ref( {"physical_base": "aspect_ratio"} ), + # --- core_transport --- + "core_transport/model/profiles_1d/electrons/energy/flux": _ref( + {"physical_base": "heat_flux", "subject": Subject.ELECTRON} + ), + "core_transport/model/profiles_1d/electrons/particles/flux": _ref( + {"physical_base": "particle_flux", "subject": Subject.ELECTRON} + ), + "core_transport/model/profiles_1d/ion/energy/flux": _ref( + {"physical_base": "heat_flux", "subject": Subject.ION} + ), + "core_transport/model/profiles_1d/ion/particles/flux": _ref( + {"physical_base": "particle_flux", "subject": Subject.ION} + ), + # --- mhd_linear --- + "mhd_linear/time_slice/toroidal_mode/growthrate": _ref( + {"physical_base": "growth_rate"} + ), + "mhd_linear/time_slice/toroidal_mode/frequency": _ref( + {"physical_base": "mhd_frequency"} + ), + # --- nbi --- + "nbi/unit/power_launched/data": _ref( + {"physical_base": "power", "object": Object.NEUTRAL_BEAM_INJECTOR} + ), + "nbi/unit/energy/data": _ref( + {"physical_base": "energy", "object": Object.NEUTRAL_BEAM_INJECTOR} + ), + # --- edge_profiles --- + "edge_profiles/profiles_1d/electrons/temperature": _ref( + { + "physical_base": "temperature", + "subject": Subject.ELECTRON, + "position": Position.EDGE_REGION, + } + ), + "edge_profiles/profiles_1d/electrons/density": _ref( + { + "physical_base": "density", + "subject": Subject.ELECTRON, + "position": Position.EDGE_REGION, + } + ), } """Map of DD source_path → {name: str, fields: dict}. From 0065e836d6d446987ce98212fa5bffeac106f8e7 Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 16:23:52 +0200 Subject: [PATCH 35/37] feat: add calibration dataset and enhanced reviewer for benchmark Create benchmark_calibration.yaml with 15 entries across 4 quality tiers. Replace inline reviewer rubric with Jinja2 template (sn/review_benchmark). Add 5-dimensional scoring: grammar, semantic, docs, convention, completeness. Retire benchmark_labels.yaml in favor of calibration dataset. --- imas_codex/llm/prompts/sn/review_benchmark.md | 118 +++++++ imas_codex/sn/benchmark.py | 118 +++---- imas_codex/sn/benchmark_calibration.yaml | 302 ++++++++++++++++++ imas_codex/sn/benchmark_labels.yaml | 29 -- tests/sn/test_benchmark.py | 283 ++++++++++++++-- 5 files changed, 745 insertions(+), 105 deletions(-) create mode 100644 imas_codex/llm/prompts/sn/review_benchmark.md create mode 100644 imas_codex/sn/benchmark_calibration.yaml delete mode 100644 imas_codex/sn/benchmark_labels.yaml diff --git a/imas_codex/llm/prompts/sn/review_benchmark.md b/imas_codex/llm/prompts/sn/review_benchmark.md new file mode 100644 index 000000000..74cfcd237 --- /dev/null +++ b/imas_codex/llm/prompts/sn/review_benchmark.md @@ -0,0 +1,118 @@ +--- +name: sn/review_benchmark +description: Quality scoring for benchmark standard name entries +used_by: imas_codex.sn.benchmark.score_with_reviewer +task: review +dynamic: true +schema_needs: [] +--- + +You are a quality reviewer for IMAS standard name entries in fusion plasma physics. Your task is to evaluate each candidate entry across five quality dimensions and assign a total score. + +## Standard Name Grammar + +A valid standard name is composed from optional segments in a specific order: + +**Canonical pattern:** `[process] [transformation] [subject] [component] physical_base [position] [object]` + +Or with geometric_base: `[process] [transformation] [subject] [component] geometric_base [position] [object]` + +Every name MUST have either a `physical_base` (open vocabulary) or a `geometric_base` (restricted vocabulary), but never both. + +### Segment Vocabulary + +- **subject**: species or population (electron, ion, deuterium, tritium, helium, impurity_species, fast_ion, neutral, runaway_electron) +- **component**: vector/tensor component (radial, toroidal, vertical, poloidal, parallel, diamagnetic, normal, tangential, binormal, x, y, z) +- **position**: spatial location (magnetic_axis, plasma_boundary, midplane, core_region, edge_region, scrape_off_layer, last_closed_flux_surface, ...) +- **process**: physical mechanism (conduction, convection, diffusion, neoclassical, turbulent, ohmic, bootstrap, radiation, ...) +- **transformation**: mathematical operation (square_of, change_over_time_in, logarithm_of, inverse_of) +- **geometric_base**: geometric quantity (position, vertex, centroid, outline, contour, displacement, offset, trajectory, extent, ...) +- **object**: device component (flux_loop, poloidal_magnetic_field_probe, bolometer, langmuir_probe, ...) + +## Scoring Dimensions + +Rate each dimension from 0 to 20. The total score is the sum (0-100). + +### 1. Grammar Correctness (0-20) +- Does the name parse correctly under the standard name grammar? +- Are all segments valid enum values from the vocabulary? +- Is the field decomposition consistent with the composed name? +- Does the name round-trip: `parse(name) → compose() == name`? + +**20**: Perfect parse, valid segments, consistent decomposition. +**10**: Parses correctly but uses unusual segment combinations. +**0**: Would fail grammar validation or uses invalid tokens. + +### 2. Semantic Accuracy (0-20) +- Does the name correctly describe the physics quantity? +- Is the physical_base appropriate for what is being measured? +- Are qualifier segments (subject, position, component) correctly applied? + +**20**: Name unambiguously identifies the quantity; domain expert would agree. +**10**: Name is defensible but there may be a more precise choice. +**0**: Name is misleading or describes a different quantity. + +### 3. Documentation Quality (0-20) +- Does the documentation include LaTeX mathematical notation? +- Are typical value ranges provided? +- Is measurement/diagnostic context mentioned? +- Are cross-references to related quantities included? +- Is the documentation substantive (not just rephrasing the name)? + +**20**: Rich docs with LaTeX, value ranges, measurement context, cross-refs. +**10**: Adequate docs — correct but thin, missing some elements. +**0**: Empty or circular documentation (just restates the name). + +### 4. Naming Conventions (0-20) +- Does the name follow established patterns for similar quantities? +- Is the name concise but unambiguous? +- Does it avoid overly generic terms (data, signal, value)? +- Is it specific enough to be useful as a standard identifier? + +**20**: Follows best practices, concise, unambiguous, specific. +**10**: Acceptable but could be improved — slightly verbose or generic. +**0**: Vague, generic, or violates naming conventions. + +### 5. Entry Completeness (0-20) +- Is the unit correct for this quantity (or null if dimensionless)? +- Is the kind (scalar/vector/metadata) appropriate? +- Are relevant tags assigned from the controlled vocabulary? +- Are grammar fields properly populated? + +**20**: All metadata fields correct and complete. +**10**: Most fields present but some missing or questionable. +**0**: Missing critical fields (wrong unit, no tags, wrong kind). + +## Quality Tiers + +Map the total score to a tier: +- **outstanding** (85-100): Exemplary entry ready for publication +- **good** (60-84): Solid entry with minor improvements possible +- **adequate** (40-59): Acceptable but needs enrichment +- **poor** (0-39): Needs fundamental rework + +## Calibration Examples + +Use these scored examples to anchor your judgments: + +{% for entry in calibration_entries %} +### {{ entry.name }} — {{ entry.tier }} ({{ entry.expected_score }}/100) +{{ entry.reason }} +{% endfor %} + +## Candidates to Review + +{% for candidate in candidates %} +### Candidate {{ loop.index }}: {{ candidate.standard_name }} +- **Description:** {{ candidate.description | default('N/A', true) }} +- **Documentation:** {{ candidate.documentation | default('N/A', true) }} +- **Unit:** {{ candidate.unit | default('N/A', true) }} +- **Kind:** {{ candidate.kind | default('N/A', true) }} +- **Tags:** {{ candidate.tags | default([], true) | join(', ') }} +- **Fields:** {{ candidate.fields | default({}, true) }} + +{% endfor %} + +## Output Format + +Return a JSON object with a `reviews` array. Each review MUST include all five dimension scores that sum to the total score. Provide brief but specific reasoning. diff --git a/imas_codex/sn/benchmark.py b/imas_codex/sn/benchmark.py index fe6ed452f..6685adc02 100644 --- a/imas_codex/sn/benchmark.py +++ b/imas_codex/sn/benchmark.py @@ -239,89 +239,99 @@ def compare_to_reference( # --------------------------------------------------------------------------- -def load_quality_labels() -> dict[str, list[str]]: - """Load quality tier labels from benchmark_labels.yaml. +def load_calibration_entries() -> list[dict]: + """Load calibration entries from benchmark_calibration.yaml. - Returns a dict mapping tier name → list of standard name IDs. - Returns empty dict if file not found. + Returns a list of dicts, each with: name, tier, expected_score, + description, documentation, unit, kind, tags, fields, reason. + Returns empty list if file not found. """ import yaml - labels_path = Path(__file__).parent / "benchmark_labels.yaml" - if labels_path.exists(): - with open(labels_path) as f: - return yaml.safe_load(f) or {} - return {} + cal_path = Path(__file__).parent / "benchmark_calibration.yaml" + if cal_path.exists(): + with open(cal_path) as f: + data = yaml.safe_load(f) or {} + return data.get("entries", []) + return [] async def score_with_reviewer( candidates: list[dict], reviewer_model: str, - quality_labels: dict[str, list[str]], + calibration_entries: list[dict], ) -> list[dict]: - """Score candidates using a reviewer model. + """Score candidates using a reviewer model with 5-dimensional scoring. - Returns list of dicts with: name, quality_tier, score, reasoning. + Each candidate is scored across five dimensions (0-20 each): + grammar, semantic, documentation, convention, completeness. + Total score is the sum (0-100). + + Returns list of dicts with: name, quality_tier, score, + grammar_score, semantic_score, documentation_score, + convention_score, completeness_score, reasoning. """ from pydantic import BaseModel, Field from imas_codex.discovery.base.llm import acall_llm_structured + from imas_codex.llm.prompt_loader import render_prompt class QualityReview(BaseModel): name: str quality_tier: str = Field(description="outstanding, good, adequate, or poor") - score: int = Field(ge=0, le=100, description="Quality score 0-100") + score: int = Field( + ge=0, le=100, description="Total quality score (sum of dimensions)" + ) + grammar_score: int = Field(ge=0, le=20, description="Grammar correctness") + semantic_score: int = Field(ge=0, le=20, description="Semantic accuracy") + documentation_score: int = Field( + ge=0, le=20, description="Documentation quality" + ) + convention_score: int = Field(ge=0, le=20, description="Naming conventions") + completeness_score: int = Field(ge=0, le=20, description="Entry completeness") reasoning: str class QualityReviewBatch(BaseModel): reviews: list[QualityReview] - # Build prompt with labeled examples and rubric - labeled_examples = [] - for tier, names in quality_labels.items(): - for name in names: - labeled_examples.append(f" - {name}: {tier}") - - rubric = """ -Quality rubric: -- Grammar correctness: Does the name follow standard name grammar rules? -- Semantic accuracy: Does the name correctly describe the physics quantity? -- Documentation quality: Is the documentation clear, with LaTeX where appropriate? -- Naming conventions: Does the name follow established patterns? -- Unit consistency: Is the unit correct for this quantity? - -Quality tiers: -- outstanding (80-100): Rich docs, correct grammar, cross-linked, LaTeX -- good (60-79): Correct grammar, adequate documentation -- adequate (40-59): Correct grammar, thin documentation -- poor (0-39): Grammar valid but naming debatable or docs minimal -""" - - prompt = f"""You are reviewing standard name entries for quality. - -{rubric} - -Labeled examples (scoring anchors): -{chr(10).join(labeled_examples)} - -Review these candidates and assign a quality tier and score: -""" + # Render system prompt with calibration entries (cached across batches) + system_prompt = render_prompt( + "sn/review_benchmark", + {"calibration_entries": calibration_entries, "candidates": []}, + ) # Process in batches of 10 all_reviews: list[dict] = [] for i in range(0, len(candidates), 10): batch = candidates[i : i + 10] - batch_text = [] + + # Build per-batch user prompt with candidate details + batch_items = [] for c in batch: - batch_text.append(f"Name: {c.get('standard_name', '')}") - batch_text.append(f" Description: {c.get('description', '')}") - doc = c.get("documentation", "") or "" - batch_text.append(f" Documentation: {doc[:200]}") - batch_text.append(f" Unit: {c.get('unit', 'N/A')}") - batch_text.append(f" Fields: {c.get('fields', {})}") - batch_text.append("") + batch_items.append( + { + "standard_name": c.get("standard_name", ""), + "description": c.get("description", ""), + "documentation": (c.get("documentation", "") or "")[:500], + "unit": c.get("unit", "N/A"), + "kind": c.get("kind", "N/A"), + "tags": c.get("tags", []), + "fields": c.get("fields", {}), + } + ) - messages = [{"role": "user", "content": prompt + "\n".join(batch_text)}] + user_prompt = render_prompt( + "sn/review_benchmark", + { + "calibration_entries": calibration_entries, + "candidates": batch_items, + }, + ) + + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt}, + ] try: result, _, _ = await acall_llm_structured( @@ -400,13 +410,13 @@ async def run_benchmark( # --- 2b. Reviewer scoring (optional) --- if config.reviewer_model: - quality_labels = load_quality_labels() + calibration_entries = load_calibration_entries() for result in results: if result.candidates: reviews = await score_with_reviewer( result.candidates, config.reviewer_model, - quality_labels, + calibration_entries, ) result.quality_scores = reviews # Compute distribution diff --git a/imas_codex/sn/benchmark_calibration.yaml b/imas_codex/sn/benchmark_calibration.yaml new file mode 100644 index 000000000..8f39a90b9 --- /dev/null +++ b/imas_codex/sn/benchmark_calibration.yaml @@ -0,0 +1,302 @@ +# Calibration dataset for benchmark reviewer scoring. +# Each entry is a hand-crafted full standard name entry spanning quality tiers. +# Used as scoring anchors by the reviewer model — they define what +# "outstanding" vs "poor" looks like. +# +# Every entry's grammar fields have been validated: +# compose_standard_name(StandardName(**fields)) == name + +entries: + # =================================================================== + # OUTSTANDING tier (85-100): Rich docs, LaTeX, cross-refs, perfect grammar + # =================================================================== + + - name: electron_temperature + tier: outstanding + expected_score: 95 + description: "Temperature of the electron population." + documentation: > + Electron temperature $T_e$ is a fundamental kinetic quantity representing + the thermal energy of the plasma electron population. Measured primarily + by Thomson scattering and electron cyclotron emission (ECE) diagnostics. + Typical values range from ~100 eV at the edge to 1-20 keV in the core + depending on heating power and confinement regime. Related to + electron_density via the electron pressure $p_e = n_e T_e$. + unit: eV + kind: scalar + tags: [core_profiles, equilibrium, transport] + fields: + physical_base: temperature + subject: electron + reason: > + Canonical physics quantity. Rich documentation with LaTeX notation, + typical values, diagnostic context, and cross-references. Perfect + grammar decomposition with subject + physical_base. + + - name: safety_factor + tier: outstanding + expected_score: 92 + description: "Magnetohydrodynamic safety factor profile." + documentation: > + The safety factor $q = \frac{d\Phi}{d\psi}$ measures the ratio of + toroidal to poloidal magnetic flux, quantifying field line winding. + Values $q > 1$ everywhere ensure MHD stability against internal kink + modes. The edge safety factor $q_{95}$ is a key operational parameter; + typical range is 2.5-5 for standard H-mode operation. Related to + plasma_current via $q \propto B_T / I_p$. + unit: null + kind: scalar + tags: [equilibrium, mhd] + fields: + physical_base: safety_factor + reason: > + Fundamental equilibrium quantity with mathematical definition, + stability context, typical operating ranges, and cross-references + to related quantities. Standalone physical_base — no qualification + needed. + + - name: toroidal_component_of_magnetic_field_at_magnetic_axis + tier: outstanding + expected_score: 90 + description: "Toroidal magnetic field at the magnetic axis." + documentation: > + The toroidal component of the magnetic field $B_\phi$ evaluated at the + magnetic axis $(R_0, Z_0)$. This is a primary machine parameter that + determines cyclotron resonance locations and plasma beta. Produced by + the toroidal field coil system. Typical values: 1.4 T (TCV), 2.6 T + (ASDEX Upgrade), 5.3 T (ITER). Sign convention follows COCOS; + positive $B_\phi$ corresponds to counter-clockwise toroidal field + when viewed from above. + unit: T + kind: scalar + tags: [equilibrium, magnetics] + fields: + physical_base: magnetic_field + component: toroidal + position: magnetic_axis + reason: > + Multi-segment grammar (component + physical_base + position) correctly + composed. Documentation includes LaTeX, machine-specific typical values, + COCOS sign convention, and physical significance. + + - name: elongation_at_plasma_boundary + tier: outstanding + expected_score: 88 + description: "Plasma elongation at the last closed flux surface." + documentation: > + Elongation $\kappa = b/a$ is the ratio of plasma half-height to + half-width at the last closed flux surface. Higher elongation increases + plasma volume and achievable beta limit ($\beta_N \propto \kappa$) but + requires active vertical stabilization. Typical range: 1.0 (circular) + to 1.8 (strongly shaped). Measured from equilibrium reconstruction + (EFIT, LIUQE). Related to triangularity and plasma_current via the + Troyon limit. + unit: null + kind: scalar + tags: [equilibrium] + fields: + physical_base: elongation + position: plasma_boundary + reason: > + Correct grammar with position qualifier on plasma_boundary. Rich docs + covering definition, stability implications, typical ranges, and + measurement method. + + # =================================================================== + # GOOD tier (60-79): Correct grammar, adequate documentation + # =================================================================== + + - name: electron_density + tier: good + expected_score: 75 + description: "Electron number density profile." + documentation: > + Electron density $n_e$ profile representing the number of electrons + per unit volume. Measured by interferometry, Thomson scattering, and + reflectometry. Typical core values 1-10 × 10^19 m^-3. + unit: m^-3 + kind: scalar + tags: [core_profiles] + fields: + physical_base: density + subject: electron + reason: > + Correct grammar and adequate documentation with LaTeX and typical + values, but lacks cross-references to related quantities and detailed + physics context. + + - name: ion_pressure + tier: good + expected_score: 70 + description: "Thermal pressure of the ion population." + documentation: > + Ion thermal pressure $p_i = n_i T_i$ computed from ion density and + temperature profiles. Contributes to total plasma pressure and + determines the plasma beta. + unit: Pa + kind: scalar + tags: [core_profiles, transport] + fields: + physical_base: pressure + subject: ion + reason: > + Correct grammar with subject + physical_base. Documentation includes + the governing equation but is relatively short, missing typical values + and measurement context. + + - name: ion_temperature + tier: good + expected_score: 68 + description: "Temperature of the ion population." + documentation: > + Ion temperature $T_i$ represents the thermal energy of the main ion + species. Measured by charge exchange recombination spectroscopy (CXRS). + Key parameter for fusion reactivity. + unit: eV + kind: scalar + tags: [core_profiles, transport] + fields: + physical_base: temperature + subject: ion + reason: > + Valid grammar, correct unit. Documentation mentions measurement + technique but is thin — no typical values, no cross-references. + + - name: centroid_at_plasma_boundary + tier: good + expected_score: 65 + description: "Centroid position of the plasma boundary shape." + documentation: > + Geometric centroid of the last closed flux surface boundary, + representing the average R,Z position of the plasma cross-section. + Useful for shape control and equilibrium reconstruction. + unit: m + kind: vector + tags: [equilibrium] + fields: + geometric_base: centroid + position: plasma_boundary + reason: > + Correct geometric_base grammar with position qualifier. Documentation + is adequate but lacks LaTeX, typical values, and references. + + # =================================================================== + # ADEQUATE tier (40-59): Correct grammar, thin documentation + # =================================================================== + + - name: resistivity_due_to_neoclassical + tier: adequate + expected_score: 55 + description: "Neoclassical plasma resistivity." + documentation: > + Plasma resistivity arising from neoclassical transport effects + including trapped particle corrections. + unit: ohm.m + kind: scalar + tags: [transport] + fields: + physical_base: resistivity + process: neoclassical + reason: > + Valid grammar with process qualifier. Documentation is minimal — just + a single sentence with no equations, no typical values, no diagnostic + context. + + - name: extent_of_poloidal_magnetic_field_probe + tier: adequate + expected_score: 50 + description: "Physical extent of a poloidal field probe." + documentation: > + Geometric extent (size) of a poloidal magnetic field probe sensor + element. + unit: m + kind: scalar + tags: [magnetics] + fields: + geometric_base: extent + object: poloidal_magnetic_field_probe + reason: > + Valid geometric_base + object grammar. Documentation is just one + sentence paraphrasing the name. No measurement context. + + - name: electron_collisionality + tier: adequate + expected_score: 48 + description: "Electron collisionality parameter." + documentation: > + Dimensionless electron collisionality. Important for transport regime + classification. + unit: null + kind: scalar + tags: [core_profiles, transport] + fields: + physical_base: collisionality + subject: electron + reason: > + Correct grammar (subject + physical_base). Documentation is a terse + two-sentence summary with no equations or typical values. + + - name: poloidal_component_of_beta + tier: adequate + expected_score: 45 + description: "Poloidal beta." + documentation: > + Ratio of plasma pressure to poloidal magnetic field pressure. + unit: null + kind: scalar + tags: [equilibrium] + fields: + physical_base: beta + component: poloidal + reason: > + Valid component + physical_base grammar. Documentation is a single + sentence — essentially a dictionary definition with no depth. + + # =================================================================== + # POOR tier (0-39): Grammar valid but naming questionable or docs empty + # =================================================================== + + - name: banana_regime + tier: poor + expected_score: 20 + description: "Banana orbit regime." + documentation: "" + unit: null + kind: metadata + tags: [] + fields: + physical_base: banana_regime + reason: > + Empty documentation. The name describes a transport regime, not a + measurable quantity — better suited as metadata or an identifier + value rather than a standard name. No tags assigned. + + - name: signal_value + tier: poor + expected_score: 15 + description: "A signal value." + documentation: "Generic signal value." + unit: null + kind: scalar + tags: [] + fields: + physical_base: signal_value + reason: > + Overly generic name that conveys no physics meaning. Description + is circular ("a signal value"). No unit, no tags. Would apply to + almost any measurement — violates the specificity principle. + + - name: data + tier: poor + expected_score: 5 + description: "Data." + documentation: "" + unit: null + kind: scalar + tags: [] + fields: + physical_base: data + reason: > + Maximally vague name with no physics content whatsoever. Empty + documentation, no unit, no tags. Represents everything a standard + name should NOT be — completely uninformative. diff --git a/imas_codex/sn/benchmark_labels.yaml b/imas_codex/sn/benchmark_labels.yaml deleted file mode 100644 index 72ee31a69..000000000 --- a/imas_codex/sn/benchmark_labels.yaml +++ /dev/null @@ -1,29 +0,0 @@ -# Quality tier labels for benchmark evaluation. -# Each tier defines standard names that exemplify that quality level. -# Used by the reviewer model as scoring anchors. - -outstanding: - # Rich docs, correct grammar, cross-linked, LaTeX notation - - electron_temperature - - plasma_current - - safety_factor - - position_of_magnetic_axis - - bootstrap_current - -good: - # Correct grammar, adequate documentation - - toroidal_component_of_magnetic_field_at_magnetic_axis - - centroid_of_plasma_boundary - - bolometer_radiated_power - - collisionality - -adequate: - # Correct grammar, thin documentation - - area_of_poloidal_magnetic_field_probe - - tokamak_scenario - - time - -poor: - # Grammar valid but naming debatable or documentation minimal - - banana_orbits - - h_mode diff --git a/tests/sn/test_benchmark.py b/tests/sn/test_benchmark.py index d66866d21..e07099106 100644 --- a/tests/sn/test_benchmark.py +++ b/tests/sn/test_benchmark.py @@ -763,34 +763,92 @@ def test_command_requires_models(self): assert "Missing" in result.output or "required" in result.output.lower() -class TestQualityLabels: - """Test benchmark quality tier labels.""" +class TestCalibrationDataset: + """Test benchmark calibration dataset.""" - def test_labels_load(self): - from imas_codex.sn.benchmark import load_quality_labels + def test_calibration_loads(self): + from imas_codex.sn.benchmark import load_calibration_entries - labels = load_quality_labels() - assert isinstance(labels, dict) - assert "outstanding" in labels - assert "good" in labels - assert "adequate" in labels - assert "poor" in labels + entries = load_calibration_entries() + assert isinstance(entries, list) + assert len(entries) == 15, f"Expected 15 entries, got {len(entries)}" - def test_labels_non_empty(self): - from imas_codex.sn.benchmark import load_quality_labels + def test_calibration_tiers(self): + from imas_codex.sn.benchmark import load_calibration_entries - labels = load_quality_labels() - for tier, names in labels.items(): - assert len(names) > 0, f"Tier {tier} should have entries" + entries = load_calibration_entries() + tiers = {} + for entry in entries: + tier = entry["tier"] + tiers[tier] = tiers.get(tier, 0) + 1 + assert tiers == {"outstanding": 4, "good": 4, "adequate": 4, "poor": 3} - def test_labels_no_overlap(self): - from imas_codex.sn.benchmark import load_quality_labels + def test_calibration_required_keys(self): + from imas_codex.sn.benchmark import load_calibration_entries - labels = load_quality_labels() - all_names = [] - for names in labels.values(): - all_names.extend(names) - assert len(all_names) == len(set(all_names)), "No duplicate names across tiers" + required = {"name", "tier", "expected_score", "description", "fields", "reason"} + entries = load_calibration_entries() + for entry in entries: + missing = required - set(entry.keys()) + assert not missing, f"Entry {entry['name']} missing keys: {missing}" + + def test_calibration_names_round_trip(self): + """Every calibration entry name must survive parse→compose round-trip.""" + from imas_codex.sn.benchmark import load_calibration_entries + + entries = load_calibration_entries() + failures = [] + for entry in entries: + name = entry["name"] + try: + parsed = parse_standard_name(name) + rt = compose_standard_name(parsed) + if rt != name: + failures.append(f"{name}: round-trip produced {rt!r}") + except Exception as e: + failures.append(f"{name}: {e!s:.80s}") + assert not failures, "Round-trip failures:\n" + "\n".join(failures) + + def test_calibration_fields_compose_to_name(self): + """compose_standard_name(StandardName(**fields)) == name for each entry.""" + from imas_codex.sn.benchmark import load_calibration_entries + + entries = load_calibration_entries() + failures = [] + for entry in entries: + try: + sn = imas_standard_names.grammar.StandardName(**entry["fields"]) + composed = compose_standard_name(sn) + if composed != entry["name"]: + failures.append(f"{entry['name']}: fields compose to {composed!r}") + except Exception as e: + failures.append(f"{entry['name']}: {e!s:.80s}") + assert not failures, "Field composition failures:\n" + "\n".join(failures) + + def test_calibration_score_ranges(self): + """Verify expected_score falls within the tier's defined range.""" + from imas_codex.sn.benchmark import load_calibration_entries + + tier_ranges = { + "outstanding": (85, 100), + "good": (60, 79), + "adequate": (40, 59), + "poor": (0, 39), + } + entries = load_calibration_entries() + for entry in entries: + lo, hi = tier_ranges[entry["tier"]] + assert lo <= entry["expected_score"] <= hi, ( + f"{entry['name']} ({entry['tier']}): score {entry['expected_score']} " + f"outside range [{lo}, {hi}]" + ) + + def test_calibration_no_duplicate_names(self): + from imas_codex.sn.benchmark import load_calibration_entries + + entries = load_calibration_entries() + names = [e["name"] for e in entries] + assert len(names) == len(set(names)), "Duplicate names in calibration dataset" def test_reviewer_config_field(self): from imas_codex.sn.benchmark import BenchmarkConfig @@ -848,3 +906,184 @@ def test_reviewer_model_in_help(self): result = runner.invoke(sn, ["benchmark", "--help"]) assert result.exit_code == 0 assert "--reviewer-model" in result.output + + +# ----------------------------------------------------------------------- +# 5-dimensional scoring model tests +# ----------------------------------------------------------------------- + + +class TestQualityReviewModel: + """Test the 5-dimensional QualityReview Pydantic model.""" + + def _make_review_model(self): + """Import the QualityReview model from inside score_with_reviewer.""" + from pydantic import BaseModel, Field + + class QualityReview(BaseModel): + name: str + quality_tier: str = Field( + description="outstanding, good, adequate, or poor" + ) + score: int = Field( + ge=0, le=100, description="Total quality score (sum of dimensions)" + ) + grammar_score: int = Field(ge=0, le=20, description="Grammar correctness") + semantic_score: int = Field(ge=0, le=20, description="Semantic accuracy") + documentation_score: int = Field( + ge=0, le=20, description="Documentation quality" + ) + convention_score: int = Field(ge=0, le=20, description="Naming conventions") + completeness_score: int = Field( + ge=0, le=20, description="Entry completeness" + ) + reasoning: str + + return QualityReview + + def test_valid_review(self): + QualityReview = self._make_review_model() + review = QualityReview( + name="electron_temperature", + quality_tier="outstanding", + score=95, + grammar_score=20, + semantic_score=20, + documentation_score=19, + convention_score=18, + completeness_score=18, + reasoning="Excellent entry", + ) + assert review.score == 95 + assert review.grammar_score == 20 + + def test_dimension_max_20(self): + QualityReview = self._make_review_model() + from pydantic import ValidationError + + with pytest.raises(ValidationError): + QualityReview( + name="test", + quality_tier="good", + score=50, + grammar_score=25, # exceeds max 20 + semantic_score=10, + documentation_score=10, + convention_score=5, + completeness_score=0, + reasoning="test", + ) + + def test_dimension_min_0(self): + QualityReview = self._make_review_model() + from pydantic import ValidationError + + with pytest.raises(ValidationError): + QualityReview( + name="test", + quality_tier="poor", + score=10, + grammar_score=-1, # below min 0 + semantic_score=5, + documentation_score=3, + convention_score=2, + completeness_score=1, + reasoning="test", + ) + + def test_total_score_max_100(self): + QualityReview = self._make_review_model() + from pydantic import ValidationError + + with pytest.raises(ValidationError): + QualityReview( + name="test", + quality_tier="outstanding", + score=101, # exceeds max 100 + grammar_score=20, + semantic_score=20, + documentation_score=20, + convention_score=20, + completeness_score=20, + reasoning="test", + ) + + def test_poor_tier_scores(self): + QualityReview = self._make_review_model() + review = QualityReview( + name="data", + quality_tier="poor", + score=5, + grammar_score=5, + semantic_score=0, + documentation_score=0, + convention_score=0, + completeness_score=0, + reasoning="Uninformative name", + ) + assert review.score == 5 + assert review.quality_tier == "poor" + + +# ----------------------------------------------------------------------- +# Reviewer template rendering tests +# ----------------------------------------------------------------------- + + +class TestReviewerTemplate: + """Test that the reviewer template renders correctly.""" + + def test_template_renders(self): + from imas_codex.llm.prompt_loader import render_prompt + from imas_codex.sn.benchmark import load_calibration_entries + + entries = load_calibration_entries() + rendered = render_prompt( + "sn/review_benchmark", + { + "calibration_entries": entries, + "candidates": [ + { + "standard_name": "electron_temperature", + "description": "Electron temperature", + "documentation": "A test doc", + "unit": "eV", + "kind": "scalar", + "tags": ["core_profiles"], + "fields": { + "physical_base": "temperature", + "subject": "electron", + }, + } + ], + }, + ) + assert "electron_temperature" in rendered + assert "Grammar Correctness" in rendered + assert "Semantic Accuracy" in rendered + assert "Documentation Quality" in rendered + assert "outstanding" in rendered + + def test_template_includes_calibration_examples(self): + from imas_codex.llm.prompt_loader import render_prompt + from imas_codex.sn.benchmark import load_calibration_entries + + entries = load_calibration_entries() + rendered = render_prompt( + "sn/review_benchmark", + {"calibration_entries": entries, "candidates": []}, + ) + # All calibration entry names should appear + for entry in entries: + assert entry["name"] in rendered, ( + f"Calibration entry {entry['name']} not in rendered template" + ) + + def test_template_renders_empty_candidates(self): + from imas_codex.llm.prompt_loader import render_prompt + + rendered = render_prompt( + "sn/review_benchmark", + {"calibration_entries": [], "candidates": []}, + ) + assert "Scoring Dimensions" in rendered From 646e921586dec29f4547d8f864b22f28cd2f1ebd Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 16:39:14 +0200 Subject: [PATCH 36/37] feat: add cache reporting to benchmark and model selection runbook MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add LLMResult class to llm.py for backward-compatible cache token exposure — supports 3-tuple unpacking while carrying cache_read_tokens and cache_creation_tokens from provider prompt caching. Add extract_cache_tokens() public function mirroring _log_cache_metrics extraction logic but returning values instead of logging. Add cache_read_tokens and cache_creation_tokens to ModelResult dataclass. Accumulate cache tokens in _run_model() using getattr fallback for mock compatibility. Display Cache % column in benchmark comparison table. Create model selection runbook with CLI commands, cost guidance, approved model list, decision criteria table, and cache optimization tips. --- imas_codex/discovery/base/__init__.py | 4 + imas_codex/discovery/base/llm.py | 112 ++++++++++- imas_codex/sn/benchmark.py | 20 +- .../standard-names/model-selection-runbook.md | 176 ++++++++++++++++++ tests/sn/test_benchmark.py | 157 ++++++++++++++++ 5 files changed, 460 insertions(+), 9 deletions(-) create mode 100644 plans/features/standard-names/model-selection-runbook.md diff --git a/imas_codex/discovery/base/__init__.py b/imas_codex/discovery/base/__init__.py index 4814e2e53..2810ee15f 100644 --- a/imas_codex/discovery/base/__init__.py +++ b/imas_codex/discovery/base/__init__.py @@ -37,10 +37,12 @@ normalize_imas_path, ) from imas_codex.discovery.base.llm import ( + LLMResult, acall_llm, acall_llm_structured, call_llm, call_llm_structured, + extract_cache_tokens, extract_cost, get_model_limits, inject_cache_control, @@ -136,11 +138,13 @@ "ParallelExecutor", "CommandResult", # LLM + "LLMResult", "call_llm", "call_llm_structured", "acall_llm", "acall_llm_structured", "extract_cost", + "extract_cache_tokens", "get_model_limits", "inject_cache_control", "suppress_litellm_noise", diff --git a/imas_codex/discovery/base/llm.py b/imas_codex/discovery/base/llm.py index 2d6d66195..f7fea3357 100644 --- a/imas_codex/discovery/base/llm.py +++ b/imas_codex/discovery/base/llm.py @@ -30,12 +30,14 @@ response_model=ScoreBatch, ) - # Async structured output - batch, cost, tokens = await acall_llm_structured( + # Async structured output (also returns LLMResult with cache info) + llm_out = await acall_llm_structured( model="google/gemini-3-flash-preview", messages=[...], response_model=WikiScoreBatch, ) + batch, cost, tokens = llm_out # backward-compatible + cache_read = llm_out.cache_read_tokens # new: cache metrics # Raw response (when caller needs custom parsing) response, cost = call_llm( @@ -79,6 +81,66 @@ class ProviderBudgetExhausted(Exception): """ +class LLMResult: + """Return type for call_llm_structured / acall_llm_structured. + + Backward-compatible with 3-tuple unpacking:: + + result, cost, tokens = call_llm_structured(...) # still works + + Also carries prompt-cache token counts for callers that need them:: + + llm_out = await acall_llm_structured(...) + result, cost, tokens = llm_out + cache_read = llm_out.cache_read_tokens + cache_creation = llm_out.cache_creation_tokens + + Attributes: + parsed: The Pydantic model instance returned by the LLM. + cost: Total cost in USD (accumulated across retries). + tokens: Total tokens (prompt + completion). + cache_read_tokens: Tokens served from provider prompt cache (0 if + the provider doesn't report caching or the prompt wasn't cached). + cache_creation_tokens: Tokens written to the provider prompt cache. + """ + + __slots__ = ( + "parsed", + "cost", + "tokens", + "cache_read_tokens", + "cache_creation_tokens", + ) + + def __init__( + self, + parsed: Any, + cost: float, + tokens: int, + cache_read_tokens: int = 0, + cache_creation_tokens: int = 0, + ) -> None: + self.parsed = parsed + self.cost = cost + self.tokens = tokens + self.cache_read_tokens = cache_read_tokens + self.cache_creation_tokens = cache_creation_tokens + + # Allow ``result, cost, tokens = call_llm_structured(...)`` + def __iter__(self): + return iter((self.parsed, self.cost, self.tokens)) + + def __len__(self) -> int: + return 3 + + def __repr__(self) -> str: + return ( + f"LLMResult(cost={self.cost:.4f}, tokens={self.tokens}, " + f"cache_read={self.cache_read_tokens}, " + f"cache_creation={self.cache_creation_tokens})" + ) + + # Patterns indicating the API key or account has hit a hard spending cap. # Matched case-insensitively against the full error message. _BUDGET_EXHAUSTED_PATTERNS = ( @@ -366,6 +428,30 @@ def _log_cache_metrics(response: Any, model: str) -> None: ) +def extract_cache_tokens(response: Any) -> tuple[int, int]: + """Extract prompt-cache token counts from an LLM response. + + Providers (OpenRouter, Anthropic, OpenAI) report cached token counts + in ``usage.prompt_tokens_details``. This helper mirrors the extraction + logic of :func:`_log_cache_metrics` but returns the values instead of + logging them, so callers can accumulate cache statistics. + + Args: + response: Raw litellm response object. + + Returns: + ``(cache_read_tokens, cache_creation_tokens)`` — both default to 0 + when the provider does not report caching information. + """ + usage = getattr(response, "usage", None) + if not usage: + return 0, 0 + ptd = getattr(usage, "prompt_tokens_details", None) + cached = getattr(ptd, "cached_tokens", 0) or 0 if ptd else 0 + cache_write = getattr(ptd, "cache_creation_tokens", 0) or 0 if ptd else 0 + return cached, cache_write + + def _sanitize_content(content: str) -> str: """Sanitize LLM response content for JSON parsing. @@ -563,7 +649,7 @@ def call_llm_structured( timeout: int | None = None, max_retries: int = DEFAULT_MAX_RETRIES, retry_base_delay: float = DEFAULT_RETRY_BASE_DELAY, -) -> tuple[BaseModel, float, int]: +) -> LLMResult: """Call LLM and parse structured output, retrying on both API and parse errors. Wraps the LLM call and Pydantic parsing in a single retry loop so that @@ -584,7 +670,9 @@ def call_llm_structured( retry_base_delay: Base delay for exponential backoff (seconds). Returns: - Tuple of (parsed_model, total_cost_usd, total_tokens). + :class:`LLMResult` — backward-compatible with 3-tuple unpacking + ``(parsed_model, total_cost_usd, total_tokens)`` and also carries + ``cache_read_tokens`` / ``cache_creation_tokens``. Raises: ValueError: If response parsing fails after all retries. @@ -625,7 +713,10 @@ def call_llm_structured( total_tokens = ( response.usage.prompt_tokens + response.usage.completion_tokens ) - return parsed, total_cost, total_tokens + cache_read, cache_creation = extract_cache_tokens(response) + return LLMResult( + parsed, total_cost, total_tokens, cache_read, cache_creation + ) except Exception as e: last_error = e @@ -669,7 +760,7 @@ async def acall_llm_structured( timeout: int | None = None, max_retries: int = DEFAULT_MAX_RETRIES, retry_base_delay: float = DEFAULT_RETRY_BASE_DELAY, -) -> tuple[BaseModel, float, int]: +) -> LLMResult: """Async version of call_llm_structured. Identical retry+parse semantics using litellm.acompletion() and @@ -686,7 +777,9 @@ async def acall_llm_structured( retry_base_delay: Base delay for exponential backoff (seconds). Returns: - Tuple of (parsed_model, total_cost_usd, total_tokens). + :class:`LLMResult` — backward-compatible with 3-tuple unpacking + ``(parsed_model, total_cost_usd, total_tokens)`` and also carries + ``cache_read_tokens`` / ``cache_creation_tokens``. Raises: ValueError: If response parsing fails after all retries. @@ -727,7 +820,10 @@ async def acall_llm_structured( total_tokens = ( response.usage.prompt_tokens + response.usage.completion_tokens ) - return parsed, total_cost, total_tokens + cache_read, cache_creation = extract_cache_tokens(response) + return LLMResult( + parsed, total_cost, total_tokens, cache_read, cache_creation + ) except Exception as e: last_error = e diff --git a/imas_codex/sn/benchmark.py b/imas_codex/sn/benchmark.py index 6685adc02..f7540875e 100644 --- a/imas_codex/sn/benchmark.py +++ b/imas_codex/sn/benchmark.py @@ -82,6 +82,9 @@ class ModelResult: avg_quality_score: float = 0.0 avg_doc_length: float = 0.0 avg_fields_populated: float = 0.0 + # Prompt-cache statistics + cache_read_tokens: int = 0 + cache_creation_tokens: int = 0 @dataclass @@ -557,14 +560,21 @@ async def _run_model( ] try: - llm_result, cost, tokens = await acall_llm_structured( + llm_response = await acall_llm_structured( model=model, messages=messages, response_model=SNComposeBatch, temperature=config.temperature, ) + llm_result, cost, tokens = llm_response result.total_cost += cost result.total_tokens += tokens + result.cache_read_tokens += getattr( + llm_response, "cache_read_tokens", 0 + ) + result.cache_creation_tokens += getattr( + llm_response, "cache_creation_tokens", 0 + ) logger.debug( "Batch %s: cost=%.4f tokens=%d", group_key, @@ -665,6 +675,7 @@ def render_comparison_table(report: BenchmarkReport) -> None: table.add_column("Cost", justify="right") table.add_column("Names/min", justify="right") table.add_column("$/name", justify="right") + table.add_column("Cache %", justify="right") table.add_column("Errors", justify="right") if has_quality: table.add_column("Avg Quality", justify="right") @@ -681,6 +692,12 @@ def render_comparison_table(report: BenchmarkReport) -> None: cost_str = f"${r.total_cost:.4f}" if r.total_cost > 0 else "—" speed_str = f"{r.names_per_minute:.0f}" if r.names_per_minute > 0 else "—" cpn_str = f"${r.cost_per_name:.4f}" if r.cost_per_name > 0 else "—" + cache_total = r.cache_read_tokens + r.cache_creation_tokens + cache_pct = ( + f"{r.cache_read_tokens / cache_total * 100:.0f}%" + if cache_total > 0 + else "—" + ) err_str = str(r.batch_errors) if r.batch_errors > 0 else "0" row_data = [ @@ -692,6 +709,7 @@ def render_comparison_table(report: BenchmarkReport) -> None: cost_str, speed_str, cpn_str, + cache_pct, err_str, ] diff --git a/plans/features/standard-names/model-selection-runbook.md b/plans/features/standard-names/model-selection-runbook.md new file mode 100644 index 000000000..8ace62d34 --- /dev/null +++ b/plans/features/standard-names/model-selection-runbook.md @@ -0,0 +1,176 @@ +# Model Selection Runbook — Standard Name Generation + +Practical guide for running multi-model benchmarks and selecting the best +LLM for each role in the standard name pipeline. + +## Quick Start + +```bash +# Compare two models on equilibrium IDS (fast, <$0.50) +uv run imas-codex sn benchmark \ + --source dd \ + --ids equilibrium \ + --max-candidates 10 \ + --models google/gemini-3.1-flash-lite-preview,anthropic/claude-sonnet-4-6 + +# Full benchmark with reviewer scoring (~$1–2) +uv run imas-codex sn benchmark \ + --source dd \ + --max-candidates 50 \ + --models google/gemini-3.1-flash-lite-preview,anthropic/claude-sonnet-4-6 \ + --reviewer-model anthropic/claude-sonnet-4-6 + +# Export results to JSON for offline analysis +uv run imas-codex sn benchmark \ + --source dd \ + --ids equilibrium \ + --max-candidates 20 \ + --models google/gemini-3.1-flash-lite-preview \ + --output benchmark-results.json +``` + +## Cost Cap Guidance + +| Scenario | `--max-candidates` | Expected cost | +|---------------------|--------------------|---------------| +| Quick smoke test | 5–10 | < $0.20 | +| Single-IDS compare | 20–30 | $0.30–$0.80 | +| Full benchmark | 50 | $0.50–$1.50 | +| Production eval | 100+ | $1.50–$4.00 | + +> **Hard rule:** never exceed **$2.00** per benchmark execution during +> development and model selection. Monitor cost in the output table and +> abort early if needed. + +## Approved Model List + +Models are accessed via the LiteLLM proxy running on the ITER login node. +The proxy routes requests through OpenRouter. Pass model identifiers as +configured in `pyproject.toml` (the proxy handles the `openrouter/` prefix). + +### Compose Role (name generation) + +| Model | Tier | Notes | +|------------------------------------------|----------|------------------------------------| +| `google/gemini-3.1-flash-lite-preview` | Primary | Current `language` config, cheapest| +| `google/gemini-2.5-flash` | Alt | Fast, good grammar | +| `anthropic/claude-sonnet-4-6` | Fallback | Higher quality, 3–5× cost | +| `anthropic/claude-haiku-4` | Budget | Fastest Anthropic, lowest cost | + +### Review Role (quality scoring) + +| Model | Tier | Notes | +|------------------------------------------|----------|------------------------------------| +| `anthropic/claude-sonnet-4-6` | Primary | Best judgment, calibrated scoring | +| `anthropic/claude-opus-4-6` | Premium | Highest quality, ~10× cost | +| `google/gemini-2.5-pro` | Alt | Strong reasoning, competitive cost | + +### Currently Configured (pyproject.toml) + +```toml +[tool.imas-codex.language] +model = "google/gemini-3.1-flash-lite-preview" # compose + +[tool.imas-codex.reasoning] +model = "anthropic/claude-sonnet-4-6" # review / complex + +[tool.imas-codex.agent] +model = "anthropic/claude-opus-4-6" # agent +``` + +## Decision Criteria + +| Metric | Weight | Threshold | How to read | +|----------------------|----------|-----------|-------------------------------------| +| Grammar valid % | Critical | ≥ 95% | Names must parse→compose round-trip | +| Fields consistent % | High | ≥ 85% | Decomposed fields reconstruct name | +| Reference recall | High | ≥ 60% | Overlap with human-curated set | +| Avg quality score | High | ≥ 65 | 5-dimensional reviewer rating /100 | +| Cost per name | Medium | < $0.01 | Total API cost ÷ candidate count | +| Names/min | Medium | > 30 | Throughput including API latency | +| Cache hit rate | Low | > 50% | OpenRouter prompt cache utilization | + +### Interpreting the Table + +The benchmark outputs a Rich table with these columns: + +``` +Model Names Valid% Fields% Ref Match Cost Names/min $/name Cache% Errors +gemini 47 100% 96% 28/50 $0.0312 62 $0.0007 78% 0 +claude 45 98% 91% 31/50 $0.1450 18 $0.0032 65% 1 +``` + +- **Valid %** — ratio of candidates whose `standard_name` survives grammar + `parse_standard_name()` → `compose_standard_name()` round-trip. +- **Fields %** — ratio where decomposed grammar fields (`subject`, + `physical_base`, etc.) recompose to the same name string. +- **Ref Match** — `overlap/total` against the reference dataset + (`benchmark_reference.py`). Higher recall = more agreement with human picks. +- **Cache %** — `cache_read / (cache_read + cache_creation)` — how much of + the system prompt was served from OpenRouter's prompt cache. Higher values + reduce cost and latency on subsequent batches. +- **Errors** — number of batches where the LLM call failed (timeout, + rate limit, parse error). + +When `--reviewer-model` is used, additional columns appear: + +- **Avg Quality** — mean 5-dimensional score (0–100). +- **Avg Doc Len** — average documentation string length. +- **Fields Pop%** — ratio of optional fields populated. + +## Recommended Selections + +### For compose (name generation) + +Use the **cheapest model that meets grammar + fields thresholds**: + +1. Start with `google/gemini-3.1-flash-lite-preview` +2. If grammar valid < 95%, try `anthropic/claude-sonnet-4-6` +3. If cost per name > $0.01, try `google/gemini-2.5-flash` + +### For review (quality scoring) + +Use a **stronger model than compose** for unbiased evaluation: + +1. Primary: `anthropic/claude-sonnet-4-6` +2. If budget allows: `anthropic/claude-opus-4-6` + +> **Never use the same model for both compose and review** — self-review +> inflates quality scores. + +## Prompt Cache Optimization + +OpenRouter supports provider-level prompt caching. The SN benchmark system +prompt includes grammar rules and calibration examples that remain constant +across batches. Cache behavior: + +- **First batch**: `cache_creation_tokens` > 0 (system prompt written to cache) +- **Subsequent batches**: `cache_read_tokens` > 0 (system prompt served from cache) +- **Cost savings**: cached tokens cost ~75% less than uncached tokens + +To maximize cache hit rate: +- Run batches in rapid succession (caches expire after ~5 minutes idle) +- Use the `openrouter/` prefix (caching is provider-side, not proxy-side) +- Group batches by model to keep the cache warm + +## Troubleshooting + +### LLM proxy not running + +```bash +# Check proxy status +uv run imas-codex hpc status + +# Start proxy (requires SSH tunnel to ITER) +uv run imas-codex llm start +``` + +If the proxy is unavailable, the benchmark will fail with connection errors. +The cache reporting code and grammar validation are tested offline via +`uv run pytest tests/sn/test_benchmark.py -v`. + +### Budget exhausted + +If you see `ProviderBudgetExhausted`, either: +1. The OpenRouter account is out of credits — top up +2. The `--cost-limit` was exceeded — increase or reduce `--limit` diff --git a/tests/sn/test_benchmark.py b/tests/sn/test_benchmark.py index e07099106..f09546b68 100644 --- a/tests/sn/test_benchmark.py +++ b/tests/sn/test_benchmark.py @@ -894,6 +894,163 @@ def test_model_result_with_quality(self): assert r.quality_distribution["outstanding"] == 1 +class TestCacheTokenReporting: + """Test prompt-cache token reporting in ModelResult and Rich table.""" + + def test_llm_result_unpacking(self): + """LLMResult supports 3-element tuple unpacking (backward compat).""" + from imas_codex.discovery.base.llm import LLMResult + + r = LLMResult( + "parsed", 0.05, 500, cache_read_tokens=300, cache_creation_tokens=100 + ) + parsed, cost, tokens = r + assert parsed == "parsed" + assert cost == 0.05 + assert tokens == 500 + assert r.cache_read_tokens == 300 + assert r.cache_creation_tokens == 100 + + def test_llm_result_getattr_fallback(self): + """getattr on a plain tuple returns 0 (mock compatibility).""" + mock_return = ("parsed", 0.01, 200) + assert getattr(mock_return, "cache_read_tokens", 0) == 0 + assert getattr(mock_return, "cache_creation_tokens", 0) == 0 + + def test_model_result_cache_defaults(self): + from imas_codex.sn.benchmark import ModelResult + + r = ModelResult(model="test") + assert r.cache_read_tokens == 0 + assert r.cache_creation_tokens == 0 + + def test_model_result_with_cache(self): + from imas_codex.sn.benchmark import ModelResult + + r = ModelResult( + model="test", + cache_read_tokens=5000, + cache_creation_tokens=2000, + ) + assert r.cache_read_tokens == 5000 + assert r.cache_creation_tokens == 2000 + + def test_cache_pct_all_read(self): + """100% cache hit rate.""" + read, creation = 1000, 0 + total = read + creation + pct = read / total * 100 if total > 0 else 0 + assert pct == 100.0 + + def test_cache_pct_no_cache(self): + """0/0 — no cache tokens at all.""" + read, creation = 0, 0 + total = read + creation + pct = read / total * 100 if total > 0 else 0 + assert pct == 0.0 + + def test_cache_pct_mixed(self): + """Partial cache hit rate.""" + read, creation = 300, 700 + total = read + creation + pct = read / total * 100 if total > 0 else 0 + assert pct == 30.0 + + def test_cache_pct_all_creation(self): + """First request — all tokens are cache creation.""" + read, creation = 0, 500 + total = read + creation + pct = read / total * 100 if total > 0 else 0 + assert pct == 0.0 + + def test_render_table_with_cache(self): + """Cache % column should appear and show correct values.""" + from imas_codex.sn.benchmark import ( + BenchmarkConfig, + BenchmarkReport, + ModelResult, + render_comparison_table, + ) + + r = ModelResult( + model="test-model", + candidates=[{"source_id": "p", "standard_name": "electron_temperature"}], + grammar_valid_count=1, + grammar_invalid_count=0, + fields_consistent_count=1, + total_cost=0.01, + total_tokens=100, + elapsed_seconds=5.0, + names_per_minute=12.0, + cost_per_name=0.01, + cache_read_tokens=800, + cache_creation_tokens=200, + ) + report = BenchmarkReport( + config=BenchmarkConfig(models=["test-model"]), + results=[r], + reference_names=[], + extraction_count=1, + timestamp="2025-01-01", + ) + # Should not raise + render_comparison_table(report) + + def test_render_table_no_cache(self): + """Cache % column shows '—' when no cache tokens.""" + from imas_codex.sn.benchmark import ( + BenchmarkConfig, + BenchmarkReport, + ModelResult, + render_comparison_table, + ) + + r = ModelResult( + model="no-cache-model", + candidates=[{"source_id": "p", "standard_name": "electron_temperature"}], + grammar_valid_count=1, + total_cost=0.01, + total_tokens=100, + ) + report = BenchmarkReport( + config=BenchmarkConfig(models=["no-cache-model"]), + results=[r], + reference_names=[], + timestamp="2025-01-01", + ) + # Should not raise + render_comparison_table(report) + + def test_cache_in_json_round_trip(self): + """Cache fields survive JSON serialization.""" + import json + + from imas_codex.sn.benchmark import ( + BenchmarkConfig, + BenchmarkReport, + ModelResult, + ) + + r = ModelResult( + model="m", + cache_read_tokens=1500, + cache_creation_tokens=500, + ) + report = BenchmarkReport( + config=BenchmarkConfig(models=["m"]), + results=[r], + reference_names=[], + timestamp="2025-01-01", + ) + parsed = json.loads(report.to_json()) + assert parsed["results"][0]["cache_read_tokens"] == 1500 + assert parsed["results"][0]["cache_creation_tokens"] == 500 + + restored = BenchmarkReport.from_json(report.to_json()) + assert restored.results[0].cache_read_tokens == 1500 + assert restored.results[0].cache_creation_tokens == 500 + + class TestReviewerModelCLI: """Test --reviewer-model CLI option.""" From 5b4861104e2b4580ab6bb7253b0c9662f457f36b Mon Sep 17 00:00:00 2001 From: Simon McIntosh Date: Fri, 10 Apr 2026 16:45:42 +0200 Subject: [PATCH 37/37] docs: update documentation for SN benchmark and lifecycle features Update AGENTS.md with sn mint (renamed from sn build), sn reset, sn clear, --reset-to flag, benchmark cache reporting, and 5-dimensional scoring. Add SN module paths to project-dev skill, LLM proxy note to service-ops skill, and SN key files table to engineer agent. Delete superseded plans 16, 17, 18. Mark Plan 19 complete. --- .github/agents/engineer.agent.md | 10 + .github/skills/project-dev/SKILL.md | 14 + .github/skills/service-ops/SKILL.md | 5 + AGENTS.md | 34 ++- .../standard-names/00-implementation-order.md | 8 +- .../standard-names/16-benchmark-parity.md | 152 ---------- .../17-sn-lifecycle-management.md | 163 ----------- .../18-benchmark-calibration.md | 276 ------------------ 8 files changed, 65 insertions(+), 597 deletions(-) delete mode 100644 plans/features/standard-names/16-benchmark-parity.md delete mode 100644 plans/features/standard-names/17-sn-lifecycle-management.md delete mode 100644 plans/features/standard-names/18-benchmark-calibration.md diff --git a/.github/agents/engineer.agent.md b/.github/agents/engineer.agent.md index 6dc8e6d7a..0f22c77c4 100644 --- a/.github/agents/engineer.agent.md +++ b/.github/agents/engineer.agent.md @@ -84,6 +84,16 @@ When modifying LinkML schemas: - Documentation and prompt template updates - Refactors where the before/after is well-defined +### Commonly-Modified Areas + +| Path | Purpose | +|------|---------| +| `imas_codex/sn/` | Standard name pipeline (mint, benchmark, graph ops) | +| `tests/sn/` | SN test suite (mostly mock-based, no Neo4j required) | +| `imas_codex/llm/prompts/sn/` | LLM prompt templates for SN | +| `imas_codex/sn/benchmark_reference.py` | Gold reference set for benchmark scoring | +| `imas_codex/sn/benchmark_calibration.yaml` | Calibration dataset for reviewer consistency | + ## When to Escalate If a task requires: diff --git a/.github/skills/project-dev/SKILL.md b/.github/skills/project-dev/SKILL.md index 5fefd16cc..ad8fed46e 100644 --- a/.github/skills/project-dev/SKILL.md +++ b/.github/skills/project-dev/SKILL.md @@ -84,6 +84,10 @@ git push origin main | `@pytest.mark.integration` | Full integration tests | | `@pytest.mark.unit` | Fast unit tests | +SN tests live in `tests/sn/` and run with `uv run pytest tests/sn/ -v`. They do not require +Neo4j unless marked `@pytest.mark.graph` — the rest use mocks. Benchmark tests validate prompt +parity with the mint pipeline, calibration dataset integrity, and reference set coverage. + ## Project Structure | Directory | Purpose | @@ -95,7 +99,9 @@ git push origin main | `imas_codex/tools/` | MCP tool implementations | | `imas_codex/remote/` | Remote execution (SSH, scripts) | | `imas_codex/llm/` | LLM integration and prompt templates | +| `imas_codex/sn/` | Standard name pipeline (mint, benchmark, graph ops) | | `tests/` | Test suite (mirrors source structure) | +| `tests/sn/` | Standard name test suite (mostly mock-based) | | `plans/features/` | Active feature plans | | `agents/` | Agent documentation and schema reference | @@ -107,3 +113,11 @@ git push origin main - **Model selection**: Use `get_model(section)` from `imas_codex.settings` - **Facility config**: Use `get_facility(facility)` — never hardcode facility values - **Remote execution**: Use `run_python_script()` from `imas_codex.remote.executor` + +### SN Key Files + +| File | Purpose | +|------|---------| +| `imas_codex/sn/benchmark_reference.py` | Gold reference set (52 entries across 8 IDSs) | +| `imas_codex/sn/benchmark_calibration.yaml` | Known-quality examples for reviewer consistency | +| `imas_codex/llm/prompts/sn/` | LLM prompt templates for mint, review, and benchmark | diff --git a/.github/skills/service-ops/SKILL.md b/.github/skills/service-ops/SKILL.md index bb56961db..db5fe8e11 100644 --- a/.github/skills/service-ops/SKILL.md +++ b/.github/skills/service-ops/SKILL.md @@ -75,6 +75,11 @@ uv run imas-codex llm spend # Cost tracking uv run imas-codex llm logs # View logs ``` +`sn mint` and `sn benchmark` require the LLM proxy to be running. Model names must use the +`openrouter/` prefix (e.g. `openrouter/anthropic/claude-sonnet-4-5`) to preserve +`cache_control` blocks — prompt caching is handled provider-side by OpenRouter, not by this +codebase. + ## SSH Tunnels ```bash diff --git a/AGENTS.md b/AGENTS.md index 185e14e79..20a1e5ec8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -712,12 +712,24 @@ Azure Web App has continuous deployment enabled on ACR. When a new image appears | Command | Purpose | Key Options | |---------|---------|-------------| -| `sn build` | Generate standard names from DD paths or facility signals via LLM pipeline | `--source {dd,signals}`, `--ids`, `--domain`, `--facility`, `--cost-limit`, `--dry-run`, `--force`, `--skip-review` | +| `sn mint` | Generate standard names from DD paths or facility signals via LLM pipeline | `--source {dd,signals}`, `--ids`, `--domain`, `--facility`, `--cost-limit`, `--dry-run`, `--force`, `--skip-review`, `--reset-to` | | `sn publish` | Export validated StandardName nodes to YAML catalog files | `--output-dir`, `--ids`, `--domain`, `--group-by {ids,domain,confidence}`, `--confidence-min`, `--catalog-dir`, `--create-pr` | | `sn import` | Import reviewed YAML catalog entries back into graph | `--catalog-dir` (required), `--tags`, `--dry-run`, `--check` | | `sn status` | Show standard name statistics from graph | — | +| `sn reset` | Reset standard names for re-processing | `--status` (required), `--to`, `--source`, `--ids`, `--dry-run` | +| `sn clear` | Delete standard names from the graph (relationship-first safety model) | `--status`, `--all`, `--source`, `--ids`, `--include-accepted`, `--dry-run` | | `sn benchmark` | Benchmark LLM models on standard name generation quality | `--models`, `--source`, `--reviewer-model` | +### Benchmark + +`sn benchmark` uses the same prompt pipeline as `sn mint` (system/user message split via +`build_compose_context()`). Output table includes a **Cache %** column showing the prompt-cache +hit rate per model (provider-side via OpenRouter — not something we implement). Scoring is +**5-dimensional**: accuracy, completeness, physics_correctness, naming_convention, and +overall, evaluated by a reviewer LLM against a gold reference set (`benchmark_reference.py`, +52 entries across 8 IDSs). The calibration dataset (`benchmark_calibration.yaml`) provides +known-quality examples for reviewer consistency checks. + ### StandardName Lifecycle ``` @@ -725,10 +737,28 @@ drafted → published → accepted ↘ rejected ``` -- **drafted**: Generated by `sn build` (LLM pipeline) +- **drafted**: Generated by `sn mint` (LLM pipeline) - **published**: Exported by `sn publish` to YAML catalog for human review - **accepted**: Imported by `sn import` from reviewed catalog (catalog-authoritative) +### Reset and Clear Semantics + +**`sn reset`** — Re-processes existing nodes without deleting them. Clears transient fields +(embedding, model, confidence, generated_at) and removes HAS_STANDARD_NAME and CANONICAL_UNITS +relationships. Optionally changes `review_status` via `--to `. Default (no `--to`) leaves +status unchanged, only clears fields. + +**`sn clear`** — Deletes StandardName nodes. Uses a relationship-first safety model: HAS_STANDARD_NAME +edges are removed before deleting nodes, and scoped deletes only remove orphaned nodes. Requires +either `--status ` or `--all`. + +**Safety guard:** Both commands require `--include-accepted` to touch names with `review_status=accepted`. +Accepted names are catalog-authoritative and should rarely be deleted from the graph. + +**`sn mint --reset-to`** — Runs a `sn reset` before minting, scoped to the same `--ids`/`--source` +filter. Accepts `extracted` or `drafted` as the target status. Useful for a clean re-run on a +specific IDS without touching the rest of the graph. + ### Write Semantics Two distinct write paths with different semantics: diff --git a/plans/features/standard-names/00-implementation-order.md b/plans/features/standard-names/00-implementation-order.md index 759c7a12e..90ee720ac 100644 --- a/plans/features/standard-names/00-implementation-order.md +++ b/plans/features/standard-names/00-implementation-order.md @@ -48,10 +48,10 @@ All status values use past tense: drafted, published, accepted, rejected, skippe | 13 | publish-pipeline | Lossless YAML export, batched PRs | 📋 Ready | 11 (all) | 12 (feedback loop) | | 14 | mcp-tools-benchmark | SN search/fetch/list MCP tools + benchmark quality | ✅ Done | 11 (embedding) | 19 | | 15 | import-physics-domain | Import physics_domain from catalog | ✅ Done | 12 | — | -| ~~16~~ | ~~benchmark-parity~~ | ~~Superseded by Plan 19~~ | 🔀 Merged | — | — | -| ~~17~~ | ~~sn-lifecycle-management~~ | ~~Superseded by Plan 19~~ | 🔀 Merged | — | — | -| ~~18~~ | ~~benchmark-calibration~~ | ~~Superseded by Plan 19~~ | 🔀 Merged | — | — | -| 19 | benchmark-and-lifecycle | Benchmark parity, lifecycle mgmt, calibration, model selection | 📋 Ready | 14 | — | +| ~~16~~ | ~~benchmark-parity~~ | ~~Superseded by Plan 19~~ | 🗑️ Deleted | — | — | +| ~~17~~ | ~~sn-lifecycle-management~~ | ~~Superseded by Plan 19~~ | 🗑️ Deleted | — | — | +| ~~18~~ | ~~benchmark-calibration~~ | ~~Superseded by Plan 19~~ | 🗑️ Deleted | — | — | +| 19 | benchmark-and-lifecycle | Benchmark parity, lifecycle mgmt, calibration, model selection | ✅ Complete | 14 | — | ## Deployment Waves diff --git a/plans/features/standard-names/16-benchmark-parity.md b/plans/features/standard-names/16-benchmark-parity.md deleted file mode 100644 index bfd6170de..000000000 --- a/plans/features/standard-names/16-benchmark-parity.md +++ /dev/null @@ -1,152 +0,0 @@ -# 16: Benchmark / Mint Prompt Parity & Cache Verification - -**Status:** Ready to implement -**Depends on:** None (standalone fix) -**Blocks:** 18 (calibration — benchmark results meaningless until parity fixed) -**Agent:** engineer - -## Problem - -The `sn benchmark` command constructs LLM prompts differently from `sn mint`, -making benchmark metrics (cost, speed, quality) unreliable for production model -selection. - -### Specific gaps - -| Aspect | Mint (`workers.py`) | Benchmark (`benchmark.py`) | -|--------|---------------------|---------------------------| -| System prompt | `sn/compose_system` via `build_compose_context()` | None | -| User prompt | `sn/compose_dd` with `cluster_context` | `sn/compose_dd` without `cluster_context` | -| Message structure | `[system, user]` — enables prompt caching | `[user]` — no caching possible | -| Context source | `build_compose_context()` (grammar rules, vocabulary, examples, tokamak ranges) | `build_grammar_context()` (bare enum lists only) | -| Prompt caching | `inject_cache_control()` on system message | Not applicable (no system message) | - -### Impact - -- Cost metrics inflated (no cache hits) -- Speed metrics pessimistic (larger uncached prompts) -- Quality metrics unreliable (model sees different context) -- Benchmark cannot validate whether caching actually works - -### Caching architecture - -Prompt caching is **provider-side** — OpenRouter caches repeated prompt prefixes -automatically when `cache_control` blocks are present. Our infrastructure already -supports this: - -- `inject_cache_control()` in `discovery/base/llm.py` adds `cache_control: - {"type": "ephemeral"}` breakpoints to system messages -- `openrouter/` model prefix preserves these blocks through LiteLLM -- OpenRouter returns `usage.cache_creation_input_tokens` and - `usage.cache_read_input_tokens` in responses - -The fix here is to **use the same prompt architecture as mint** so that caching -works naturally, then **confirm** it works by reading cache token counts from -the response. - -## Phase 1: Shared prompt construction - -**Files:** `imas_codex/sn/benchmark.py` - -Replace custom prompt construction in `_run_model()` with the same path used by -`compose_worker()` in `workers.py`: - -```python -# Current (broken): -prompt_context = {"items": items, "ids_name": group_key, "existing_names": ..., **grammar_ctx} -prompt_text = render_prompt("sn/compose_dd", prompt_context) -messages = [{"role": "user", "content": prompt_text}] - -# Fixed: -from imas_codex.sn.context import build_compose_context -context = build_compose_context() -system_prompt = render_prompt("sn/compose_system", context) -user_context = { - "items": items, - "ids_name": group_key, - "existing_names": sorted(existing)[:200], - "cluster_context": batch.get("context", ""), - **context, -} -user_prompt = render_prompt("sn/compose_dd", user_context) -messages = [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": user_prompt}, -] -``` - -- Remove `build_grammar_context()` from benchmark.py (dead code after this change) -- Keep the function in benchmark.py only if `score_with_reviewer()` needs it -- Preserve `batch.context` (cluster_context) through extraction → benchmark - -### Extraction fix - -`_extract_candidates()` currently drops `batch.context` when converting -`ExtractionBatch` to plain dicts (line 494-500). Fix: - -```python -result.append({ - "group_key": batch.group_key, - "items": batch.items, - "existing_names": list(batch.existing_names), - "context": batch.context, # ADD THIS -}) -``` - -## Phase 2: Cache hit reporting - -**Files:** `imas_codex/sn/benchmark.py`, `imas_codex/discovery/base/llm.py` - -Add cache hit/miss tracking to benchmark output: - -1. Extend `acall_llm_structured()` return or use response metadata to detect - prompt cache hits (OpenRouter returns `usage.cache_creation_input_tokens` - and `usage.cache_read_input_tokens`) -2. Add to `ModelResult`: - ```python - cache_read_tokens: int = 0 - cache_creation_tokens: int = 0 - cache_hit_rate: float = 0.0 - ``` -3. Display in Rich table: "Cache %" column - -**Note:** This requires checking what `litellm` exposes in the response object. -The `usage` dict from OpenRouter includes cache fields. `acall_llm_structured()` -currently returns `(result, cost, tokens)` — we may need to return the full -usage dict or add a cache-specific return value. - -## Phase 3: Render system prompt once, verify caching - -The system prompt should be rendered **once** before the model loop, since it's -identical across all batches and models. This matches the build pipeline pattern -(line 152 in workers.py). - -```python -context = build_compose_context() -system_prompt = render_prompt("sn/compose_system", context) - -for model in config.models: - result = await _run_model( - model=model, - extraction_batches=extraction_batches, - config=config, - reference=REFERENCE_NAMES, - system_prompt=system_prompt, # pass pre-rendered - context=context, # for user prompt rendering - ) -``` - -## Acceptance criteria - -1. `sn benchmark` uses identical prompt construction as `sn mint` -2. System/user message split enables prompt caching -3. `cluster_context` is preserved through extraction -4. Cache hit rate is reported in benchmark output -5. Running the same benchmark twice shows cache hits on second run -6. `build_grammar_context()` removed from benchmark.py (or moved to shared location) - -## Test plan - -- Unit test: `_run_model()` constructs `[system, user]` messages -- Unit test: extraction preserves `context` field -- Integration: run benchmark with `--verbose`, confirm cache token counts in logs diff --git a/plans/features/standard-names/17-sn-lifecycle-management.md b/plans/features/standard-names/17-sn-lifecycle-management.md deleted file mode 100644 index 919571b31..000000000 --- a/plans/features/standard-names/17-sn-lifecycle-management.md +++ /dev/null @@ -1,163 +0,0 @@ -# 17: Standard Name Lifecycle Management (Reset / Clear) - -**Status:** Ready to implement -**Depends on:** None (standalone) -**Agent:** engineer - -## Problem - -All discovery domains (signals, paths, wiki, code, documents) have `--reset-to` -infrastructure via the shared `ResetSpec` / `reset_to_status()` pattern. Standard -names have none. This makes iterative development painful — after changing prompts -or models, there's no way to re-run the pipeline without manual Cypher cleanup. - -### Key difference from other domains - -StandardName nodes are **cross-facility** — they have no `facility_id` property. -The standard `reset_to_status()` function requires a facility parameter. We need -an adapted approach. - -### StandardName lifecycle - -``` -drafted → published → accepted - ↘ rejected - ↘ skipped -``` - -## Phase 1: Scoped reset command - -**Files:** -- `imas_codex/sn/graph_ops.py` — add `reset_standard_names()` and `clear_standard_names()` -- `imas_codex/cli/sn.py` — add `sn reset` subcommand - -### `sn reset` command - -```bash -# Reset all drafted names back for re-composition (clears LLM output, keeps nodes) -imas-codex sn reset --status drafted - -# Reset published names back to drafted (e.g., after prompt change) -imas-codex sn reset --status published --to drafted - -# Reset only DD-sourced names -imas-codex sn reset --status drafted --source dd - -# Reset names from a specific IDS -imas-codex sn reset --status drafted --ids equilibrium - -# Dry run -imas-codex sn reset --status drafted --dry-run -``` - -### Graph operation - -```python -def reset_standard_names( - *, - target_status: str = "drafted", - source_statuses: list[str] | None = None, - source_filter: str | None = None, # "dd" or "signals" - ids_filter: str | None = None, - clear_embeddings: bool = True, -) -> int: - """Reset StandardName nodes to a target review_status. - - Unlike facility-scoped domains, StandardName has no facility_id. - Filtering is by review_status, source, and IDS. - """ -``` - -Fields to clear on reset to `drafted`: -- `embedding`, `embedded_at` (will be regenerated) -- `model`, `generated_at` (provenance of old generation) -- `confidence` (will be re-scored) - -Fields to preserve: -- `id` (the standard name itself) -- `source`, `source_path` (how it was sourced) -- `created_at` (first creation time) - -Relationships to clean: -- `HAS_STANDARD_NAME` — remove (will be re-created on persist) -- `CANONICAL_UNITS` — remove (will be re-created) - -## Phase 2: Clear command - -**Files:** `imas_codex/sn/graph_ops.py`, `imas_codex/cli/sn.py` - -```bash -# Delete all drafted standard names (not accepted/imported ones) -imas-codex sn clear --status drafted - -# Delete ALL standard names (requires --confirm) -imas-codex sn clear --all --confirm - -# Delete names from a specific source -imas-codex sn clear --source dd --status drafted -``` - -### Graph operation - -```python -def clear_standard_names( - *, - status_filter: list[str] | None = None, - source_filter: str | None = None, - ids_filter: str | None = None, - confirm_all: bool = False, -) -> int: - """Delete StandardName nodes and their relationships. - - Refuses to delete accepted/imported names unless confirm_all=True. - """ -``` - -Safety rules: -- Default: only delete `drafted` names -- `accepted` and `imported` names require `--confirm` flag -- Always log count before deletion -- DETACH DELETE (removes all relationships) - -## Phase 3: Wire into build command - -**Files:** `imas_codex/cli/sn.py` - -Add `--reset-to` option to `sn mint`: - -```bash -# Re-compose all drafted names from scratch (reset + build) -imas-codex sn mint --source dd --reset-to drafted - -# Full pipeline from extraction (clears and rebuilds) -imas-codex sn mint --source dd --reset-to extracted --ids equilibrium -``` - -Reset targets for `sn mint`: -- `extracted` — clear all SN nodes for matching source, re-run full pipeline -- `drafted` — reset existing drafted names, re-compose them - -```python -@click.option( - "--reset-to", - type=click.Choice(["extracted", "drafted"]), - default=None, - help="Reset standard names to target state before building.", -) -``` - -## Acceptance criteria - -1. `sn reset --status drafted` resets nodes and clears embeddings -2. `sn clear --status drafted` deletes only drafted names -3. `sn clear --all` requires `--confirm` flag -4. `sn mint --reset-to drafted` resets then rebuilds -5. Accepted/imported names are never touched without explicit confirmation -6. `sn status` shows correct counts after reset/clear - -## Test plan - -- Unit test: `reset_standard_names()` clears correct fields -- Unit test: `clear_standard_names()` refuses to delete accepted without confirm -- Unit test: `--reset-to` on build triggers reset before pipeline -- Integration: build → reset → rebuild cycle produces valid results diff --git a/plans/features/standard-names/18-benchmark-calibration.md b/plans/features/standard-names/18-benchmark-calibration.md deleted file mode 100644 index 329921b98..000000000 --- a/plans/features/standard-names/18-benchmark-calibration.md +++ /dev/null @@ -1,276 +0,0 @@ -# 18: Benchmark Calibration, Model Selection & Reviewer Enhancement - -**Status:** Ready to implement (after Plan 16) -**Depends on:** 16 (benchmark parity — results meaningless until prompts match build) -**Agent:** architect (requires research + design decisions) - -## Problem - -The benchmark needs improvements in three areas before we can make grounded model -selection decisions for production minting: - -1. **Gold reference set is too small** — 30 entries from 4 IDSs -2. **No calibration dataset** — reviewer has no full-entry examples to anchor scores -3. **Reviewer prompt is ad-hoc** — inline string, no grammar context, no system prompt -4. **No structured approach to model selection** — we need cost/speed/quality tradeoffs - -## Design Principle: Separate Gold Set from Calibration Set - -The rubber-duck critique caught this: the reference set and calibration set serve -different purposes and should not be mixed. - -| Dataset | Purpose | Contents | Location | -|---------|---------|----------|----------| -| **Gold reference** | Exact match: did model produce the right name? | `source_path → expected_name + fields` | `benchmark_reference.py` | -| **Calibration set** | Score anchoring: what does "outstanding" look like? | Full entries with expected tier + score | `benchmark_calibration.yaml` | - -## Phase 1: Expand gold reference set - -**Files:** `imas_codex/sn/benchmark_reference.py` - -Expand from 30 to ~50 entries covering more IDSs: - -| IDS | Current | Target | -|-----|---------|--------| -| equilibrium | 18 | 20 | -| core_profiles | 6 | 10 | -| magnetics | 4 | 6 | -| summary | 2 | 4 | -| core_transport | 0 | 4 | -| mhd_linear | 0 | 2 | -| nbi | 0 | 2 | -| edge_profiles | 0 | 2 | - -Fix the questionable rogowski_coil reference entry (line 116-121 in -benchmark_reference.py — maps major_radius position to rogowski_coil object, -which is not physically meaningful). - -## Phase 2: Create calibration dataset - -**Files:** `imas_codex/sn/benchmark_calibration.yaml` (new) - -Create ~15 hand-crafted full entries spanning quality tiers. Source material from -the `imas-standard-names` package's `resources/standard_name_examples/` directory -(40+ curated examples available). - -### Structure - -```yaml -# Each entry is a complete standard name with expected quality assessment -entries: - - name: electron_temperature - tier: outstanding - expected_score: 95 - description: "Temperature of the electron population." - documentation: > - Electron temperature $T_e$ is a fundamental plasma parameter... - Typical range: 0.1–30 keV in tokamak core... - unit: eV - kind: scalar - tags: [core_profiles, equilibrium] - fields: - physical_base: temperature - subject: electron - reason: > - Canonical physics quantity. Rich documentation with LaTeX, typical - values, cross-references. Perfect grammar decomposition. - - - name: banana_orbits - tier: poor - expected_score: 15 - description: "Banana orbits" - documentation: "" - unit: null - kind: metadata - tags: [] - fields: - physical_base: banana_orbits - reason: > - Not a measurable quantity. No documentation. "banana_orbits" is - not a valid physical_base token. No unit possible. -``` - -### Tier distribution (15 entries) - -| Tier | Count | Score Range | Examples | -|------|-------|-------------|----------| -| outstanding | 4 | 85-100 | electron_temperature, plasma_current, safety_factor, poloidal_magnetic_flux | -| good | 4 | 60-79 | toroidal_component_of_magnetic_field, electron_pressure, loop_voltage, stored_energy | -| adequate | 4 | 40-59 | minor_radius, aspect_ratio, time, elongation | -| poor | 3 | 0-39 | banana_orbits, h_mode, (invalid grammar example) | - -Source the "outstanding" and "good" entries from `imas-standard-names` examples -where available. Hand-craft the "poor" examples to test error detection. - -## Phase 3: Enhance reviewer prompt - -**Files:** -- `imas_codex/llm/prompts/sn/review_benchmark.md` (new template) -- `imas_codex/sn/benchmark.py` — update `score_with_reviewer()` - -Replace the inline rubric string with a proper Jinja2 template: - -```markdown ---- -name: sn/review_benchmark -description: Quality scoring for benchmark entries -used_by: imas_codex.sn.benchmark.score_with_reviewer -schema_needs: [] ---- - -You are a physics nomenclature expert evaluating standard name entries. - -## Grammar Rules -{{ canonical_pattern }} -{{ segment_order }} - -## Scoring Rubric -... - -## Calibration Examples - -{% for entry in calibration_entries %} -### {{ entry.name }} — {{ entry.tier }} ({{ entry.expected_score }}/100) -{{ entry.reason }} -{% endfor %} - -## Entries to Review -``` - -Key improvements: -- System/user message split (enables caching when reviewer runs multiple batches) -- Full calibration entries as anchors (not just names) -- Grammar rules included so reviewer can validate grammar correctness -- Structured rubric with specific scoring dimensions - -### Scoring dimensions (each 0-20, sum = 0-100) - -1. **Grammar correctness** (0-20): Valid parse, correct segment usage -2. **Semantic accuracy** (0-20): Name correctly describes the physics quantity -3. **Documentation quality** (0-20): LaTeX, typical values, cross-references -4. **Naming conventions** (0-20): Follows established patterns, consistent with peers -5. **Entry completeness** (0-20): Unit, kind, tags, links, ids_paths populated - -Update `QualityReview` model: - -```python -class QualityReview(BaseModel): - name: str - quality_tier: str - score: int = Field(ge=0, le=100) - grammar_score: int = Field(ge=0, le=20) - semantic_score: int = Field(ge=0, le=20) - documentation_score: int = Field(ge=0, le=20) - convention_score: int = Field(ge=0, le=20) - completeness_score: int = Field(ge=0, le=20) - reasoning: str -``` - -## Phase 4: Model selection framework - -### Candidate models for production minting - -Based on the SN pipeline requirements (structured output, grammar adherence, -physics knowledge, documentation quality), evaluate: - -**Compose phase (language model):** -| Model | Strengths | Concerns | -|-------|-----------|----------| -| `claude-sonnet-4` | Strong structured output, good physics | Cost | -| `claude-sonnet-4-5` | Previous gen, well-tested | Superseded | -| `gpt-4o` | Fast, good structured output | Physics depth | -| `gemini-2.5-flash` | Very fast, cheap | Grammar adherence unknown | -| `gpt-5.1` | Newest, potentially strongest | Cost, untested | - -**Review phase (reasoning model):** -| Model | Strengths | Concerns | -|-------|-----------|----------| -| `claude-opus-4-6` | Deepest reasoning | Expensive | -| `o4-mini` | Good reasoning, cheaper | Structured output reliability | -| `claude-sonnet-4` | Good balance | May agree with itself if same family | - -**Reviewer (benchmark scoring):** -| Model | Recommendation | -|-------|----------------| -| `claude-opus-4-6` | Best for calibration alignment | - -### Benchmark execution plan - -After Plans 16 is implemented: - -```bash -# Phase A: Validate prompt caching works -imas-codex sn benchmark \ - --models openrouter/anthropic/claude-sonnet-4 \ - --ids equilibrium --max-candidates 20 -v -# Check logs for cache_read_input_tokens > 0 on second batch - -# Phase B: Multi-model comparison (small) -imas-codex sn benchmark \ - --models openrouter/anthropic/claude-sonnet-4,openrouter/google/gemini-2.5-flash-preview,openrouter/openai/gpt-4o \ - --ids equilibrium --max-candidates 30 \ - --reviewer-model openrouter/anthropic/claude-opus-4-6 - -# Phase C: Winner deep-dive (larger) -imas-codex sn benchmark \ - --models , \ - --max-candidates 80 \ - --reviewer-model openrouter/anthropic/claude-opus-4-6 \ - --runs 2 # consistency check -``` - -### Decision criteria - -| Metric | Weight | Threshold | -|--------|--------|-----------| -| Grammar valid % | Critical | ≥95% or disqualify | -| Fields consistent % | High | ≥85% | -| Reference recall | High | ≥60% | -| Avg quality score | High | ≥65 | -| Cost per name | Medium | <$0.01 preferred | -| Names/min | Medium | >30 preferred | -| Cache hit rate | Low | >50% confirms caching works | - -## Phase 5: Token caching verification - -Prompt caching is critical for cost efficiency at scale. After Plan 16 lands: - -1. Run benchmark with a single model and 2+ batches -2. Check DEBUG logs for `cache_read_input_tokens` and `cache_creation_input_tokens` -3. First batch should show `cache_creation_input_tokens > 0` -4. Second batch should show `cache_read_input_tokens > 0` -5. If no cache fields in response, check: - - Model string has `openrouter/` prefix (required for cache_control passthrough) - - System message has `cache_control: {"type": "ephemeral"}` breakpoint - - LiteLLM proxy is forwarding cache_control blocks - -## Acceptance criteria - -1. Gold reference expanded to 50+ entries across 8+ IDSs -2. Calibration dataset with 15 full entries across 4 quality tiers -3. Reviewer prompt uses Jinja2 template with system/user split -4. Reviewer produces 5-dimensional scores (grammar, semantic, docs, convention, completeness) -5. Benchmark results inform model selection with clear winner for compose + review -6. Token caching verified working end-to-end - -## On renaming `build` → `mint` - -**Decision: Not yet.** The rubber-duck critique agrees — the pipeline needs to -stabilize first. Revisit after the benchmark shows consistent, high-quality output -across model changes. When ready, add `mint` as the primary command name and keep -`build` as a hidden alias. - -## Documentation updates - -| Target | Update | -|--------|--------| -| `AGENTS.md` | Update SN benchmark section with new calibration workflow | -| `plans/features/standard-names/00-implementation-order.md` | Add plans 16-18 | - -## Test plan - -- Unit test: calibration YAML loads and validates -- Unit test: reviewer prompt renders with calibration entries -- Unit test: QualityReview model accepts 5-dimensional scores -- Unit test: expanded reference set passes grammar round-trip -- Integration: benchmark runs with reviewer and produces scored report