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/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 8cfbb70e5..a3a5a96a2 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -1,8 +1,5 @@ name: Benchmark -env: - FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true - on: push: tags: ["v*"] @@ -54,7 +51,7 @@ jobs: steps: # ── Setup ────────────────────────────────────────────────────────── - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: fetch-depth: 0 @@ -94,7 +91,7 @@ jobs: df -h - name: Install UV + ASV - uses: astral-sh/setup-uv@v5 + uses: astral-sh/setup-uv@v8.0.0 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 bb3644279..0046463b3 100644 --- a/.github/workflows/docker-build-push.yml +++ b/.github/workflows/docker-build-push.yml @@ -11,11 +11,12 @@ 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 }} - IMAS_DD_VERSION: "4.1.0" + # 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 # 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. @@ -52,7 +53,7 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Install oras CLI run: | @@ -207,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.0.0 with: enable-cache: true github-token: ${{ secrets.GITHUB_TOKEN }} @@ -257,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 @@ -270,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 @@ -316,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 @@ -436,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 @@ -528,7 +529,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 }} @@ -555,7 +556,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 }} @@ -583,7 +583,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 }} @@ -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..0883540c6 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.0.0 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..923f12eea 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -1,8 +1,5 @@ name: Release -env: - FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true - on: push: tags: @@ -23,12 +20,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.0.0 with: github-token: ${{ secrets.GITHUB_TOKEN }} @@ -90,7 +87,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..5b73ac0bf 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -1,8 +1,5 @@ name: Test -env: - FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true - on: push: branches: [main] @@ -23,7 +20,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 +34,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 +48,7 @@ jobs: run: playwright install --with-deps chromium - name: Install uv - uses: astral-sh/setup-uv@v5 + uses: astral-sh/setup-uv@v8.0.0 with: enable-cache: true github-token: ${{ secrets.GITHUB_TOKEN }} @@ -85,7 +82,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 +113,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.0.0 with: enable-cache: true github-token: ${{ secrets.GITHUB_TOKEN }} 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/.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 diff --git a/AGENTS.md b/AGENTS.md index a688e127d..20a1e5ec8 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. @@ -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: @@ -650,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 @@ -689,18 +693,95 @@ 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 +## Standard Names + +### CLI Commands + +| Command | Purpose | Key Options | +|---------|---------|-------------| +| `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 + +``` +drafted → published → accepted + ↘ rejected +``` + +- **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: + +- **`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`. @@ -739,11 +820,11 @@ 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, 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 d8d07a13d..f1947b3df 100644 --- a/hatch_build_hooks.py +++ b/hatch_build_hooks.py @@ -92,50 +92,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" @@ -281,14 +237,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/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/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) diff --git a/imas_codex/cli/sn.py b/imas_codex/cli/sn.py index e35487465..984fdb89d 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,17 @@ 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( +@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, domain_filter: str | None, @@ -92,19 +102,41 @@ def sn_build( skip_review: bool, verbose: bool, quiet: bool, + reset_to: str | None, ) -> 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: 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() @@ -149,7 +181,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 +219,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, @@ -290,6 +322,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 +339,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 +377,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 +390,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 +615,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: @@ -596,3 +647,286 @@ def sn_publish( console.print( f" [yellow]Would create PR for {batch.group_key}[/yellow]" ) + + +@sn.command("import") +@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( + "--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_dir: str, + tags: str | None, + dry_run: bool, + check_mode: 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 + Use --check to compare catalog vs graph without importing. + + \b + Examples: + 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 + + if verbose: + logging.basicConfig(level=logging.DEBUG) + + 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: + 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.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]: + 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") + + +@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/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/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/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/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/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/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/llm/search_formatters.py b/imas_codex/llm/search_formatters.py index b122b2c44..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: @@ -931,11 +934,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 +996,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]}") @@ -1287,75 +1296,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']}" @@ -1376,6 +1317,14 @@ def format_ids_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( @@ -1403,12 +1352,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 +1379,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 +1424,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", "") @@ -1597,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 472bdddae..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: @@ -1139,30 +1142,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 +1164,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 +1182,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 +1192,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 +1220,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 +1558,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 +1632,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, @@ -1950,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. - return schema_for(task=scope) + 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) if not self.read_only: # ===================================================================== @@ -2516,11 +2514,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(): @@ -2532,7 +2525,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, ) ) @@ -2716,33 +2710,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 +2817,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 +2827,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 +2843,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( @@ -3106,6 +3034,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..da7b203b0 --- /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/models/result_models.py b/imas_codex/models/result_models.py index a98735a16..b5b6a2703 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) @@ -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/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..b63dc389b 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 @@ -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: >- @@ -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/imas_codex/schemas/standard_name.yaml b/imas_codex/schemas/standard_name.yaml index 43511b580..a691b36e0 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,50 @@ 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 + imas_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.) + physics_domain: + 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`` 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/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/sn/benchmark.py b/imas_codex/sn/benchmark.py index 37997c383..f7540875e 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,15 @@ 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 + # Prompt-cache statistics + cache_read_tokens: int = 0 + cache_creation_tokens: int = 0 @dataclass @@ -105,31 +116,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 # --------------------------------------------------------------------------- @@ -251,6 +237,118 @@ def compare_to_reference( return overlap, ref_total, precision, recall +# --------------------------------------------------------------------------- +# Quality tier labels +# --------------------------------------------------------------------------- + + +def load_calibration_entries() -> list[dict]: + """Load calibration entries from benchmark_calibration.yaml. + + 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 + + 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, + calibration_entries: list[dict], +) -> list[dict]: + """Score candidates using a reviewer model with 5-dimensional scoring. + + 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="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] + + # 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] + + # Build per-batch user prompt with candidate details + batch_items = [] + for c in batch: + 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", {}), + } + ) + + 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( + 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 +393,12 @@ async def run_benchmark( # --- 2. Run each model --- results: list[ModelResult] = [] + 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( @@ -302,9 +406,57 @@ async def run_benchmark( extraction_batches=extraction_batches, config=config, reference=REFERENCE_NAMES, + system_prompt=system_prompt, + context=context, ) results.append(model_result) + # --- 2b. Reviewer scoring (optional) --- + if config.reviewer_model: + calibration_entries = load_calibration_entries() + for result in results: + if result.candidates: + reviews = await score_with_reviewer( + result.candidates, + config.reviewer_model, + calibration_entries, + ) + 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, @@ -319,7 +471,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 @@ -337,6 +489,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 @@ -363,12 +516,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] = [] @@ -383,32 +537,50 @@ 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( + 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, + cost, + tokens, + ) # Collect candidates for c in llm_result.candidates: @@ -486,6 +658,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, @@ -500,7 +675,12 @@ 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") + table.add_column("Avg Doc Len", justify="right") + table.add_column("Fields Pop%", justify="right") for r in report.results: n = len(r.candidates) @@ -512,9 +692,15 @@ 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" - table.add_row( + row_data = [ r.model, str(n), valid_pct, @@ -523,15 +709,66 @@ def render_comparison_table(report: BenchmarkReport) -> None: cost_str, speed_str, cpn_str, + cache_pct, 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_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_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}. diff --git a/imas_codex/sn/catalog_import.py b/imas_codex/sn/catalog_import.py new file mode 100644 index 000000000..7e69f2930 --- /dev/null +++ b/imas_codex/sn/catalog_import.py @@ -0,0 +1,476 @@ +"""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 +import subprocess +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) + 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]: + """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]], + catalog_commit_sha: str | None = None, +) -> 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 + + # 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( + """ + 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.catalog_commit_sha = b.catalog_commit_sha, + 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) + # 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( + 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, 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/imas_codex/sn/graph_ops.py b/imas_codex/sn/graph_ops.py index 6c74b63e2..c4084e043 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``, ``imas_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.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), + 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, + "imas_paths": n.get("imas_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 @@ -249,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 ---------- @@ -266,32 +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. """ - from imas_codex.graph.client import GraphClient - 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 @@ -299,19 +334,328 @@ 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 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. + + 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 412a64678..f5d023085 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): @@ -48,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="candidate", description="Entry status") + 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/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/imas_codex/sn/publish.py b/imas_codex/sn/publish.py index 97823c678..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="candidate", + 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/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..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() @@ -657,7 +778,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 +802,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/imas_codex/tools/graph_search.py b/imas_codex/tools/graph_search.py index de6d6cec3..f2b361bc3 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,12 @@ 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: + 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] @@ -596,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, @@ -620,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 @@ -996,6 +1005,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: @@ -1050,6 +1060,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 ] @@ -1082,25 +1093,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,55 +1147,14 @@ 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"], "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"]) @@ -1205,28 +1173,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 +1192,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 +1827,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 +1936,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( - 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( + @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]: - """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 +1986,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 +2009,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) @@ -2243,9 +2041,25 @@ async def get_ids_structure( **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) - 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 +2071,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 +2080,17 @@ 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 [])}, + "lifecycle_distribution": { + r["status"]: r["count"] for r in (lifecycle_dist 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 +2158,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 +2206,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/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 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) diff --git a/plans/features/standard-names/00-implementation-order.md b/plans/features/standard-names/00-implementation-order.md index 323417bfe..90ee720ac 100644 --- a/plans/features/standard-names/00-implementation-order.md +++ b/plans/features/standard-names/00-implementation-order.md @@ -46,7 +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 | 📋 Ready | 11 (embedding) | — | +| 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~~ | 🗑️ 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 @@ -77,9 +82,19 @@ 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 (Plan 19) -- End-to-end: build → publish → review → import +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 + +- 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/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). 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) | 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/pyproject.toml b/pyproject.toml index ee95fd2da..45b22b869 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -155,6 +155,7 @@ dev = [ # --- Serve (embedding + LLM proxy) --- "fastapi>=0.115.0", "uvicorn>=0.31.1", + "imas-standard-names>=0.7.0rc2", ] [project.urls] @@ -432,7 +433,9 @@ torch = [ { index = "pytorch-gpu", extra = "gpu" }, ] + [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 c46cbbc04..f687caebe 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/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/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.""" 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/sn/conftest.py b/tests/sn/conftest.py new file mode 100644 index 000000000..e15ebb9cb --- /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"], + "imas_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_benchmark.py b/tests/sn/test_benchmark.py index 3c0157706..f09546b68 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 @@ -661,3 +761,486 @@ 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 TestCalibrationDataset: + """Test benchmark calibration dataset.""" + + def test_calibration_loads(self): + from imas_codex.sn.benchmark import load_calibration_entries + + entries = load_calibration_entries() + assert isinstance(entries, list) + assert len(entries) == 15, f"Expected 15 entries, got {len(entries)}" + + def test_calibration_tiers(self): + from imas_codex.sn.benchmark import load_calibration_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_calibration_required_keys(self): + from imas_codex.sn.benchmark import load_calibration_entries + + 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 + + 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 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.""" + + 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 + + +# ----------------------------------------------------------------------- +# 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 diff --git a/tests/sn/test_catalog_import.py b/tests/sn/test_catalog_import.py new file mode 100644 index 000000000..bcf94b3cf --- /dev/null +++ b/tests/sn/test_catalog_import.py @@ -0,0 +1,1253 @@ +"""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 + + +# ============================================================================= +# 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 + + +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." diff --git a/tests/sn/test_graph_ops.py b/tests/sn/test_graph_ops.py new file mode 100644 index 000000000..740d06cf8 --- /dev/null +++ b/tests/sn/test_graph_ops.py @@ -0,0 +1,560 @@ +"""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.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 + 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": [], + "imas_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["imas_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 + + +# ============================================================================= +# 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 diff --git a/tests/sn/test_integration.py b/tests/sn/test_integration.py new file mode 100644 index 000000000..bdbb24d9d --- /dev/null +++ b/tests/sn/test_integration.py @@ -0,0 +1,1040 @@ +"""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). +6. Full E2E lifecycle: write_standard_names → get_validated_standard_names + → graph_records_to_entries → generate_catalog_files → import_catalog. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock, call, patch + +import pytest +import yaml + +imas_sn = pytest.importorskip("imas_standard_names") + +# ============================================================================= +# 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}" + ) + + +# ============================================================================= +# 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" diff --git a/tests/sn/test_publish.py b/tests/sn/test_publish.py index 73be92b30..bfb79ac35 100644 --- a/tests/sn/test_publish.py +++ b/tests/sn/test_publish.py @@ -43,10 +43,10 @@ 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="candidate", + status="drafted", description="Electron temperature profile", provenance=sample_provenance, ) @@ -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.status == "candidate" + 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,9 +204,9 @@ 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"] == "candidate" + assert doc["status"] == "drafted" assert doc["description"] == "Electron temperature profile" assert doc["provenance"]["source"] == "dd" assert doc["provenance"]["confidence"] == 0.95 @@ -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"}] 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 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 diff --git a/uv.lock b/uv.lock index ed77a5c55..ce808559c 100644 --- a/uv.lock +++ b/uv.lock @@ -37,6 +37,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9f/d2/c581486aa6c4fbd7394c23c47b83fa1a919d34194e16944241daf9e762dd/accelerate-1.12.0-py3-none-any.whl", hash = "sha256:3e2091cd341423207e2f084a6654b1efcd250dc326f2a37d6dde446e07cabb11", size = 380935, upload-time = "2025-11-21T11:27:44.522Z" }, ] +[[package]] +name = "ag-ui-protocol" +version = "0.1.15" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pydantic" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/57/71/96c21ae7e2fb9b610c1a90d38bd2de8b6e5b2900a63001f3882f43e519af/ag_ui_protocol-0.1.15.tar.gz", hash = "sha256:5e23c1042c7d4e364d685e68d2fb74d37c16bc83c66d270102d8eaedce56ad82", size = 6269, upload-time = "2026-04-01T15:44:33.136Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e4/a0/a73398d30bb0f9ad70cd70426151a4a19527a7296e48a3a16a50e1d5db05/ag_ui_protocol-0.1.15-py3-none-any.whl", hash = "sha256:85cde077023ccbc37b5ce2ad953537883c262d210320f201fc2ec4e85408b06a", size = 8661, upload-time = "2026-04-01T15:44:32.079Z" }, +] + [[package]] name = "aiofile" version = "3.9.0" @@ -132,6 +144,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" }, ] +[[package]] +name = "anthropic" +version = "0.92.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "distro" }, + { name = "docstring-parser" }, + { name = "httpx" }, + { name = "jiter" }, + { name = "pydantic" }, + { name = "sniffio" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/01/2d/fc5c5a369db977efbaa646d77ba42b38a6de4e95789884032b0e2e3fc834/anthropic-0.92.0.tar.gz", hash = "sha256:d1e792ed0692379452a1af6b266df495e973c3695cd0aace2a108b838393cbc4", size = 652420, upload-time = "2026-04-08T16:55:35.37Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c3/21/bf5b5ab10b6932c5c43eaa66b6e3f256de569cf0323d89f9cc281a0d0f39/anthropic-0.92.0-py3-none-any.whl", hash = "sha256:f92a4bd065d5cab90a96b65bb44e473bf7c6fe731a743cd156e9ad1d245c381e", size = 621195, upload-time = "2026-04-08T16:55:33.639Z" }, +] + [[package]] name = "antlr4-python3-runtime" version = "4.9.3" @@ -160,6 +191,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/81/29/5ecc3a15d5a33e31b26c11426c45c501e439cb865d0bff96315d86443b78/appnope-0.1.4-py2.py3-none-any.whl", hash = "sha256:502575ee11cd7a28c0205f379b525beefebab9d161b7c964670864014ed7213c", size = 4321, upload-time = "2024-02-06T09:43:09.663Z" }, ] +[[package]] +name = "argcomplete" +version = "3.6.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/38/61/0b9ae6399dd4a58d8c1b1dc5a27d6f2808023d0b5dd3104bb99f45a33ff6/argcomplete-3.6.3.tar.gz", hash = "sha256:62e8ed4fd6a45864acc8235409461b72c9a28ee785a2011cc5eb78318786c89c", size = 73754, upload-time = "2025-10-20T03:33:34.741Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/74/f5/9373290775639cb67a2fce7f629a1c240dce9f12fe927bc32b2736e16dfc/argcomplete-3.6.3-py3-none-any.whl", hash = "sha256:f5007b3a600ccac5d25bbce33089211dfd49eab4a7718da3f10e3082525a92ce", size = 43846, upload-time = "2025-10-20T03:33:33.021Z" }, +] + [[package]] name = "arrow" version = "1.4.0" @@ -311,6 +351,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1a/39/47f9197bdd44df24d67ac8893641e16f386c984a0619ef2ee4c51fbbc019/beautifulsoup4-4.14.3-py3-none-any.whl", hash = "sha256:0918bfe44902e6ad8d57732ba310582e98da931428d231a5ecb9e7c703a735bb", size = 107721, upload-time = "2025-11-30T15:08:24.087Z" }, ] +[[package]] +name = "boto3" +version = "1.42.85" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "botocore" }, + { name = "jmespath" }, + { name = "s3transfer" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/31/9d/a9a7b5a9351e3ff0baae01136f71ba6fc4652fe0dc2da3b0a8ebdfc1be44/boto3-1.42.85.tar.gz", hash = "sha256:1cd3dcbfaba85c6071ba9397c1804b6a94a1a97031b8f1993fdba27c0c5d6eba", size = 112769, upload-time = "2026-04-07T19:40:53.834Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/78/ab/3167b8ec3cf1d87ad08d2ad5f15823a22945cae7870798274c283c3a18f1/boto3-1.42.85-py3-none-any.whl", hash = "sha256:4f6ac066e41d18ec33f532253fac0f35e0fdca373724458f983ce3d531340b7a", size = 140556, upload-time = "2026-04-07T19:40:52.186Z" }, +] + +[[package]] +name = "botocore" +version = "1.42.85" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jmespath" }, + { name = "python-dateutil" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0a/ac/7f14b05cf43e4baae99f4570b02e10b2aebf242dfd86245523340390c834/botocore-1.42.85.tar.gz", hash = "sha256:2ee61f80b7724a143e16d0a85408ef5fa20b99dce7a3c8ec5d25cc8dced164c1", size = 15159562, upload-time = "2026-04-07T19:40:43.831Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/16/f3/c1fbaff4c509c616fd01f44357283a8992f10b3a05d932b22e602aa3a221/botocore-1.42.85-py3-none-any.whl", hash = "sha256:828b67722caeb7e240eefedee74050e803d1fa102958ead9c4009101eefd5381", size = 14839741, upload-time = "2026-04-07T19:40:40.733Z" }, +] + [[package]] name = "build" version = "1.4.0" @@ -443,6 +511,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/98/78/01c019cdb5d6498122777c1a43056ebb3ebfeef2076d9d026bfe15583b2b/click-8.3.1-py3-none-any.whl", hash = "sha256:981153a64e25f12d547d3426c367a4857371575ee7ad18df2a6183ab0545b2a6", size = 108274, upload-time = "2025-11-15T20:45:41.139Z" }, ] +[[package]] +name = "cohere" +version = "5.21.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "fastavro" }, + { name = "httpx" }, + { name = "pydantic" }, + { name = "pydantic-core" }, + { name = "requests" }, + { name = "tokenizers" }, + { name = "types-requests" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d2/75/4c346f6e2322e545f8452692304bd4eca15a2a0209ab9af6a0d1a7810b67/cohere-5.21.1.tar.gz", hash = "sha256:e5ade4423b928b01ff2038980e1b62b2a5bb412c8ab83e30882753b810a5509f", size = 191272, upload-time = "2026-03-26T15:09:27.857Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0a/50/5538f02ec6d10fbb84f29c1b18c68ff2a03d7877926a80275efdf8755a9f/cohere-5.21.1-py3-none-any.whl", hash = "sha256:f15592ec60d8cf12f01563db94ec28c388c61269d9617f23c2d6d910e505344e", size = 334262, upload-time = "2026-03-26T15:09:26.284Z" }, +] + [[package]] name = "colorama" version = "0.4.6" @@ -631,6 +718,17 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/02/10/5da547df7a391dcde17f59520a231527b8571e6f46fc8efb02ccb370ab12/docutils-0.22.4-py3-none-any.whl", hash = "sha256:d0013f540772d1420576855455d050a2180186c91c15779301ac2ccb3eeb68de", size = 633196, upload-time = "2025-12-18T19:00:18.077Z" }, ] +[[package]] +name = "dotenv" +version = "0.9.9" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "python-dotenv" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/b2/b7/545d2c10c1fc15e48653c91efde329a790f2eecfbbf2bd16003b5db2bab0/dotenv-0.9.9-py2.py3-none-any.whl", hash = "sha256:29cf74a087b31dafdb5a446b6d7e11cbce8ed2741540e2339c69fbef92c94ce9", size = 1892, upload-time = "2025-02-19T22:15:01.647Z" }, +] + [[package]] name = "email-validator" version = "2.3.0" @@ -653,6 +751,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c1/8b/5fe2cc11fee489817272089c4203e679c63b570a5aaeb18d852ae3cbba6a/et_xmlfile-2.0.0-py3-none-any.whl", hash = "sha256:7a91720bc756843502c3b7504c77b8fe44217c85c537d85037f0f536151b2caa", size = 18059, upload-time = "2024-10-25T17:25:39.051Z" }, ] +[[package]] +name = "eval-type-backport" +version = "0.3.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fb/a3/cafafb4558fd638aadfe4121dc6cefb8d743368c085acb2f521df0f3d9d7/eval_type_backport-0.3.1.tar.gz", hash = "sha256:57e993f7b5b69d271e37482e62f74e76a0276c82490cf8e4f0dffeb6b332d5ed", size = 9445, upload-time = "2025-12-02T11:51:42.987Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cf/22/fdc2e30d43ff853720042fa15baa3e6122722be1a7950a98233ebb55cd71/eval_type_backport-0.3.1-py3-none-any.whl", hash = "sha256:279ab641905e9f11129f56a8a78f493518515b83402b860f6f06dd7c011fdfa8", size = 6063, upload-time = "2025-12-02T11:51:41.665Z" }, +] + [[package]] name = "exceptiongroup" version = "1.3.1" @@ -714,6 +821,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9e/dd/d0ee25348ac58245ee9f90b6f3cbb666bf01f69be7e0911f9851bddbda16/fastapi-0.129.0-py3-none-any.whl", hash = "sha256:b4946880e48f462692b31c083be0432275cbfb6e2274566b1be91479cc1a84ec", size = 102950, upload-time = "2026-02-12T13:54:54.528Z" }, ] +[[package]] +name = "fastavro" +version = "1.12.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/65/8b/fa2d3287fd2267be6261d0177c6809a7fa12c5600ddb33490c8dc29e77b2/fastavro-1.12.1.tar.gz", hash = "sha256:2f285be49e45bc047ab2f6bed040bb349da85db3f3c87880e4b92595ea093b2b", size = 1025661, upload-time = "2025-10-10T15:40:55.41Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7c/f0/10bd1a3d08667fa0739e2b451fe90e06df575ec8b8ba5d3135c70555c9bd/fastavro-1.12.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:509818cb24b98a804fc80be9c5fed90f660310ae3d59382fc811bfa187122167", size = 1009057, upload-time = "2025-10-10T15:41:24.556Z" }, + { url = "https://files.pythonhosted.org/packages/78/ad/0d985bc99e1fa9e74c636658000ba38a5cd7f5ab2708e9c62eaf736ecf1a/fastavro-1.12.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:089e155c0c76e0d418d7e79144ce000524dd345eab3bc1e9c5ae69d500f71b14", size = 3391866, upload-time = "2025-10-10T15:41:26.882Z" }, + { url = "https://files.pythonhosted.org/packages/0d/9e/b4951dc84ebc34aac69afcbfbb22ea4a91080422ec2bfd2c06076ff1d419/fastavro-1.12.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:44cbff7518901c91a82aab476fcab13d102e4999499df219d481b9e15f61af34", size = 3458005, upload-time = "2025-10-10T15:41:29.017Z" }, + { url = "https://files.pythonhosted.org/packages/af/f8/5a8df450a9f55ca8441f22ea0351d8c77809fc121498b6970daaaf667a21/fastavro-1.12.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:a275e48df0b1701bb764b18a8a21900b24cf882263cb03d35ecdba636bbc830b", size = 3295258, upload-time = "2025-10-10T15:41:31.564Z" }, + { url = "https://files.pythonhosted.org/packages/99/b2/40f25299111d737e58b85696e91138a66c25b7334f5357e7ac2b0e8966f8/fastavro-1.12.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:2de72d786eb38be6b16d556b27232b1bf1b2797ea09599507938cdb7a9fe3e7c", size = 3430328, upload-time = "2025-10-10T15:41:33.689Z" }, + { url = "https://files.pythonhosted.org/packages/e0/07/85157a7c57c5f8b95507d7829b5946561e5ee656ff80e9dd9a757f53ddaf/fastavro-1.12.1-cp312-cp312-win_amd64.whl", hash = "sha256:9090f0dee63fe022ee9cc5147483366cc4171c821644c22da020d6b48f576b4f", size = 444140, upload-time = "2025-10-10T15:41:34.902Z" }, +] + [[package]] name = "fastjsonschema" version = "2.21.2" @@ -850,6 +971,70 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e6/ab/fb21f4c939bb440104cc2b396d3be1d9b7a9fd3c6c2a53d98c45b3d7c954/fsspec-2026.2.0-py3-none-any.whl", hash = "sha256:98de475b5cb3bd66bedd5c4679e87b4fdfe1a3bf4d707b151b3c07e58c9a2437", size = 202505, upload-time = "2026-02-05T21:50:51.819Z" }, ] +[[package]] +name = "genai-prices" +version = "0.0.56" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "httpx" }, + { name = "pydantic" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/44/6b/94b3018a672c7775edfb485f0fed8f6068fba75e49b067e8a1ac5eb96764/genai_prices-0.0.56.tar.gz", hash = "sha256:ac24b16a84d0ab97539bfa48dfa4649689de8e3ce71c12ebacef29efb1998045", size = 65872, upload-time = "2026-03-20T20:33:00.732Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a3/f6/8ef7e4c286deb2709d11ca96a5237caae3ef4876ab3c48095856cfd2df30/genai_prices-0.0.56-py3-none-any.whl", hash = "sha256:dbe86be8f3f556bed1b72209ed36851fec8b01793b3b220f42921a4e7da945f6", size = 68966, upload-time = "2026-03-20T20:33:02.555Z" }, +] + +[[package]] +name = "google-auth" +version = "2.49.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cryptography" }, + { name = "pyasn1-modules" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ea/80/6a696a07d3d3b0a92488933532f03dbefa4a24ab80fb231395b9a2a1be77/google_auth-2.49.1.tar.gz", hash = "sha256:16d40da1c3c5a0533f57d268fe72e0ebb0ae1cc3b567024122651c045d879b64", size = 333825, upload-time = "2026-03-12T19:30:58.135Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e9/eb/c6c2478d8a8d633460be40e2a8a6f8f429171997a35a96f81d3b680dec83/google_auth-2.49.1-py3-none-any.whl", hash = "sha256:195ebe3dca18eddd1b3db5edc5189b76c13e96f29e73043b923ebcf3f1a860f7", size = 240737, upload-time = "2026-03-12T19:30:53.159Z" }, +] + +[package.optional-dependencies] +requests = [ + { name = "requests" }, +] + +[[package]] +name = "google-genai" +version = "1.71.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "distro" }, + { name = "google-auth", extra = ["requests"] }, + { name = "httpx" }, + { name = "pydantic" }, + { name = "requests" }, + { name = "sniffio" }, + { name = "tenacity" }, + { name = "typing-extensions" }, + { name = "websockets" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/88/49/a13e9cf4d963691fc79d661f2d78f041bc1f2e7287d41ef0f831b82462f0/google_genai-1.71.0.tar.gz", hash = "sha256:044f7ac453437d5d380ec192f823dba64e001c478d7878c5a2d327432f4a28ac", size = 520044, upload-time = "2026-04-08T17:55:51.084Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/55/d4/63c97d487f0b4861a6f530628a9e56a380f9c9704d2e207f0bed9d16e31a/google_genai-1.71.0-py3-none-any.whl", hash = "sha256:6213ebfee7fc8e6a21692c2c340309e463322cc35c2c603d1ea59e8ea34ac240", size = 760561, upload-time = "2026-04-08T17:55:49.107Z" }, +] + +[[package]] +name = "googleapis-common-protos" +version = "1.74.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "protobuf" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/20/18/a746c8344152d368a5aac738d4c857012f2c5d1fd2eac7e17b647a7861bd/googleapis_common_protos-1.74.0.tar.gz", hash = "sha256:57971e4eeeba6aad1163c1f0fc88543f965bb49129b8bb55b2b7b26ecab084f1", size = 151254, upload-time = "2026-04-02T21:23:26.679Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b6/b0/be5d3329badb9230b765de6eea66b73abd5944bdeb5afb3562ddcd80ae84/googleapis_common_protos-1.74.0-py3-none-any.whl", hash = "sha256:702216f78610bb510e3f12ac3cafd281b7ac45cc5d86e90ad87e4d301a3426b5", size = 300743, upload-time = "2026-04-02T21:22:49.108Z" }, +] + [[package]] name = "graphviz" version = "0.21" @@ -876,6 +1061,53 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c8/ab/717c58343cf02c5265b531384b248787e04d8160b8afe53d9eec053d7b44/greenlet-3.3.1-cp312-cp312-win_arm64.whl", hash = "sha256:bfb2d1763d777de5ee495c85309460f6fd8146e50ec9d0ae0183dbf6f0a829d1", size = 226403, upload-time = "2026-01-23T15:31:39.372Z" }, ] +[[package]] +name = "griffelib" +version = "2.0.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9d/82/74f4a3310cdabfbb10da554c3a672847f1ed33c6f61dd472681ce7f1fe67/griffelib-2.0.2.tar.gz", hash = "sha256:3cf20b3bc470e83763ffbf236e0076b1211bac1bc67de13daf494640f2de707e", size = 166461, upload-time = "2026-03-27T11:34:51.091Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/11/8c/c9138d881c79aa0ea9ed83cbd58d5ca75624378b38cee225dcf5c42cc91f/griffelib-2.0.2-py3-none-any.whl", hash = "sha256:925c857658fb1ba40c0772c37acbc2ab650bd794d9c1b9726922e36ea4117ea1", size = 142357, upload-time = "2026-03-27T11:34:46.275Z" }, +] + +[[package]] +name = "groq" +version = "1.1.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "distro" }, + { name = "httpx" }, + { name = "pydantic" }, + { name = "sniffio" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d3/c7/a2153b639062f59f9bc93a1b5507c0c4a6b654b8a9edbf432ec2f4a62d2d/groq-1.1.2.tar.gz", hash = "sha256:9ec2b5b6a1c4856a8c6c38741353c5ab37472a4e3fded02af783750d849cc988", size = 154033, upload-time = "2026-03-25T23:16:10.313Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/34/b0/83e3892a4597a4b8ebf8a662aeaf314765c4c2340516eb1d049b459b24fc/groq-1.1.2-py3-none-any.whl", hash = "sha256:348cb7a674b6aa7105719b533f6fc48fd32b503bc9256924aaed6dc186f778b5", size = 141700, upload-time = "2026-03-25T23:16:08.998Z" }, +] + +[[package]] +name = "grpcio" +version = "1.80.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b7/48/af6173dbca4454f4637a4678b67f52ca7e0c1ed7d5894d89d434fecede05/grpcio-1.80.0.tar.gz", hash = "sha256:29aca15edd0688c22ba01d7cc01cb000d72b2033f4a3c72a81a19b56fd143257", size = 12978905, upload-time = "2026-03-30T08:49:10.502Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5c/e8/a2b749265eb3415abc94f2e619bbd9e9707bebdda787e61c593004ec927a/grpcio-1.80.0-cp312-cp312-linux_armv7l.whl", hash = "sha256:c624cc9f1008361014378c9d776de7182b11fe8b2e5a81bc69f23a295f2a1ad0", size = 6015616, upload-time = "2026-03-30T08:47:13.428Z" }, + { url = "https://files.pythonhosted.org/packages/3e/97/b1282161a15d699d1e90c360df18d19165a045ce1c343c7f313f5e8a0b77/grpcio-1.80.0-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:f49eddcac43c3bf350c0385366a58f36bed8cc2c0ec35ef7b74b49e56552c0c2", size = 12014204, upload-time = "2026-03-30T08:47:15.873Z" }, + { url = "https://files.pythonhosted.org/packages/6e/5e/d319c6e997b50c155ac5a8cb12f5173d5b42677510e886d250d50264949d/grpcio-1.80.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d334591df610ab94714048e0d5b4f3dd5ad1bee74dfec11eee344220077a79de", size = 6563866, upload-time = "2026-03-30T08:47:18.588Z" }, + { url = "https://files.pythonhosted.org/packages/ae/f6/fdd975a2cb4d78eb67769a7b3b3830970bfa2e919f1decf724ae4445f42c/grpcio-1.80.0-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:0cb517eb1d0d0aaf1d87af7cc5b801d686557c1d88b2619f5e31fab3c2315921", size = 7273060, upload-time = "2026-03-30T08:47:21.113Z" }, + { url = "https://files.pythonhosted.org/packages/db/f0/a3deb5feba60d9538a962913e37bd2e69a195f1c3376a3dd44fe0427e996/grpcio-1.80.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4e78c4ac0d97dc2e569b2f4bcbbb447491167cb358d1a389fc4af71ab6f70411", size = 6782121, upload-time = "2026-03-30T08:47:23.827Z" }, + { url = "https://files.pythonhosted.org/packages/ca/84/36c6dcfddc093e108141f757c407902a05085e0c328007cb090d56646cdf/grpcio-1.80.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2ed770b4c06984f3b47eb0517b1c69ad0b84ef3f40128f51448433be904634cd", size = 7383811, upload-time = "2026-03-30T08:47:26.517Z" }, + { url = "https://files.pythonhosted.org/packages/7c/ef/f3a77e3dc5b471a0ec86c564c98d6adfa3510d38f8ee99010410858d591e/grpcio-1.80.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:256507e2f524092f1473071a05e65a5b10d84b82e3ff24c5b571513cfaa61e2f", size = 8393860, upload-time = "2026-03-30T08:47:29.439Z" }, + { url = "https://files.pythonhosted.org/packages/9b/8d/9d4d27ed7f33d109c50d6b5ce578a9914aa68edab75d65869a17e630a8d1/grpcio-1.80.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:9a6284a5d907c37db53350645567c522be314bac859a64a7a5ca63b77bb7958f", size = 7830132, upload-time = "2026-03-30T08:47:33.254Z" }, + { url = "https://files.pythonhosted.org/packages/14/e4/9990b41c6d7a44e1e9dee8ac11d7a9802ba1378b40d77468a7761d1ad288/grpcio-1.80.0-cp312-cp312-win32.whl", hash = "sha256:c71309cfce2f22be26aa4a847357c502db6c621f1a49825ae98aa0907595b193", size = 4140904, upload-time = "2026-03-30T08:47:35.319Z" }, + { url = "https://files.pythonhosted.org/packages/2f/2c/296f6138caca1f4b92a31ace4ae1b87dab692fc16a7a3417af3bb3c805bf/grpcio-1.80.0-cp312-cp312-win_amd64.whl", hash = "sha256:9fe648599c0e37594c4809d81a9e77bd138cc82eb8baa71b6a86af65426723ff", size = 4880944, upload-time = "2026-03-30T08:47:37.831Z" }, +] + [[package]] name = "h11" version = "0.16.0" @@ -1078,6 +1310,7 @@ dev = [ { name = "fastapi" }, { name = "hdbscan" }, { name = "imas-python" }, + { name = "imas-standard-names" }, { name = "ipykernel" }, { name = "ipython" }, { name = "jellyfish" }, @@ -1166,6 +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", specifier = ">=0.7.0rc2" }, { name = "ipykernel", specifier = ">=6.29.5" }, { name = "ipython", specifier = ">=9.2.0" }, { name = "jellyfish", specifier = ">=1.2.1" }, @@ -1242,6 +1476,29 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/80/c3/51724c1ba79aa3f34566750de5a0ad41176a272b19e0585dc62aea3a987b/imas_python-2.2.0-py3-none-any.whl", hash = "sha256:52a16cd13d7756413ff918c0cf754d42ab9ac61ae2524ab7f72a9df00a70637c", size = 2405647, upload-time = "2026-02-12T15:32:16.657Z" }, ] +[[package]] +name = "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" }, + { name = "pyyaml" }, + { name = "requests" }, + { name = "strictyaml" }, + { name = "textual" }, +] +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]] name = "importlib-metadata" version = "8.7.1" @@ -1468,6 +1725,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/67/8a/a342b2f0251f3dac4ca17618265d93bf244a2a4d089126e81e4c1056ac50/jiter-0.13.0-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7bb00b6d26db67a05fe3e12c76edc75f32077fb51deed13822dc648fa373bc19", size = 343768, upload-time = "2026-02-02T12:37:55.055Z" }, ] +[[package]] +name = "jmespath" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d3/59/322338183ecda247fb5d1763a6cbe46eff7222eaeebafd9fa65d4bf5cb11/jmespath-1.1.0.tar.gz", hash = "sha256:472c87d80f36026ae83c6ddd0f1d05d4e510134ed462851fd5f754c8c3cbb88d", size = 27377, upload-time = "2026-01-22T16:35:26.279Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl", hash = "sha256:a5663118de4908c91729bea0acadca56526eb2698e83de10cd116ae0f4e97c64", size = 20419, upload-time = "2026-01-22T16:35:24.919Z" }, +] + [[package]] name = "joblib" version = "1.5.3" @@ -1520,6 +1786,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e5/90/0d93963711f811efe528e3cead2f2bfb78c196df74d8a24fe8d655288e50/jsonasobj2-1.0.4-py3-none-any.whl", hash = "sha256:12e86f86324d54fcf60632db94ea74488d5314e3da554c994fe1e2c6f29acb79", size = 6324, upload-time = "2021-06-02T17:43:27.126Z" }, ] +[[package]] +name = "jsonpath-python" +version = "1.1.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2d/db/2f4ecc24da35c6142b39c353d5b7c16eef955cc94b35a48d3fa47996d7c3/jsonpath_python-1.1.5.tar.gz", hash = "sha256:ceea2efd9e56add09330a2c9631ea3d55297b9619348c1055e5bfb9cb0b8c538", size = 87352, upload-time = "2026-03-17T06:16:40.597Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/28/50/1a313fb700526b134c71eb8a225d8b83be0385dbb0204337b4379c698cef/jsonpath_python-1.1.5-py3-none-any.whl", hash = "sha256:a60315404d70a65e76c9a782c84e50600480221d94a58af47b7b4d437351cb4b", size = 14090, upload-time = "2026-03-17T06:16:39.152Z" }, +] + [[package]] name = "jsonpointer" version = "3.0.0" @@ -1659,6 +1934,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c7/7c/c614252f9acda59b01a66e2ddfd243ed1c7e1deab0293332dfbccf862808/librt-0.8.1-cp312-cp312-win_arm64.whl", hash = "sha256:0f2ae3725904f7377e11cc37722d5d401e8b3d5851fb9273d7f4fe04f6b3d37d", size = 52441, upload-time = "2026-02-17T16:11:56.801Z" }, ] +[[package]] +name = "linkify-it-py" +version = "2.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "uc-micro-py" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/2e/c9/06ea13676ef354f0af6169587ae292d3e2406e212876a413bf9eece4eb23/linkify_it_py-2.1.0.tar.gz", hash = "sha256:43360231720999c10e9328dc3691160e27a718e280673d444c38d7d3aaa3b98b", size = 29158, upload-time = "2026-03-01T07:48:47.683Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b4/de/88b3be5c31b22333b3ca2f6ff1de4e863d8fe45aaea7485f591970ec1d3e/linkify_it_py-2.1.0-py3-none-any.whl", hash = "sha256:0d252c1594ecba2ecedc444053db5d3a9b7ec1b0dd929c8f1d74dce89f86c05e", size = 19878, upload-time = "2026-03-01T07:48:46.098Z" }, +] + [[package]] name = "linkml" version = "1.9.3" @@ -1741,6 +2028,38 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/be/f3/fffb7932870163cea7addc392165647a9a8a5489967de486c854226f1141/litellm-1.81.13-py3-none-any.whl", hash = "sha256:ae4aea2a55e85993f5f6dd36d036519422d24812a1a3e8540d9e987f2d7a4304", size = 14587505, upload-time = "2026-02-17T02:00:44.22Z" }, ] +[[package]] +name = "logfire" +version = "4.31.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "executing" }, + { name = "opentelemetry-exporter-otlp-proto-http" }, + { name = "opentelemetry-instrumentation" }, + { name = "opentelemetry-sdk" }, + { name = "protobuf" }, + { name = "rich" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/61/fc/21f923243d8c3ca2ebfa97de46970ced734e66ac634c1c35b6abb41300f1/logfire-4.31.0.tar.gz", hash = "sha256:361bfda17c9d70ada5d220211033bae06b871ddac9d5b06978bc0ceca6b8e658", size = 1080609, upload-time = "2026-03-27T19:00:46.339Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/1a/8c860e35bf847ac0d647d94bad89dccbb66cbcafdd61d8334f8cc7cfdd58/logfire-4.31.0-py3-none-any.whl", hash = "sha256:49fad38b5e6f199a98e9c8814e860c8a42595bb81479b52a20413e53ee475b72", size = 308896, upload-time = "2026-03-27T19:00:43.107Z" }, +] + +[package.optional-dependencies] +httpx = [ + { name = "opentelemetry-instrumentation-httpx" }, +] + +[[package]] +name = "logfire-api" +version = "4.31.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/08/a2/8d5a3c1c282d5f2bd9f5e9ddd5288d1414a53301ce389af9016b6d82bd50/logfire_api-4.31.0.tar.gz", hash = "sha256:fc4b01257ebd4ce297ad374ed201eb1a9213b999f6ae6df45cfca5bd0ef378f8", size = 77838, upload-time = "2026-03-27T19:00:47.545Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/26/27/9372b7492b3e146908d520f8599909311cd930175801ad219171fafc6f3e/logfire_api-4.31.0-py3-none-any.whl", hash = "sha256:3c1f502fd4eb8ef0996427a5cf275fd8f327f38600650a1f53071a8171c812db", size = 123402, upload-time = "2026-03-27T19:00:44.952Z" }, +] + [[package]] name = "lxml" version = "6.0.2" @@ -1767,6 +2086,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ea/7b/93c73c67db235931527301ed3785f849c78991e2e34f3fd9a6663ffda4c5/lxml-6.0.2-cp312-cp312-win_arm64.whl", hash = "sha256:61cb10eeb95570153e0c0e554f58df92ecf5109f75eacad4a95baa709e26c3d6", size = 3672836, upload-time = "2025-09-22T04:01:52.145Z" }, ] +[[package]] +name = "markdown" +version = "3.10.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2b/f4/69fa6ed85ae003c2378ffa8f6d2e3234662abd02c10d216c0ba96081a238/markdown-3.10.2.tar.gz", hash = "sha256:994d51325d25ad8aa7ce4ebaec003febcce822c3f8c911e3b17c52f7f589f950", size = 368805, upload-time = "2026-02-09T14:57:26.942Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/de/1f/77fa3081e4f66ca3576c896ae5d31c3002ac6607f9747d2e3aa49227e464/markdown-3.10.2-py3-none-any.whl", hash = "sha256:e91464b71ae3ee7afd3017d9f358ef0baf158fd9a298db92f1d4761133824c36", size = 108180, upload-time = "2026-02-09T14:57:25.787Z" }, +] + [[package]] name = "markdown-it-py" version = "4.0.0" @@ -1779,6 +2107,14 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl", hash = "sha256:87327c59b172c5011896038353a81343b6754500a08cd7a4973bb48c6d578147", size = 87321, upload-time = "2025-08-11T12:57:51.923Z" }, ] +[package.optional-dependencies] +linkify = [ + { name = "linkify-it-py" }, +] +plugins = [ + { name = "mdit-py-plugins" }, +] + [[package]] name = "markupsafe" version = "3.0.3" @@ -1835,6 +2171,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fd/d9/eaa1f80170d2b7c5ba23f3b59f766f3a0bb41155fbc32a69adfa1adaaef9/mcp-1.26.0-py3-none-any.whl", hash = "sha256:904a21c33c25aa98ddbeb47273033c435e595bbacfdb177f4bd87f6dceebe1ca", size = 233615, upload-time = "2026-01-24T19:40:30.652Z" }, ] +[[package]] +name = "mdit-py-plugins" +version = "0.5.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b2/fd/a756d36c0bfba5f6e39a1cdbdbfdd448dc02692467d83816dff4592a1ebc/mdit_py_plugins-0.5.0.tar.gz", hash = "sha256:f4918cb50119f50446560513a8e311d574ff6aaed72606ddae6d35716fe809c6", size = 44655, upload-time = "2025-08-11T07:25:49.083Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/86/dd6e5db36df29e76c7a7699123569a4a18c1623ce68d826ed96c62643cae/mdit_py_plugins-0.5.0-py3-none-any.whl", hash = "sha256:07a08422fc1936a5d26d146759e9155ea466e842f5ab2f7d2266dd084c8dab1f", size = 57205, upload-time = "2025-08-11T07:25:47.597Z" }, +] + [[package]] name = "mdurl" version = "0.1.2" @@ -1844,6 +2192,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, ] +[[package]] +name = "mistralai" +version = "2.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "eval-type-backport" }, + { name = "httpx" }, + { name = "jsonpath-python" }, + { name = "opentelemetry-api" }, + { name = "opentelemetry-semantic-conventions" }, + { name = "pydantic" }, + { name = "python-dateutil" }, + { name = "typing-inspection" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/9f/96/b8ab9bbdcefda9803cf3b51e11548730ac94303850028fb86c163472aac3/mistralai-2.3.1.tar.gz", hash = "sha256:02989e509124cb28aaffd92660bf7511b3f8f5c215e1de8d49d0c8276bacc72a", size = 390323, upload-time = "2026-04-07T14:49:18.38Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c3/1d/b0da235154e9c7039c27b91785ec81f01ba1b3c48092924c6770ba2da22a/mistralai-2.3.1-py3-none-any.whl", hash = "sha256:8f4f783cb7603f6060490105f55b16a5d0a7e854c05e96fed316efcc4b393fe3", size = 930912, upload-time = "2026-04-07T14:49:16.863Z" }, +] + [[package]] name = "more-itertools" version = "10.8.0" @@ -1964,6 +2331,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl", hash = "sha256:d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762", size = 2068504, upload-time = "2025-12-08T17:02:38.159Z" }, ] +[[package]] +name = "nexus-rpc" +version = "1.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/35/d5/cd1ffb202b76ebc1b33c1332a3416e55a39929006982adc2b1eb069aaa9b/nexus_rpc-1.4.0.tar.gz", hash = "sha256:3b8b373d4865671789cc43623e3dc0bcbf192562e40e13727e17f1c149050fba", size = 82367, upload-time = "2026-02-25T22:01:34.053Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/11/52/6327a5f4fda01207205038a106a99848a41c83e933cd23ea2cab3d2ebc6c/nexus_rpc-1.4.0-py3-none-any.whl", hash = "sha256:14c953d3519113f8ccec533a9efdb6b10c28afef75d11cdd6d422640c40b3a49", size = 29645, upload-time = "2026-02-25T22:01:33.122Z" }, +] + [[package]] name = "nodeenv" version = "1.10.0" @@ -2167,6 +2546,115 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cf/df/d3f1ddf4bb4cb50ed9b1139cc7b1c54c34a1e7ce8fd1b9a37c0d1551a6bd/opentelemetry_api-1.39.1-py3-none-any.whl", hash = "sha256:2edd8463432a7f8443edce90972169b195e7d6a05500cd29e6d13898187c9950", size = 66356, upload-time = "2025-12-11T13:32:17.304Z" }, ] +[[package]] +name = "opentelemetry-exporter-otlp-proto-common" +version = "1.39.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-proto" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e9/9d/22d241b66f7bbde88a3bfa6847a351d2c46b84de23e71222c6aae25c7050/opentelemetry_exporter_otlp_proto_common-1.39.1.tar.gz", hash = "sha256:763370d4737a59741c89a67b50f9e39271639ee4afc999dadfe768541c027464", size = 20409, upload-time = "2025-12-11T13:32:40.885Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8c/02/ffc3e143d89a27ac21fd557365b98bd0653b98de8a101151d5805b5d4c33/opentelemetry_exporter_otlp_proto_common-1.39.1-py3-none-any.whl", hash = "sha256:08f8a5862d64cc3435105686d0216c1365dc5701f86844a8cd56597d0c764fde", size = 18366, upload-time = "2025-12-11T13:32:20.2Z" }, +] + +[[package]] +name = "opentelemetry-exporter-otlp-proto-http" +version = "1.39.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "googleapis-common-protos" }, + { name = "opentelemetry-api" }, + { name = "opentelemetry-exporter-otlp-proto-common" }, + { name = "opentelemetry-proto" }, + { name = "opentelemetry-sdk" }, + { name = "requests" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/80/04/2a08fa9c0214ae38880df01e8bfae12b067ec0793446578575e5080d6545/opentelemetry_exporter_otlp_proto_http-1.39.1.tar.gz", hash = "sha256:31bdab9745c709ce90a49a0624c2bd445d31a28ba34275951a6a362d16a0b9cb", size = 17288, upload-time = "2025-12-11T13:32:42.029Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/95/f1/b27d3e2e003cd9a3592c43d099d2ed8d0a947c15281bf8463a256db0b46c/opentelemetry_exporter_otlp_proto_http-1.39.1-py3-none-any.whl", hash = "sha256:d9f5207183dd752a412c4cd564ca8875ececba13be6e9c6c370ffb752fd59985", size = 19641, upload-time = "2025-12-11T13:32:22.248Z" }, +] + +[[package]] +name = "opentelemetry-instrumentation" +version = "0.60b1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "opentelemetry-semantic-conventions" }, + { name = "packaging" }, + { name = "wrapt" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/41/0f/7e6b713ac117c1f5e4e3300748af699b9902a2e5e34c9cf443dde25a01fa/opentelemetry_instrumentation-0.60b1.tar.gz", hash = "sha256:57ddc7974c6eb35865af0426d1a17132b88b2ed8586897fee187fd5b8944bd6a", size = 31706, upload-time = "2025-12-11T13:36:42.515Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/77/d2/6788e83c5c86a2690101681aeef27eeb2a6bf22df52d3f263a22cee20915/opentelemetry_instrumentation-0.60b1-py3-none-any.whl", hash = "sha256:04480db952b48fb1ed0073f822f0ee26012b7be7c3eac1a3793122737c78632d", size = 33096, upload-time = "2025-12-11T13:35:33.067Z" }, +] + +[[package]] +name = "opentelemetry-instrumentation-httpx" +version = "0.60b1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "opentelemetry-instrumentation" }, + { name = "opentelemetry-semantic-conventions" }, + { name = "opentelemetry-util-http" }, + { name = "wrapt" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/86/08/11208bcfcab4fc2023252c3f322aa397fd9ad948355fea60f5fc98648603/opentelemetry_instrumentation_httpx-0.60b1.tar.gz", hash = "sha256:a506ebaf28c60112cbe70ad4f0338f8603f148938cb7b6794ce1051cd2b270ae", size = 20611, upload-time = "2025-12-11T13:37:01.661Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/59/b98e84eebf745ffc75397eaad4763795bff8a30cbf2373a50ed4e70646c5/opentelemetry_instrumentation_httpx-0.60b1-py3-none-any.whl", hash = "sha256:f37636dd742ad2af83d896ba69601ed28da51fa4e25d1ab62fde89ce413e275b", size = 15701, upload-time = "2025-12-11T13:36:04.56Z" }, +] + +[[package]] +name = "opentelemetry-proto" +version = "1.39.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "protobuf" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/49/1d/f25d76d8260c156c40c97c9ed4511ec0f9ce353f8108ca6e7561f82a06b2/opentelemetry_proto-1.39.1.tar.gz", hash = "sha256:6c8e05144fc0d3ed4d22c2289c6b126e03bcd0e6a7da0f16cedd2e1c2772e2c8", size = 46152, upload-time = "2025-12-11T13:32:48.681Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/51/95/b40c96a7b5203005a0b03d8ce8cd212ff23f1793d5ba289c87a097571b18/opentelemetry_proto-1.39.1-py3-none-any.whl", hash = "sha256:22cdc78efd3b3765d09e68bfbd010d4fc254c9818afd0b6b423387d9dee46007", size = 72535, upload-time = "2025-12-11T13:32:33.866Z" }, +] + +[[package]] +name = "opentelemetry-sdk" +version = "1.39.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "opentelemetry-semantic-conventions" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/eb/fb/c76080c9ba07e1e8235d24cdcc4d125ef7aa3edf23eb4e497c2e50889adc/opentelemetry_sdk-1.39.1.tar.gz", hash = "sha256:cf4d4563caf7bff906c9f7967e2be22d0d6b349b908be0d90fb21c8e9c995cc6", size = 171460, upload-time = "2025-12-11T13:32:49.369Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7c/98/e91cf858f203d86f4eccdf763dcf01cf03f1dae80c3750f7e635bfa206b6/opentelemetry_sdk-1.39.1-py3-none-any.whl", hash = "sha256:4d5482c478513ecb0a5d938dcc61394e647066e0cc2676bee9f3af3f3f45f01c", size = 132565, upload-time = "2025-12-11T13:32:35.069Z" }, +] + +[[package]] +name = "opentelemetry-semantic-conventions" +version = "0.60b1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/91/df/553f93ed38bf22f4b999d9be9c185adb558982214f33eae539d3b5cd0858/opentelemetry_semantic_conventions-0.60b1.tar.gz", hash = "sha256:87c228b5a0669b748c76d76df6c364c369c28f1c465e50f661e39737e84bc953", size = 137935, upload-time = "2025-12-11T13:32:50.487Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7a/5e/5958555e09635d09b75de3c4f8b9cae7335ca545d77392ffe7331534c402/opentelemetry_semantic_conventions-0.60b1-py3-none-any.whl", hash = "sha256:9fa8c8b0c110da289809292b0591220d3a7b53c1526a23021e977d68597893fb", size = 219982, upload-time = "2025-12-11T13:32:36.955Z" }, +] + +[[package]] +name = "opentelemetry-util-http" +version = "0.60b1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/50/fc/c47bb04a1d8a941a4061307e1eddfa331ed4d0ab13d8a9781e6db256940a/opentelemetry_util_http-0.60b1.tar.gz", hash = "sha256:0d97152ca8c8a41ced7172d29d3622a219317f74ae6bb3027cfbdcf22c3cc0d6", size = 11053, upload-time = "2025-12-11T13:37:25.115Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/16/5c/d3f1733665f7cd582ef0842fb1d2ed0bc1fba10875160593342d22bba375/opentelemetry_util_http-0.60b1-py3-none-any.whl", hash = "sha256:66381ba28550c91bee14dcba8979ace443444af1ed609226634596b4b0faf199", size = 8947, upload-time = "2025-12-11T13:36:37.151Z" }, +] + [[package]] name = "packaging" version = "25.0" @@ -2371,6 +2859,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5b/5a/bc7b4a4ef808fa59a816c17b20c4bef6884daebbdf627ff2a161da67da19/propcache-0.4.1-py3-none-any.whl", hash = "sha256:af2a6052aeb6cf17d3e46ee169099044fd8224cbaf75c76a2ef596e8163e2237", size = 13305, upload-time = "2025-10-08T19:49:00.792Z" }, ] +[[package]] +name = "protobuf" +version = "6.33.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/66/70/e908e9c5e52ef7c3a6c7902c9dfbb34c7e29c25d2f81ade3856445fd5c94/protobuf-6.33.6.tar.gz", hash = "sha256:a6768d25248312c297558af96a9f9c929e8c4cee0659cb07e780731095f38135", size = 444531, upload-time = "2026-03-18T19:05:00.988Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fc/9f/2f509339e89cfa6f6a4c4ff50438db9ca488dec341f7e454adad60150b00/protobuf-6.33.6-cp310-abi3-win32.whl", hash = "sha256:7d29d9b65f8afef196f8334e80d6bc1d5d4adedb449971fefd3723824e6e77d3", size = 425739, upload-time = "2026-03-18T19:04:48.373Z" }, + { url = "https://files.pythonhosted.org/packages/76/5d/683efcd4798e0030c1bab27374fd13a89f7c2515fb1f3123efdfaa5eab57/protobuf-6.33.6-cp310-abi3-win_amd64.whl", hash = "sha256:0cd27b587afca21b7cfa59a74dcbd48a50f0a6400cfb59391340ad729d91d326", size = 437089, upload-time = "2026-03-18T19:04:50.381Z" }, + { url = "https://files.pythonhosted.org/packages/5c/01/a3c3ed5cd186f39e7880f8303cc51385a198a81469d53d0fdecf1f64d929/protobuf-6.33.6-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:9720e6961b251bde64edfdab7d500725a2af5280f3f4c87e57c0208376aa8c3a", size = 427737, upload-time = "2026-03-18T19:04:51.866Z" }, + { url = "https://files.pythonhosted.org/packages/ee/90/b3c01fdec7d2f627b3a6884243ba328c1217ed2d978def5c12dc50d328a3/protobuf-6.33.6-cp39-abi3-manylinux2014_aarch64.whl", hash = "sha256:e2afbae9b8e1825e3529f88d514754e094278bb95eadc0e199751cdd9a2e82a2", size = 324610, upload-time = "2026-03-18T19:04:53.096Z" }, + { url = "https://files.pythonhosted.org/packages/9b/ca/25afc144934014700c52e05103c2421997482d561f3101ff352e1292fb81/protobuf-6.33.6-cp39-abi3-manylinux2014_s390x.whl", hash = "sha256:c96c37eec15086b79762ed265d59ab204dabc53056e3443e702d2681f4b39ce3", size = 339381, upload-time = "2026-03-18T19:04:54.616Z" }, + { url = "https://files.pythonhosted.org/packages/16/92/d1e32e3e0d894fe00b15ce28ad4944ab692713f2e7f0a99787405e43533a/protobuf-6.33.6-cp39-abi3-manylinux2014_x86_64.whl", hash = "sha256:e9db7e292e0ab79dd108d7f1a94fe31601ce1ee3f7b79e0692043423020b0593", size = 323436, upload-time = "2026-03-18T19:04:55.768Z" }, + { url = "https://files.pythonhosted.org/packages/c4/72/02445137af02769918a93807b2b7890047c32bfb9f90371cbc12688819eb/protobuf-6.33.6-py3-none-any.whl", hash = "sha256:77179e006c476e69bf8e8ce866640091ec42e1beb80b213c3900006ecfba6901", size = 170656, upload-time = "2026-03-18T19:04:59.826Z" }, +] + [[package]] name = "psutil" version = "7.2.2" @@ -2439,6 +2942,27 @@ memory = [ { name = "cachetools" }, ] +[[package]] +name = "pyasn1" +version = "0.6.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5c/5f/6583902b6f79b399c9c40674ac384fd9cd77805f9e6205075f828ef11fb2/pyasn1-0.6.3.tar.gz", hash = "sha256:697a8ecd6d98891189184ca1fa05d1bb00e2f84b5977c481452050549c8a72cf", size = 148685, upload-time = "2026-03-17T01:06:53.382Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5d/a0/7d793dce3fa811fe047d6ae2431c672364b462850c6235ae306c0efd025f/pyasn1-0.6.3-py3-none-any.whl", hash = "sha256:a80184d120f0864a52a073acc6fc642847d0be408e7c7252f31390c0f4eadcde", size = 83997, upload-time = "2026-03-17T01:06:52.036Z" }, +] + +[[package]] +name = "pyasn1-modules" +version = "0.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyasn1" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e9/e6/78ebbb10a8c8e4b61a59249394a4a594c1a7af95593dc933a349c8d00964/pyasn1_modules-0.4.2.tar.gz", hash = "sha256:677091de870a80aae844b1ca6134f54652fa2c8c5a52aa396440ac3106e941e6", size = 307892, upload-time = "2025-03-28T02:41:22.17Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/47/8d/d529b5d697919ba8c11ad626e835d4039be708a35b0d22de83a269a6682c/pyasn1_modules-0.4.2-py3-none-any.whl", hash = "sha256:29253a9207ce32b64c3ac6600edc75368f98473906e8fd1043bd6b5b1de2c14a", size = 181259, upload-time = "2025-03-28T02:41:19.028Z" }, +] + [[package]] name = "pycparser" version = "3.0" @@ -2468,6 +2992,106 @@ email = [ { name = "email-validator" }, ] +[[package]] +name = "pydantic-ai" +version = "1.78.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pydantic-ai-slim", extra = ["ag-ui", "anthropic", "bedrock", "cli", "cohere", "evals", "fastmcp", "google", "groq", "huggingface", "logfire", "mcp", "mistral", "openai", "retries", "spec", "temporal", "ui", "vertexai", "xai"] }, +] +sdist = { url = "https://files.pythonhosted.org/packages/04/e1/ced6f04f60accb11deb1a8ca4fc576270022e646f011a9e1674695420710/pydantic_ai-1.78.0.tar.gz", hash = "sha256:dd3f56306c671f7785126e78d72924e5a80c30bca27460081941ad22b63fcc8d", size = 12645, upload-time = "2026-04-08T05:20:34.096Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/1a/497fdf8224aed69752559350e437395266ff0f3107ca3bff63b6925358cc/pydantic_ai-1.78.0-py3-none-any.whl", hash = "sha256:aa0fdacec813fa457243206a9dae4d5152e0814bf17fc7eee75045d3469f5ca8", size = 7551, upload-time = "2026-04-08T05:20:24.209Z" }, +] + +[[package]] +name = "pydantic-ai-slim" +version = "1.78.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "genai-prices" }, + { name = "griffelib" }, + { name = "httpx" }, + { name = "opentelemetry-api" }, + { name = "pydantic" }, + { name = "pydantic-graph" }, + { name = "typing-inspection" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/85/84/4cf98c41a2019a5c5ed6aa5d5fa2bbc8b70b152b527c56d27dabbeaeb75c/pydantic_ai_slim-1.78.0.tar.gz", hash = "sha256:97c6467a6bb09f61fd48cd828db066204ae77419d14a4edf47f90f72d06ab11f", size = 531385, upload-time = "2026-04-08T05:20:36.291Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/91/487992a441c03f16525885ee541083f66f579500edd7aa8b838f3296455f/pydantic_ai_slim-1.78.0-py3-none-any.whl", hash = "sha256:36f88bab6016186b958363ca1254741999df276322d7066dd69793dcb2134b66", size = 680002, upload-time = "2026-04-08T05:20:27.508Z" }, +] + +[package.optional-dependencies] +ag-ui = [ + { name = "ag-ui-protocol" }, + { name = "starlette" }, +] +anthropic = [ + { name = "anthropic" }, +] +bedrock = [ + { name = "boto3" }, +] +cli = [ + { name = "argcomplete" }, + { name = "prompt-toolkit" }, + { name = "pyperclip" }, + { name = "pyyaml" }, + { name = "rich" }, +] +cohere = [ + { name = "cohere", marker = "sys_platform != 'emscripten' or (extra == 'extra-10-imas-codex-cpu' and extra == 'extra-10-imas-codex-gpu') or (extra == 'extra-10-imas-codex-gpu' and extra == 'extra-10-imas-codex-test')" }, +] +evals = [ + { name = "pydantic-evals" }, +] +fastmcp = [ + { name = "fastmcp" }, +] +google = [ + { name = "google-genai" }, +] +groq = [ + { name = "groq" }, +] +huggingface = [ + { name = "huggingface-hub" }, +] +logfire = [ + { name = "logfire", extra = ["httpx"] }, +] +mcp = [ + { name = "mcp" }, +] +mistral = [ + { name = "mistralai" }, +] +openai = [ + { name = "openai" }, + { name = "tiktoken" }, +] +retries = [ + { name = "tenacity" }, +] +spec = [ + { name = "pydantic-handlebars" }, + { name = "pyyaml" }, +] +temporal = [ + { name = "temporalio" }, +] +ui = [ + { name = "starlette" }, +] +vertexai = [ + { name = "google-auth" }, + { name = "requests" }, +] +xai = [ + { name = "xai-sdk" }, +] + [[package]] name = "pydantic-core" version = "2.41.5" @@ -2497,6 +3121,50 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f7/07/34573da085946b6a313d7c42f82f16e8920bfd730665de2d11c0c37a74b5/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:76d0819de158cd855d1cbb8fcafdf6f5cf1eb8e470abe056d5d161106e38062b", size = 2139017, upload-time = "2025-11-04T13:42:59.471Z" }, ] +[[package]] +name = "pydantic-evals" +version = "1.78.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "logfire-api" }, + { name = "pydantic" }, + { name = "pydantic-ai-slim" }, + { name = "pyyaml" }, + { name = "rich" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3b/80/539238f284f6a4fdef5737a0cb7efa746e4259bb27526af271898784f2fa/pydantic_evals-1.78.0.tar.gz", hash = "sha256:8608068c2569a0169977526a93ddea45e924331115233eb1f297ab19653e14e0", size = 65818, upload-time = "2026-04-08T05:20:37.844Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f7/4c/e6bcf445ff3f15211c217975a11d401e78b3f2d1e989c50d859527e81050/pydantic_evals-1.78.0-py3-none-any.whl", hash = "sha256:b00cc22e2d24a0771f40fc7e3b2b4afeec7c99bda8cfce37e6bee58e7a29fe3c", size = 77739, upload-time = "2026-04-08T05:20:29.577Z" }, +] + +[[package]] +name = "pydantic-graph" +version = "1.78.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "httpx" }, + { name = "logfire-api" }, + { name = "pydantic" }, + { name = "typing-inspection" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/30/e4/eb52021f43f2ac495955af19219c1ab261d707bbd15f4dc229079c2276d4/pydantic_graph-1.78.0.tar.gz", hash = "sha256:dd627e37cb3adaf8c95cca6a4b33e0d1b7fc9bed075dc3b8ad5df2c2a3cb432b", size = 58682, upload-time = "2026-04-08T05:20:39.288Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/47/97/e3dd1a1c6f6b9c104c844c1a2383361a831d1b027c99c0ac1a20544f0b13/pydantic_graph-1.78.0-py3-none-any.whl", hash = "sha256:0302835f46da3ee70ba3602a4d886c41a76fa8750f23f4257b968163ba4bb89f", size = 72500, upload-time = "2026-04-08T05:20:31.237Z" }, +] + +[[package]] +name = "pydantic-handlebars" +version = "0.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pydantic" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/90/16/d41768bd3fd77e6250c20be11a3e68fee5fff07c3356455e6708f6a60f2a/pydantic_handlebars-0.1.0.tar.gz", hash = "sha256:1931c54946add1b5e3796c9bf6a005ed7662cef0109bb05c352f0b3d031a1260", size = 159826, upload-time = "2026-03-01T20:00:17.497Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/99/5f/86b1630be61bdebf253c2f953a6c3f073ec21bb0725565ea3896802e1ca3/pydantic_handlebars-0.1.0-py3-none-any.whl", hash = "sha256:8a436fe8bc607295eb04bec58bd6e2c9498c9e069c557ff0b505e3d568c783bc", size = 40890, upload-time = "2026-03-01T20:00:16.106Z" }, +] + [[package]] name = "pydantic-settings" version = "2.13.1" @@ -3170,6 +3838,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2a/07/5bda6a85b220c64c65686bc85bd0bbb23b29c62b3a9f9433fa55f17cda93/ruff-0.15.1-py3-none-win_arm64.whl", hash = "sha256:5ff7d5f0f88567850f45081fac8f4ec212be8d0b963e385c3f7d0d2eb4899416", size = 10874604, upload-time = "2026-02-12T23:09:05.515Z" }, ] +[[package]] +name = "s3transfer" +version = "0.16.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "botocore" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/05/04/74127fc843314818edfa81b5540e26dd537353b123a4edc563109d8f17dd/s3transfer-0.16.0.tar.gz", hash = "sha256:8e990f13268025792229cd52fa10cb7163744bf56e719e0b9cb925ab79abf920", size = 153827, upload-time = "2025-12-01T02:30:59.114Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fc/51/727abb13f44c1fcf6d145979e1535a35794db0f6e450a0cb46aa24732fe2/s3transfer-0.16.0-py3-none-any.whl", hash = "sha256:18e25d66fed509e3868dc1572b3f427ff947dd2c56f844a5bf09481ad3f3b2fe", size = 86830, upload-time = "2025-12-01T02:30:57.729Z" }, +] + [[package]] name = "safetensors" version = "0.7.0" @@ -3515,6 +4195,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/81/0d/13d1d239a25cbfb19e740db83143e95c772a1fe10202dda4b76792b114dd/starlette-0.52.1-py3-none-any.whl", hash = "sha256:0029d43eb3d273bc4f83a08720b4912ea4b071087a3b48db01b7c839f7954d74", size = 74272, upload-time = "2026-01-18T13:34:09.188Z" }, ] +[[package]] +name = "strictyaml" +version = "1.7.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "python-dateutil" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b3/08/efd28d49162ce89c2ad61a88bd80e11fb77bc9f6c145402589112d38f8af/strictyaml-1.7.3.tar.gz", hash = "sha256:22f854a5fcab42b5ddba8030a0e4be51ca89af0267961c8d6cfa86395586c407", size = 115206, upload-time = "2023-03-10T12:50:27.062Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/96/7c/a81ef5ef10978dd073a854e0fa93b5d8021d0594b639cc8f6453c3c78a1d/strictyaml-1.7.3-py3-none-any.whl", hash = "sha256:fb5c8a4edb43bebb765959e420f9b3978d7f1af88c80606c03fb420888f5d1c7", size = 123917, upload-time = "2023-03-10T12:50:17.242Z" }, +] + [[package]] name = "sympy" version = "1.13.1" @@ -3552,6 +4244,50 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/40/44/4a5f08c96eb108af5cb50b41f76142f0afa346dfa99d5296fe7202a11854/tabulate-0.9.0-py3-none-any.whl", hash = "sha256:024ca478df22e9340661486f85298cff5f6dcdba14f3813e8830015b9ed1948f", size = 35252, upload-time = "2022-10-06T17:21:44.262Z" }, ] +[[package]] +name = "temporalio" +version = "1.25.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nexus-rpc" }, + { name = "protobuf" }, + { name = "types-protobuf" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/de/9c/3782bab0bf11a40b550147c19a5d1a476c17405391751982408902d9f138/temporalio-1.25.0.tar.gz", hash = "sha256:a3bbec1dcc904f674402cfa4faae480fda490b1c53ea5440c1f1996c562016fb", size = 2152534, upload-time = "2026-04-08T18:53:55.388Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/19/e3/5676dd10d1164b6d6ca8752314054097b89c5da931e936af402a7b15236c/temporalio-1.25.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:6dc1bc8e1773b1a833d86a7ede2dd90ef4e031ced5b748b59e7f09a5bf9b327d", size = 13943906, upload-time = "2026-04-08T18:53:30.022Z" }, + { url = "https://files.pythonhosted.org/packages/89/50/7cbf7f845973be986ec165348f72f7a409750842a04d554965a39be5cb4f/temporalio-1.25.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:3c8fdcf79ea5ae8ae2cf6f48072e4a86c3e0f4778f6a8a066c6ff1d336587db4", size = 13298719, upload-time = "2026-04-08T18:53:35.95Z" }, + { url = "https://files.pythonhosted.org/packages/d2/31/d474bab8535552add6ed289911bf1ffae5d7071823ece1069842190fcaed/temporalio-1.25.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:141f37aaafd7d090ba5c8776e4e9bc60df1fbc64b9f50c8f00e905a436588ddc", size = 13555435, upload-time = "2026-04-08T18:53:41.36Z" }, + { url = "https://files.pythonhosted.org/packages/2a/c8/e7dc053d6107bf2a037a3c9fe7b86639a25dcb888bde0e1ca366901ee47f/temporalio-1.25.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ff7ca5bb80264976477d4dc7a839b3d22af8577ae92306526a061481db49bf92", size = 14052050, upload-time = "2026-04-08T18:53:46.44Z" }, + { url = "https://files.pythonhosted.org/packages/08/70/9340ed3a578321cbc153041d34834bb1ec3f1f3e3d9cded47cd1b7c3e403/temporalio-1.25.0-cp310-abi3-win_amd64.whl", hash = "sha256:9411534279a2e64847231b6059c214bff4d57cfd1532bd09f333d0b1603daa7f", size = 14299684, upload-time = "2026-04-08T18:53:52.482Z" }, +] + +[[package]] +name = "tenacity" +version = "9.1.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/47/c6/ee486fd809e357697ee8a44d3d69222b344920433d3b6666ccd9b374630c/tenacity-9.1.4.tar.gz", hash = "sha256:adb31d4c263f2bd041081ab33b498309a57c77f9acf2db65aadf0898179cf93a", size = 49413, upload-time = "2026-02-07T10:45:33.841Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d7/c1/eb8f9debc45d3b7918a32ab756658a0904732f75e555402972246b0b8e71/tenacity-9.1.4-py3-none-any.whl", hash = "sha256:6095a360c919085f28c6527de529e76a06ad89b23659fa881ae0649b867a9d55", size = 28926, upload-time = "2026-02-07T10:45:32.24Z" }, +] + +[[package]] +name = "textual" +version = "6.2.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py", extra = ["linkify", "plugins"] }, + { name = "platformdirs" }, + { name = "pygments" }, + { name = "rich" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a2/30/38b615f7d4b16f6fdd73e4dcd8913e2d880bbb655e68a076e3d91181a7ee/textual-6.2.1.tar.gz", hash = "sha256:4699d8dfae43503b9c417bd2a6fb0da1c89e323fe91c4baa012f9298acaa83e1", size = 1570645, upload-time = "2025-10-01T16:11:24.467Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c5/93/02c7adec57a594af28388d85da9972703a4af94ae1399542555cd9581952/textual-6.2.1-py3-none-any.whl", hash = "sha256:3c7190633cd4d8bfe6049ae66808b98da91ded2edb85cef54e82bf77b03d2a54", size = 710702, upload-time = "2025-10-01T16:11:22.161Z" }, +] + [[package]] name = "threadpoolctl" version = "3.6.0" @@ -3873,6 +4609,27 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4a/91/48db081e7a63bb37284f9fbcefda7c44c277b18b0e13fbc36ea2335b71e6/typer-0.24.1-py3-none-any.whl", hash = "sha256:112c1f0ce578bfb4cab9ffdabc68f031416ebcc216536611ba21f04e9aa84c9e", size = 56085, upload-time = "2026-02-21T16:54:41.616Z" }, ] +[[package]] +name = "types-protobuf" +version = "6.32.1.20260221" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5f/e2/9aa4a3b2469508bd7b4e2ae11cbedaf419222a09a1b94daffcd5efca4023/types_protobuf-6.32.1.20260221.tar.gz", hash = "sha256:6d5fb060a616bfb076cbb61b4b3c3969f5fc8bec5810f9a2f7e648ee5cbcbf6e", size = 64408, upload-time = "2026-02-21T03:55:13.916Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2e/e8/1fd38926f9cf031188fbc5a96694203ea6f24b0e34bd64a225ec6f6291ba/types_protobuf-6.32.1.20260221-py3-none-any.whl", hash = "sha256:da7cdd947975964a93c30bfbcc2c6841ee646b318d3816b033adc2c4eb6448e4", size = 77956, upload-time = "2026-02-21T03:55:12.894Z" }, +] + +[[package]] +name = "types-requests" +version = "2.33.0.20260408" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/69/6a/749dc53a54a3f35842c1f8197b3ca6b54af6d7458a1bfc75f6629b6da666/types_requests-2.33.0.20260408.tar.gz", hash = "sha256:95b9a86376807a216b2fb412b47617b202091c3ea7c078f47cc358d5528ccb7b", size = 23882, upload-time = "2026-04-08T04:34:49.33Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/90/b8/78fd6c037de4788c040fdd323b3369804400351b7827473920f6c1d03c10/types_requests-2.33.0.20260408-py3-none-any.whl", hash = "sha256:81f31d5ea4acb39f03be7bc8bed569ba6d5a9c5d97e89f45ac43d819b68ca50f", size = 20739, upload-time = "2026-04-08T04:34:48.325Z" }, +] + [[package]] name = "typing-extensions" version = "4.15.0" @@ -3903,6 +4660,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c7/b0/003792df09decd6849a5e39c28b513c06e84436a54440380862b5aeff25d/tzdata-2025.3-py2.py3-none-any.whl", hash = "sha256:06a47e5700f3081aab02b2e513160914ff0694bce9947d6b76ebd6bf57cfc5d1", size = 348521, upload-time = "2025-12-13T17:45:33.889Z" }, ] +[[package]] +name = "uc-micro-py" +version = "2.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/78/67/9a363818028526e2d4579334460df777115bdec1bb77c08f9db88f6389f2/uc_micro_py-2.0.0.tar.gz", hash = "sha256:c53691e495c8db60e16ffc4861a35469b0ba0821fe409a8a7a0a71864d33a811", size = 6611, upload-time = "2026-03-01T06:31:27.526Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/61/73/d21edf5b204d1467e06500080a50f79d49ef2b997c79123a536d4a17d97c/uc_micro_py-2.0.0-py3-none-any.whl", hash = "sha256:3603a3859af53e5a39bc7677713c78ea6589ff188d70f4fee165db88e22b242c", size = 6383, upload-time = "2026-03-01T06:31:26.257Z" }, +] + [[package]] name = "uncalled-for" version = "0.3.1" @@ -4071,6 +4837,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1f/f6/a933bd70f98e9cf3e08167fc5cd7aaaca49147e48411c0bd5ae701bb2194/wrapt-1.17.3-py3-none-any.whl", hash = "sha256:7171ae35d2c33d326ac19dd8facb1e82e5fd04ef8c6c0e394d7af55a55051c22", size = 23591, upload-time = "2025-08-12T05:53:20.674Z" }, ] +[[package]] +name = "xai-sdk" +version = "1.11.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiohttp" }, + { name = "googleapis-common-protos" }, + { name = "grpcio" }, + { name = "opentelemetry-sdk" }, + { name = "packaging" }, + { name = "protobuf" }, + { name = "pydantic" }, + { name = "requests" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/49/32/bb8385f7a3b05ce406b689aa000c9a34289caa1526f1c093a1cefc0d9695/xai_sdk-1.11.0.tar.gz", hash = "sha256:ca87a830d310fb8e06fba44fb2a8c5cdf0d9f716b61126eddd51b7f416a63932", size = 404313, upload-time = "2026-03-27T18:23:10.091Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/76/86d9a3589c725ce825d2ed3e7cb3ecf7f956d3fd015353d52197bb341bcd/xai_sdk-1.11.0-py3-none-any.whl", hash = "sha256:fe58ce6d8f8115ae8bd57ded57bcd847d0bb7cb28bb7b236abefd4626df1ed8d", size = 251388, upload-time = "2026-03-27T18:23:08.573Z" }, +] + [[package]] name = "xlrd" version = "2.0.2"