diff --git a/.github/workflows/dependency_tests.yml b/.github/workflows/dependency_tests.yml index 7c16d628..6c8bf8be 100644 --- a/.github/workflows/dependency_tests.yml +++ b/.github/workflows/dependency_tests.yml @@ -14,7 +14,7 @@ jobs: fail-fast: false matrix: python-version: ["3.10", "3.11", "3.12", "3.13"] - install_extras: [core, calculators, uma, ui, parsl, viz] + install_extras: [core, calculators, uma, ui, parsl, xanes, rag] steps: - uses: actions/checkout@v4 @@ -74,9 +74,19 @@ jobs: 'uma': ['fairchem.core', 'e3nn'], 'ui': ['streamlit', 'stmol'], 'parsl': ['parsl'], - 'viz': ['pyppeteer', 'grandalf', 'nest_asyncio'] + 'xanes': ['parsl'], + 'rag': [ + 'faiss', + 'langchain_text_splitters', + 'langchain_huggingface', + 'sentence_transformers', + 'fitz', + ], } + if extras == 'xanes' and sys.version_info >= (3, 11): + extra_packages['xanes'].append('mp_api') + if extras != 'core': # Check packages for specific extra if extras in extra_packages: diff --git a/.github/workflows/ghcr-publish.yml b/.github/workflows/ghcr-publish.yml index 6a5b79bb..b325e796 100644 --- a/.github/workflows/ghcr-publish.yml +++ b/.github/workflows/ghcr-publish.yml @@ -4,6 +4,9 @@ on: push: tags: - "v*" + branches: + - feature_k8s + - dev workflow_dispatch: permissions: @@ -20,9 +23,6 @@ jobs: - arch: amd64 platform: linux/amd64 dockerfile: Dockerfile - - arch: arm64 - platform: linux/arm64 - dockerfile: Dockerfile.arm steps: - name: Checkout code @@ -84,10 +84,11 @@ jobs: images: ${{ steps.prep.outputs.image }} tags: | type=ref,event=tag + type=ref,event=branch type=sha,prefix=sha- - type=raw,value=latest + type=raw,value=latest,enable=${{ startsWith(github.ref, 'refs/tags/') }} - - name: Create and push multi-arch manifests + - name: Create and push manifests env: IMAGE: ${{ steps.prep.outputs.image }} TAGS: ${{ steps.meta.outputs.tags }} @@ -96,11 +97,10 @@ jobs: [ -z "$tag" ] && continue docker buildx imagetools create \ -t "$tag" \ - "${IMAGE}:build-${GITHUB_SHA}-amd64" \ - "${IMAGE}:build-${GITHUB_SHA}-arm64" + "${IMAGE}:build-${GITHUB_SHA}-amd64" done <<< "$TAGS" - - name: Inspect latest manifest + - name: Inspect manifest env: IMAGE: ${{ steps.prep.outputs.image }} - run: docker buildx imagetools inspect "${IMAGE}:latest" + run: docker buildx imagetools inspect "${IMAGE}:${GITHUB_REF_NAME}" diff --git a/.github/workflows/test-pypi-package.yml b/.github/workflows/test-pypi-package.yml index b8f07170..8c9bee93 100644 --- a/.github/workflows/test-pypi-package.yml +++ b/.github/workflows/test-pypi-package.yml @@ -34,8 +34,7 @@ jobs: run: | python -c "import chemgraph; print('ChemGraph imported successfully')" python -c "from chemgraph.cli import main; print('CLI module imported successfully')" - # Test that CLI command is available - chemgraph --help || echo "CLI help command executed" + chemgraph --help - name: Run basic import tests run: | diff --git a/.gitignore b/.gitignore index 957eeda0..307e2461 100644 --- a/.gitignore +++ b/.gitignore @@ -57,5 +57,8 @@ test.csv nwchem/ nwchem.nwi nwchem.nwo - +config.toml vib*.traj + +# Kubernetes secrets (keep secrets.yaml.template, ignore actual secrets) +k8s/secrets.yaml diff --git a/.opencode/plans/refactoring-plan.md b/.opencode/plans/refactoring-plan.md new file mode 100644 index 00000000..14f862cd --- /dev/null +++ b/.opencode/plans/refactoring-plan.md @@ -0,0 +1,426 @@ +# ChemGraph Codebase Reorganization Plan + +**Date:** 2025-04-25 +**Status:** Completed (2026-04-25) + +## Goal + +Establish a consistent `*_core.py` + thin-wrapper pattern across all tool +domains, eliminate all code duplication, and consolidate schemas into the +`schemas/` package. After this refactoring every domain follows: + +``` +schemas/_schema.py # Pydantic input/output models +tools/_core.py # Pure-Python logic (no decorators) +tools/_tools.py # LangChain @tool wrappers (thin) +mcp/_mcp.py or *_parsl.py # MCP @mcp.tool wrappers (thin) +``` + +--- + +## Context: What Was Already Done + +In the first refactoring pass we unified the ASE simulation logic: + +- Created `tools/core.py` — single `run_ase_core()` + shared helpers + (`atoms_to_atomsdata`, `is_linear_molecule`, `get_symmetry_number`, + `load_calculator`, `_resolve_path`, etc.) +- Simplified `tools/ase_tools.py` → thin `@tool` wrappers (737 → 187 lines) +- Simplified `mcp/mcp_tools.py` → thin `@mcp.tool` wrappers (554 → 206 lines) +- Simplified `tools/parsl_tools.py` → `run_mace_core()` delegates to + `run_ase_core()` via schema conversion (406 → 204 lines) +- Reduced `tools/mcp_helper.py` to a backward-compat re-export shim (158 → 13 lines) + +All 19 relevant tests pass. The 59 failures in the full suite are pre-existing +(missing `vision_test_graph`, `langchain_huggingface`, `architector_tools`, +`_detect_executor_failure`). + +--- + +## Current File Inventory (with line counts) + +### tools/ +| File | Lines | Status | +|----------------------------|------:|-------------------------------------------| +| `core.py` | 611 | To rename → `ase_core.py` | +| `ase_tools.py` | 187 | Already refactored (thin wrappers) | +| `cheminformatics_tools.py` | 152 | Needs core extraction | +| `graspa_tools.py` | 307 | Has core+wrapper in one file | +| `xanes_tools.py` | 616 | Has core+wrapper in one file | +| `parsl_tools.py` | 204 | Schemas need to move to `schemas/` | +| `report_tools.py` | 776 | Has duplicated element_map (lines 347,415)| +| `mcp_helper.py` | 13 | Delete (re-export shim) | +| `rag_tools.py` | 343 | No changes needed | +| `generic_tools.py` | 101 | No changes needed | + +### mcp/ +| File | Lines | Status | +|-------------------------|------:|-------------------------------------------| +| `mcp_tools.py` | 206 | Needs cheminformatics core extraction | +| `mace_mcp_parsl.py` | 256 | Update schema import path | +| `graspa_mcp_parsl.py` | 201 | Extract `load_parsl_config` to server_utils| +| `xanes_mcp_parsl.py` | 290 | Merge with `xanes_mcp.py`, extract config | +| `xanes_mcp.py` | 97 | Delete (merge into `xanes_mcp_parsl.py`) | +| `data_analysis_mcp.py` | 304 | No changes needed | +| `server_utils.py` | 79 | Add `load_parsl_config()` | + +### schemas/ +| File | Status | +|-------------------------|---------------------------------------------| +| `ase_input.py` | No changes | +| `atomsdata.py` | No changes | +| `graspa_schema.py` | No changes | +| `graspa_input.py` | No changes (separate schema, used elsewhere)| +| `xanes_schema.py` | No changes | +| `agent_response.py` | No changes | +| `multi_agent_response.py`| No changes | +| `calculators/` | Add `__init__.py` | +| **`mace_parsl_schema.py`** | **NEW** — moved from `parsl_tools.py` | + +--- + +## Phases + +### Phase 1: Rename and Relocate (mechanical, low risk) + +#### 1A. Rename `tools/core.py` → `tools/ase_core.py` + +**Why:** The name `core.py` is ambiguous once we add `cheminformatics_core.py`, +`graspa_core.py`, `xanes_core.py`. + +**Files to update (imports):** +- `tools/ase_tools.py` — `from chemgraph.tools.core import …` → `from chemgraph.tools.ase_core import …` +- `tools/mcp_helper.py` — same (will be deleted in Phase 4A anyway, but update + first so we can delete cleanly) +- `tools/parsl_tools.py` — `from chemgraph.tools.core import run_ase_core` → `from chemgraph.tools.ase_core import run_ase_core` +- `mcp/mcp_tools.py` — `from chemgraph.tools.core import …` → `from chemgraph.tools.ase_core import …` + +No other files import from `tools/core.py` directly. + +**Verification:** `pytest tests/test_tools.py tests/test_mace.py tests/test_mcp.py tests/test_calculators.py -v` + +#### 1B. Move MACE schemas to `schemas/mace_parsl_schema.py` + +**What moves:** +- `mace_input_schema` (class) — from `parsl_tools.py` +- `mace_input_schema_ensemble` (class) — from `parsl_tools.py` +- `mace_output_schema` (class) — from `parsl_tools.py` + +**Files to update (imports):** +- `tools/parsl_tools.py` — add `from chemgraph.schemas.mace_parsl_schema import mace_input_schema, mace_output_schema` +- `mcp/mace_mcp_parsl.py` line 15-19 — `from chemgraph.tools.parsl_tools import (mace_input_schema, mace_input_schema_ensemble, run_mace_core)` → split into: + - `from chemgraph.schemas.mace_parsl_schema import mace_input_schema, mace_input_schema_ensemble` + - `from chemgraph.tools.parsl_tools import run_mace_core` +- `mcp/mace_mcp_parsl.py` line 39 (deferred import inside `run_mace_parsl_app`) — update similarly + +**Verification:** `pytest tests/test_mace.py tests/test_mcp.py -v` + +#### 1C. Add missing `__init__.py` files + +Create empty `__init__.py` in: +- `src/chemgraph/mcp/__init__.py` +- `src/chemgraph/schemas/calculators/__init__.py` + +These are currently implicit namespace packages. Adding `__init__.py` makes them +consistent with `tools/` and `utils/`. + +--- + +### Phase 2: Extract Core Modules (eliminates duplication) + +#### 2A. Create `tools/cheminformatics_core.py` + +**Currently duplicated logic:** +- RDKit SMILES→3D pipeline appears **3 times**: + - `cheminformatics_tools.py:smiles_to_atomsdata()` (lines 37-83) + - `cheminformatics_tools.py:smiles_to_coordinate_file()` (lines 86-152) + - `mcp/mcp_tools.py:smiles_to_coordinate_file()` (lines 82-159) +- PubChem name→SMILES lookup appears **2 times** with **different return types**: + - `cheminformatics_tools.py:molecule_name_to_smiles()` → returns `{"name": ..., "smiles": ...}` + - `mcp/mcp_tools.py:molecule_name_to_smiles()` → returns just the SMILES string + +**New file: `tools/cheminformatics_core.py`** +```python +def smiles_to_3d(smiles: str, seed: int = 2025) -> tuple[list[int], list[list[float]]]: + """SMILES → (atomic_numbers, positions) via RDKit. Single implementation.""" + +def molecule_name_to_smiles_core(name: str) -> str: + """PubChem name → canonical SMILES. Returns the raw SMILES string.""" + +def smiles_to_coordinate_file_core(smiles: str, output_file: str, seed: int) -> dict: + """SMILES → coordinate file on disk. Returns result dict.""" + +def smiles_to_atomsdata_core(smiles: str, seed: int) -> AtomsData: + """SMILES → AtomsData object.""" +``` + +**Files to simplify:** +- `cheminformatics_tools.py` → thin `@tool` wrappers calling `cheminformatics_core.*` + - `molecule_name_to_smiles` wraps `molecule_name_to_smiles_core`, adds the dict return + - `smiles_to_atomsdata` wraps `smiles_to_atomsdata_core` + - `smiles_to_coordinate_file` wraps `smiles_to_coordinate_file_core` +- `mcp/mcp_tools.py` → `molecule_name_to_smiles` and `smiles_to_coordinate_file` become + thin `@mcp.tool` wrappers calling `cheminformatics_core.*` + - Remove `_resolve_path` local definition (use from `ase_core` or `cheminformatics_core`) + +**Verification:** `pytest tests/test_tools.py tests/test_mcp.py -v -k "smiles or molecule"` + +#### 2B. Create `tools/xanes_core.py` (higher priority — full file duplication) + +**Currently duplicated:** +- `xanes_mcp.py` is an exact subset of `xanes_mcp_parsl.py` — `run_xanes_single`, + `fetch_mp_structures`, `plot_xanes` are identical in both files. +- Both use default port 9007 (would conflict if run simultaneously). + +**New file: `tools/xanes_core.py`** + +Extract from `xanes_tools.py`: +```python +def write_fdmnes_input(...): ... +def get_normalized_xanes(...): ... +def extract_conv(...): ... +def _get_data_dir(): ... +def run_xanes_core(params): ... +def fetch_materials_project_data(params, db_path): ... +def create_fdmnes_inputs(...): ... +def expand_database_results(...): ... +def plot_xanes_results(...): ... +``` + +**Files to simplify:** +- `xanes_tools.py` → thin `@tool` wrappers: + - `run_xanes` wraps `run_xanes_core` + - `fetch_xanes_data` wraps `fetch_materials_project_data` + - `plot_xanes_data` wraps `plot_xanes_results` +- `xanes_mcp_parsl.py` → imports from `xanes_core` instead of deferred imports + from `xanes_tools` +- **Delete `xanes_mcp.py`** — merge its tools into `xanes_mcp_parsl.py` with + conditional Parsl support (only expose `run_xanes_ensemble` if Parsl is loaded) + +**Verification:** Manual (XANES tests require FDMNES executable which is HPC-only) + +#### 2C. Create `tools/graspa_core.py` (lower priority) + +**Current state:** `graspa_tools.py` already has `run_graspa_core()` as a plain +function with `@tool run_graspa()` as the wrapper — the pattern is there, just +in a single file. + +**New file: `tools/graspa_core.py`** + +Extract from `graspa_tools.py`: +```python +def _read_graspa_sycl_output(output_dir): ... +def mock_graspa(): ... +def run_graspa_core(params: graspa_input_schema): ... +``` + +**Files to simplify:** +- `graspa_tools.py` → only `@tool run_graspa()` wrapper +- `graspa_mcp_parsl.py` → deferred import changes from + `from chemgraph.tools.graspa_tools import run_graspa_core` to + `from chemgraph.tools.graspa_core import run_graspa_core` + +**Verification:** Manual (gRASPA requires SYCL runtime) + +--- + +### Phase 3: Consolidate Shared Utilities + +#### 3A. Move `extract_output_json` to `ase_core.py` + +**Currently duplicated 3 times** (identical `json.load` logic): +- `tools/ase_tools.py` — `@tool extract_output_json` +- `mcp/mcp_tools.py` — `@mcp.tool extract_output_json` +- `mcp/mace_mcp_parsl.py` — `@mcp.tool extract_output_json` + +**Action:** +- Add `extract_output_json_core(json_file: str) -> dict` to `tools/ase_core.py` +- All three wrappers call it + +#### 3B. Consolidate `load_parsl_config()` into `mcp/server_utils.py` + +**Currently duplicated 2 times** (identical function): +- `mcp/graspa_mcp_parsl.py` lines 44-66 +- `mcp/xanes_mcp_parsl.py` lines 39-65 + +**Action:** +- Add `load_parsl_config(compute_system: str) -> Config` to `mcp/server_utils.py` +- Both MCP servers import from there + +#### 3C. Fix `report_tools.py` duplicated element map + +**Currently:** Two identical `element_map` dicts (lines 347 and 415) covering +only H(1) through Ar(18). Heavier elements fall through to `X{num}`. + +**Action:** +- Replace with `from ase.data import chemical_symbols` +- Use `chemical_symbols[num]` instead of `element_map.get(num, f"X{num}")` +- Handles all elements automatically + +--- + +### Phase 4: Cleanup + +#### 4A. Delete `tools/mcp_helper.py` + +**Current state:** 13-line re-export shim. Only one remaining consumer imports +from `mcp_helper` instead of `ase_core`: +- `tools/cheminformatics_tools.py` line 10: `from chemgraph.tools.mcp_helper import _resolve_path` + +The `new_eval/scripts/mcp_example/` files also import from `mcp_helper`, but +those are separate example scripts (see 4B). + +**Action:** +1. Update `cheminformatics_tools.py` to import `_resolve_path` from + `chemgraph.tools.ase_core` (or from the new `cheminformatics_core.py` if + Phase 2A is done first) +2. Delete `tools/mcp_helper.py` + +#### 4B. Refactor `new_eval/scripts/mcp_example/` scripts + +**Currently:** Both `mcp_http/start_mcp_server.py` and +`mcp_stdio/mcp_tools_stdio.py` contain **inline copies** of `run_ase()` +(~200+ lines each) plus import helpers from `mcp_helper`. + +**Action:** +- Replace inline `run_ase` with `from chemgraph.tools.ase_core import run_ase_core` +- Update helper imports from `mcp_helper` → `ase_core` +- Each file shrinks from ~200+ lines of duplicated simulation logic to ~20 lines + +#### 4C. Fix stale import in `utils/tool_call_eval.py` + +**Line 4:** `from chemgraph.models.ase_input import ASEInputSchema` +**Should be:** `from chemgraph.schemas.ase_input import ASEInputSchema` + +(The `chemgraph.models` path does not contain `ase_input` — this is either dead +code or a bug.) + +--- + +## Execution Order (with dependencies) + +``` +Phase 1A: Rename core.py → ase_core.py + └── Phase 1B: Move MACE schemas to schemas/mace_parsl_schema.py + └── Phase 1C: Add missing __init__.py files + └── Phase 4A: Delete mcp_helper.py + └── Phase 2A: Create cheminformatics_core.py + └── Phase 3A: Move extract_output_json to ase_core.py + └── Phase 2B: Create xanes_core.py + merge xanes_mcp.py + └── Phase 3B: Consolidate load_parsl_config to server_utils.py + └── Phase 2C: Create graspa_core.py (optional, lower priority) + └── Phase 3C: Fix report_tools.py element map + └── Phase 4B: Refactor new_eval example scripts + └── Phase 4C: Fix stale import in tool_call_eval.py +``` + +Suggested implementation order (each step leaves the repo in a passing state): + +1. **Phase 1A** — rename `core.py` → `ase_core.py`, update 4 import sites +2. **Phase 1B** — move MACE schemas, update 2 files +3. **Phase 1C** — add 2 empty `__init__.py` files +4. **Phase 4A** — fix 1 import in `cheminformatics_tools.py`, delete `mcp_helper.py` +5. **Phase 3A** — add `extract_output_json_core` to `ase_core.py`, update 3 wrappers +6. **Phase 2A** — create `cheminformatics_core.py`, simplify 2 files +7. **Phase 3C** — fix `report_tools.py` element map (2 occurrences) +8. **Phase 4C** — fix stale import (1 line) +9. **Phase 2B** — create `xanes_core.py`, merge `xanes_mcp.py` into `xanes_mcp_parsl.py` +10. **Phase 3B** — extract `load_parsl_config` to `server_utils.py` +11. **Phase 2C** — create `graspa_core.py` (optional) +12. **Phase 4B** — refactor `new_eval` example scripts + +--- + +## Target File Structure (after all phases) + +``` +src/chemgraph/ + schemas/ + __init__.py + ase_input.py # ASEInputSchema, ASEOutputSchema + atomsdata.py # AtomsData + graspa_schema.py # graspa_input_schema, graspa_input_schema_ensemble + graspa_input.py # GRASPAInputSchema (separate, pre-existing) + xanes_schema.py # xanes_input_schema, xanes_input_schema_ensemble + mace_parsl_schema.py # NEW: mace_input/output/ensemble schemas + agent_response.py # ResponseFormatter, etc. + multi_agent_response.py # MultiAgentResponse + calculators/ + __init__.py # NEW (empty) + mace_calc.py + emt_calc.py + tblite_calc.py + orca_calc.py + nwchem_calc.py + fairchem_calc.py + aimnet2_calc.py + mopac_calc.py + psi4_calc.py + + tools/ + __init__.py + ase_core.py # RENAMED from core.py + ase_tools.py # @tool wrappers → ase_core + cheminformatics_core.py # NEW: smiles_to_3d, name→SMILES, etc. + cheminformatics_tools.py # @tool wrappers → cheminformatics_core + graspa_core.py # NEW: run_graspa_core extracted + graspa_tools.py # @tool wrapper → graspa_core + xanes_core.py # NEW: run_xanes_core, helpers extracted + xanes_tools.py # @tool wrappers → xanes_core + parsl_tools.py # run_mace_core (schemas imported from schemas/) + report_tools.py # fixed element_map → ase.data.chemical_symbols + rag_tools.py # unchanged + generic_tools.py # unchanged + # mcp_helper.py # DELETED + + mcp/ + __init__.py # NEW (empty) + server_utils.py # + load_parsl_config() consolidated + mcp_tools.py # @mcp.tool → ase_core + cheminformatics_core + mace_mcp_parsl.py # schema import updated + graspa_mcp_parsl.py # → graspa_core, config from server_utils + xanes_mcp_parsl.py # → xanes_core, absorbed xanes_mcp.py + # xanes_mcp.py # DELETED (merged into xanes_mcp_parsl.py) + data_analysis_mcp.py # unchanged + + utils/ + __init__.py + async_utils.py # unchanged + config_utils.py # unchanged + get_workflow_from_llm.py # unchanged + logging_config.py # unchanged + parsing.py # unchanged + tool_call_eval.py # fixed stale import +``` + +--- + +## Test Commands + +After each phase, run the relevant subset: + +```bash +# Core ASE tests (phases 1A, 3A, 4A) +pytest tests/test_tools.py tests/test_mace.py tests/test_mcp.py tests/test_calculators.py -v + +# After phase 2A (cheminformatics) +pytest tests/test_tools.py tests/test_mcp.py -v -k "smiles or molecule" + +# Full suite (excluding pre-existing failures) +pytest tests/ -v \ + --ignore=tests/test_architector.py \ + --ignore=tests/test_multi_agent_retry.py \ + -k "not test_real_new_evaluation_ground_truth" +``` + +--- + +## Risks and Mitigations + +| Risk | Mitigation | +|------|------------| +| Breaking `new_eval/` scripts | Phase 4B is last; scripts are not part of CI | +| Parsl deferred imports break after rename | Update the deferred import inside `run_mace_parsl_app` in Phase 1B | +| `xanes_mcp.py` deletion breaks HPC users | Phase 2B merges functionality into `xanes_mcp_parsl.py` with conditional Parsl | +| Circular imports from `ase_core.py` | `ase_core.py` has no LangChain/MCP deps — no risk | +| Calculator schemas missing `__init__.py` | Phase 1C — adding empty file, no functional change | diff --git a/Dockerfile b/Dockerfile index 0d238836..3064e96d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -34,7 +34,7 @@ RUN pip install --no-cache-dir . && \ # Install Python tblite from source with conservative flags to avoid ABI/symbol issues on ARM. RUN CFLAGS="-O2 -fno-tree-vectorize" \ FFLAGS="-O2 -fno-tree-vectorize" \ - pip install --no-cache-dir --no-binary=tblite --force-reinstall "tblite==0.5.0" + pip install --no-cache-dir --no-binary=tblite --force-reinstall "tblite==0.4.0" # Validate calculator runtimes at build time after package install. RUN which nwchem && python -c "from tblite.ase import TBLite" diff --git a/README.md b/README.md index e5e46b35..d1c6df52 100644 --- a/README.md +++ b/README.md @@ -158,7 +158,6 @@ If you need to install from source for the latest version: 2. Create and activate a virtual environment using uv: ```bash uv venv --python 3.11 chemgraph-env - # uv venv --python 3.11 chemgraph-env # For specific python version source chemgraph-env/bin/activate # Unix/macos # OR @@ -218,9 +217,9 @@ If you need to install from source for the latest version: - **[Single-Agent System with UMA](notebooks/Demo_single_agent_UMA.ipynb)**: This notebook demonstrates how a single agent can utilize multiple tools with UMA support. - - **[Multi-Agent System](notebooks/2_Demo-multi_agent.ipynb)**: This notebook demonstrates a multi-agent setup where different agents (Planner, Executor and Aggregator) handle various tasks exemplifying the collaborative potential of ChemGraph. + - **[Multi-Agent System](notebooks/2_Demo-multi_agent.ipynb)**: This notebook demonstrates a multi-agent setup where planner and executor agents decompose and run computational chemistry tasks. - - **[Model Context Protocol (MCP) Server](notebooks/3_MCP_server.ipynb)**: This notebook demonstrates how to run an MCP server and connect to ChemGraph. + - **[Model Context Protocol (MCP) Server](notebooks/3_Demo_using_MCP.ipynb)**: This notebook demonstrates how to run an MCP server and connect to ChemGraph. - **[Single-Agent System with gRASPA](notebooks/Demo_graspa_agent.ipynb)**: This notebook provides a sample guide on executing a gRASPA simulation using a single agent. For gRASPA-related installation instructions, visit the [gRASPA GitHub repository](https://github.com/snurr-group/gRASPA). The notebook's functionality has been validated on a single compute node at ALCF Polaris. - **[Infrared absorption spectrum prediction](notebooks/Demo_infrared_spectrum.ipynb)**: This notebook demonstrates how to calculate an infrared absorption spectrum. @@ -231,16 +230,18 @@ If you need to install from source for the latest version:
Streamlit Web Interface -ChemGraph includes a **Streamlit web interface** that provides an intuitive, chat-based UI for interacting with computational chemistry agents. The interface supports 3D molecular visualization, conversation history, and easy access to various ChemGraph workflows. +ChemGraph includes a **Streamlit web interface** for chat-driven computational chemistry workflows. The UI auto-initializes the selected agent, streams tool-call progress while a query runs, shows generated structures and reports, and stores conversations in the same local session database used by the CLI. ### Features - **🧪 Interactive Chat Interface**: Natural language queries for computational chemistry tasks - **🧬 3D Molecular Visualization**: Interactive molecular structure display using `stmol` and `py3Dmol` -- **📊 Report Integration**: Embedded HTML reports from computational calculations +- **📊 Report Integration**: Embedded and downloadable HTML reports from computational calculations - **💾 Data Export**: Download molecular structures as XYZ or JSON files +- **🧮 Math Rendering**: Display LaTeX-style equations and reaction arrows in assistant responses - **🔧 Multiple Workflows**: Support for single-agent, multi-agent, Python REPL, and gRASPA workflows -- **🎨 Modern UI**: Clean, responsive interface with conversation bubbles and molecular properties display +- **💬 Session Memory**: Browse, load, and delete saved conversations from `~/.chemgraph/sessions.db` +- **👤 Human Supervision**: Optional follow-up prompts when the agent needs confirmation or missing inputs ### Installation Requirements @@ -278,16 +279,18 @@ pip install -e ".[uma]" ### Using the Interface #### Configuration -- **Model Selection**: Choose from GPT-4o, GPT-4o-mini, or Claude models -- **Workflow Type**: Select single-agent, multi-agent, Python REPL, or gRASPA workflows +- Use the **Configuration** page to edit `config.toml`, provider base URLs, API timeouts, workflow, recursion limit, report generation, and human supervision. +- API keys entered in the UI are applied only to the current Streamlit process and are not written to `config.toml`. +- The main sidebar shows calculators detected during ChemGraph initialization and marks the default calculator used when a query does not specify one. +- To change model, workflow, thread, or report settings, edit them on the **Configuration** page, save, then use **Reload Config** or **Refresh Agents**. #### Interaction -1. **Initialize Agent**: Click "Initialize Agent" in the sidebar to set up your ChemGraph instance -2. **Ask Questions**: Use the text area to enter computational chemistry queries -3. **View Results**: See responses in chat bubbles with automatic structure detection -4. **3D Visualization**: When molecular structures are detected, they're automatically displayed in 3D -5. **Download Data**: Export structures and calculation results directly from the interface +1. **Open the main page**: The agent initializes automatically from the active configuration. +2. **Ask Questions**: Use the chat input to enter computational chemistry queries. +3. **Monitor Tools**: Tool calls and completions stream in the assistant response while the workflow runs. +4. **Respond to Prompts**: If human supervision is enabled and the agent pauses, answer in the same chat input. +5. **View and Export Results**: Structures, IR artifacts, HTML reports, and download controls appear with the response when available. #### Example Queries - "What is the SMILES string for caffeine?" @@ -304,8 +307,9 @@ The interface automatically detects molecular structure data in agent responses #### Conversation Management - **History Display**: All queries and responses are preserved in conversation bubbles +- **Saved Sessions**: Recent sessions can be loaded or deleted from the sidebar - **Structure Detection**: Molecular structures are automatically extracted and visualized -- **Report Integration**: HTML reports from calculations are embedded directly in the interface +- **Report Integration**: HTML reports and run artifacts are embedded directly in the interface - **Debug Information**: Expandable sections show detailed message processing information ### Troubleshooting @@ -317,13 +321,14 @@ The interface automatically detects molecular structure data in agent responses **Agent Initialization:** - Verify API keys are set correctly +- Verify provider base URLs and local model endpoints on the Configuration page - Check that ChemGraph package is installed: `pip install -e .` - Ensure all dependencies are available in your environment **Performance:** - For large molecular systems, visualization may take longer to load -- Use the refresh button if the interface becomes unresponsive -- Clear conversation history to improve performance with many queries +- Start a new chat or load a smaller saved session if rendering many prior structures becomes slow +- Use **Refresh Agents** after changing credentials or external model services
@@ -343,7 +348,8 @@ Create a `config.toml` file in your project directory to configure ChemGraph beh [general] # Default model to use for queries model = "gpt-4o-mini" -# Workflow type: single_agent, multi_agent, python_repl, graspa +# Workflow type: single_agent, multi_agent, python_relp, graspa, mock_agent +# Alias accepted by CLI/UI: python_repl -> python_relp workflow = "single_agent" # Output format: state, last_message output = "state" @@ -351,9 +357,13 @@ output = "state" structured = false # Generate detailed reports report = true +# Default LangGraph thread ID +thread = 1 # Recursion limit for agent workflows recursion_limit = 20 +# Allow the agent to pause and ask for human input +human_supervised = false # Enable verbose output verbose = false @@ -460,6 +470,11 @@ rate_limit = true max_requests_per_minute = 60 ``` +The core CLI and UI currently consume `[general]`, `[api]`, `[chemistry]`, and +`[output]` directly. The agent uses deterministic LLM defaults internally +(`temperature=0.0`, fixed token limits); `[llm]` entries are kept for +documentation/forward compatibility rather than active runtime tuning. + ### Using Configuration Files #### With the Command Line Interface @@ -495,20 +510,17 @@ export OPENAI_API_KEY="" export ARGO_USER="" ``` -3. Use an Argo model ID (from `supported_argo_models` in `src/chemgraph/models/supported_models.py`): +3. Use an Argo model ID with the `argo:` prefix (from `supported_argo_models` in `src/chemgraph/models/supported_models.py`), for example: ```text -gpt4o, gpt4olatest, gpto3mini, gpto1, gpto3, gpto4mini, -gpt41, gpt41mini, gpt41nano, gpt5, gpt5mini, gpt5nano, gpt51, gpt52, -gemini25pro, gemini25flash, -claudeopus46, claudeopus45, claudeopus41, claudeopus4, -claudehaiku45, claudesonnet45, claudesonnet4, claudesonnet35v2, claudehaiku35 +argo:gpt-4o, argo:gpt-4o-latest, argo:gpt-5, argo:gpt-5-mini, +argo:gemini-2.5-flash, argo:claude-sonnet-4.5 ``` 4. Run with config: ```bash -chemgraph --config config.toml -m gpt4olatest -q "calculate the energy for water molecule using mace_mp" +chemgraph --config config.toml -m argo:gpt-4o-latest -q "calculate the energy for water molecule using mace_mp" ``` Notes: @@ -602,7 +614,7 @@ For Groq, the `groq:` prefix is stripped before sending to the Groq API. Any mod | Section | Description | | ------------- | ------------------------------------------------------- | | `[general]` | Basic settings like model, workflow, and output format | -| `[llm]` | LLM-specific parameters (temperature, max_tokens, etc.) | +| `[llm]` | Reserved/legacy LLM parameter documentation | | `[api]` | API endpoints and timeouts for different providers | | `[chemistry]` | Chemistry-specific calculation settings | | `[output]` | Output file formats and visualization settings | @@ -1058,7 +1070,7 @@ recursion_limit = 50 To generate a new ground-truth dataset from custom molecules and reactions: ```bash -cd scripts/new_evaluation +cd scripts/evaluations # Full execution (runs tool chains, captures actual results) python generate_ground_truth.py --input_file input_data.json @@ -1170,6 +1182,19 @@ If you use [Colmena](https://github.com/exalearn/colmena), run ChemGraph service 3. Mount the same project/output volume if Colmena workers and Docker run on the same host. Use this as an orchestration layer: Colmena schedules tasks; ChemGraph handles chemistry execution. + +### Kubernetes Deployment + +ChemGraph Streamlit can also be deployed on Kubernetes clusters. See the [`k8s/`](k8s/) directory for deployment manifests and instructions. + +**Quick deployment:** + +```bash +cd k8s +./deploy.sh deploy +``` + +For detailed instructions, see [`k8s/README.md`](k8s/README.md).
@@ -1190,26 +1215,27 @@ The servers are located in `src/chemgraph/mcp/`: * Supports **ensemble** calculations (directories of structures) using **Parsl** for parallel execution on HPC systems (e.g., Polaris, Aurora). * **`graspa_mcp_parsl.py`**: Tools for **gRASPA** simulations (Gas Adsorption in MOFs). * Supports single and ensemble runs via Parsl. +* **`xanes_mcp_parsl.py`**: Tools for running XANES/FDMNES ensembles with Parsl. * **`data_analysis_mcp.py`**: Tools for analyzing simulation results. * Aggregating JSONL logs from ensemble runs into CSV/DataFrames. * Plotting isotherms and other data. ### Running a Server -You can run the servers using Python. They support both `stdio` (default) and `streamable_http` (SSE) transports. +You can run the servers using Python. They support both `stdio` (default) and `streamable_http` transports. **Basic Usage (stdio)** Connect this directly to your MCP client (e.g., Claude Desktop config): ```bash -python src/chemgraph/mcp/mcp_tools.py +python -m chemgraph.mcp.mcp_tools ``` -**Using HTTP/SSE** +**Using streamable HTTP** To run a server that listens for HTTP connections (useful for remote deployment or debugging): ```bash -python src/chemgraph/mcp/mcp_tools.py --transport streamable_http --port 8000 +python -m chemgraph.mcp.mcp_tools --transport streamable_http --host 0.0.0.0 --port 8000 ``` **Configuration via Arguments** @@ -1219,7 +1245,7 @@ All servers in `src/chemgraph/mcp/` support the following arguments: * `--host`: Host address (default: 127.0.0.1). **Note on HPC Servers:** -For `mace_mcp_parsl.py` and `graspa_mcp_parsl.py`, ensure your environment is configured for the target HPC system if running actual parallel jobs. They leverage `chemgraph.hpc_configs` to load system-specific Parsl configurations (like `polaris` or `aurora`). +For `graspa_mcp_parsl.py` and `xanes_mcp_parsl.py`, set `COMPUTE_SYSTEM=polaris` or `COMPUTE_SYSTEM=aurora` before launch so the server loads the matching Parsl configuration from `chemgraph.hpc_configs`. `mace_mcp_parsl.py` currently contains site-specific `worker_init` settings; review and edit those paths/modules before production use.
diff --git a/config.toml b/config.toml deleted file mode 100644 index f4100a3e..00000000 --- a/config.toml +++ /dev/null @@ -1,109 +0,0 @@ -[general] -model = "argo:gpt-4o" -workflow = "single_agent" -output = "state" -structured = true -report = true -thread = 1 -recursion_limit = 20 -verbose = false - -[logging] -level = "WARNING" -file = "./chemgraph.log" -console = true -format = "%(asctime)s - %(name)s - %(levelname)s - %(message)s" - -[features] -enable_experimental = false -enable_cache = true -cache_dir = "./cache" -cache_expiry = 24 - -[security] -validate_keys = true -rate_limit = true -max_requests_per_minute = 60 - -[eval] -default_profile = "standard" - -[api.openai] -base_url = "https://apps.inside.anl.gov/argoapi/api/v1/resource/chat/" -timeout = 30 -argo_user = "" - -[api.groq] -base_url = "https://api.groq.com/openai/v1" -timeout = 30 - -[api.anthropic] -base_url = "https://api.anthropic.com" -timeout = 30 - -[api.google] -base_url = "https://generativelanguage.googleapis.com/v1beta" -timeout = 30 - -[api.alcf] -base_url = "https://inference-api.alcf.anl.gov/resource_server/sophia/vllm/v1" -timeout = 30 - -[api.local] -base_url = "http://localhost:11434" -timeout = 60 - -[chemistry.optimization] -method = "BFGS" -fmax = 0.05 -steps = 200 - -[chemistry.frequencies] -displacement = 0.01 -nprocs = 1 - -[chemistry.calculators] -default = "mace_mp" -fallback = "emt" - -[output.files] -directory = "./chemgraph_output" -pattern = "{timestamp}_{query_hash}" -formats = [ "xyz", "json", "html",] - -[output.visualization] -enable_3d = true -viewer = "py3dmol" -dpi = 300 - -[advanced.agent] -custom_system_prompt = "" -max_memory_tokens = 8000 -enable_function_calling = true - -[advanced.parallel] -enable_parallel = false -num_workers = 2 - -[environments.development] -model = "gpt-4o-mini" -verbose = true -enable_cache = false - -[environments.production] -model = "gpt-4o" -verbose = false -enable_cache = true -rate_limit = true - -[environments.testing] -model = "gpt-4o-mini" -verbose = true -enable_cache = false - -[eval.profiles.standard] -workflow_types = [ "single_agent",] -judge_model = "gpt4o" -recursion_limit = 50 -structured_output = true -max_queries = 0 diff --git a/docker-compose.yml b/docker-compose.yml index f8c52ec2..1e2180dd 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -13,6 +13,8 @@ x-chemgraph-common: &chemgraph-common GROQ_API_KEY: ${GROQ_API_KEY:-} CHEMGRAPH_LOG_DIR: /app/cg_logs PYTHONPATH: /app/src + ARGO_USER: ${ARGO_USER:-} + ALCF_ACCESS_TOKEN: ${ALCF_ACCESS_TOKEN:-} services: jupyter: diff --git a/docs/configuration_with_toml.md b/docs/configuration_with_toml.md index 31db7104..b5908d8d 100644 --- a/docs/configuration_with_toml.md +++ b/docs/configuration_with_toml.md @@ -12,7 +12,8 @@ Create a `config.toml` file in your project directory to configure ChemGraph beh [general] # Default model to use for queries model = "gpt-4o-mini" -# Workflow type: single_agent, multi_agent, python_repl, graspa +# Workflow type: single_agent, multi_agent, python_relp, graspa, mock_agent +# Alias accepted by CLI/UI: python_repl -> python_relp workflow = "single_agent" # Output format: state, last_message output = "state" @@ -20,9 +21,13 @@ output = "state" structured = false # Generate detailed reports report = true +# Default LangGraph thread ID +thread = 1 # Recursion limit for agent workflows recursion_limit = 20 +# Allow the agent to pause and ask for human input +human_supervised = false # Enable verbose output verbose = false @@ -128,6 +133,11 @@ rate_limit = true max_requests_per_minute = 60 ``` +The core CLI and UI currently consume `[general]`, `[api]`, `[chemistry]`, and +`[output]` directly. The agent uses deterministic LLM defaults internally +(`temperature=0.0`, fixed token limits); `[llm]` entries are kept for +documentation/forward compatibility rather than active runtime tuning. + ### Using Configuration Files #### With the Command Line Interface @@ -222,7 +232,7 @@ Direct model names (no prefix) are used for OpenAI, Anthropic, Google, ALCF, and | Section | Description | | ---------------- | ------------------------------------------------------- | | `[general]` | Basic settings like model, workflow, and output format | -| `[llm]` | LLM-specific parameters (temperature, max_tokens, etc.) | +| `[llm]` | Reserved/legacy LLM parameter documentation | | `[api]` | API endpoints and timeouts for different providers | | `[chemistry]` | Chemistry-specific calculation settings | | `[output]` | Output file formats and visualization settings | diff --git a/docs/example_usage.md b/docs/example_usage.md index 23359968..db7003a5 100644 --- a/docs/example_usage.md +++ b/docs/example_usage.md @@ -64,8 +64,8 @@ - **[Single-Agent System with UMA](https://github.com/argonne-lcf/ChemGraph/blob/main/notebooks/Demo_single_agent_UMA.ipynb)**: This notebook demonstrates how a single agent can utilize multiple tools with UMA support. - - **[Multi-Agent System](https://github.com/argonne-lcf/ChemGraph/blob/main/notebooks/Demo-multi_agent.ipynb)**: This notebook demonstrates a multi-agent setup where different agents (Planner, Executor and Aggregator) handle various tasks exemplifying the collaborative potential of ChemGraph. + - **[Multi-Agent System](https://github.com/argonne-lcf/ChemGraph/blob/main/notebooks/2_Demo-multi_agent.ipynb)**: This notebook demonstrates a multi-agent setup where planner and executor agents decompose and run computational chemistry tasks. - - **[Model Context Protocol (MCP) Server](https://github.com/argonne-lcf/ChemGraph/blob/main/notebooks/3_MCP_server.ipynb)**: This notebook shows how to run and connect to ChemGraph MCP tooling. + - **[Model Context Protocol (MCP) Server](https://github.com/argonne-lcf/ChemGraph/blob/main/notebooks/3_Demo_using_MCP.ipynb)**: This notebook shows how to run and connect to ChemGraph MCP tooling. - **[Single-Agent System with gRASPA](https://github.com/argonne-lcf/ChemGraph/blob/main/notebooks/Demo_graspa_agent.ipynb)**: This notebook provides a sample guide on executing a gRASPA simulation using a single agent. For gRASPA-related installation instructions, visit the [gRASPA GitHub repository](https://github.com/snurr-group/gRASPA). The notebook's functionality has been validated on a single compute node at ALCF Polaris. diff --git a/docs/mcp_servers.md b/docs/mcp_servers.md index dcf81005..85167655 100644 --- a/docs/mcp_servers.md +++ b/docs/mcp_servers.md @@ -6,6 +6,7 @@ - `mcp_tools.py`: general ASE-powered chemistry tools - `mace_mcp_parsl.py`: MACE + Parsl workflows - `graspa_mcp_parsl.py`: gRASPA + Parsl workflows +- `xanes_mcp_parsl.py`: XANES/FDMNES + Parsl workflows - `data_analysis_mcp.py`: analysis utilities for generated results ## Run a server @@ -85,6 +86,7 @@ The example config (`.opencode/opencode.example.jsonc`) includes all servers. En | `chemgraph` | `chemgraph.mcp.mcp_tools` | molecule_name_to_smiles, smiles_to_coordinate_file, run_ase, extract_output_json | Stable | `chemgraph-mace-parsl` | `chemgraph.mcp.mace_mcp_parsl` | MACE ensemble calculations via Parsl (HPC) | Experimental | `chemgraph-graspa-parsl` | `chemgraph.mcp.graspa_mcp_parsl` | gRASPA gas adsorption via Parsl (HPC) | Experimental +| `chemgraph-xanes-parsl` | `chemgraph.mcp.xanes_mcp_parsl` | XANES/FDMNES ensembles via Parsl (HPC) | Experimental | `chemgraph-data-analysis` | `chemgraph.mcp.data_analysis_mcp` | CIF splitting, JSONL aggregation, isotherm plotting | Experimental ### How it works @@ -93,4 +95,17 @@ OpenCode spawns the MCP server as a local child process using stdio transport. T ## Notes for Parsl-based servers -`mace_mcp_parsl.py` and `graspa_mcp_parsl.py` rely on Parsl and HPC-specific configuration. Ensure your environment is prepared for the target system before running production jobs. +Install the Parsl optional dependency when using HPC-backed servers: + +```bash +pip install -e ".[parsl]" +``` + +`graspa_mcp_parsl.py` and `xanes_mcp_parsl.py` load system-specific Parsl configuration through `COMPUTE_SYSTEM`: + +```bash +export COMPUTE_SYSTEM=polaris # or aurora +python -m chemgraph.mcp.graspa_mcp_parsl --transport streamable_http --host 0.0.0.0 --port 9001 +``` + +`mace_mcp_parsl.py` also uses Parsl, but currently contains site-specific `worker_init` settings in the module. Review the module loads, conda environment path, and filesystem paths before running production jobs. diff --git a/docs/streamlit_web_interface.md b/docs/streamlit_web_interface.md index 20b041e3..1a858b19 100644 --- a/docs/streamlit_web_interface.md +++ b/docs/streamlit_web_interface.md @@ -1,14 +1,17 @@ !!! note - ChemGraph includes a Streamlit web UI for chat-driven chemistry workflows, structure visualization, and report viewing. + ChemGraph includes a Streamlit web UI for chat-driven chemistry workflows, live tool progress, structure visualization, report viewing, and saved-session management. ## Run the app -Set provider keys as needed: +Install ChemGraph, then set the provider credentials required by the model you plan to use: ```bash export OPENAI_API_KEY="..." export ANTHROPIC_API_KEY="..." export GEMINI_API_KEY="..." +export GROQ_API_KEY="..." +# ALCF inference endpoints: +export ALCF_ACCESS_TOKEN="..." ``` Launch: @@ -21,11 +24,43 @@ Then open `http://localhost:8501`. ## Features -- Chat interface for single-agent and multi-agent workflows -- Model selection across supported providers -- 3D molecular visualization with `stmol`/`py3Dmol` -- Embedded report display and structure export -- Config editor for `config.toml` +- Chat input for single-agent, multi-agent, Python REPL, gRASPA, and mock-agent workflows exposed in the UI. +- Automatic agent initialization from the active `config.toml`. +- Sidebar calculator availability panel showing calculators detected at startup and the selected default. +- Live tool-call status while workflows run. +- Optional human-supervised pauses through the `ask_human` tool. +- 3D molecular visualization with `stmol` and `py3Dmol`, with table/XYZ fallback when the viewer is unavailable. +- Math-aware assistant rendering for LaTeX-style equations, reaction arrows, and thermochemistry expressions. +- Embedded and downloadable HTML reports, IR spectrum artifacts, normal-mode trajectory controls, and structure export. +- Session browser backed by `~/.chemgraph/sessions.db`. +- Configuration editor for `config.toml` plus session-only API key entry. + +## Configuration + +The UI reads `config.toml` from the working directory where Streamlit is launched. If the file is missing, the app creates one with defaults. + +Use the Configuration page for persistent settings: + +- `general.model`: default model. +- `general.workflow`: workflow type. The UI accepts `single_agent`, `multi_agent`, `python_relp`, `graspa`, and `mock_agent`; `python_repl` is accepted as an alias for `python_relp`. +- `general.thread`: default LangGraph thread ID. +- `general.recursion_limit`: workflow recursion limit. +- `general.report`: generate HTML reports when supported. +- `general.human_supervised`: allow the agent to pause and request human input. +- `api.*.base_url` and `api.*.timeout`: provider endpoint settings. +- `api.openai.argo_user`: optional Argo username; `ARGO_USER` is used only as a fallback. + +API keys entered in the UI are applied as process environment variables for the current Streamlit process and are not saved to `config.toml`. For shared deployments, prefer server-side environment variables. + +## Sessions + +The main sidebar lists recent saved sessions. Loading a session rebuilds the visible conversation history from `~/.chemgraph/sessions.db`; deleting a session removes it from that database. A new chat clears the visible conversation and starts a new saved session on the next successful exchange. + +The UI uses the active saved configuration for model, workflow, thread, report generation, and human-supervision settings. To change these settings, use the Configuration page, save the configuration, then click **Reload Config** or **Refresh Agents** on the main page. + +## Artifacts + +The UI detects structures and reports from agent messages. For IR calculations, it looks in the run directory referenced by the result message for files such as `ir_spectrum_.png`, `frequencies_.csv`, and `_vib..traj`. ## Troubleshooting @@ -33,3 +68,5 @@ Then open `http://localhost:8501`. `pip install stmol` - If model calls fail, verify API keys and endpoint settings in `config.toml`. - If Argo is used, ensure `api.openai.base_url` and optional `api.openai.argo_user` are configured. +- If a local model endpoint is selected, the UI probes `/models` and blocks queries when the local endpoint is unreachable. +- If the UI still shows an old model, workflow, or calculator default after editing configuration, click **Reload Config** or **Refresh Agents**. diff --git a/environment.yml b/environment.yml index 7c50dc6c..1a2cbc80 100644 --- a/environment.yml +++ b/environment.yml @@ -16,6 +16,9 @@ dependencies: - ase==3.25.0 - rdkit==2025.3.3 - langgraph==0.4.7 + - langgraph-checkpoint==2.1.0 + - langgraph-prebuilt==0.5.2 + - langgraph-sdk==0.1.72 - langchain - langchain-openai==0.3.27 - langchain-ollama==0.3.4 @@ -24,7 +27,6 @@ dependencies: - langchain-groq - langchain-experimental==0.3.4 - langchain-mcp-adapters - - langgraph-prebuilt - pydantic==2.11.7 - pubchempy==1.0.5 - pyppeteer==2.0.0 diff --git a/k8s/DEPLOYMENT_GUIDE.md b/k8s/DEPLOYMENT_GUIDE.md new file mode 100644 index 00000000..8f64b803 --- /dev/null +++ b/k8s/DEPLOYMENT_GUIDE.md @@ -0,0 +1,456 @@ +# ChemGraph Streamlit Kubernetes Deployment Guide + +This guide provides step-by-step instructions for deploying the ChemGraph Streamlit application on a Kubernetes cluster. + +## Table of Contents + +1. [Prerequisites](#prerequisites) +2. [Quick Start](#quick-start) +3. [Detailed Setup](#detailed-setup) +4. [Configuration Options](#configuration-options) +5. [Accessing the Application](#accessing-the-application) +6. [Troubleshooting](#troubleshooting) +7. [Production Considerations](#production-considerations) + +## Prerequisites + +Before you begin, ensure you have: + +- A Kubernetes cluster (v1.19+) +- `kubectl` installed and configured +- Access to push/pull Docker images (or use the public `ghcr.io/argonne-lcf/chemgraph:latest`) +- API keys for at least one LLM provider (OpenAI, Anthropic, Google, or Groq) + +## Quick Start + +The fastest way to deploy ChemGraph Streamlit: + +```bash +cd k8s + +# 1. Create and configure secrets +cp secrets.yaml.template secrets.yaml +# Edit secrets.yaml with your API keys + +# 2. Deploy using the deployment script +./deploy.sh deploy + +# 3. Check status +./deploy.sh status + +# 4. Access via port-forward (for testing) +./deploy.sh port-forward +# Then open http://localhost:8501 in your browser +``` + +## Detailed Setup + +### Step 1: Create Secrets + +Kubernetes Secrets are used to securely store your API keys. + +```bash +# Copy the template +cp secrets.yaml.template secrets.yaml + +# Edit the file with your actual API keys +vim secrets.yaml +``` + +Your `secrets.yaml` should look like: + +```yaml +apiVersion: v1 +kind: Secret +metadata: + name: chemgraph-secrets +type: Opaque +stringData: + openai-api-key: "sk-proj-..." + anthropic-api-key: "sk-ant-..." + gemini-api-key: "..." + groq-api-key: "..." + argo-user: "..." + alcf-access-token: "..." +``` + +**Important:** Do NOT commit this file! It's already in `.gitignore`. + +Apply the secret: + +```bash +kubectl apply -f secrets.yaml +``` + +### Step 2: Deploy the Application + +Apply the deployment manifest: + +```bash +kubectl apply -f deployment.yaml +``` + +This creates a Deployment with: +- 1 replica (can be scaled) +- 2Gi memory request, 4Gi limit +- 1 CPU request, 2 CPU limit +- Health checks (liveness and readiness probes) +- Environment variables for API keys + +### Step 3: Create the Service + +Apply the service manifest: + +```bash +kubectl apply -f service.yaml +``` + +This creates a LoadBalancer service that exposes the Streamlit app on port 8501. + +### Step 4: Verify Deployment + +Check the status: + +```bash +# Check pods +kubectl get pods -l app=chemgraph + +# Check deployment +kubectl get deployment chemgraph-streamlit + +# Check service +kubectl get svc chemgraph-streamlit + +# View logs +kubectl logs -l app=chemgraph,component=streamlit -f +``` + +## Configuration Options + +### Using a Custom Namespace + +Create and use a dedicated namespace: + +```bash +kubectl create namespace chemgraph +kubectl apply -f secrets.yaml -n chemgraph +kubectl apply -f deployment.yaml -n chemgraph +kubectl apply -f service.yaml -n chemgraph +``` + +Or use the deploy script: + +```bash +NAMESPACE=chemgraph ./deploy.sh deploy +``` + +### Scaling the Deployment + +To run multiple replicas: + +```bash +kubectl scale deployment chemgraph-streamlit --replicas=3 +``` + +Or edit `deployment.yaml`: + +```yaml +spec: + replicas: 3 +``` + +### Adjusting Resources + +Edit `deployment.yaml` to change resource allocation: + +```yaml +resources: + requests: + memory: "4Gi" # Increase for larger workloads + cpu: "2000m" + limits: + memory: "8Gi" + cpu: "4000m" +``` + +### Using a Custom Docker Image + +If you've built your own image: + +1. Build and push to your registry: + ```bash + docker build -t your-registry/chemgraph:v1.0 . + docker push your-registry/chemgraph:v1.0 + ``` + +2. Update `deployment.yaml`: + ```yaml + spec: + containers: + - name: streamlit + image: your-registry/chemgraph:v1.0 + ``` + +## Accessing the Application + +### Method 1: LoadBalancer (Cloud Providers) + +If your cluster supports LoadBalancer services: + +```bash +# Get the external IP +kubectl get svc chemgraph-streamlit + +# Access the app +# http://:8501 +``` + +### Method 2: NodePort (On-Premise) + +Edit `service.yaml`: + +```yaml +spec: + type: NodePort +``` + +Apply and get the NodePort: + +```bash +kubectl apply -f service.yaml +kubectl get svc chemgraph-streamlit + +# Access at http://: +``` + +### Method 3: Port Forwarding (Development) + +For local development/testing: + +```bash +kubectl port-forward svc/chemgraph-streamlit 8501:8501 + +# Access at http://localhost:8501 +``` + +Or use the deploy script: + +```bash +./deploy.sh port-forward +``` + +### Method 4: Ingress (Production) + +For production with a domain name: + +1. Install an Ingress controller (e.g., nginx-ingress) +2. Edit `ingress.yaml` with your domain +3. Apply: + ```bash + kubectl apply -f ingress.yaml + ``` + +## Troubleshooting + +### Pods Not Starting + +Check pod events: + +```bash +kubectl describe pod +``` + +Common issues: +- **ImagePullBackOff**: Check image name/tag and registry access +- **CrashLoopBackOff**: Check logs with `kubectl logs ` +- **Pending**: Check if nodes have sufficient resources + +### Application Not Accessible + +1. Verify the pod is running: + ```bash + kubectl get pods -l app=chemgraph + ``` + +2. Check service endpoints: + ```bash + kubectl get endpoints chemgraph-streamlit + ``` + +3. Test from inside the cluster: + ```bash + kubectl run -it --rm debug --image=curlimages/curl --restart=Never -- \ + curl http://chemgraph-streamlit:8501/_stcore/health + ``` + +### API Key Issues + +If the app can't access LLMs: + +1. Verify secrets are created: + ```bash + kubectl get secret chemgraph-secrets -o yaml + ``` + +2. Check if environment variables are set in pod: + ```bash + kubectl exec -- env | grep API_KEY + ``` + +3. View application logs: + ```bash + kubectl logs + ``` + +### Performance Issues + +1. Check resource usage: + ```bash + kubectl top pod + ``` + +2. Increase resource limits in `deployment.yaml` + +3. Scale horizontally: + ```bash + kubectl scale deployment chemgraph-streamlit --replicas=3 + ``` + +## Production Considerations + +### Security + +1. **Use Secrets Management**: Consider using external secrets (Vault, AWS Secrets Manager) + + Example with External Secrets Operator: + ```yaml + apiVersion: external-secrets.io/v1beta1 + kind: ExternalSecret + metadata: + name: chemgraph-secrets + spec: + secretStoreRef: + name: vault-backend + kind: SecretStore + target: + name: chemgraph-secrets + data: + - secretKey: openai-api-key + remoteRef: + key: chemgraph/openai + ``` + +2. **Enable RBAC**: Create a service account with minimal permissions + +3. **Network Policies**: Restrict pod-to-pod communication + + ```yaml + apiVersion: networking.k8s.io/v1 + kind: NetworkPolicy + metadata: + name: chemgraph-network-policy + spec: + podSelector: + matchLabels: + app: chemgraph + policyTypes: + - Ingress + - Egress + ingress: + - from: + - podSelector: {} + ports: + - protocol: TCP + port: 8501 + egress: + - to: + - namespaceSelector: {} + ``` + +4. **Enable TLS**: Use Ingress with cert-manager for HTTPS + +### High Availability + +1. **Multiple Replicas**: Run at least 3 replicas + +2. **Pod Disruption Budget**: + ```yaml + apiVersion: policy/v1 + kind: PodDisruptionBudget + metadata: + name: chemgraph-pdb + spec: + minAvailable: 1 + selector: + matchLabels: + app: chemgraph + ``` + +3. **Anti-Affinity**: Spread pods across nodes + ```yaml + affinity: + podAntiAffinity: + preferredDuringSchedulingIgnoredDuringExecution: + - weight: 100 + podAffinityTerm: + labelSelector: + matchExpressions: + - key: app + operator: In + values: + - chemgraph + topologyKey: kubernetes.io/hostname + ``` + +### Monitoring + +1. **Prometheus + Grafana**: Monitor resource usage and application metrics + +2. **Logging**: Use Fluentd/Fluent Bit to collect logs + +3. **Health Checks**: Already configured in `deployment.yaml` + +### Backup and Recovery + +1. **Persistent Storage**: If needed, mount a PVC for data persistence + + ```yaml + volumes: + - name: data + persistentVolumeClaim: + claimName: chemgraph-data-pvc + ``` + +2. **Disaster Recovery**: Keep deployment manifests in version control + +## Helper Commands + +```bash +# Deploy +./deploy.sh deploy + +# Check status +./deploy.sh status + +# View logs +./deploy.sh logs + +# Port forward +./deploy.sh port-forward + +# Delete +./deploy.sh delete + +# Manual commands +kubectl get all -l app=chemgraph +kubectl describe pod +kubectl logs -f +kubectl exec -it -- /bin/bash +``` + +## Additional Resources + +- [Kubernetes Documentation](https://kubernetes.io/docs/) +- [ChemGraph Documentation](../README.md) +- [Streamlit Deployment Guide](https://docs.streamlit.io/deploy) +- [Docker Guide](../docs/docker_support.md) diff --git a/k8s/README.md b/k8s/README.md new file mode 100644 index 00000000..d5ceefa0 --- /dev/null +++ b/k8s/README.md @@ -0,0 +1,216 @@ +# Kubernetes Deployment for ChemGraph Streamlit App + +This directory contains Kubernetes manifests to deploy the ChemGraph Streamlit application on a Kubernetes cluster. + +## Prerequisites + +- A running Kubernetes cluster +- `kubectl` configured to communicate with your cluster +- Docker image available at `ghcr.io/argonne-lcf/chemgraph:latest` (or build your own) +- API keys for the LLM providers you want to use + +## Files + +- `deployment.yaml` - Deployment manifest for the Streamlit app +- `service.yaml` - Service manifest to expose the app +- `secrets.yaml.template` - Template for storing API keys securely +- `ingress.yaml` - (Optional) Ingress configuration for external access + +## Quick Start + +### 1. Create the Secrets + +First, copy the secrets template and fill in your API keys: + +```bash +cp secrets.yaml.template secrets.yaml +``` + +Edit `secrets.yaml` and replace the placeholder values with your actual API keys: + +```yaml +stringData: + openai-api-key: "sk-..." + anthropic-api-key: "sk-ant-..." + gemini-api-key: "..." + groq-api-key: "..." + argo-user: "..." + alcf-access-token: "..." +``` + +**Important:** Never commit `secrets.yaml` to version control! Add it to `.gitignore`. + +Apply the secret: + +```bash +kubectl create namespace chemgraph # Optional: create a dedicated namespace +kubectl apply -f secrets.yaml +``` + +### 2. Deploy the Application + +Apply the deployment and service manifests: + +```bash +kubectl apply -f deployment.yaml +kubectl apply -f service.yaml +``` + +### 3. Access the Application + +Check the status of your deployment: + +```bash +kubectl get pods -l app=chemgraph +kubectl get svc chemgraph-streamlit +``` + +Get the external IP (if using LoadBalancer type): + +```bash +kubectl get svc chemgraph-streamlit +``` + +Once the service has an external IP, access the Streamlit app at: +``` +http://:8501 +``` + +If using a cloud provider, it may take a few minutes for the LoadBalancer to provision an external IP. + +## Alternative Access Methods + +### Using NodePort + +If your cluster doesn't support LoadBalancer, edit `service.yaml` and change the service type: + +```yaml +spec: + type: NodePort +``` + +Then access the app using any node IP and the assigned NodePort: + +```bash +kubectl get svc chemgraph-streamlit +# Access at http://: +``` + +### Using Port Forwarding (Development) + +For local testing, use port forwarding: + +```bash +kubectl port-forward svc/chemgraph-streamlit 8501:8501 +``` + +Then access at `http://localhost:8501` + +### Using Ingress (Production) + +For production deployments with a domain name, use the included ingress configuration: + +```bash +kubectl apply -f ingress.yaml +``` + +Make sure you have an Ingress controller installed (like nginx-ingress or traefik) and update the host in `ingress.yaml`. + +## Customization + +### Resource Limits + +Edit `deployment.yaml` to adjust resource requests and limits: + +```yaml +resources: + requests: + memory: "2Gi" + cpu: "1000m" + limits: + memory: "4Gi" + cpu: "2000m" +``` + +### Replicas + +To run multiple replicas for high availability: + +```yaml +spec: + replicas: 3 +``` + +### Using a Custom Image + +If you've built your own ChemGraph image: + +```yaml +containers: +- name: streamlit + image: your-registry/chemgraph:your-tag +``` + +## Monitoring and Troubleshooting + +### Check Pod Status + +```bash +kubectl get pods -l app=chemgraph +kubectl describe pod +``` + +### View Logs + +```bash +kubectl logs -f deployment/chemgraph-streamlit +``` + +### Check Health Probes + +The deployment includes liveness and readiness probes that check the Streamlit health endpoint: + +```bash +kubectl describe pod | grep -A 10 "Liveness\|Readiness" +``` + +### Common Issues + +**Pod not starting:** +- Check if the image can be pulled: `kubectl describe pod ` +- Verify secrets are created: `kubectl get secrets` + +**Application crashes:** +- Check logs: `kubectl logs ` +- Verify API keys are correct +- Check resource limits are sufficient + +**Cannot access the service:** +- Verify service is running: `kubectl get svc` +- Check if LoadBalancer has an external IP assigned +- Verify firewall rules allow traffic on port 8501 + +## Cleanup + +To remove all ChemGraph resources: + +```bash +kubectl delete -f deployment.yaml +kubectl delete -f service.yaml +kubectl delete secret chemgraph-secrets +``` + +## Security Considerations + +1. **Never commit secrets to version control** +2. **Use RBAC** to limit access to the namespace +3. **Enable TLS** for production deployments (use Ingress with cert-manager) +4. **Rotate API keys** regularly +5. **Use network policies** to restrict pod-to-pod communication if needed +6. **Consider using a secrets management solution** like HashiCorp Vault or Sealed Secrets + +## Additional Resources + +- [Kubernetes Documentation](https://kubernetes.io/docs/) +- [ChemGraph README](../README.md) +- [Streamlit Documentation](https://docs.streamlit.io/) diff --git a/k8s/deploy.sh b/k8s/deploy.sh new file mode 100755 index 00000000..8a283a5a --- /dev/null +++ b/k8s/deploy.sh @@ -0,0 +1,215 @@ +#!/bin/bash +# ChemGraph Kubernetes Deployment Script +# Deploys both Streamlit UI and MCP server + +set -e + +# Colors for output +RED='\033[0;31m' +GREEN='\033[0;32m' +YELLOW='\033[1;33m' +NC='\033[0m' # No Color + +# Function to print colored messages +print_info() { + echo -e "${GREEN}[INFO]${NC} $1" +} + +print_warn() { + echo -e "${YELLOW}[WARN]${NC} $1" +} + +print_error() { + echo -e "${RED}[ERROR]${NC} $1" +} + +# Check if kubectl is installed +if ! command -v kubectl &> /dev/null; then + print_error "kubectl is not installed. Please install it first." + exit 1 +fi + +# Default namespace (set early so the connection check can use it) +NAMESPACE="${NAMESPACE:-chemgraph}" + +# Check kubectl connection using the target namespace +# (cluster-info requires kube-system access which RBAC-limited users may not have) +if ! kubectl get namespace "$NAMESPACE" &> /dev/null; then + print_error "Cannot access namespace '$NAMESPACE'. Check your kubeconfig and permissions." + exit 1 +fi + +print_info "Connected to Kubernetes cluster (namespace: $NAMESPACE)" + +# Get the directory where this script is located +SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )" + +# Image tag override (defaults to whatever is in the YAML manifests) +IMAGE_TAG="${IMAGE_TAG:-}" +IMAGE_REPO="ghcr.io/argonne-lcf/chemgraph" + +# Parse command line arguments +ACTION="${1:-deploy}" +TARGET="${2:-}" + +# Helper: wait for LoadBalancer IP and print access URL +wait_for_lb() { + local svc_name="$1" + local port="$2" + local label="$3" + + print_info "Waiting for $label LoadBalancer IP..." + local external_ip="" + for i in {1..30}; do + external_ip=$(kubectl get svc "$svc_name" -n "$NAMESPACE" \ + -o jsonpath='{.status.loadBalancer.ingress[0].ip}' 2>/dev/null || echo "") + if [ -n "$external_ip" ]; then + print_info "$label is available at: http://$external_ip:$port" + return + fi + sleep 5 + done + print_warn "$label LoadBalancer IP not assigned yet." + print_info " Check: kubectl get svc $svc_name -n $NAMESPACE" + print_info " Or use: $0 port-forward ${svc_name#chemgraph-}" +} + +case "$ACTION" in + deploy) + print_info "Deploying ChemGraph (Streamlit + MCP) to namespace: $NAMESPACE" + + # Check if secrets file exists + if [ ! -f "$SCRIPT_DIR/secrets.yaml" ]; then + print_warn "secrets.yaml not found. Creating from template..." + cp "$SCRIPT_DIR/secrets.yaml.template" "$SCRIPT_DIR/secrets.yaml" + print_warn "Please edit k8s/secrets.yaml with your API keys before proceeding." + print_warn "Press Enter to continue after editing secrets.yaml, or Ctrl+C to cancel..." + read -r + fi + + # Apply secrets (shared by both services) + print_info "Creating secrets..." + kubectl apply -f "$SCRIPT_DIR/secrets.yaml" -n "$NAMESPACE" + + # Deploy Streamlit + print_info "Creating Streamlit deployment..." + kubectl apply -f "$SCRIPT_DIR/deployment.yaml" -n "$NAMESPACE" + kubectl apply -f "$SCRIPT_DIR/service.yaml" -n "$NAMESPACE" + + # Deploy MCP + print_info "Creating MCP deployment..." + kubectl apply -f "$SCRIPT_DIR/mcp-deployment.yaml" -n "$NAMESPACE" + kubectl apply -f "$SCRIPT_DIR/mcp-service.yaml" -n "$NAMESPACE" + + # Override image tag if IMAGE_TAG is set + if [ -n "$IMAGE_TAG" ]; then + print_info "Overriding image tag to: $IMAGE_REPO:$IMAGE_TAG" + kubectl set image deployment/chemgraph-streamlit streamlit="$IMAGE_REPO:$IMAGE_TAG" -n "$NAMESPACE" + kubectl set image deployment/chemgraph-mcp mcp="$IMAGE_REPO:$IMAGE_TAG" -n "$NAMESPACE" + fi + + # Wait for both rollouts + print_info "Waiting for deployments to be ready..." + kubectl rollout status deployment/chemgraph-streamlit -n "$NAMESPACE" --timeout=5m + kubectl rollout status deployment/chemgraph-mcp -n "$NAMESPACE" --timeout=5m + + print_info "Deployment successful!" + echo "" + + # Show service info + print_info "Service information:" + kubectl get svc -l app=chemgraph -n "$NAMESPACE" + echo "" + + # Wait for LoadBalancer IPs + wait_for_lb "chemgraph-streamlit" 8501 "Streamlit" + wait_for_lb "chemgraph-mcp" 9003 "MCP server" + ;; + + status) + print_info "Checking deployment status in namespace: $NAMESPACE" + echo "" + print_info "Pods:" + kubectl get pods -l app=chemgraph -n "$NAMESPACE" + echo "" + print_info "Services:" + kubectl get svc -l app=chemgraph -n "$NAMESPACE" + echo "" + print_info "Deployments:" + kubectl get deployment -l app=chemgraph -n "$NAMESPACE" + ;; + + logs) + case "$TARGET" in + streamlit) + print_info "Fetching Streamlit logs..." + kubectl logs -l app=chemgraph,component=streamlit -n "$NAMESPACE" --tail=100 -f + ;; + mcp) + print_info "Fetching MCP server logs..." + kubectl logs -l app=chemgraph,component=mcp -n "$NAMESPACE" --tail=100 -f + ;; + *) + print_info "Fetching logs from all ChemGraph pods..." + kubectl logs -l app=chemgraph -n "$NAMESPACE" --tail=100 -f --prefix + ;; + esac + ;; + + delete) + print_warn "This will delete all ChemGraph deployments from namespace: $NAMESPACE" + read -p "Are you sure? (y/N) " -n 1 -r + echo + if [[ $REPLY =~ ^[Yy]$ ]]; then + print_info "Deleting resources..." + kubectl delete -f "$SCRIPT_DIR/mcp-deployment.yaml" -n "$NAMESPACE" || true + kubectl delete -f "$SCRIPT_DIR/mcp-service.yaml" -n "$NAMESPACE" || true + kubectl delete -f "$SCRIPT_DIR/deployment.yaml" -n "$NAMESPACE" || true + kubectl delete -f "$SCRIPT_DIR/service.yaml" -n "$NAMESPACE" || true + print_info "Keeping secrets (delete manually if needed)" + print_info "Cleanup complete!" + else + print_info "Deletion cancelled" + fi + ;; + + port-forward) + case "$TARGET" in + streamlit) + print_info "Port forwarding Streamlit to localhost:8501" + kubectl port-forward svc/chemgraph-streamlit 8501:8501 -n "$NAMESPACE" + ;; + mcp) + print_info "Port forwarding MCP server to localhost:9003" + kubectl port-forward svc/chemgraph-mcp 9003:9003 -n "$NAMESPACE" + ;; + *) + print_info "Port forwarding both services (Ctrl+C to stop)..." + print_info " Streamlit: http://localhost:8501" + print_info " MCP: http://localhost:9003/mcp/" + kubectl port-forward svc/chemgraph-streamlit 8501:8501 -n "$NAMESPACE" & + PF_PID_STREAMLIT=$! + kubectl port-forward svc/chemgraph-mcp 9003:9003 -n "$NAMESPACE" & + PF_PID_MCP=$! + trap "kill $PF_PID_STREAMLIT $PF_PID_MCP 2>/dev/null; exit" INT TERM + wait + ;; + esac + ;; + + *) + echo "Usage: $0 {deploy|status|logs|delete|port-forward} [target]" + echo "" + echo "Commands:" + echo " deploy - Deploy Streamlit and MCP server to Kubernetes" + echo " status - Check deployment status" + echo " logs [streamlit|mcp] - View logs (default: all pods)" + echo " delete - Remove all ChemGraph deployments" + echo " port-forward [streamlit|mcp] - Forward ports to localhost (default: both)" + echo "" + echo "Environment variables:" + echo " NAMESPACE - Kubernetes namespace (default: chemgraph)" + echo " IMAGE_TAG - Override image tag (e.g. dev, sha-abc123, v1.0.0)" + exit 1 + ;; +esac diff --git a/k8s/deployment.yaml b/k8s/deployment.yaml new file mode 100644 index 00000000..9d2379d7 --- /dev/null +++ b/k8s/deployment.yaml @@ -0,0 +1,105 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: chemgraph-streamlit + labels: + app: chemgraph + component: streamlit +spec: + replicas: 1 + selector: + matchLabels: + app: chemgraph + component: streamlit + template: + metadata: + labels: + app: chemgraph + component: streamlit + spec: + containers: + - name: streamlit + image: ghcr.io/argonne-lcf/chemgraph:dev + imagePullPolicy: Always + ports: + - containerPort: 8501 + name: http + protocol: TCP + env: + - name: OPENAI_API_KEY + valueFrom: + secretKeyRef: + name: chemgraph-secrets + key: openai-api-key + optional: true + - name: ANTHROPIC_API_KEY + valueFrom: + secretKeyRef: + name: chemgraph-secrets + key: anthropic-api-key + optional: true + - name: GEMINI_API_KEY + valueFrom: + secretKeyRef: + name: chemgraph-secrets + key: gemini-api-key + optional: true + - name: GROQ_API_KEY + valueFrom: + secretKeyRef: + name: chemgraph-secrets + key: groq-api-key + optional: true + - name: ARGO_USER + valueFrom: + secretKeyRef: + name: chemgraph-secrets + key: argo-user + optional: true + - name: ALCF_ACCESS_TOKEN + valueFrom: + secretKeyRef: + name: chemgraph-secrets + key: alcf-access-token + optional: true + - name: CHEMGRAPH_LOG_DIR + value: /app/cg_logs + - name: PYTHONPATH + value: /app/src + - name: HTTP_PROXY + value: http://proxy.alcf.anl.gov:3128 + - name: HTTPS_PROXY + value: http://proxy.alcf.anl.gov:3128 + - name: http_proxy + value: http://proxy.alcf.anl.gov:3128 + - name: https_proxy + value: http://proxy.alcf.anl.gov:3128 + command: + - streamlit + - run + - src/ui/app.py + - --server.address=0.0.0.0 + - --server.port=8501 + resources: + requests: + memory: "2Gi" + cpu: "1000m" + limits: + memory: "4Gi" + cpu: "2000m" + livenessProbe: + httpGet: + path: /_stcore/health + port: 8501 + initialDelaySeconds: 30 + periodSeconds: 10 + timeoutSeconds: 5 + failureThreshold: 3 + readinessProbe: + httpGet: + path: /_stcore/health + port: 8501 + initialDelaySeconds: 10 + periodSeconds: 5 + timeoutSeconds: 3 + failureThreshold: 3 diff --git a/k8s/ingress.yaml b/k8s/ingress.yaml new file mode 100644 index 00000000..28b6e04c --- /dev/null +++ b/k8s/ingress.yaml @@ -0,0 +1,32 @@ +apiVersion: networking.k8s.io/v1 +kind: Ingress +metadata: + name: chemgraph-streamlit + labels: + app: chemgraph + component: streamlit + annotations: + # Adjust these based on your ingress controller + # Example for nginx-ingress: + nginx.ingress.kubernetes.io/rewrite-target: / + nginx.ingress.kubernetes.io/ssl-redirect: "true" + # Example for cert-manager (automatic HTTPS): + # cert-manager.io/cluster-issuer: letsencrypt-prod +spec: + ingressClassName: nginx # Adjust based on your ingress controller + rules: + - host: chemgraph.example.com # Replace with your domain + http: + paths: + - path: / + pathType: Prefix + backend: + service: + name: chemgraph-streamlit + port: + number: 8501 + # Uncomment for HTTPS with cert-manager: + # tls: + # - hosts: + # - chemgraph.example.com + # secretName: chemgraph-tls-cert diff --git a/k8s/mcp-deployment.yaml b/k8s/mcp-deployment.yaml new file mode 100644 index 00000000..2ba24853 --- /dev/null +++ b/k8s/mcp-deployment.yaml @@ -0,0 +1,107 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: chemgraph-mcp + labels: + app: chemgraph + component: mcp +spec: + replicas: 1 + selector: + matchLabels: + app: chemgraph + component: mcp + template: + metadata: + labels: + app: chemgraph + component: mcp + spec: + containers: + - name: mcp + image: ghcr.io/argonne-lcf/chemgraph:dev + imagePullPolicy: Always + ports: + - containerPort: 9003 + name: http + protocol: TCP + env: + - name: OPENAI_API_KEY + valueFrom: + secretKeyRef: + name: chemgraph-secrets + key: openai-api-key + optional: true + - name: ANTHROPIC_API_KEY + valueFrom: + secretKeyRef: + name: chemgraph-secrets + key: anthropic-api-key + optional: true + - name: GEMINI_API_KEY + valueFrom: + secretKeyRef: + name: chemgraph-secrets + key: gemini-api-key + optional: true + - name: GROQ_API_KEY + valueFrom: + secretKeyRef: + name: chemgraph-secrets + key: groq-api-key + optional: true + - name: ARGO_USER + valueFrom: + secretKeyRef: + name: chemgraph-secrets + key: argo-user + optional: true + - name: ALCF_ACCESS_TOKEN + valueFrom: + secretKeyRef: + name: chemgraph-secrets + key: alcf-access-token + optional: true + - name: CHEMGRAPH_LOG_DIR + value: /app/cg_logs + - name: PYTHONPATH + value: /app/src + - name: HTTP_PROXY + value: http://proxy.alcf.anl.gov:3128 + - name: HTTPS_PROXY + value: http://proxy.alcf.anl.gov:3128 + - name: http_proxy + value: http://proxy.alcf.anl.gov:3128 + - name: https_proxy + value: http://proxy.alcf.anl.gov:3128 + command: + - python + - -m + - chemgraph.mcp.mcp_tools + - --transport + - streamable_http + - --host + - 0.0.0.0 + - --port + - "9003" + resources: + requests: + memory: "2Gi" + cpu: "1000m" + limits: + memory: "4Gi" + cpu: "2000m" + livenessProbe: + tcpSocket: + port: 9003 + initialDelaySeconds: 30 + periodSeconds: 10 + timeoutSeconds: 5 + failureThreshold: 3 + readinessProbe: + tcpSocket: + port: 9003 + initialDelaySeconds: 10 + periodSeconds: 5 + timeoutSeconds: 3 + failureThreshold: 3 diff --git a/k8s/mcp-service.yaml b/k8s/mcp-service.yaml new file mode 100644 index 00000000..b5eb7554 --- /dev/null +++ b/k8s/mcp-service.yaml @@ -0,0 +1,17 @@ +apiVersion: v1 +kind: Service +metadata: + name: chemgraph-mcp + labels: + app: chemgraph + component: mcp +spec: + type: LoadBalancer + selector: + app: chemgraph + component: mcp + ports: + - name: http + protocol: TCP + port: 9003 + targetPort: 9003 diff --git a/k8s/secrets.yaml.template b/k8s/secrets.yaml.template new file mode 100644 index 00000000..27844e2f --- /dev/null +++ b/k8s/secrets.yaml.template @@ -0,0 +1,13 @@ +apiVersion: v1 +kind: Secret +metadata: + name: chemgraph-secrets +type: Opaque +stringData: + # Replace with your actual API keys + openai-api-key: "YOUR_OPENAI_API_KEY" + anthropic-api-key: "YOUR_ANTHROPIC_API_KEY" + gemini-api-key: "YOUR_GEMINI_API_KEY" + groq-api-key: "YOUR_GROQ_API_KEY" + argo-user: "YOUR_ARGO_USER" + alcf-access-token: "YOUR_ALCF_ACCESS_TOKEN" diff --git a/k8s/service.yaml b/k8s/service.yaml new file mode 100644 index 00000000..da65e084 --- /dev/null +++ b/k8s/service.yaml @@ -0,0 +1,17 @@ +apiVersion: v1 +kind: Service +metadata: + name: chemgraph-streamlit + labels: + app: chemgraph + component: streamlit +spec: + type: LoadBalancer + selector: + app: chemgraph + component: streamlit + ports: + - name: http + protocol: TCP + port: 8501 + targetPort: 8501 diff --git a/new_eval/scripts/mcp_example/mcp_http/start_mcp_server.py b/new_eval/scripts/mcp_example/mcp_http/start_mcp_server.py new file mode 100755 index 00000000..8fe84e8b --- /dev/null +++ b/new_eval/scripts/mcp_example/mcp_http/start_mcp_server.py @@ -0,0 +1,86 @@ +"""HTTP-based MCP server for ChemGraph chemistry tools. + +This is a thin wrapper that delegates to the core implementations +in :mod:`chemgraph.tools.ase_core` and :mod:`chemgraph.tools.cheminformatics_core`. +""" + +from __future__ import annotations + +import io +from contextlib import redirect_stdout +from typing import Literal + +import uvicorn +from mcp.server.fastmcp import FastMCP + +from chemgraph.tools.ase_core import extract_output_json_core, run_ase_core +from chemgraph.tools.cheminformatics_core import ( + molecule_name_to_smiles_core, + smiles_to_coordinate_file_core, +) +from chemgraph.schemas.ase_input import ASEInputSchema + +mcp = FastMCP( + name="Chemistry Tools MCP", + instructions=( + "You provide chemistry tools for converting molecule names to SMILES, " + "building 3D coordinates, running ASE simulations, and reading results. " + "Each tool has its own description — follow those to decide when to use them.\n\n" + "General guidance:\n" + "• Keep outputs compact; large results are written to files.\n" + "• Do not invent data. If a tool raises an error, report it as-is.\n" + "• Use absolute file paths when returning artifacts.\n" + "• Energies are in eV, vibrational frequencies in cm-1, wall times in seconds.\n" + ), +) + + +@mcp.tool( + name="molecule_name_to_smiles", + description="Convert a molecule name to a canonical SMILES string using PubChem.", +) +async def molecule_name_to_smiles(name: str) -> str: + """Resolve a molecule name to its canonical SMILES via PubChem.""" + return molecule_name_to_smiles_core(name) + + +@mcp.tool( + name="smiles_to_coordinate_file", + description="Convert a SMILES string to a coordinate file", +) +async def smiles_to_coordinate_file( + smiles: str, + output_file: str = "molecule.xyz", + randomSeed: int = 2025, + fmt: Literal["xyz"] = "xyz", +) -> dict: + """Convert a SMILES string to a coordinate file on disk.""" + return smiles_to_coordinate_file_core( + smiles, output_file=output_file, seed=randomSeed, fmt=fmt + ) + + +@mcp.tool( + name="extract_output_json", + description="Load simulation results from a JSON file produced by run_ase.", +) +def extract_output_json(json_file: str) -> dict: + """Load simulation results from a JSON file produced by run_ase.""" + return extract_output_json_core(json_file) + + +@mcp.tool( + name="run_ase", + description="Run ASE calculations using specified input parameters.", +) +async def run_ase(params: ASEInputSchema) -> dict: + """Run ASE calculations using specified input parameters.""" + f = io.StringIO() + with redirect_stdout(f): + return run_ase_core(params) + + +app = mcp.streamable_http_app() # exposes endpoints under /mcp + +if __name__ == "__main__": + uvicorn.run(app, host="127.0.0.1", port=9001) diff --git a/new_eval/scripts/mcp_example/mcp_stdio/mcp_tools_stdio.py b/new_eval/scripts/mcp_example/mcp_stdio/mcp_tools_stdio.py new file mode 100755 index 00000000..5315232e --- /dev/null +++ b/new_eval/scripts/mcp_example/mcp_stdio/mcp_tools_stdio.py @@ -0,0 +1,83 @@ +"""stdio-based MCP server for ChemGraph chemistry tools. + +This is a thin wrapper that delegates to the core implementations +in :mod:`chemgraph.tools.ase_core` and :mod:`chemgraph.tools.cheminformatics_core`. +""" + +from __future__ import annotations + +import io +from contextlib import redirect_stdout +from typing import Literal + +from mcp.server.fastmcp import FastMCP + +from chemgraph.tools.ase_core import extract_output_json_core, run_ase_core +from chemgraph.tools.cheminformatics_core import ( + molecule_name_to_smiles_core, + smiles_to_coordinate_file_core, +) +from chemgraph.schemas.ase_input import ASEInputSchema + +mcp = FastMCP( + name="Chemistry Tools MCP", + instructions=( + "You provide chemistry tools for converting molecule names to SMILES, " + "building 3D coordinates, running ASE simulations, and reading results. " + "Each tool has its own description — follow those to decide when to use them.\n\n" + "General guidance:\n" + "• Keep outputs compact; large results are written to files.\n" + "• Do not invent data. If a tool raises an error, report it as-is.\n" + "• Use absolute file paths when returning artifacts.\n" + "• Energies are in eV, vibrational frequencies in cm-1, wall times in seconds.\n" + ), +) + + +@mcp.tool( + name="molecule_name_to_smiles", + description="Convert a molecule name to a canonical SMILES string using PubChem.", +) +async def molecule_name_to_smiles(name: str) -> str: + """Resolve a molecule name to its canonical SMILES via PubChem.""" + return molecule_name_to_smiles_core(name) + + +@mcp.tool( + name="smiles_to_coordinate_file", + description="Convert a SMILES string to a coordinate file", +) +async def smiles_to_coordinate_file( + smiles: str, + output_file: str = "molecule.xyz", + randomSeed: int = 2025, + fmt: Literal["xyz"] = "xyz", +) -> dict: + """Convert a SMILES string to a coordinate file on disk.""" + return smiles_to_coordinate_file_core( + smiles, output_file=output_file, seed=randomSeed, fmt=fmt + ) + + +@mcp.tool( + name="extract_output_json", + description="Load simulation results from a JSON file produced by run_ase.", +) +def extract_output_json(json_file: str) -> dict: + """Load simulation results from a JSON file produced by run_ase.""" + return extract_output_json_core(json_file) + + +@mcp.tool( + name="run_ase", + description="Run ASE calculations using specified input parameters.", +) +async def run_ase(params: ASEInputSchema) -> dict: + """Run ASE calculations using specified input parameters.""" + f = io.StringIO() + with redirect_stdout(f): + return run_ase_core(params) + + +if __name__ == "__main__": + mcp.run() diff --git a/pyproject.toml b/pyproject.toml index 40ccb6a0..d61f3918 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "chemgraphagent" -version = "0.4.0" +version = "0.5.0" description = "A computational chemistry agent for molecular simulation tasks." authors = [ { name = "Thang Pham", email = "tpham@anl.gov" }, @@ -13,7 +13,10 @@ authors = [ ] requires-python = ">=3.10" dependencies = [ - "langgraph", + "langgraph==0.4.7", + "langgraph-checkpoint==2.1.0", + "langgraph-prebuilt==0.5.2", + "langgraph-sdk==0.1.72", "langchain", "langchain-openai", "langchain-ollama", @@ -22,7 +25,6 @@ dependencies = [ "langchain-groq", "langchain-experimental", "langchain-mcp-adapters", - "langgraph-prebuilt", "pydantic", "pandas==2.2.3", "pubchempy==1.0.5", @@ -30,6 +32,7 @@ dependencies = [ "numpy==2.2.6", "numexpr==2.11.0", "pytest==8.4.1", + "deepdiff==8.5.0", "ase", "rdkit", @@ -50,7 +53,7 @@ dependencies = [ [project.optional-dependencies] calculators = [ - "tblite==0.5.0; python_version < '3.13' or platform_system != 'Darwin'", + "tblite==0.4.0; python_version < '3.13' or platform_system != 'Darwin'", ] uma = [ "fairchem-core==2.3.0", @@ -102,6 +105,7 @@ indent-style = "space" # Use spaces for indentation skip-magic-trailing-comma = false # Ensure Black-style formatting [tool.pytest.ini_options] +testpaths = ["tests"] markers = [ "llm: marks tests as requiring LLM API access (run with --run-llm)", "asyncio: marks async tests", @@ -109,4 +113,7 @@ markers = [ filterwarnings = [ "ignore:Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected.*:UserWarning", "ignore:In future, it will be an error for 'np.bool' scalars to be interpreted as an index:DeprecationWarning", + "ignore:authlib.jose module is deprecated.*:authlib.deprecate.AuthlibDeprecationWarning", + "ignore:`torch.jit.script` is deprecated.*:DeprecationWarning:torch.jit._script", + "ignore:`torch.jit.load` is deprecated.*:DeprecationWarning:torch.jit._serialization", ] diff --git a/scripts/mcp_parsl_example/README.md b/scripts/mcp_parsl_example/README.md index ac65be63..eebc64a2 100644 --- a/scripts/mcp_parsl_example/README.md +++ b/scripts/mcp_parsl_example/README.md @@ -1,86 +1,127 @@ -# Using MCP+Parsl via Port-Forwarding on Aurora (ALCF) +# Using MCP + Parsl on Aurora or Polaris -This directory provides an example of how to use **MCP (Model Control Protocol)** with port-forwarding on **Aurora at ALCF**. The instructions below guide you through launching the MCP server and connecting ChemGraph to it. +This example shows how to run a ChemGraph Parsl-backed MCP server on an HPC +compute node and connect a ChemGraph client to it through port forwarding. + +Use this pattern when the LLM client should stay on a login node or workstation, +while simulation tools run on allocated compute resources. ## Prerequisites -- ChemGraph installed in your environment -- `OPENAI_API_KEY` set (or enter interactively when running ChemGraph) +- ChemGraph installed in the Python environment used on the compute node. +- `parsl` installed, either through the optional dependency or directly: + + ```bash + pip install -e ".[parsl]" + ``` -## Step-by-Step Instructions +- Any workflow-specific runtime installed, such as gRASPA, MACE, or FDMNES. +- Provider credentials for the LLM client, such as `OPENAI_API_KEY` or the + endpoint-specific variables used by your deployment. -### 1. Secure a Compute Node +## 1. Request a compute node -Request an interactive job on a compute node: +For a PBS-based ALCF system: ```bash -qsub -I -q debug -l select=1,walltime=60:00 -A your_account_name -l filesystems=flare +qsub -I -q debug -l select=1,walltime=01:00:00 -A your_account -l filesystems=home:flare ``` -### 2. SSH to the Compute Node + +After the allocation starts, connect to the compute node if your site requires a +separate SSH hop: + ```bash ssh YOUR_COMPUTE_NODE_ID ``` -### 3. Launch the MCP Server -Navigate to this directory, activate the environment and start the MCP server: + +## 2. Start the Parsl-backed MCP server + +Activate your environment and select the target system: + ```bash -# Set proxy for tools that query external databases +module load frameworks +source /path/to/venv/bin/activate + +export COMPUTE_SYSTEM=aurora # or polaris +export CHEMGRAPH_LOG_DIR="$PWD/chemgraph_mcp_logs" export http_proxy="proxy.alcf.anl.gov:3128" export https_proxy="proxy.alcf.anl.gov:3128" +export NO_PROXY=127.0.0.1,localhost,::1 +``` -# Load environment modules and activate your Python environment -module load frameworks -source /path/to/venv/bin/activate -pip install parsl # Run this to install Parsl (not yet included in ChemGraph pyproject.toml) +Start one of the Parsl-backed MCP servers. For gRASPA: + +```bash +python -m chemgraph.mcp.graspa_mcp_parsl \ + --transport streamable_http \ + --host 0.0.0.0 \ + --port 9001 +``` + +For XANES/FDMNES: -# Start MCP server -python -m chemgraph.tools.mcp_parsl +```bash +python -m chemgraph.mcp.xanes_mcp_parsl \ + --transport streamable_http \ + --host 0.0.0.0 \ + --port 9007 ``` -The server will run on port 9001 by default. -You can also launch the MCP server as a batch job using the `start_mcp_server.sub` script. -First, open the script and update the placeholders for your account name and path to your virtual environment. -Then submit the job with: +`mace_mcp_parsl.py` is also available, but it contains site-specific +`worker_init` settings in the module. Review module loads, conda environment +paths, and filesystem paths before using it in production. + +## 3. Or submit the server as a batch job + +Edit `start_mcp_server.sub` and update: + +- `#PBS -A your_account` +- the environment activation path +- `COMPUTE_SYSTEM` +- the MCP module and port if you want a server other than gRASPA + +Then submit: + ```bash qsub start_mcp_server.sub ``` -Once the job is running, you can find the compute node ID with: + +Find the compute node assigned to the job: + ```bash -qstat -f | awk -F'=' '/exec_host =/ {gsub(/^[ \t]+/,"",$2); sub(/\/.*/,"",$2); print $2}' +qstat -f JOB_ID | awk -F'=' '/exec_host =/ {gsub(/^[ \t]+/,"",$2); sub(/\/.*/,"",$2); print $2}' ``` -### 4. Set Up Port Forwarding -Open a new terminal on the login node, forwarding port 9001 so you can access the MCP server running on the compute node: +## 4. Forward the MCP port + +From the login node, forward the server port from the compute node: + ```bash ssh -N -L 9001:localhost:9001 YOUR_COMPUTE_NODE_ID ``` -Keep this terminal open while using ChemGraph. This ensures that all traffic from Aurora compute node to login node is routed through port 9001. -### 5. Launch ChemGraph -In another terminal session on the same login node used in Step 4, run ChemGraph and connect it to the MCP server (listening on port 9001 by default): +Keep this terminal open while the client runs. + +## 5. Run the ChemGraph client + +In another terminal on the login node: + ```bash -# Load environment modules and activate your Python environment module load frameworks source /path/to/venv/bin/activate +export NO_PROXY=127.0.0.1,localhost,::1 +export no_proxy=127.0.0.1,localhost,::1 python run_mcp_parsl.py ``` -### Troubleshooting +The example client connects to `http://127.0.0.1:9001/mcp/`. -If you get an error like this: -``` -httpx.HTTPStatusError: Server error '503 Service Unavailable' for url 'http://127.0.0.1:9001/mcp/' -``` -Try: -``` -export NO_PROXY=127.0.0.1,localhost,::1 -export no_proxy=127.0.0.1,localhost,::1 -``` -And run ChemGraph again. -======= -Finally, in another terminal, activate the environment and run ChemGraph to connect to the MCP server, which listens on port 9001. -```bash -python scripts/run_mcp_parsl/run_mcp_parsl.py -``` -python run_chemgraph.py -``` \ No newline at end of file +## Troubleshooting + +- If the client gets `503 Service Unavailable`, verify that the MCP server is + still running and that the SSH tunnel points to the correct compute node. +- If localhost requests go through the site proxy, set both `NO_PROXY` and + `no_proxy` for `127.0.0.1,localhost,::1`. +- If Parsl fails at startup, confirm that `PBS_NODEFILE` exists inside the + allocation and that `COMPUTE_SYSTEM` matches a supported config. diff --git a/scripts/mcp_parsl_example/start_mcp_server.sub b/scripts/mcp_parsl_example/start_mcp_server.sub index e77e1987..0e0c7c25 100644 --- a/scripts/mcp_parsl_example/start_mcp_server.sub +++ b/scripts/mcp_parsl_example/start_mcp_server.sub @@ -4,14 +4,21 @@ #PBS -l filesystems=home:flare #PBS -q debug #PBS -A your_account -#PBE -N ChemGraph +#PBS -N ChemGraphMCP cd $PBS_O_WORKDIR export http_proxy="proxy.alcf.anl.gov:3128" export https_proxy="proxy.alcf.anl.gov:3128" +export NO_PROXY=127.0.0.1,localhost,::1 +export no_proxy=127.0.0.1,localhost,::1 +export COMPUTE_SYSTEM=aurora +export CHEMGRAPH_LOG_DIR="$PBS_O_WORKDIR/chemgraph_mcp_logs" module load frameworks source /path/to/venv/bin/activate -python -m chemgraph.tools.mcp_parsl +python -m chemgraph.mcp.graspa_mcp_parsl \ + --transport streamable_http \ + --host 0.0.0.0 \ + --port 9001 diff --git a/src/chemgraph/__init__.py b/src/chemgraph/__init__.py index 216a3532..d1e0118f 100644 --- a/src/chemgraph/__init__.py +++ b/src/chemgraph/__init__.py @@ -1,11 +1,74 @@ - """ChemGraph package metadata.""" -from importlib.metadata import PackageNotFoundError, packages_distributions, version +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path +from typing import Any -try: - dist_names = packages_distributions().get("chemgraph", []) - __version__ = version(dist_names[0]) if dist_names else "unknown" -except (PackageNotFoundError, IndexError): - # Local source tree without installed package metadata. - __version__ = "unknown" +_DISTRIBUTION_NAMES = ("chemgraphagent", "chemgraph") + + +def _load_toml(path: Path) -> dict[str, Any]: + """Load TOML data from a file path. + + Parameters + ---------- + path : pathlib.Path + TOML file to read. + + Returns + ------- + dict[str, Any] + Parsed TOML data. + """ + text = path.read_text(encoding="utf-8") + try: + import tomllib + + return tomllib.loads(text) + except ModuleNotFoundError: + import toml + + return toml.loads(text) + + +def _version_from_distribution() -> str | None: + """Return the installed package version when distribution metadata exists. + + Returns + ------- + str or None + Installed version string, or ``None`` when ChemGraph is running from a + source checkout without package metadata. + """ + for dist_name in _DISTRIBUTION_NAMES: + try: + return version(dist_name) + except PackageNotFoundError: + continue + return None + + +def _version_from_pyproject() -> str | None: + """Return the project version from a source checkout's ``pyproject.toml``. + + Returns + ------- + str or None + Project version string, or ``None`` when the file cannot be found or + parsed. + """ + start = Path(__file__).resolve() + for directory in start.parents: + pyproject = directory / "pyproject.toml" + if not pyproject.exists(): + continue + try: + project = _load_toml(pyproject).get("project", {}) + except Exception: + return None + project_version = project.get("version") + return str(project_version) if project_version else None + return None + + +__version__ = _version_from_pyproject() or _version_from_distribution() or "unknown" diff --git a/src/chemgraph/agent/llm_agent.py b/src/chemgraph/agent/llm_agent.py index a397197a..d1f5b373 100644 --- a/src/chemgraph/agent/llm_agent.py +++ b/src/chemgraph/agent/llm_agent.py @@ -1,6 +1,8 @@ +import asyncio import datetime +import dataclasses import os -from typing import List, Optional +from typing import Callable, List, Optional import uuid from chemgraph.memory.store import SessionStore @@ -20,9 +22,15 @@ supported_gemini_models, ) +from chemgraph.schemas.ase_input import ( + get_available_calculator_names, + get_calculator_selection_context, + get_default_calculator_name, +) from chemgraph.prompt.single_agent_prompt import ( single_agent_prompt, + get_single_agent_prompt, formatter_prompt as default_formatter_prompt, report_prompt as default_report_prompt, ) @@ -32,7 +40,12 @@ aggregator_prompt as default_aggregator_prompt, planner_prompt as default_planner_prompt, ) +from langgraph.types import Command +from langgraph.errors import GraphInterrupt + from chemgraph.graphs.single_agent import construct_single_agent_graph + + from chemgraph.graphs.python_relp_agent import construct_relp_graph from chemgraph.graphs.multi_agent import construct_multi_agent_graph from chemgraph.graphs.graspa_agent import construct_graspa_graph @@ -52,30 +65,121 @@ logger = logging.getLogger(__name__) -def serialize_state(state): +def _is_mock_object(value) -> bool: + """Return True for unittest.mock objects without importing test-only APIs. + + Parameters + ---------- + value : Any + Object to inspect. + + Returns + ------- + bool + ``True`` when the object comes from ``unittest.mock``. + """ + return value.__class__.__module__.startswith("unittest.mock") + + +def serialize_state(state, *, max_depth: int = 50, _seen: set[int] | None = None): """Convert non-serializable objects in state to a JSON-friendly format. Parameters ---------- state : Any The state object to be serialized. Can be a list, dict, or object with __dict__ + max_depth : int, optional + Maximum object nesting depth to serialize before falling back to a + placeholder. This prevents runaway recursion for complex graph objects. Returns ------- Any A JSON-serializable version of the input state """ - if isinstance(state, (int, float, bool)) or state is None: + if _seen is None: + _seen = set() + + if max_depth < 0: + return f"" + + if isinstance(state, (str, int, float, bool)) or state is None: return state - elif isinstance(state, list): - return [serialize_state(item) for item in state] - elif isinstance(state, dict): - return {key: serialize_state(value) for key, value in state.items()} - elif hasattr(state, "__dict__"): - return {key: serialize_state(value) for key, value in state.__dict__.items()} - else: + + if isinstance(state, (datetime.datetime, datetime.date)): + return state.isoformat() + + if _is_mock_object(state): return str(state) + state_id = id(state) + if state_id in _seen: + return f"" + + if isinstance(state, dict): + _seen.add(state_id) + try: + return { + str(key): serialize_state( + value, max_depth=max_depth - 1, _seen=_seen + ) + for key, value in state.items() + } + finally: + _seen.remove(state_id) + + if isinstance(state, (list, tuple, set, frozenset)): + _seen.add(state_id) + try: + return [ + serialize_state(item, max_depth=max_depth - 1, _seen=_seen) + for item in state + ] + finally: + _seen.remove(state_id) + + model_dump = getattr(state, "model_dump", None) + if callable(model_dump): + _seen.add(state_id) + try: + try: + dumped = model_dump(mode="json") + except TypeError: + dumped = model_dump() + return serialize_state(dumped, max_depth=max_depth - 1, _seen=_seen) + except Exception: + return str(state) + finally: + _seen.remove(state_id) + + if dataclasses.is_dataclass(state) and not isinstance(state, type): + _seen.add(state_id) + try: + return { + field.name: serialize_state( + getattr(state, field.name), + max_depth=max_depth - 1, + _seen=_seen, + ) + for field in dataclasses.fields(state) + } + finally: + _seen.remove(state_id) + + if hasattr(state, "__dict__"): + _seen.add(state_id) + try: + return { + str(key): serialize_state( + value, max_depth=max_depth - 1, _seen=_seen + ) + for key, value in vars(state).items() + } + finally: + _seen.remove(state_id) + + return str(state) + class ChemGraph: """A graph-based workflow for LLM-powered computational chemistry tasks. @@ -114,6 +218,18 @@ class ChemGraph: max_retries : int, optional Maximum number of LLM retry attempts when an agent fails to parse its output, by default 1 + human_input_handler : callable, optional + A callback ``f(question: str) -> str`` invoked when the graph + pauses for human input (via ``interrupt()``). Receives the + question text and must return the human's answer as a string. + If ``None`` (default), interrupts will propagate as + ``GraphInterrupt`` exceptions. The handler may also be an + ``async`` callable. + human_supervised : bool, optional + Whether to include the ``ask_human`` tool so the agent can + pause and request human input. When ``False`` the tool is + excluded from the tool list and the corresponding instruction + is removed from the default system prompt, by default False. Raises ------ @@ -149,7 +265,66 @@ def __init__( memory_db_path: Optional[str] = None, log_dir: Optional[str] = None, max_retries: int = 1, + human_input_handler: Optional[Callable[[str], str]] = None, + human_supervised: bool = False, ): + """Initialize a ChemGraph workflow instance. + + Parameters + ---------- + model_name : str, optional + LLM model identifier. + workflow_type : str, optional + Workflow constructor key. + base_url : str, optional + Custom provider endpoint URL. + api_key : str, optional + API key passed to compatible model loaders. + argo_user : str, optional + Argo username for Argo-hosted models. + system_prompt : str, optional + System prompt for single-agent-style workflows. + formatter_prompt : str, optional + Prompt used to format single-agent final output. + structured_output : bool, optional + Whether structured final output is requested. + return_option : str, optional + Return mode, such as ``"last_message"`` or ``"state"``. + recursion_limit : int, optional + LangGraph recursion limit. + planner_prompt : str, optional + Planner prompt for multi-agent workflows. + executor_prompt : str, optional + Executor prompt for multi-agent workflows. + aggregator_prompt : str, optional + Aggregator prompt retained for compatibility. + formatter_multi_prompt : str, optional + Formatter prompt for multi-agent workflows. + generate_report : bool, optional + Whether report generation is enabled. + report_prompt : str, optional + Prompt used by the report-generation workflow. + support_structured_output : bool, optional + Whether the selected model supports structured output. + tools : list, optional + Custom tool list for applicable workflows. + data_tools : list, optional + Additional data-analysis tools for MCP workflows. + session_store : SessionStore, optional + Existing session store instance. + enable_memory : bool, optional + Whether persistent session memory is enabled. + memory_db_path : str, optional + SQLite path for the session store. + log_dir : str, optional + Directory for run logs and artifacts. + max_retries : int, optional + LLM parse-retry limit for formatter/planner nodes. + human_input_handler : Callable[[str], str], optional + Callback used to answer graph human-interrupt prompts. + human_supervised : bool, optional + Whether to expose human-supervision tools to the agent. + """ # Always generate a unique identifier for this instance self.uuid = str(uuid.uuid4())[:8] @@ -277,6 +452,41 @@ def __init__( self.tools = tools self.data_tools = data_tools self.max_retries = max_retries + self.human_input_handler = human_input_handler + self.human_supervised = human_supervised + + # When human supervision is disabled and the caller is using the + # default system prompt, strip the ask_human instructions so the + # LLM is not told to call a tool that is unavailable. + if not self.human_supervised and self.system_prompt == single_agent_prompt: + self.system_prompt = get_single_agent_prompt(human_supervised=False) + + self.available_calculators = get_available_calculator_names() + self.default_calculator = get_default_calculator_name() + self.calculator_selection_context = get_calculator_selection_context() + + def append_calculator_context(prompt: str) -> str: + """Append calculator availability guidance to a prompt once. + + Parameters + ---------- + prompt : str + Prompt text to augment. + + Returns + ------- + str + Prompt with calculator-selection context appended. + """ + if self.calculator_selection_context in prompt: + return prompt + return f"{prompt}{self.calculator_selection_context}" + + if self.workflow_type in {"single_agent", "mock_agent", "single_agent_mcp"}: + self.system_prompt = append_calculator_context(self.system_prompt) + elif self.workflow_type == "multi_agent": + self.planner_prompt = append_calculator_context(self.planner_prompt) + self.executor_prompt = append_calculator_context(self.executor_prompt) if model_name in supported_argo_models: self.support_structured_output = False @@ -310,6 +520,7 @@ def __init__( self.report_prompt, self.tools, max_retries=self.max_retries, + human_supervised=self.human_supervised, ) elif self.workflow_type == "multi_agent": self.workflow = self.workflow_map[workflow_type]["constructor"]( @@ -377,6 +588,18 @@ def visualize(self, method: str = "ascii"): This method creates and displays a visual representation of the workflow graph using Mermaid diagrams. The visualization is shown in Jupyter notebooks. + Parameters + ---------- + method : str, optional + Visualization backend. ``"ascii"`` returns an ASCII graph; + any other value renders a Mermaid PNG in the active notebook. + + Returns + ------- + str or None + ASCII graph text when ``method`` is ``"ascii"``; otherwise + displays an image and returns ``None``. + Notes ----- Requires IPython and nest_asyncio to be installed. @@ -547,7 +770,13 @@ def session_id(self) -> str: return self.uuid def _ensure_session(self, query: str) -> None: - """Create a session record on first run if memory is enabled.""" + """Create a session record on first run if memory is enabled. + + Parameters + ---------- + query : str + User query used to generate the session title. + """ if self.session_store is None: return if self._session_created: @@ -565,7 +794,15 @@ def _ensure_session(self, query: str) -> None: logger.info(f"Created session {self.uuid}: {self._session_title}") def _save_messages_to_store(self, last_state: dict, query: str) -> None: - """Extract messages from workflow state and persist to session store.""" + """Extract messages from workflow state and persist to session store. + + Parameters + ---------- + last_state : dict + Latest LangGraph state containing a ``messages`` sequence. + query : str + Original user query associated with the saved messages. + """ if self.session_store is None or not self._session_created: return @@ -593,6 +830,19 @@ def _save_messages_to_store(self, last_state: dict, query: str) -> None: content = msg.get("content", "") tool_name = msg.get("name") + # MCP tool messages may return content as a list of + # content blocks (e.g. [{'type': 'text', 'text': '...'}]) + # instead of a plain string. Normalize to str. + if isinstance(content, list): + content = "\n".join( + block.get("text", str(block)) + if isinstance(block, dict) + else str(block) + for block in content + ) + elif not isinstance(content, str): + content = str(content) + if role and content: messages_to_save.append( SessionMessage( @@ -640,12 +890,42 @@ def load_previous_context( return "" return self.session_store.build_context_summary(session_id) + async def _call_human_input_handler(self, question: str) -> str: + """Invoke the human_input_handler, supporting both sync and async callables. + + Raises :class:`HumanInputRequired` when no handler is configured, + allowing external callers (CLI, UI) to catch it, prompt the user, + and resume the graph. + + Parameters + ---------- + question : str + Prompt emitted by the graph for a human response. + + Returns + ------- + str + Human response returned by the configured handler. + """ + handler = self.human_input_handler + if handler is None: + raise HumanInputRequired(question) + if asyncio.iscoroutinefunction(handler): + return await handler(question) + return handler(question) + async def run(self, query: str, config=None, resume_from: Optional[str] = None): """ Async-only runner. Requires `self.workflow.astream(...)`. Streams values, logs new messages, writes state, and returns according to `self.return_option` ("last_message" or "state"). + When the graph pauses for human input (via ``interrupt()``), the + ``human_input_handler`` callback is invoked to obtain the user's + response, and the graph is automatically resumed. If no handler + is configured, the ``GraphInterrupt`` exception propagates to the + caller. + Parameters ---------- query : str @@ -658,6 +938,19 @@ async def run(self, query: str, config=None, resume_from: Optional[str] = None): """ def _validate_config(cfg): + """Normalize and validate the LangGraph run configuration. + + Parameters + ---------- + cfg : dict or None + User-provided configuration, optionally with top-level + ``thread_id``. + + Returns + ------- + dict + Config with ``configurable.thread_id`` and recursion limit set. + """ if cfg is None: cfg = {} if not isinstance(cfg, dict): @@ -676,6 +969,21 @@ def _validate_config(cfg): return cfg def _save_state_and_select_return(last_state, cfg): + """Persist the final state and apply the configured return option. + + Parameters + ---------- + last_state : dict + Final streamed graph state. + cfg : dict + LangGraph run configuration used to retrieve/write state. + + Returns + ------- + Any + Final message or serialized state, depending on + ``self.return_option``. + """ log_dir = self.log_dir if not log_dir: log_dir = "cg_logs" @@ -693,6 +1001,95 @@ def _save_state_and_select_return(last_state, cfg): f"Unsupported return_option: {self.return_option}. Use 'last_message' or 'state'." ) + async def _stream_until_interrupt(stream_input, cfg): + """Stream the workflow until completion or an interrupt. + + Parameters + ---------- + stream_input : dict or Command + Initial graph input or resume command to stream. + cfg : dict + LangGraph run configuration. + + Returns + ------- + tuple + ``(last_state, interrupt_value)`` where ``interrupt_value`` is + ``None`` when the graph completed normally. + + LangGraph's ``astream(stream_mode="values")`` does **not** + raise ``GraphInterrupt``. Instead the stream emits a state + containing an ``__interrupt__`` key and then ends. We + detect this in two ways: + + 1. Check for the ``__interrupt__`` key in streamed states. + 2. After the stream ends, inspect the checkpoint snapshot + for pending interrupt tasks. + """ + prev_msgs: list = [] + last_st = None + interrupt_val = None + try: + async for s in self.workflow.astream( + stream_input, stream_mode="values", config=cfg + ): + # Detect inline interrupt marker emitted by astream. + if "__interrupt__" in s: + int_data = s["__interrupt__"] + if isinstance(int_data, (list, tuple)) and int_data: + interrupt_val = int_data[0].value + elif hasattr(int_data, "value"): + interrupt_val = int_data.value + else: + interrupt_val = { + "question": "The workflow needs your input." + } + + if "messages" in s and s["messages"] != prev_msgs: + new_message = s["messages"][-1] + try: + new_message.pretty_print() + except Exception: + pass + logger.info(new_message) + prev_msgs = s["messages"] + last_st = s + except GraphInterrupt as gi: + # Fallback: some LangGraph versions may still raise. + interrupts = gi.args[0] if gi.args else [] + if interrupts: + interrupt_val = interrupts[0].value + else: + interrupt_val = { + "question": "The workflow needs your input." + } + + # Double-check the checkpoint for pending interrupts that + # the stream may not have surfaced explicitly. + if interrupt_val is None: + try: + snapshot = self.workflow.get_state(cfg) + if snapshot and snapshot.tasks: + for t in snapshot.tasks: + t_interrupts = getattr(t, "interrupts", None) + if t_interrupts: + interrupt_val = t_interrupts[0].value + break + except Exception: + pass + + if interrupt_val is not None: + logger.info("Graph interrupted: %s", interrupt_val) + # Refresh state from checkpoint for consistency. + try: + snapshot = self.workflow.get_state(cfg) + if snapshot: + last_st = snapshot.values + except Exception: + pass + + return last_st, interrupt_val + logger.debug("run called with config=%s", config) config = _validate_config(config) logger.debug("validated config=%s", config) @@ -718,21 +1115,47 @@ def _save_state_and_select_return(last_state, cfg): inputs = {"messages": query} - prev_messages = [] - last_state = None try: - async for s in self.workflow.astream( - inputs, stream_mode="values", config=config - ): - if "messages" in s and s["messages"] != prev_messages: - new_message = s["messages"][-1] - try: - new_message.pretty_print() - except Exception: - pass - logger.info(new_message) - prev_messages = s["messages"] - last_state = s + last_state, interrupt_value = await _stream_until_interrupt(inputs, config) + + # --- Human-in-the-loop resume loop --- + # When the graph pauses with an interrupt, ask the human and + # resume. This loop handles chains of multiple interrupts + # (e.g., the agent asks a follow-up question after receiving + # the first answer). + max_interrupts = 10 # safety guard against infinite interrupt loops + interrupt_count = 0 + while interrupt_value is not None: + interrupt_count += 1 + if interrupt_count > max_interrupts: + logger.error( + "Exceeded maximum number of human interrupts (%d); " + "aborting workflow.", + max_interrupts, + ) + raise RuntimeError( + f"Workflow exceeded maximum of {max_interrupts} " + f"human interrupts." + ) + + # Extract the question text from the interrupt value. + if isinstance(interrupt_value, dict): + question = interrupt_value.get( + "question", + interrupt_value.get("message", str(interrupt_value)), + ) + else: + question = str(interrupt_value) + + logger.info("Requesting human input: %s", question) + human_answer = await self._call_human_input_handler(question) + logger.info("Human responded: %s", human_answer) + + # Resume the graph from the checkpoint with the human's answer. + resume_cmd = Command(resume=human_answer) + last_state, interrupt_value = await _stream_until_interrupt( + resume_cmd, config + ) if last_state is None: raise RuntimeError("Workflow produced no states.") @@ -742,6 +1165,29 @@ def _save_state_and_select_return(last_state, cfg): return _save_state_and_select_return(last_state, config) + except HumanInputRequired: + # No human_input_handler configured — propagate so the + # caller (CLI / UI) can prompt the user and resume. + raise except Exception as e: logger.error(f"Error running workflow {self.workflow_type}: {e}") raise + +class HumanInputRequired(Exception): + """Raised when the graph needs human input but no handler is configured. + + Carries the question text so that external callers (CLI, UI) can + present it to the user and resume the graph with + ``Command(resume=answer)``. + """ + + def __init__(self, question: str): + """Initialize the exception with the pending human question. + + Parameters + ---------- + question : str + Question that should be presented to the user. + """ + self.question = question + super().__init__(question) diff --git a/src/chemgraph/cli/commands.py b/src/chemgraph/cli/commands.py index 6f490f38..abbd0fff 100644 --- a/src/chemgraph/cli/commands.py +++ b/src/chemgraph/cli/commands.py @@ -58,7 +58,18 @@ def resolve_workflow(name: str) -> str: - """Resolve a workflow name, applying aliases.""" + """Resolve a workflow name, applying aliases. + + Parameters + ---------- + name : str + Workflow name or supported alias. + + Returns + ------- + str + Canonical workflow name. + """ return WORKFLOW_ALIASES.get(name, name) @@ -70,7 +81,16 @@ def resolve_workflow(name: str) -> str: def check_api_keys(model_name: str) -> tuple[bool, str]: """Check if required API keys are available for *model_name*. - Returns ``(is_available, error_message)``. + Parameters + ---------- + model_name : str + Model identifier selected for a run. + + Returns + ------- + tuple[bool, str] + ``(is_available, error_message)``. The message is empty when the + required credentials are available or not required. """ model_lower = model_name.lower() @@ -155,11 +175,44 @@ def initialize_agent( base_url: Optional[str] = None, argo_user: Optional[str] = None, verbose: bool = False, + human_supervised: bool = False, + tools: Optional[list] = None, ) -> Any: """Initialize a ChemGraph agent with progress indication. Uses a thread-pool executor for the timeout so it works on all platforms. + + Parameters + ---------- + model_name : str + LLM model identifier. + workflow_type : str + ChemGraph workflow name or alias. + structured_output : bool + Whether to request structured final output. + return_option : str + Agent return mode, such as ``"state"`` or ``"last_message"``. + generate_report : bool + Whether the agent should generate an HTML report. + recursion_limit : int + LangGraph recursion limit for the run. + base_url : str, optional + Custom model endpoint URL. + argo_user : str, optional + Argo username for Argo-hosted models. + verbose : bool, optional + Whether to print initialization details. + human_supervised : bool, optional + Whether to enable human-interrupt tooling. + tools : list, optional + Custom tool list for MCP-backed workflows. + + Returns + ------- + Any + Initialized ``ChemGraph`` instance, or ``None`` when initialization + fails. """ # Resolve workflow alias before initializing. workflow_type = resolve_workflow(workflow_type) @@ -171,11 +224,14 @@ def initialize_agent( console.print(f" Structured Output: {structured_output}") console.print(f" Return Option: {return_option}") console.print(f" Generate Report: {generate_report}") + console.print(f" Human Supervised: {human_supervised}") console.print(f" Recursion Limit: {recursion_limit}") if base_url: console.print(f" Base URL: {base_url}") if argo_user: console.print(f" Argo User: {argo_user}") + if tools: + console.print(f" MCP Tools: {len(tools)} loaded") # Check API keys before attempting initialization api_key_available, error_msg = check_api_keys(model_name) @@ -203,6 +259,13 @@ def initialize_agent( task = progress.add_task("Initializing ChemGraph agent...", total=None) def _create_agent() -> Any: + """Create the ChemGraph agent inside the initialization worker. + + Returns + ------- + Any + Initialized ``ChemGraph`` instance. + """ from chemgraph.agent.llm_agent import ChemGraph return ChemGraph( @@ -215,6 +278,8 @@ def _create_agent() -> Any: return_option=return_option, recursion_limit=recursion_limit, structured_output=structured_output, + human_supervised=human_supervised, + tools=tools, ) try: @@ -260,6 +325,13 @@ def _create_agent() -> Any: def _next_thread_id() -> int: + """Return the next interactive-mode thread ID. + + Returns + ------- + int + Incremented thread ID. + """ global _thread_counter _thread_counter += 1 return _thread_counter @@ -272,7 +344,35 @@ def run_query( verbose: bool = False, resume_from: Optional[str] = None, ) -> Any: - """Execute a query with the agent.""" + """Execute a query with the agent. + + When the graph pauses for human input (``HumanInputRequired``), the + spinner is stopped, the question is shown in a Rich panel, and the + user is prompted for a response. The graph is then resumed with the + user's answer and the spinner restarts. This loop repeats until the + graph completes or a non-interrupt error occurs. + + Parameters + ---------- + agent : Any + Initialized ChemGraph-like agent with ``run`` and ``workflow`` methods. + query : str + User query to execute. + thread_id : int, optional + LangGraph thread identifier. A new ID is allocated when omitted. + verbose : bool, optional + Whether to print execution details. + resume_from : str, optional + Previous ChemGraph session ID to load as context. + + Returns + ------- + Any + Agent result, resumed graph result, or ``None`` on failure. + """ + from langgraph.types import Command + from chemgraph.agent.llm_agent import HumanInputRequired + if thread_id is None: thread_id = _next_thread_id() @@ -282,6 +382,11 @@ def run_query( if resume_from: console.print(f"[blue]Resuming from session:[/blue] {resume_from}") + config = {"configurable": {"thread_id": thread_id}} + max_interrupts = 10 # safety guard + interrupt_count = 0 + + # --- First invocation: run the full agent.run() --- with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), @@ -289,22 +394,94 @@ def run_query( transient=True, ) as progress: task = progress.add_task("Processing query...", total=None) - try: - config = {"configurable": {"thread_id": thread_id}} result = run_async_callable( lambda: agent.run(query, config=config, resume_from=resume_from) ) - progress.update(task, description="[green]Query completed!") - time.sleep(0.5) + time.sleep(0.3) return result - + except HumanInputRequired as hir: + progress.update(task, description="[yellow]Agent needs your input") + time.sleep(0.2) + question = hir.question except Exception as e: progress.update(task, description="[red]Query failed!") console.print(f"[red]Error processing query: {e}[/red]") return None + # --- Interrupt-resume loop --- + # The spinner's `with` block has exited, so the terminal is free + # for interactive user input. + while question is not None: + interrupt_count += 1 + if interrupt_count > max_interrupts: + console.print( + "[red]Exceeded maximum number of human interrupts. Aborting.[/red]" + ) + return None + + console.print( + Panel( + question, + title="[bold yellow]Agent needs your input[/bold yellow]", + style="yellow", + ) + ) + human_answer = Prompt.ask("[bold cyan]Your response[/bold cyan]") + + # Resume the graph, streaming messages so tool-call parameters + # are printed just like the initial invocation. + resume_config = dict(config) + resume_config["recursion_limit"] = agent.recursion_limit + + async def _resume_stream(): + """Resume an interrupted graph and stream updates until completion. + + Returns + ------- + dict or None + Final streamed graph state. + """ + prev_msgs: list = [] + last_st = None + async for s in agent.workflow.astream( + Command(resume=human_answer), + stream_mode="values", + config=resume_config, + ): + if "messages" in s and s["messages"] != prev_msgs: + new_message = s["messages"][-1] + try: + new_message.pretty_print() + except Exception: + pass + prev_msgs = s["messages"] + last_st = s + return last_st + + try: + result = run_async_callable(_resume_stream) + + if result is None: + console.print("[red]Resume produced no output.[/red]") + return None + + if agent.return_option == "last_message": + return result["messages"][-1] if result else None + elif agent.return_option == "state": + from chemgraph.agent.llm_agent import serialize_state + + return serialize_state(agent.get_state(config=config)) + return result + except HumanInputRequired as hir: + question = hir.question + except Exception as e: + console.print(f"[red]Error processing query: {e}[/red]") + return None + + return None + # --------------------------------------------------------------------------- # Session management @@ -312,7 +489,15 @@ def run_query( def list_sessions(limit: int = 20, db_path: Optional[str] = None) -> None: - """Display recent sessions in a formatted table.""" + """Display recent sessions in a formatted table. + + Parameters + ---------- + limit : int, optional + Maximum number of sessions to display. + db_path : str, optional + Path to the session SQLite database. + """ store = SessionStore(db_path=db_path) sessions = store.list_sessions(limit=limit) @@ -354,7 +539,17 @@ def show_session( db_path: Optional[str] = None, max_content: int = 500, ) -> None: - """Display a session's full conversation.""" + """Display a session's full conversation. + + Parameters + ---------- + session_id : str + Session ID or unique session prefix. + db_path : str, optional + Path to the session SQLite database. + max_content : int, optional + Maximum number of characters displayed for each message. + """ store = SessionStore(db_path=db_path) session = store.get_session(session_id) @@ -414,7 +609,15 @@ def show_session( def delete_session_cmd(session_id: str, db_path: Optional[str] = None) -> None: - """Delete a session from the database.""" + """Delete a session from the database. + + Parameters + ---------- + session_id : str + Session ID or unique session prefix to delete. + db_path : str, optional + Path to the session SQLite database. + """ store = SessionStore(db_path=db_path) # Show session info before deleting @@ -440,7 +643,15 @@ def delete_session_cmd(session_id: str, db_path: Optional[str] = None) -> None: def save_output(content: str, output_file: str) -> None: - """Save output to a file.""" + """Save output to a file. + + Parameters + ---------- + content : str + Text content to write. + output_file : str + Destination file path. + """ try: with open(output_file, "w") as f: f.write(content) @@ -460,16 +671,43 @@ def interactive_mode( structured: bool = False, return_option: str = "state", generate_report: bool = True, + human_supervised: bool = False, recursion_limit: int = 20, base_url: Optional[str] = None, argo_user: Optional[str] = None, verbose: bool = False, + tools: Optional[list] = None, ) -> None: """Start interactive REPL mode for ChemGraph CLI. Accepts the same configuration parameters as a normal run so that ``--config`` and CLI flags are honoured when entering interactive mode. + + Parameters + ---------- + model : str, optional + Initial model selection. + workflow : str, optional + Initial workflow selection. + structured : bool, optional + Whether structured output is requested. + return_option : str, optional + Agent return mode. + generate_report : bool, optional + Whether report generation is enabled. + human_supervised : bool, optional + Whether human supervision tools are enabled. + recursion_limit : int, optional + LangGraph recursion limit. + base_url : str, optional + Custom model endpoint URL. + argo_user : str, optional + Argo username for Argo-hosted models. + verbose : bool, optional + Whether to print diagnostic output. + tools : list, optional + Custom tool list for MCP-backed workflows. """ console.print(create_banner()) console.print("[bold green]Welcome to ChemGraph Interactive Mode![/bold green]") @@ -501,6 +739,8 @@ def interactive_mode( base_url=base_url, argo_user=argo_user, verbose=verbose, + human_supervised=human_supervised, + tools=tools, ) if not agent: return @@ -593,6 +833,8 @@ def interactive_mode( recursion_limit, base_url=base_url, argo_user=argo_user, + human_supervised=human_supervised, + tools=tools, ) if agent: console.print(f"[green]Model changed to: {model}[/green]") @@ -610,6 +852,8 @@ def interactive_mode( recursion_limit, base_url=base_url, argo_user=argo_user, + human_supervised=human_supervised, + tools=tools, ) if agent: console.print( diff --git a/src/chemgraph/cli/formatting.py b/src/chemgraph/cli/formatting.py index 27f4b433..be3b8cce 100644 --- a/src/chemgraph/cli/formatting.py +++ b/src/chemgraph/cli/formatting.py @@ -167,6 +167,16 @@ def _is_atomic_json(content: str) -> bool: This replaces the old fragile substring check (Bug 10) with a proper parse attempt. + + Parameters + ---------- + content : str + Candidate JSON text. + + Returns + ------- + bool + ``True`` when the parsed object contains atomic-structure keys. """ try: data = json.loads(content.strip()) @@ -179,7 +189,15 @@ def _is_atomic_json(content: str) -> bool: def format_response(result: Any, verbose: bool = False) -> None: - """Format the agent response for display.""" + """Format the agent response for display. + + Parameters + ---------- + result : Any + Agent result, message list, state dictionary, or message object. + verbose : bool, optional + Whether to include raw message details. + """ if not result: console.print("[red]No response received from agent.[/red]") return diff --git a/src/chemgraph/cli/main.py b/src/chemgraph/cli/main.py index 3f4baf15..badf4168 100644 --- a/src/chemgraph/cli/main.py +++ b/src/chemgraph/cli/main.py @@ -59,6 +59,11 @@ def _add_run_args(parser: argparse.ArgumentParser) -> None: Used by both the ``run`` subcommand and the legacy (no subcommand) argument parser for backward compatibility. + + Parameters + ---------- + parser : argparse.ArgumentParser + Parser or subparser to receive query/run arguments. """ parser.add_argument( "-q", "--query", type=str, help="The computational chemistry query to execute" @@ -92,6 +97,11 @@ def _add_run_args(parser: argparse.ArgumentParser) -> None: parser.add_argument( "-r", "--report", action="store_true", help="Generate detailed report" ) + parser.add_argument( + "--human-supervised", + action="store_true", + help="Enable the ask_human tool for human-in-the-loop interaction", + ) parser.add_argument( "--recursion-limit", type=int, @@ -145,6 +155,24 @@ def _add_run_args(parser: argparse.ArgumentParser) -> None: default=None, help="Base URL for the LLM API endpoint (overrides config file)", ) + parser.add_argument( + "--mcp-url", + type=str, + default=None, + help="MCP server URL for streamable_http transport (e.g. http://localhost:9003/mcp/)", + ) + parser.add_argument( + "--mcp-command", + type=str, + default=None, + help="MCP server command for stdio transport (e.g. 'python -m chemgraph.mcp.mcp_tools')", + ) + parser.add_argument( + "--mcp-server-name", + type=str, + default="ChemGraph General Tools", + help="Display name for the MCP server connection (default: 'ChemGraph General Tools')", + ) def create_argument_parser() -> argparse.ArgumentParser: @@ -228,6 +256,16 @@ def load_config(config_file: str) -> Dict[str, Any]: Merges missing keys from a sensible default so that partial config files don't crash the CLI (addresses Bug 4 -- parity with the Streamlit config loader). + + Parameters + ---------- + config_file : str + Path to a TOML configuration file. + + Returns + ------- + dict[str, Any] + Flattened configuration dictionary with defaults filled in. """ try: with open(config_file, "r") as f: @@ -244,11 +282,13 @@ def load_config(config_file: str) -> Dict[str, Any]: "report": False, "thread": 1, "recursion_limit": 20, + "human_supervised": False, "verbose": False, }, "api": {}, "chemistry": {}, "output": {}, + "mcp": {}, } for section, defaults in _DEFAULT_SECTIONS.items(): @@ -258,7 +298,14 @@ def load_config(config_file: str) -> Dict[str, Any]: for key, value in defaults.items(): raw_config[section].setdefault(key, value) - return flatten_config(raw_config) + flat = flatten_config(raw_config) + + # Inject MCP config keys (not handled by flatten_config). + if "mcp" in raw_config and isinstance(raw_config["mcp"], dict): + for key, value in raw_config["mcp"].items(): + flat[f"mcp_{key}"] = value + + return flat except FileNotFoundError: console.print(f"[red]Configuration file not found: {config_file}[/red]") @@ -274,7 +321,13 @@ def load_config(config_file: str) -> Dict[str, Any]: def _handle_run(args: argparse.Namespace) -> None: - """Handle the ``run`` subcommand (and legacy no-subcommand mode).""" + """Handle the ``run`` subcommand and legacy no-subcommand mode. + + Parameters + ---------- + args : argparse.Namespace + Parsed CLI arguments. + """ # Handle special commands first if getattr(args, "list_models", False): list_models() @@ -334,6 +387,27 @@ def _handle_run(args: argparse.Namespace) -> None: # Resolve workflow alias (e.g. python_repl -> python_relp) args.workflow = resolve_workflow(args.workflow) + # ---- MCP tool loading ---------------------------------------------- + mcp_tools = None + mcp_url = getattr(args, "mcp_url", None) or config.get("mcp_url") + mcp_command = getattr(args, "mcp_command", None) or config.get("mcp_command") + mcp_server_name = ( + getattr(args, "mcp_server_name", None) + or config.get("mcp_server_name", "ChemGraph General Tools") + ) + + if mcp_url or mcp_command: + from chemgraph.cli.mcp_utils import load_mcp_tools_from_config + + mcp_tools = load_mcp_tools_from_config( + url=mcp_url, + command=mcp_command, + server_name=mcp_server_name, + verbose=(args.verbose > 0), + ) + if mcp_tools is None: + sys.exit(1) + if getattr(args, "interactive", False): interactive_mode( model=args.model, @@ -341,10 +415,12 @@ def _handle_run(args: argparse.Namespace) -> None: structured=args.structured, return_option=args.output, generate_report=args.report, + human_supervised=args.human_supervised, recursion_limit=args.recursion_limit, base_url=base_url, argo_user=argo_user, verbose=(args.verbose > 0), + tools=mcp_tools, ) return @@ -375,6 +451,8 @@ def _handle_run(args: argparse.Namespace) -> None: base_url=base_url, argo_user=argo_user, verbose=(args.verbose > 0), + human_supervised=args.human_supervised, + tools=mcp_tools, ) if not agent: diff --git a/src/chemgraph/cli/mcp_utils.py b/src/chemgraph/cli/mcp_utils.py new file mode 100644 index 00000000..ce287af9 --- /dev/null +++ b/src/chemgraph/cli/mcp_utils.py @@ -0,0 +1,124 @@ +"""MCP client utilities for the ChemGraph CLI. + +Handles connecting to MCP servers and loading tools for use with +MCP-enabled workflows (single_agent, multi_agent, etc.). +""" + +from __future__ import annotations + +import shlex +import time +from typing import List, Optional + +from rich.progress import Progress, SpinnerColumn, TextColumn + +from chemgraph.cli.formatting import console +from chemgraph.utils.async_utils import run_async_callable + + +def load_mcp_tools_from_config( + url: Optional[str] = None, + command: Optional[str] = None, + server_name: str = "ChemGraph General Tools", + verbose: bool = False, +) -> Optional[List]: + """Connect to an MCP server and return loaded tools. + + Supports two transports: + + - **streamable_http**: specify *url* (e.g. ``http://localhost:9003/mcp/``) + - **stdio**: specify *command* as a shell command string + + Parameters + ---------- + url : str, optional + MCP server URL for streamable_http transport. + command : str, optional + Shell command to launch an MCP server via stdio transport. + The first token is the executable; the rest are arguments. + server_name : str + Display name for the MCP server connection. + verbose : bool + Print extra diagnostic information. + + Returns + ------- + list or None + List of loaded MCP tools, or ``None`` on failure. + """ + from langchain_mcp_adapters.client import MultiServerMCPClient + + if url and command: + console.print( + "[yellow]Both --mcp-url and --mcp-command specified; " + "using --mcp-url (streamable_http).[/yellow]" + ) + + # Build connection config + if url: + connections = { + server_name: { + "transport": "streamable_http", + "url": url, + } + } + transport_label = f"streamable_http @ {url}" + elif command: + parts = shlex.split(command) + connections = { + server_name: { + "command": parts[0], + "args": parts[1:], + "transport": "stdio", + } + } + transport_label = f"stdio: {command}" + else: + console.print("[red]No MCP server URL or command provided.[/red]") + return None + + if verbose: + console.print(f"[blue]Connecting to MCP server: {transport_label}[/blue]") + + client = MultiServerMCPClient(connections) + + with Progress( + SpinnerColumn(), + TextColumn("[progress.description]{task.description}"), + console=console, + transient=True, + ) as progress: + task = progress.add_task("Loading MCP tools...", total=None) + try: + tools = run_async_callable(lambda: client.get_tools()) + progress.update( + task, + description=f"[green]Loaded {len(tools)} MCP tools!", + ) + time.sleep(0.3) + + if verbose: + tool_names = [t.name for t in tools] + console.print(f"[blue]MCP tools: {tool_names}[/blue]") + + if not tools: + console.print( + "[yellow]Warning: MCP server returned zero tools.[/yellow]" + ) + + return tools + + except Exception as e: + progress.update(task, description="[red]MCP connection failed!") + console.print(f"[red]Failed to load MCP tools: {e}[/red]") + + err_str = str(e).lower() + if "connection" in err_str or "refused" in err_str: + console.print( + "[dim]Check that the MCP server is running and reachable.[/dim]" + ) + elif "timeout" in err_str: + console.print( + "[dim]The MCP server did not respond in time.[/dim]" + ) + return None diff --git a/src/chemgraph/eval/cli.py b/src/chemgraph/eval/cli.py index cebe9154..5b587997 100644 --- a/src/chemgraph/eval/cli.py +++ b/src/chemgraph/eval/cli.py @@ -38,6 +38,11 @@ def add_eval_args(parser: argparse.ArgumentParser) -> None: This function is used by both the standalone ``chemgraph-eval`` entry point and the ``chemgraph eval`` subcommand so that the argument interface is consistent. + + Parameters + ---------- + parser : argparse.ArgumentParser + Parser or subparser to receive evaluation arguments. """ parser.add_argument( "--models", @@ -152,6 +157,16 @@ def _resolve_profile(args: argparse.Namespace) -> Optional[str]: ``[eval] default_profile``, use that as the profile name. Returns ``None`` if no profile should be used. + + Parameters + ---------- + args : argparse.Namespace + Parsed evaluation arguments. + + Returns + ------- + str or None + Selected profile name, or ``None`` when no profile applies. """ if args.profile: return args.profile @@ -180,6 +195,16 @@ def build_config_from_args(args: argparse.Namespace) -> BenchmarkConfig: When ``--config`` is provided without ``--profile``, the ``[eval] default_profile`` from the config file is used automatically if it exists. + + Parameters + ---------- + args : argparse.Namespace + Parsed evaluation arguments. + + Returns + ------- + BenchmarkConfig + Validated benchmark configuration. """ profile = _resolve_profile(args) @@ -247,7 +272,13 @@ def build_config_from_args(args: argparse.Namespace) -> BenchmarkConfig: def run_eval(args: argparse.Namespace) -> None: - """Execute an evaluation benchmark from parsed CLI arguments.""" + """Execute an evaluation benchmark from parsed CLI arguments. + + Parameters + ---------- + args : argparse.Namespace + Parsed evaluation arguments. + """ config = build_config_from_args(args) runner = ModelBenchmarkRunner(config) @@ -274,7 +305,18 @@ def run_eval(args: argparse.Namespace) -> None: def parse_args(argv=None) -> argparse.Namespace: - """Parse arguments for the standalone ``chemgraph-eval`` command.""" + """Parse arguments for the standalone ``chemgraph-eval`` command. + + Parameters + ---------- + argv : list[str], optional + Argument list to parse. Uses ``sys.argv`` when omitted. + + Returns + ------- + argparse.Namespace + Parsed command-line arguments. + """ parser = argparse.ArgumentParser( prog="chemgraph-eval", description="Run ChemGraph multi-model evaluation benchmarks.", @@ -284,7 +326,13 @@ def parse_args(argv=None) -> argparse.Namespace: def main(argv=None) -> None: - """Standalone entry point for ``chemgraph-eval``.""" + """Standalone entry point for ``chemgraph-eval``. + + Parameters + ---------- + argv : list[str], optional + Argument list to parse. Uses ``sys.argv`` when omitted. + """ args = parse_args(argv) run_eval(args) diff --git a/src/chemgraph/eval/config.py b/src/chemgraph/eval/config.py index 38a7ce2b..8c2fa23b 100644 --- a/src/chemgraph/eval/config.py +++ b/src/chemgraph/eval/config.py @@ -139,6 +139,18 @@ class BenchmarkConfig(BaseModel): @field_validator("dataset") @classmethod def dataset_must_exist(cls, v: str) -> str: + """Validate that the dataset path exists and points to JSON. + + Parameters + ---------- + v : str + Dataset path supplied to the benchmark config. + + Returns + ------- + str + Absolute resolved dataset path. + """ p = Path(v) if not p.exists(): raise ValueError(f"Dataset file does not exist: {v}") @@ -173,6 +185,18 @@ def validate_judge_model_required(self): @field_validator("judge_type") @classmethod def validate_judge_type(cls, v: str) -> str: + """Validate the requested judge strategy. + + Parameters + ---------- + v : str + Judge strategy name. + + Returns + ------- + str + Validated judge strategy. + """ valid = {"llm", "structured", "both"} if v not in valid: raise ValueError(f"Unknown judge_type: {v!r}. Valid: {sorted(valid)}") @@ -181,6 +205,18 @@ def validate_judge_type(cls, v: str) -> str: @field_validator("workflow_types") @classmethod def validate_workflow_types(cls, v: List[str]) -> List[str]: + """Validate benchmark workflow names. + + Parameters + ---------- + v : list[str] + Workflow names requested for the benchmark. + + Returns + ------- + list[str] + Validated workflow names. + """ valid = { "single_agent", "multi_agent", @@ -202,6 +238,16 @@ def get_base_url(self, model_name: str) -> Optional[str]: Returns ``None`` when no config file was provided (the provider loaders will fall back to their defaults / environment variables). + + Parameters + ---------- + model_name : str + Model identifier whose provider URL should be resolved. + + Returns + ------- + str or None + Configured base URL, or ``None`` when no override is available. """ if not self._flat_config: return None diff --git a/src/chemgraph/eval/reporter.py b/src/chemgraph/eval/reporter.py index 011fa703..2ae3e305 100644 --- a/src/chemgraph/eval/reporter.py +++ b/src/chemgraph/eval/reporter.py @@ -14,7 +14,18 @@ def _safe_pct(value: float) -> str: - """Format a 0-1 fraction as a percentage string.""" + """Format a 0-1 fraction as a percentage string. + + Parameters + ---------- + value : float + Fractional value to format. + + Returns + ------- + str + Percentage string with one decimal place. + """ return f"{value * 100:.1f}%" @@ -313,7 +324,18 @@ def print_summary_table(results: Dict[str, Dict[str, dict]]) -> None: def _make_serializable(obj): - """Recursively convert non-serializable objects to strings.""" + """Recursively convert non-serializable objects to strings. + + Parameters + ---------- + obj : Any + Object to convert. + + Returns + ------- + Any + JSON-serializable object. + """ if isinstance(obj, dict): return {k: _make_serializable(v) for k, v in obj.items()} elif isinstance(obj, (list, tuple)): diff --git a/src/chemgraph/eval/runner.py b/src/chemgraph/eval/runner.py index fb58257a..6c85b9e8 100644 --- a/src/chemgraph/eval/runner.py +++ b/src/chemgraph/eval/runner.py @@ -61,6 +61,13 @@ class ModelBenchmarkRunner: """ def __init__(self, config: BenchmarkConfig): + """Initialize the benchmark runner. + + Parameters + ---------- + config : BenchmarkConfig + Validated benchmark configuration. + """ self.config = config full_dataset: List[GroundTruthItem] = load_dataset(config.dataset) # Apply max_queries limit if configured (0 = no limit). @@ -109,7 +116,20 @@ def _checkpoint_dir(self) -> str: return d def _checkpoint_path(self, model_name: str, workflow_type: str) -> str: - """Return the JSONL checkpoint file path for a (model, workflow) pair.""" + """Return the JSONL checkpoint file path for a model/workflow pair. + + Parameters + ---------- + model_name : str + Model identifier being evaluated. + workflow_type : str + Workflow type being evaluated. + + Returns + ------- + str + Checkpoint JSONL file path. + """ safe_name = model_name.replace("/", "_").replace(":", "_") return os.path.join( self._checkpoint_dir(), @@ -130,6 +150,19 @@ def _save_query_checkpoint( the query ID, index, and full result (raw output + judge scores). Append-only writes make this crash-safe: at worst the last line may be truncated (one query lost, not all). + + Parameters + ---------- + model_name : str + Model identifier being evaluated. + workflow_type : str + Workflow type being evaluated. + query_id : str + Ground-truth query identifier. + query_idx : int + Query index used as the LangGraph thread ID. + query_result : dict + Result payload to checkpoint. """ record = { "query_id": query_id, @@ -143,6 +176,13 @@ def _save_query_checkpoint( def _load_checkpoint(self, model_name: str, workflow_type: str) -> Dict[str, dict]: """Load completed query results from a checkpoint file. + Parameters + ---------- + model_name : str + Model identifier being evaluated. + workflow_type : str + Workflow type being evaluated. + Returns ------- dict @@ -186,6 +226,13 @@ def _clear_checkpoint(self, model_name: str, workflow_type: str) -> None: Called when *not* resuming, so that stale checkpoint data from a previous run does not leak into the current run. + + Parameters + ---------- + model_name : str + Model identifier being evaluated. + workflow_type : str + Workflow type being evaluated. """ path = self._checkpoint_path(model_name, workflow_type) if os.path.exists(path): @@ -203,6 +250,13 @@ async def _run_single_model_workflow( ) -> dict: """Run all queries for one (model, workflow) pair. + Parameters + ---------- + model_name : str + Model identifier to evaluate. + workflow_type : str + Workflow type to evaluate. + Returns ------- dict @@ -342,6 +396,24 @@ async def _run_single_query( """Execute and evaluate a single query. Returns ``{"raw": ..., "judge": ..., "structured_judge": ...}``. + + Parameters + ---------- + cg : ChemGraph + Initialized ChemGraph agent. + item : GroundTruthItem + Ground-truth query item. + idx : int + Query index used as the LangGraph thread ID. + model_name : str + Model identifier being evaluated. + workflow_type : str + Workflow type being evaluated. + + Returns + ------- + dict + Query result containing raw output and judge results. """ try: config = {"configurable": {"thread_id": str(idx)}} @@ -540,7 +612,20 @@ def report(self, format: str = "all") -> None: @staticmethod def _make_error_result(error_msg: str, n_queries: int) -> dict: - """Build an error placeholder result for a failed model init.""" + """Build an error placeholder result for a failed model init. + + Parameters + ---------- + error_msg : str + Error message to store in the aggregate result. + n_queries : int + Number of benchmark queries that were skipped. + + Returns + ------- + dict + Placeholder aggregate result. + """ return { "judge_aggregate": { "n_queries": n_queries, diff --git a/src/chemgraph/eval/structured_output_judge.py b/src/chemgraph/eval/structured_output_judge.py index e9570ef6..aacb7b4e 100644 --- a/src/chemgraph/eval/structured_output_judge.py +++ b/src/chemgraph/eval/structured_output_judge.py @@ -63,6 +63,20 @@ def _relative_close(a: float, b: float, tol: float = 0.05) -> bool: """Return True if *a* and *b* are within *tol* relative tolerance. Falls back to absolute comparison when *b* is near zero. + + Parameters + ---------- + a : float + Actual value. + b : float + Expected value. + tol : float, optional + Relative tolerance. + + Returns + ------- + bool + ``True`` when the values are close enough. """ if b == 0: return abs(a) < 1e-8 @@ -70,7 +84,18 @@ def _relative_close(a: float, b: float, tol: float = 0.05) -> bool: def _parse_numeric(val: Any) -> Optional[float]: - """Try to parse *val* as a float, returning None on failure.""" + """Try to parse a value as a float. + + Parameters + ---------- + val : Any + Candidate numeric value. + + Returns + ------- + float or None + Parsed float, or ``None`` on failure. + """ if isinstance(val, (int, float)): return float(val) if isinstance(val, str): @@ -84,12 +109,34 @@ def _parse_numeric(val: Any) -> Optional[float]: def _is_imaginary_freq(val: str) -> bool: - """Return True if *val* represents an imaginary frequency.""" + """Return True if a value represents an imaginary frequency. + + Parameters + ---------- + val : str + Frequency value to inspect. + + Returns + ------- + bool + ``True`` when the value ends with the imaginary-frequency marker. + """ return isinstance(val, str) and val.strip().endswith("i") def _canonicalise_smiles(smiles: str) -> Optional[str]: - """Return the RDKit canonical SMILES, or None if RDKit is unavailable.""" + """Return the RDKit canonical SMILES. + + Parameters + ---------- + smiles : str + Input SMILES string. + + Returns + ------- + str or None + Canonical SMILES, or ``None`` if RDKit is unavailable/invalid. + """ try: from rdkit import Chem @@ -114,6 +161,20 @@ def _compare_scalar( """Compare two ``ScalarResult`` dicts. Returns ``(passed, reason)``. + + Parameters + ---------- + expected : dict[str, Any] + Expected scalar result. + actual : dict[str, Any] + Actual scalar result. + tolerance : float + Relative tolerance for value comparison. + + Returns + ------- + tuple[bool, str] + Pass/fail flag and explanation. """ reasons: List[str] = [] @@ -158,6 +219,18 @@ def _compare_smiles( string comparison. Returns ``(passed, reason)``. + + Parameters + ---------- + expected : list[str] + Expected SMILES strings. + actual : list[str] + Actual SMILES strings. + + Returns + ------- + tuple[bool, str] + Pass/fail flag and explanation. """ if not expected: return True, "expected smiles list is empty (skipped)" @@ -167,6 +240,18 @@ def _compare_smiles( # Build canonical sets. def _canon_set(smiles_list: List[str]) -> set[str]: + """Canonicalize a SMILES list into a set. + + Parameters + ---------- + smiles_list : list[str] + SMILES strings to canonicalize. + + Returns + ------- + set[str] + Canonicalized SMILES strings. + """ result: set[str] = set() for s in smiles_list: canon = _canonicalise_smiles(s) @@ -197,6 +282,20 @@ def _compare_vibrational( """Compare two ``VibrationalFrequency`` dicts. Filters imaginary frequencies and compares real ones element-wise. + + Parameters + ---------- + expected : dict[str, Any] + Expected vibrational data. + actual : dict[str, Any] + Actual vibrational data. + tolerance : float + Relative tolerance for frequency comparison. + + Returns + ------- + tuple[bool, str] + Pass/fail flag and explanation. """ exp_freqs = expected.get("frequency_cm1", []) act_freqs = actual.get("frequency_cm1", []) @@ -230,7 +329,22 @@ def _compare_ir_spectrum( actual: Dict[str, Any], tolerance: float, ) -> tuple[bool, str]: - """Compare two ``IRSpectrum`` dicts (frequencies + intensities).""" + """Compare two ``IRSpectrum`` dicts. + + Parameters + ---------- + expected : dict[str, Any] + Expected IR spectrum data. + actual : dict[str, Any] + Actual IR spectrum data. + tolerance : float + Relative tolerance for frequency/intensity comparison. + + Returns + ------- + tuple[bool, str] + Pass/fail flag and explanation. + """ # Compare frequencies. freq_ok, freq_reason = _compare_vibrational( {"frequency_cm1": expected.get("frequency_cm1", [])}, diff --git a/src/chemgraph/graphs/graspa_mcp.py b/src/chemgraph/graphs/graspa_mcp.py index a43ac6e4..b82fc04d 100644 --- a/src/chemgraph/graphs/graspa_mcp.py +++ b/src/chemgraph/graphs/graspa_mcp.py @@ -27,6 +27,22 @@ def planner_agent( llm: ChatOpenAI, system_prompt: str, ): + """Plan the next gRASPA MCP workflow step. + + Parameters + ---------- + state : PlannerState + Current planner state containing messages and executor results. + llm : ChatOpenAI + Chat model used for planning. + system_prompt : str + Planner system prompt. + + Returns + ------- + dict + Planner state update containing messages, next step, and tasks. + """ executor_outputs = state.get("executor_results", []) content_block = f"Current Conversation History: {state['messages']}" if executor_outputs: @@ -51,8 +67,17 @@ def planner_agent( def unified_planner_router(state: PlannerState) -> Union[str, list[Send]]: - """ - Routes based on the Planner's structured 'next_step'. + """Route based on the planner's structured ``next_step``. + + Parameters + ---------- + state : PlannerState + Current planner state. + + Returns + ------- + str or list[Send] + Next node name, ``END``, or fan-out executor sends. """ next_step = state.get("next_step") @@ -81,9 +106,23 @@ async def executor_model_node( system_prompt: str, tools: list, ): - """ - The reasoning engine for a single executor. - It sees its own 'task_prompt' and its own 'messages' history. + """Run the reasoning step for a single gRASPA executor. + + Parameters + ---------- + state : ExecutorState + Local executor state. + llm : ChatOpenAI + Chat model used by the executor. + system_prompt : str + Executor system prompt. + tools : list + Tools available to the executor. + + Returns + ------- + dict + Executor state update containing the model response. """ messages = [{"role": "system", "content": system_prompt}] + state["messages"] @@ -106,7 +145,18 @@ async def executor_model_node( return {"messages": [response]} def route_executor(state: ExecutorState): - """Standard ReAct routing: Tool vs End.""" + """Route executor output to tools or completion. + + Parameters + ---------- + state : ExecutorState + Local executor state. + + Returns + ------- + str + ``"tools"`` when tool calls are present, otherwise ``"done"``. + """ messages = state["messages"] last_message = messages[-1] if hasattr(last_message, "tool_calls") and last_message.tool_calls: @@ -115,9 +165,17 @@ def route_executor(state: ExecutorState): def format_executor_output(state: ExecutorState) -> PlannerState: - """ - Bridge function: - Converts the Local ExecutorState into an update for the Global PlannerState. + """Convert local executor state into a global planner update. + + Parameters + ---------- + state : ExecutorState + Local executor state at subgraph completion. + + Returns + ------- + PlannerState + Planner update containing executor results and logs. """ executor_id = state["executor_id"] final_message = state["messages"][-1].content @@ -130,7 +188,22 @@ def format_executor_output(state: ExecutorState) -> PlannerState: def construct_executor_subgraph(llm: ChatOpenAI, tools: list, system_prompt: str): - """Builds the reusable executor subgraph (Agent -> Tools -> Agent).""" + """Build the reusable executor subgraph. + + Parameters + ---------- + llm : ChatOpenAI + Chat model used by executor agents. + tools : list + Tools available to executor agents. + system_prompt : str + Executor system prompt. + + Returns + ------- + CompiledStateGraph + Compiled executor subgraph. + """ workflow = StateGraph(ExecutorState) workflow.add_node( "executor_agent", @@ -160,7 +233,24 @@ def insight_analyst_node( tools: list, system_prompt: str, ): - """Analyzes the gathered results.""" + """Analyze gathered executor results. + + Parameters + ---------- + state : PlannerState + Planner state containing executor results. + llm : ChatOpenAI + Chat model used by the analyst. + tools : list + Analysis tools available to the analyst. + system_prompt : str + Analyst system prompt. + + Returns + ------- + dict + Planner state update containing the analyst response. + """ results_text = "\n".join(state["executor_results"]) messages = [ {"role": "system", "content": system_prompt}, @@ -176,8 +266,18 @@ def insight_analyst_node( def route_analyst(state: PlannerState): - """ - Determines if the Analyst is calling a tool or giving the final answer. + """Route analyst output to tools or back to the planner. + + Parameters + ---------- + state : PlannerState + Planner state containing the analyst's latest message. + + Returns + ------- + str + ``"analyst_tools"`` when tool calls are present, otherwise + ``"Planner"``. """ last_msg = state["messages"][-1] @@ -197,8 +297,27 @@ def construct_graspa_mcp_graph( executor_tools: list = None, analysis_tools: list = None, ): - """ - Constructs the Main Graph using the Map-Reduce (Send) pattern. + """Construct the gRASPA MCP map-reduce graph. + + Parameters + ---------- + llm : ChatOpenAI + Chat model shared by planner, executors, and analyst. + planner_prompt : str, optional + Planner system prompt. + executor_prompt : str, optional + Executor system prompt. + analyst_prompt : str, optional + Analyst system prompt. + executor_tools : list, optional + Tools available to executor subgraphs. + analysis_tools : list, optional + Tools available to the analyst node. + + Returns + ------- + CompiledStateGraph + Compiled gRASPA MCP graph. """ checkpointer = MemorySaver() diff --git a/src/chemgraph/graphs/multi_agent.py b/src/chemgraph/graphs/multi_agent.py index d2074573..507e54bf 100644 --- a/src/chemgraph/graphs/multi_agent.py +++ b/src/chemgraph/graphs/multi_agent.py @@ -24,7 +24,7 @@ from langchain_core.messages import AIMessage, BaseMessage, ToolMessage from langchain_openai import ChatOpenAI from langgraph.checkpoint.memory import MemorySaver -from langgraph.types import Send +from langgraph.types import Send, interrupt from chemgraph.utils.logging_config import setup_logger from chemgraph.utils.parsing import extract_json_block, parse_response_formatter @@ -45,7 +45,18 @@ def _to_jsonable(obj: Any) -> Any: - """Recursively convert Pydantic models to plain dicts.""" + """Recursively convert Pydantic models to plain dictionaries. + + Parameters + ---------- + obj : Any + Object to convert. + + Returns + ------- + Any + JSON-compatible version of the object where possible. + """ if isinstance(obj, BaseModel): return obj.model_dump() elif isinstance(obj, dict): @@ -68,6 +79,16 @@ def sanitize_tool_calls(messages: list[BaseMessage]) -> list[BaseMessage]: This function walks every ``AIMessage.tool_calls`` entry and recursively converts Pydantic models back to plain dicts. + + Parameters + ---------- + messages : list[BaseMessage] + LangChain messages that may contain tool-call arguments. + + Returns + ------- + list[BaseMessage] + Messages with JSON-serializable tool-call arguments. """ for m in messages: if isinstance(m, AIMessage) and getattr(m, "tool_calls", None): @@ -92,6 +113,16 @@ def _parse_planner_response( Returns ``(parsed_response, None)`` on success, or ``(None, error_msg)`` on failure. + + Parameters + ---------- + raw_text : str + Raw planner model output. + + Returns + ------- + tuple[PlannerResponse | None, str | None] + Parsed response and optional parse error. """ # 1. Direct validation try: @@ -135,6 +166,23 @@ def planner_agent( The LLM is prompted to return a JSON object matching the ``PlannerResponse`` schema. If parsing fails, the LLM is retried up to ``max_retries`` times with error feedback. + + Parameters + ---------- + state : PlannerState + Current global planner state. + llm : ChatOpenAI + Chat model used for planning. + system_prompt : str + Planner system prompt. + max_retries : int, optional + Number of parse-retry attempts after invalid planner output. + + Returns + ------- + dict + Planner state update containing messages, next step, tasks, and + iteration count. """ executor_outputs = state.get("executor_results", []) failed_tasks = state.get("failed_tasks", []) @@ -204,12 +252,66 @@ def planner_agent( logger.info("PLANNER: %s", response_obj.model_dump_json()) current_iterations = state.get("planner_iterations", 0) - return { + result = { "messages": [AIMessage(content=response_obj.thought_process)], "next_step": response_obj.next_step, "tasks": response_obj.tasks if response_obj.tasks else [], "planner_iterations": current_iterations + 1, } + if response_obj.next_step == "ask_human" and response_obj.clarification: + result["clarification"] = response_obj.clarification + return result + + +# --------------------------------------------------------------------------- +# Human review node (interrupt for human-in-the-loop) +# --------------------------------------------------------------------------- + + +def human_review_node(state: PlannerState): + """Pause the graph and ask the human for clarification. + + This node calls ``interrupt()`` with the planner's clarification + question. Execution halts until a human provides a response via + ``Command(resume=...)``. The human's answer is injected back into + the conversation as an ``AIMessage`` summarising what was asked and + what the human replied, then control returns to the Planner. + + Parameters + ---------- + state : PlannerState + Current planner state containing the clarification question. + + Returns + ------- + dict + Planner state update containing the human clarification message. + """ + question = state.get("clarification", "Could you please provide more details?") + logger.info("HUMAN_REVIEW: interrupting with question: %s", question) + + human_response = interrupt({"question": question}) + + # Normalise the response to a plain string. + if isinstance(human_response, dict): + answer = human_response.get( + "answer", human_response.get("response", str(human_response)) + ) + else: + answer = str(human_response) + + logger.info("HUMAN_REVIEW: received response: %s", answer) + return { + "messages": [ + AIMessage( + content=( + f"Human clarification received.\n" + f"Question: {question}\n" + f"Answer: {answer}" + ) + ) + ], + } # --------------------------------------------------------------------------- @@ -226,6 +328,7 @@ def unified_planner_router( """Route based on the planner's ``next_step`` decision. * ``executor_subgraph`` -- fan-out tasks via ``Send()`` + * ``ask_human`` -- pause for human clarification via ``human_review`` * ``FINISH`` -- go to ``ResponseAgent`` (if structured_output) or ``END`` A cycle guard forces ``FINISH`` when the planner has dispatched @@ -234,10 +337,29 @@ def unified_planner_router( For retried tasks, the ``retry_count`` from the ``WorkerTask`` is checked against ``max_task_retries``. Tasks that have exceeded the retry limit are skipped and logged as permanently failed. + + Parameters + ---------- + state : PlannerState + Current planner state. + structured_output : bool, optional + Whether to route final output through the response formatter. + max_planner_iterations : int, optional + Maximum planner dispatch iterations before forcing completion. + max_task_retries : int, optional + Maximum retry count for individual executor tasks. + + Returns + ------- + str or list[Send] + Next graph node name, ``END``, or fan-out ``Send`` instructions. """ next_step = state.get("next_step") iterations = state.get("planner_iterations", 0) + if next_step == "ask_human": + return "human_review" + if next_step == "executor_subgraph": if iterations > max_planner_iterations: logger.warning( @@ -321,6 +443,22 @@ async def executor_model_node( Reads its own ``messages`` history, calls the LLM with bound tools, and returns the response. + + Parameters + ---------- + state : ExecutorState + Local executor state. + llm : ChatOpenAI + Chat model used by the executor. + system_prompt : str + Executor system prompt. + tools : list + Tools available to the executor. + + Returns + ------- + dict + Executor state update containing the new model message. """ sanitized = sanitize_tool_calls(list(state["messages"])) messages = [{"role": "system", "content": system_prompt}] + sanitized @@ -348,7 +486,19 @@ async def executor_model_node( def route_executor(state: ExecutorState): - """Standard ReAct routing: tool calls -> ``tools``, else -> ``done``.""" + """Route executor output to tools or completion. + + Parameters + ---------- + state : ExecutorState + Local executor state. + + Returns + ------- + str + ``"tools"`` when the last message has tool calls, otherwise + ``"done"``. + """ last_message = state["messages"][-1] if hasattr(last_message, "tool_calls") and last_message.tool_calls: return "tools" @@ -382,6 +532,16 @@ def _detect_executor_failure(messages: list) -> tuple[bool, str | None]: 2. Error markers in the final assistant message content. Returns ``(is_failed, error_summary)``. + + Parameters + ---------- + messages : list + Executor message history. + + Returns + ------- + tuple[bool, str | None] + Failure flag and optional error summary. """ # Collect all tool-level errors tool_errors = [] @@ -421,6 +581,16 @@ def format_executor_output(state: ExecutorState) -> dict: Detects executor failures by scanning the message history for tool errors and error markers. When a failure is detected, populates ``failed_tasks`` so the planner can decide whether to retry. + + Parameters + ---------- + state : ExecutorState + Local executor state at subgraph completion. + + Returns + ------- + dict + Planner-state update with executor results, logs, and failure data. """ executor_id = state["executor_id"] task_index = state.get("task_index", -1) @@ -473,6 +643,20 @@ def construct_executor_subgraph( The subgraph is compiled and used as a node in the main graph. Each ``Send()`` invocation creates an independent copy with its own ``ExecutorState``. + + Parameters + ---------- + llm : ChatOpenAI + Chat model used by executor agents. + tools : list + Tools available to executor agents. + system_prompt : str + Executor system prompt. + + Returns + ------- + CompiledStateGraph + Compiled executor subgraph. """ workflow = StateGraph(ExecutorState) workflow.add_node( @@ -512,6 +696,22 @@ def response_agent( Mirrors the ``ResponseAgent`` from ``single_agent.py``: invokes the LLM with a formatter prompt and manually parses the response into a ``ResponseFormatter`` with retry logic on parse failure. + + Parameters + ---------- + state : PlannerState + Final planner state to summarize. + llm : ChatOpenAI + Chat model used for response formatting. + formatter_prompt : str + Prompt instructing the model how to format the final answer. + max_retries : int, optional + Number of parse-retry attempts after invalid formatter output. + + Returns + ------- + dict + State update containing the formatted response message. """ messages = [ {"role": "system", "content": formatter_prompt}, @@ -639,9 +839,10 @@ def construct_multi_agent_graph( ), ) graph_builder.add_node("executor_subgraph", executor_subgraph) + graph_builder.add_node("human_review", human_review_node) # Conditional destinations list for the planner router - conditional_targets = ["executor_subgraph", END] + conditional_targets = ["executor_subgraph", "human_review", END] if structured_output: graph_builder.add_node( @@ -671,6 +872,9 @@ def construct_multi_agent_graph( # Executors feed results back to the planner graph_builder.add_edge("executor_subgraph", "Planner") + # After human clarification, return to the planner for re-planning + graph_builder.add_edge("human_review", "Planner") + if structured_output: graph_builder.add_edge("ResponseAgent", END) diff --git a/src/chemgraph/graphs/python_relp_agent.py b/src/chemgraph/graphs/python_relp_agent.py index 1cfc2bc6..dd8edf98 100644 --- a/src/chemgraph/graphs/python_relp_agent.py +++ b/src/chemgraph/graphs/python_relp_agent.py @@ -45,6 +45,13 @@ class BasicToolNode: """ def __init__(self, tools: list) -> None: + """Initialize the tool node. + + Parameters + ---------- + tools : list + Tool objects keyed by their ``name`` attribute. + """ self.tools_by_name = {tool.name: tool for tool in tools} def __call__(self, inputs: State) -> State: diff --git a/src/chemgraph/graphs/rag_agent.py b/src/chemgraph/graphs/rag_agent.py index e7c1e9f3..91611166 100644 --- a/src/chemgraph/graphs/rag_agent.py +++ b/src/chemgraph/graphs/rag_agent.py @@ -51,7 +51,18 @@ # Helpers (reuse the repeated-tool-call detection from single_agent) # --------------------------------------------------------------------------- def _tool_call_signature(tool_calls) -> tuple: - """Create a comparable signature for a list of tool calls.""" + """Create a comparable signature for a list of tool calls. + + Parameters + ---------- + tool_calls : list + Tool-call dictionaries from an AI message. + + Returns + ------- + tuple + Deterministic signature of tool names and arguments. + """ signature = [] for call in tool_calls or []: name = call.get("name") if isinstance(call, dict) else None @@ -65,7 +76,18 @@ def _tool_call_signature(tool_calls) -> tuple: def _is_repeated_tool_cycle(messages) -> bool: - """Detect if the most recent AI tool-call set repeats the previous one.""" + """Detect if the most recent AI tool-call set repeats the previous one. + + Parameters + ---------- + messages : list + Message history to inspect. + + Returns + ------- + bool + ``True`` when the last two AI tool-call sets are identical. + """ ai_with_calls = [ m for m in messages diff --git a/src/chemgraph/graphs/single_agent.py b/src/chemgraph/graphs/single_agent.py index 8b8a1e2e..7be83d71 100644 --- a/src/chemgraph/graphs/single_agent.py +++ b/src/chemgraph/graphs/single_agent.py @@ -13,7 +13,7 @@ smiles_to_coordinate_file, ) from chemgraph.tools.report_tools import generate_html -from chemgraph.tools.generic_tools import calculator +from chemgraph.tools.generic_tools import calculator, ask_human from chemgraph.prompt.single_agent_prompt import ( single_agent_prompt, formatter_prompt, @@ -27,7 +27,18 @@ def _tool_call_signature(tool_calls) -> tuple: - """Create a comparable signature for a list of tool calls.""" + """Create a comparable signature for a list of tool calls. + + Parameters + ---------- + tool_calls : list + Tool-call dictionaries from an AI message. + + Returns + ------- + tuple + Deterministic signature of tool names and arguments. + """ signature = [] for call in tool_calls or []: name = call.get("name") if isinstance(call, dict) else None @@ -42,7 +53,18 @@ def _tool_call_signature(tool_calls) -> tuple: def _is_repeated_tool_cycle(messages) -> bool: - """Detect if the most recent AI tool-call set repeats the previous AI tool-call set.""" + """Detect if the most recent AI tool-call set repeats the previous one. + + Parameters + ---------- + messages : list + Message history to inspect. + + Returns + ------- + bool + ``True`` when the last two AI tool-call sets are identical. + """ ai_with_calls = [] for message in messages: if hasattr(message, "tool_calls") and getattr(message, "tool_calls", None): @@ -57,21 +79,54 @@ def _is_repeated_tool_cycle(messages) -> bool: def _tool_message_name(message): - """Extract tool name from a message-like object.""" + """Extract tool name from a message-like object. + + Parameters + ---------- + message : Any + Message dictionary or object. + + Returns + ------- + str or None + Tool name when present. + """ if isinstance(message, dict): return message.get("name") return getattr(message, "name", None) def _tool_message_content(message): - """Extract content text from a message-like object.""" + """Extract content text from a message-like object. + + Parameters + ---------- + message : Any + Message dictionary or object. + + Returns + ------- + Any + Message content, or an empty string when unavailable. + """ if isinstance(message, dict): return message.get("content", "") return getattr(message, "content", "") def _is_successful_report_message(message) -> bool: - """Return True when message indicates successful generate_html execution.""" + """Return True when a message indicates successful report generation. + + Parameters + ---------- + message : Any + Tool message dictionary or object. + + Returns + ------- + bool + ``True`` for non-error ``generate_html`` tool output. + """ if _tool_message_name(message) != "generate_html": return False @@ -111,7 +166,18 @@ def route_tools(state: State): def route_report_tools(state: State): - """Route report tool execution and stop if a report was already generated.""" + """Route report tool execution and stop if a report was already generated. + + Parameters + ---------- + state : State + Current graph state or message list. + + Returns + ------- + str + ``"tools"`` when ``generate_html`` should run, otherwise ``"done"``. + """ if isinstance(state, list): messages = state ai_message = state[-1] if state else None @@ -140,7 +206,18 @@ def route_report_tools(state: State): def route_after_report_tools(state: State): - """After report tool execution, stop on success; otherwise retry report generation.""" + """After report tool execution, stop on success or retry on failure. + + Parameters + ---------- + state : State + Current graph state or message list after report tool execution. + + Returns + ------- + str + ``"done"`` after a successful report message, otherwise ``"retry"``. + """ if isinstance(state, list): messages = state elif messages := state.get("messages", []): @@ -151,7 +228,13 @@ def route_after_report_tools(state: State): return "done" if _is_successful_report_message(messages[-1]) else "retry" -def ChemGraphAgent(state: State, llm: ChatOpenAI, system_prompt: str, tools=None): +def ChemGraphAgent( + state: State, + llm: ChatOpenAI, + system_prompt: str, + tools=None, + human_supervised: bool = False, +): """LLM node that processes messages and decides next actions. Parameters @@ -164,6 +247,8 @@ def ChemGraphAgent(state: State, llm: ChatOpenAI, system_prompt: str, tools=None The system prompt to guide the LLM's behavior tools : list, optional List of tools available to the agent, by default None + human_supervised : bool, optional + Whether to include the ``ask_human`` tool, by default False Returns ------- @@ -180,6 +265,12 @@ def ChemGraphAgent(state: State, llm: ChatOpenAI, system_prompt: str, tools=None extract_output_json, calculator, ] + if human_supervised: + tools.append(ask_human) + elif human_supervised and ask_human not in tools: + # Ensure ask_human is available when custom tools are provided + # and human supervision is enabled. + tools = list(tools) + [ask_human] messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": f"{state['messages']}"}, @@ -319,6 +410,7 @@ def construct_single_agent_graph( report_prompt: str = report_prompt, tools: list = None, max_retries: int = 1, + human_supervised: bool = False, ): """Construct a geometry optimization graph. @@ -341,6 +433,10 @@ def construct_single_agent_graph( max_retries : int, optional Maximum number of LLM retry attempts when the ResponseAgent fails to parse the formatter output, by default 1 + human_supervised : bool, optional + Whether to include the ``ask_human`` tool so the agent can + pause and request human input, by default False + Returns ------- StateGraph @@ -357,6 +453,12 @@ def construct_single_agent_graph( extract_output_json, calculator, ] + if human_supervised: + tools.append(ask_human) + elif human_supervised and ask_human not in tools: + # Ensure ask_human is available when custom tools are provided + # and human supervision is enabled. + tools = list(tools) + [ask_human] tool_node = ToolNode(tools=tools) graph_builder = StateGraph(State) @@ -364,7 +466,11 @@ def construct_single_agent_graph( graph_builder.add_node( "ChemGraphAgent", lambda state: ChemGraphAgent( - state, llm, system_prompt=system_prompt, tools=tools + state, + llm, + system_prompt=system_prompt, + tools=tools, + human_supervised=human_supervised, ), ) graph_builder.add_node("tools", tool_node) @@ -403,7 +509,6 @@ def construct_single_agent_graph( {"tools": "tools", "done": END}, ) graph_builder.add_edge("tools", "ChemGraphAgent") - graph_builder.add_edge("ChemGraphAgent", END) graph = graph_builder.compile(checkpointer=checkpointer) logger.info("Graph construction completed") @@ -412,7 +517,11 @@ def construct_single_agent_graph( graph_builder.add_node( "ChemGraphAgent", lambda state: ChemGraphAgent( - state, llm, system_prompt=system_prompt, tools=tools + state, + llm, + system_prompt=system_prompt, + tools=tools, + human_supervised=human_supervised, ), ) graph_builder.add_node("tools", tool_node) diff --git a/src/chemgraph/graphs/single_agent_xanes.py b/src/chemgraph/graphs/single_agent_xanes.py index 9fe40cf0..1c3935d8 100644 --- a/src/chemgraph/graphs/single_agent_xanes.py +++ b/src/chemgraph/graphs/single_agent_xanes.py @@ -25,7 +25,18 @@ def _tool_call_signature(tool_calls) -> tuple: - """Create a comparable signature for a list of tool calls.""" + """Create a comparable signature for a list of tool calls. + + Parameters + ---------- + tool_calls : list + Tool-call dictionaries from an AI message. + + Returns + ------- + tuple + Deterministic signature of tool names and arguments. + """ signature = [] for call in tool_calls or []: name = call.get("name") if isinstance(call, dict) else None @@ -39,7 +50,18 @@ def _tool_call_signature(tool_calls) -> tuple: def _is_repeated_tool_cycle(messages) -> bool: - """Detect if the most recent AI tool-call set repeats the previous AI tool-call set.""" + """Detect if the most recent AI tool-call set repeats the previous one. + + Parameters + ---------- + messages : list + Message history to inspect. + + Returns + ------- + bool + ``True`` when the last two AI tool-call sets are identical. + """ ai_with_calls = [] for message in messages: if hasattr(message, "tool_calls") and getattr(message, "tool_calls", None): diff --git a/src/chemgraph/hpc_configs/aurora_parsl.py b/src/chemgraph/hpc_configs/aurora_parsl.py index f2a0f761..61793aaf 100644 --- a/src/chemgraph/hpc_configs/aurora_parsl.py +++ b/src/chemgraph/hpc_configs/aurora_parsl.py @@ -9,6 +9,18 @@ def get_aurora_config( run_dir=None, ): + """Create a Parsl configuration for Aurora PBS jobs. + + Parameters + ---------- + run_dir : str, optional + Directory used as Parsl's run directory and worker working directory. + + Returns + ------- + parsl.config.Config + Configured Parsl ``Config`` for Aurora. + """ if run_dir is None: run_dir = os.getcwd() diff --git a/src/chemgraph/hpc_configs/polaris_parsl.py b/src/chemgraph/hpc_configs/polaris_parsl.py index d90277c4..ef60f207 100644 --- a/src/chemgraph/hpc_configs/polaris_parsl.py +++ b/src/chemgraph/hpc_configs/polaris_parsl.py @@ -9,8 +9,19 @@ def get_polaris_config( run_dir=None, worker_init: str = "export TMPDIR=/tmp", ): - """ - Generates the Parsl configuration for the Polaris supercomputer. + """Generate the Parsl configuration for the Polaris supercomputer. + + Parameters + ---------- + run_dir : str, optional + Directory used as Parsl's run directory. + worker_init : str, optional + Shell initialization snippet run by each Parsl worker. + + Returns + ------- + parsl.config.Config + Configured Parsl ``Config`` for Polaris. """ if run_dir is None: run_dir = os.getcwd() diff --git a/src/chemgraph/mcp/__init__.py b/src/chemgraph/mcp/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/chemgraph/mcp/data_analysis_mcp.py b/src/chemgraph/mcp/data_analysis_mcp.py index f85d1900..0b360ad7 100644 --- a/src/chemgraph/mcp/data_analysis_mcp.py +++ b/src/chemgraph/mcp/data_analysis_mcp.py @@ -98,9 +98,19 @@ def aggregate_simulation_results( file_paths: list[str], output_csv_path: str, ) -> str: - """ - Reads a provided list of specific JSONL simulation file paths and combines them into a CSV. - Splits the absolute 'cif_path' into 'cif_base_path' and 'cif_filename'. + """Aggregate JSONL simulation records into a CSV summary. + + Parameters + ---------- + file_paths : list[str] + JSONL files to read. Each line should contain one simulation result. + output_csv_path : str + Destination CSV path. + + Returns + ------- + str + Human-readable success or error message. """ all_data = [] @@ -218,6 +228,20 @@ def rank_mofs_performance( for cif_name, group in grouped: # Helper: Robust lookup with tolerances def get_uptake(target_p, target_t): + """Return the uptake matching target pressure and temperature. + + Parameters + ---------- + target_p : float or None + Target pressure in Pa. + target_t : float or None + Target temperature in K. + + Returns + ------- + float or None + Mean uptake for matching rows, or ``None`` when no match exists. + """ if target_p is None or target_t is None: return None diff --git a/src/chemgraph/mcp/graspa_mcp_parsl.py b/src/chemgraph/mcp/graspa_mcp_parsl.py index 9efeba86..3b55690a 100644 --- a/src/chemgraph/mcp/graspa_mcp_parsl.py +++ b/src/chemgraph/mcp/graspa_mcp_parsl.py @@ -1,13 +1,12 @@ import asyncio import json -import logging import os from pathlib import Path from mcp.server.fastmcp import FastMCP import parsl -from chemgraph.mcp.server_utils import run_mcp_server +from chemgraph.mcp.server_utils import load_parsl_config, run_mcp_server from chemgraph.schemas.graspa_schema import ( graspa_input_schema_ensemble, ) @@ -27,7 +26,7 @@ def run_graspa_parsl_app(job: dict): from chemgraph.schemas.graspa_schema import ( graspa_input_schema, ) - from chemgraph.tools.graspa_tools import run_graspa_core + from chemgraph.tools.graspa_core import run_graspa_core if isinstance(job, dict): params = graspa_input_schema(**job) @@ -41,31 +40,6 @@ def run_graspa_parsl_app(job: dict): return run_graspa_core(params) -def load_parsl_config(system_name: str): - """ - Dynamically imports and returns the Parsl config based on the system name. - """ - system_name = system_name.lower() - run_dir = os.getcwd() - - logging.info("Initializing Parsl for system: %s", system_name) - - if system_name == "polaris": - from chemgraph.hpc_configs.polaris_parsl import get_polaris_config - - return get_polaris_config(run_dir=run_dir) - - elif system_name == "aurora": - from chemgraph.hpc_configs.aurora_parsl import get_aurora_config - - return get_aurora_config(run_dir=run_dir) - - else: - raise ValueError( - f"Unknown system specified: '{system_name}'. Supported: polaris, aurora" - ) - - # Load Parsl Config target_system = os.getenv("COMPUTE_SYSTEM", "polaris") parsl.load(load_parsl_config(target_system)) @@ -166,6 +140,20 @@ async def run_graspa_ensemble( pending_tasks.append((task_meta, fut)) async def wait_for_task(struct_name, parsl_future): + """Await a Parsl gRASPA task and normalize failures. + + Parameters + ---------- + struct_name : str + Structure identifier associated with the task. + parsl_future : concurrent.futures.Future + Parsl future returned by the submitted app. + + Returns + ------- + dict + Task result dictionary, or a normalized failure dictionary. + """ try: # Wrap the Parsl/Concurrent future so it becomes Awaitable res = await asyncio.wrap_future(parsl_future) diff --git a/src/chemgraph/mcp/mace_mcp_parsl.py b/src/chemgraph/mcp/mace_mcp_parsl.py index ce6d24c9..4b3f03fc 100644 --- a/src/chemgraph/mcp/mace_mcp_parsl.py +++ b/src/chemgraph/mcp/mace_mcp_parsl.py @@ -1,4 +1,3 @@ -import json import os from pathlib import Path @@ -12,11 +11,11 @@ import parsl from chemgraph.mcp.server_utils import run_mcp_server -from chemgraph.tools.parsl_tools import ( +from chemgraph.schemas.mace_parsl_schema import ( mace_input_schema, mace_input_schema_ensemble, - run_mace_core, ) +from chemgraph.tools.parsl_tools import run_mace_core from parsl import python_app @@ -36,7 +35,8 @@ def run_mace_parsl_app(job: dict): dict The result of `run_mace_core(job)`. """ - from chemgraph.tools.parsl_tools import mace_input_schema, run_mace_core + from chemgraph.schemas.mace_parsl_schema import mace_input_schema + from chemgraph.tools.parsl_tools import run_mace_core if isinstance(job, dict): params = mace_input_schema(**job) @@ -73,6 +73,18 @@ def run_mace_parsl_app(job: dict): description="Run a single MACE calculation", ) def run_mace_single(params: mace_input_schema): + """Run one MACE calculation. + + Parameters + ---------- + params : mace_input_schema + Input parameters for the MACE calculation. + + Returns + ------- + dict + MACE calculation result. + """ return run_mace_core(params) @@ -177,29 +189,21 @@ def run_mace_ensemble(params: mace_input_schema_ensemble): description="Load output from a JSON file.", ) def extract_output_json(json_file: str) -> dict: - """ - Load simulation results from a JSON file produced by run_ase. + """Load simulation results from a JSON file produced by run_ase. Parameters ---------- json_file : str - Path to the JSON file containing ASE simulation results. + Path to a JSON output file. Returns ------- - Dict[str, Any] - Parsed results from the JSON file as a Python dictionary. - - Raises - ------ - FileNotFoundError - If the specified file does not exist. - json.JSONDecodeError - If the file is not valid JSON. + dict + Parsed simulation results. """ - with open(json_file, "r") as f: - data = json.load(f) - return data + from chemgraph.tools.ase_core import extract_output_json_core + + return extract_output_json_core(json_file) # User-specific paths and settings diff --git a/src/chemgraph/mcp/mcp_tools.py b/src/chemgraph/mcp/mcp_tools.py index 4109e892..42650253 100644 --- a/src/chemgraph/mcp/mcp_tools.py +++ b/src/chemgraph/mcp/mcp_tools.py @@ -1,33 +1,21 @@ +"""FastMCP server exposing general chemistry tools. + +The ``run_ase`` tool delegates to :func:`chemgraph.tools.ase_core.run_ase_core` +so that the simulation logic lives in a single place. +""" + from __future__ import annotations -import os -import glob -import json -import time -from pathlib import Path + from typing import Literal from mcp.server.fastmcp import FastMCP - -import pubchempy as pcp -from ase import Atoms - -from chemgraph.tools.mcp_helper import ( - load_calculator, - atoms_to_atomsdata, - is_linear_molecule, - get_symmetry_number, +from chemgraph.tools.ase_core import extract_output_json_core, run_ase_core +from chemgraph.tools.cheminformatics_core import ( + molecule_name_to_smiles_core, + smiles_to_coordinate_file_core, ) -from chemgraph.schemas.ase_input import ASEInputSchema, ASEOutputSchema - - -def _resolve_path(path: str) -> str: - """If CHEMGRAPH_LOG_DIR is set and path is relative, prepend it.""" - log_dir = os.environ.get("CHEMGRAPH_LOG_DIR") - if log_dir and not os.path.isabs(path): - os.makedirs(log_dir, exist_ok=True) - return os.path.join(log_dir, path) - return path +from chemgraph.schemas.ase_input import ASEInputSchema mcp = FastMCP( @@ -50,33 +38,19 @@ def _resolve_path(path: str) -> str: description="Convert a molecule name to a canonical SMILES string using PubChem.", ) async def molecule_name_to_smiles(name: str) -> str: - """ + """Resolve a molecule name to its canonical SMILES via PubChem. + Parameters ---------- name : str - The molecule/common name to resolve. + Molecule name to resolve. Returns ------- str Canonical SMILES string. - - Raises - ------ - ValueError - If no match is found on PubChem. """ - if not name or not str(name).strip(): - raise ValueError("Parameter 'name' must be a non-empty string.") - - comps = pcp.get_compounds(str(name).strip(), "name") - if not comps: - raise ValueError(f"No PubChem compound found for name: {name!r}") - - smiles = comps[0].canonical_smiles - if not smiles: - raise ValueError(f"PubChem returned an empty SMILES for {name!r}.") - return smiles + return molecule_name_to_smiles_core(name) @mcp.tool( @@ -89,104 +63,47 @@ async def smiles_to_coordinate_file( seed: int = 2025, fmt: Literal["xyz"] = "xyz", ) -> dict: - """Convert a SMILES string to a coordinate file. + """Convert a SMILES string to a coordinate file on disk. Parameters ---------- smiles : str - SMILES string representation of the molecule. + Input SMILES string. output_file : str, optional - Path to save the output coordinate file (currently XYZ only). + Coordinate file path to write. seed : int, optional - Random seed for RDKit 3D structure generation, by default 2025. + Random seed used for conformer generation. fmt : {"xyz"}, optional - Output format. Only "xyz" supported for now. + Output coordinate format. Returns ------- - str - A single-line JSON string LLMs can parse, e.g. - { - "ok": true, - "artifact": "coordinate_file", - "format": "xyz", - "path": "...", - "smiles": "...", - "natoms": 12 - } - - Raises - ------ - ValueError - If the SMILES string is invalid or if 3D structure generation fails. + dict + Coordinate-generation result metadata. """ - from rdkit import Chem - from rdkit.Chem import AllChem - from ase.io import write as ase_write - - # Generate the molecule object - mol = Chem.MolFromSmiles(smiles) - if mol is None: - raise ValueError("Invalid SMILES string.") - - # Add hydrogens and optimize 3D structure - mol = Chem.AddHs(mol) - if AllChem.EmbedMolecule(mol, randomSeed=seed) != 0: - raise ValueError("Failed to generate 3D coordinates.") - if AllChem.UFFOptimizeMolecule(mol) != 0: - raise ValueError("Failed to optimize 3D geometry.") - # Extract atomic information - conf = mol.GetConformer() - numbers = [atom.GetAtomicNum() for atom in mol.GetAtoms()] - positions = [list(conf.GetAtomPosition(i)) for i in range(mol.GetNumAtoms())] - - # Create Atoms object - atoms = Atoms(numbers=numbers, positions=positions) - - final_output_file = _resolve_path(output_file) - ase_write( - final_output_file, - atoms, + return smiles_to_coordinate_file_core( + smiles, output_file=output_file, seed=seed, fmt=fmt ) - # Return dict for LLM/tool chaining - return { - "ok": True, - "artifact": "coordinate_file", - "path": os.path.abspath(final_output_file), - "smiles": smiles, - "natoms": len(numbers), - } - @mcp.tool( name="extract_output_json", description="Load simulation results from a JSON file produced by run_ase.", ) def extract_output_json(json_file: str) -> dict: - """ - Load simulation results from a JSON file produced by run_ase. + """Load simulation results from a JSON file produced by run_ase. Parameters ---------- json_file : str - Path to the JSON file containing ASE simulation results. + Path to the JSON output file. Returns ------- dict - Parsed results from the JSON file as a Python dictionary. - - Raises - ------ - FileNotFoundError - If the specified file does not exist. - json.JSONDecodeError - If the file is not valid JSON. + Parsed simulation results. """ - with open(json_file, "r", encoding="utf-8") as f: - data = json.load(f) - return data + return extract_output_json_core(json_file) @mcp.tool( @@ -213,339 +130,10 @@ async def run_ase(params: ASEInputSchema) -> dict: """ import io from contextlib import redirect_stdout - from ase.io import read - from ase.optimize import BFGS, LBFGS, GPMin, FIRE, MDMin - from chemgraph.schemas.atomsdata import AtomsData f = io.StringIO() with redirect_stdout(f): - try: - calculator = params.calculator.model_dump() - except Exception as e: - return f"Missing calculator parameter for the simulation. Raised exception: {str(e)}" - - # Calculate wall time. - start_time = time.time() - - input_structure_file = params.input_structure_file - output_results_file = _resolve_path(params.output_results_file) - optimizer = params.optimizer - fmax = params.fmax - steps = params.steps - driver = params.driver - temperature = params.temperature - pressure = params.pressure - - # # Validate that the input structure file exists - if not os.path.isfile(input_structure_file): - err = f"Input structure file {input_structure_file} does not exist." - raise ValueError(err) - - # Validate the output results file (if provided) - if not output_results_file.endswith(".json"): - err = f"Output results file must end with '.json', got: {params.output_results_file}" - raise ValueError(err) - - calc, system_info, calc_model = load_calculator(calculator) - params.calculator = calc_model - - if calc is None: - err = ( - f"Unsupported calculator: {calculator}. Available calculators are MACE" - "(mace_mp, mace_off, mace_anicc), EMT, TBLite (GFN2-xTB, GFN1-xTB), NWChem and Orca" - ) - raise ValueError(err) - - try: - atoms = read(input_structure_file) - except Exception as e: - err = ( - f"Cannot read {input_structure_file} using ASE. Exception from ASE: {e}" - ) - raise ValueError(err) from e - - atoms.info.update(system_info) - atoms.calc = calc - - if driver == "energy" or driver == "dipole": - energy = atoms.get_potential_energy() - final_structure = atoms_to_atomsdata(atoms=atoms) - - dipole = [None, None, None] - if driver == "dipole": - # Catch exception if calculator doesn't have get_dipole_moment() - try: - dipole = list(atoms.get_dipole_moment()) - except Exception: - pass - - end_time = time.time() - wall_time = end_time - start_time - - simulation_output = ASEOutputSchema( - input_structure_file=input_structure_file, - converged=True, - final_structure=final_structure, - simulation_input=params, - success=True, - dipole_value=dipole, - single_point_energy=energy, - wall_time=wall_time, - ) - with open(output_results_file, "w", encoding="utf-8") as wf: - wf.write(simulation_output.model_dump_json(indent=4)) - return { - "status": "success", - "message": f"Simulation completed. Results saved to {os.path.abspath(output_results_file)}", - "single_point_energy": energy, - "unit": "eV", - } - - OPTIMIZERS = { - "bfgs": BFGS, - "lbfgs": LBFGS, - "gpmin": GPMin, - "fire": FIRE, - "mdmin": MDMin, - } - try: - optimizer_class = OPTIMIZERS.get(optimizer.lower()) - if optimizer_class is None: - raise ValueError(f"Unsupported optimizer: {optimizer_class}") - - # Do optimization only if number of atoms > 1 to avoid error. - if len(atoms) > 1: - dyn = optimizer_class(atoms) - converged = dyn.run(fmax=fmax, steps=steps) - else: - converged = True - - single_point_energy = float(atoms.get_potential_energy()) - final_structure = AtomsData( - numbers=atoms.numbers, - positions=atoms.positions, - cell=atoms.cell, - pbc=atoms.pbc, - ) - thermo_data = {} - vib_data = {} - ir_data = {} - - if driver in {"vib", "thermo", "ir"}: - from ase.vibrations import Vibrations - from ase import units - import tempfile - import shutil - - ir_plot_path = None # Will be set inside tmpdir block if driver == "ir" - # Use a temporary directory to isolate parallel vibration runs. - # ASE's Vibrations class writes cache files (vib/cache.*.json) and - # trajectory files (vib.*.traj) using the `name` parameter. Without - # isolation, parallel calls for different molecules write to the same - # files, causing shape-mismatch errors and corrupted thermochemistry. - mol_stem = ( - Path(input_structure_file).stem if input_structure_file else "mol" - ) - - with tempfile.TemporaryDirectory( - prefix=f"chemgraph_vib_{mol_stem}_" - ) as tmpdir: - vib_name = os.path.join(tmpdir, "vib") - vib = Vibrations(atoms, name=vib_name) - - vib.clean() - vib.run() - - vib_data = { - "energies": [], - "energy_unit": "meV", - "frequencies": [], - "frequency_unit": "cm-1", - } - - energies = vib.get_energies() - linear = is_linear_molecule(atomsdata=final_structure) - - for idx, e in enumerate(energies): - is_imag = abs(e.imag) > 1e-8 - e_val = e.imag if is_imag else e.real - energy_meV = 1e3 * e_val - freq_cm1 = e_val / units.invcm - suffix = "i" if is_imag else "" - vib_data["energies"].append(f"{energy_meV}{suffix}") - vib_data["frequencies"].append(f"{freq_cm1}{suffix}") - - # Write frequencies.csv to the resolved output directory - freq_file_path = _resolve_path(f"frequencies_{mol_stem}.csv") - freq_file = Path(freq_file_path) - if freq_file.exists(): - freq_file.unlink() - - with freq_file.open("w", encoding="utf-8") as f: - for i, freq in enumerate(vib_data["frequencies"], start=0): - f.write(f"{mol_stem}_vib.{i}.traj,{freq}\n") - - # Write normal modes .traj files inside tmpdir, then copy out - for i in range(len(energies)): - vib.write_mode(n=i, kT=units.kB * 300, nimages=30) - - # Copy .traj files to the resolved output directory with molecule prefix - traj_dest_dir = _resolve_path("") - if traj_dest_dir: - os.makedirs(traj_dest_dir, exist_ok=True) - for traj_file in glob.glob(os.path.join(tmpdir, "vib.*.traj")): - dest_name = f"{mol_stem}_{Path(traj_file).name}" - dest_path = ( - os.path.join(traj_dest_dir, dest_name) - if traj_dest_dir - else dest_name - ) - shutil.copy2(traj_file, dest_path) - - if driver == "ir": - from ase.vibrations import Infrared - import matplotlib.pyplot as plt - - ir_data["spectrum_frequencies"] = [] - ir_data["spectrum_frequencies_units"] = "cm-1" - - ir_data["spectrum_intensities"] = [] - ir_data["spectrum_intensities_units"] = "D/Å^2 amu^-1" - - ir_name = os.path.join(tmpdir, "ir") - ir = Infrared(atoms, name=ir_name) - ir.clean() - ir.run() - - IR_SPECTRUM_START = 500 # Start of IR spectrum range - IR_SPECTRUM_END = 4000 # End of IR spectrum range - freq_intensity = ir.get_spectrum( - start=IR_SPECTRUM_START, end=IR_SPECTRUM_END - ) - # Generate IR spectrum plot - fig, ax = plt.subplots() - ax.plot(freq_intensity[0], freq_intensity[1]) - ax.set_xlabel("Frequency (cm⁻¹)") - ax.set_ylabel("Intensity (a.u.)") - ax.set_title("Infrared Spectrum") - ax.grid(True) - ir_plot_path = _resolve_path(f"ir_spectrum_{mol_stem}.png") - fig.savefig(ir_plot_path, format="png", dpi=300) - plt.close(fig) - - ir_data["IR Plot"] = f"Saved to {os.path.abspath(ir_plot_path)}" - ir_data["Normal mode data"] = ( - f"Normal modes saved as individual .traj files with prefix {mol_stem}_" - ) - - if driver == "thermo": - # Approximation for a single atom system. - if len(atoms) == 1: - thermo_data = { - "enthalpy": single_point_energy, - "entropy": 0.0, - "gibbs_free_energy": single_point_energy, - "unit": "eV", - } - else: - from ase.thermochemistry import IdealGasThermo - - linear = is_linear_molecule(atomsdata=final_structure) - geometry = "linear" if linear else "nonlinear" - symmetrynumber = get_symmetry_number( - atomsdata=final_structure - ) - - thermo = IdealGasThermo( - vib_energies=energies, - potentialenergy=single_point_energy, - atoms=atoms, - geometry=geometry, - symmetrynumber=symmetrynumber, - spin=0, # Only support spin=0 - ) - thermo_data = { - "enthalpy": float( - thermo.get_enthalpy(temperature=temperature) - ), - "entropy": float( - thermo.get_entropy( - temperature=temperature, pressure=pressure - ) - ), - "gibbs_free_energy": float( - thermo.get_gibbs_energy( - temperature=temperature, pressure=pressure - ) - ), - "unit": "eV", - } - - end_time = time.time() - wall_time = end_time - start_time - - simulation_output = ASEOutputSchema( - input_structure_file=input_structure_file, - converged=converged, - final_structure=final_structure, - simulation_input=params, - vibrational_frequencies=vib_data, - thermochemistry=thermo_data, - success=True, - ir_data=ir_data, - single_point_energy=single_point_energy, - wall_time=wall_time, - ) - - with open(output_results_file, "w", encoding="utf-8") as wf: - wf.write(simulation_output.model_dump_json(indent=4)) - - # Return message based on driver. Keep the return output minimal. - if driver == "opt": - return { - "status": "success", - "message": f"Simulation completed. Results saved to {os.path.abspath(output_results_file)}", - "single_point_energy": single_point_energy, # small payload for LLMs - "unit": "eV", - } - elif driver == "vib": - return { - "status": "success", - "result": { - "vibrational_frequencies": vib_data, - }, # small payload for LLMs - "message": ( - "Vibrational analysis completed; frequencies returned. " - f"Full results (structure, vibrations and metadata) saved to {os.path.abspath(output_results_file)}." - ), - } - elif driver == "thermo": - return { - "status": "success", - "result": { - "thermochemistry": thermo_data - }, # small payload for LLMs - "message": ( - "Thermochemistry computed and returned. " - f"Full results (structure, vibrations, thermochemistry and metadata) saved to {os.path.abspath(output_results_file)}" - ), - } - elif driver == "ir": - return { - "status": "success", - "result": { - "vibrational_frequencies": vib_data - }, # small payload for LLMs - "message": ( - "Infrared computed and returned. " - f"Full results (structure, vibrations, thermochemistry and metadata) saved to {os.path.abspath(output_results_file)}. " - f"IR plot saved to {os.path.abspath(ir_plot_path) if ir_plot_path else 'N/A'}. Normal modes saved as individual .traj files" - ), - } - - except Exception as e: - err = f"ASE simulation gave an exception:{e}" - raise ValueError(err) from e + return run_ase_core(params) if __name__ == "__main__": diff --git a/src/chemgraph/mcp/server_utils.py b/src/chemgraph/mcp/server_utils.py index b285fe3f..91fce11e 100644 --- a/src/chemgraph/mcp/server_utils.py +++ b/src/chemgraph/mcp/server_utils.py @@ -1,6 +1,8 @@ import argparse import logging +import os import sys + import uvicorn from mcp.server.fastmcp import FastMCP @@ -8,10 +10,19 @@ def run_mcp_server( mcp: FastMCP, default_port: int = 8000, default_host: str = "127.0.0.1" ): - """ - Standardizes the startup process for ChemGraph MCP servers. + """Standardize the startup process for ChemGraph MCP servers. + Supports 'stdio' (default) and 'sse' (HTTP) transports. Ensures logging is correctly routed to stderr to avoid corrupting stdio transport. + + Parameters + ---------- + mcp : FastMCP + MCP server instance to run. + default_port : int, optional + Default port for streamable HTTP transport. + default_host : str, optional + Default host for streamable HTTP transport. """ parser = argparse.ArgumentParser(description=f"Run {mcp.name} MCP Server") parser.add_argument( @@ -49,8 +60,6 @@ def run_mcp_server( logger.addHandler(stderr_handler) # If CHEMGRAPH_LOG_DIR is set, also log to a file - import os - log_dir = os.environ.get("CHEMGRAPH_LOG_DIR") if log_dir: try: @@ -77,3 +86,47 @@ def run_mcp_server( logging.info("Starting %s via stdio transport...", mcp.name) # FastMCP.run(transport='stdio') handles the stdio loop mcp.run(transport="stdio") + + +# --------------------------------------------------------------------------- +# Parsl HPC configuration loader +# --------------------------------------------------------------------------- + + +def load_parsl_config(system_name: str): + """Dynamically import and return a Parsl config for the given HPC system. + + Parameters + ---------- + system_name : str + Target system name. Supported: ``polaris``, ``aurora``. + + Returns + ------- + parsl.config.Config + Parsl configuration object for the requested system. + + Raises + ------ + ValueError + If the system name is not recognised. + """ + system_name = system_name.lower() + run_dir = os.getcwd() + + logging.info("Initializing Parsl for system: %s", system_name) + + if system_name == "polaris": + from chemgraph.hpc_configs.polaris_parsl import get_polaris_config + + return get_polaris_config(run_dir=run_dir) + + elif system_name == "aurora": + from chemgraph.hpc_configs.aurora_parsl import get_aurora_config + + return get_aurora_config(run_dir=run_dir) + + else: + raise ValueError( + f"Unknown system specified: '{system_name}'. Supported: polaris, aurora" + ) diff --git a/src/chemgraph/mcp/xanes_mcp.py b/src/chemgraph/mcp/xanes_mcp.py deleted file mode 100644 index d8ba9ead..00000000 --- a/src/chemgraph/mcp/xanes_mcp.py +++ /dev/null @@ -1,97 +0,0 @@ -from pathlib import Path - -from mcp.server.fastmcp import FastMCP - -from chemgraph.mcp.server_utils import run_mcp_server -from chemgraph.schemas.xanes_schema import xanes_input_schema, mp_query_schema - -# Start MCP server -mcp = FastMCP( - name="ChemGraph XANES Tools", - instructions=""" - You expose tools for running XANES/FDMNES simulations. - The available tools are: - 1. run_xanes_single: run a single FDMNES calculation for one structure. - 2. fetch_mp_structures: fetch optimized structures from Materials Project. - 3. plot_xanes: generate normalized XANES plots for completed calculations. - - Guidelines: - - Use each tool only when its input schema matches the user request. - - Do not guess numerical values; report tool errors exactly as they occur. - - Keep responses compact -- full results are in the output directories. - - When returning paths, use absolute paths. - - Energies are in eV. - - The FDMNES executable path is read from the FDMNES_EXE environment variable. - """, -) - - -@mcp.tool( - name="run_xanes_single", - description="Run a single XANES/FDMNES calculation for one input structure.", -) -def run_xanes_single(params: xanes_input_schema): - """Run a single FDMNES calculation using the core engine.""" - from chemgraph.tools.xanes_tools import run_xanes_core - - return run_xanes_core(params) - - -@mcp.tool( - name="fetch_mp_structures", - description="Fetch optimized structures from Materials Project.", -) -def fetch_mp_structures(params: mp_query_schema): - """Fetch structures from Materials Project and save as CIF files and pickle database.""" - from chemgraph.tools.xanes_tools import ( - fetch_materials_project_data, - _get_data_dir, - ) - - data_dir = _get_data_dir() - result = fetch_materials_project_data(params, data_dir) - return { - "status": "success", - "n_structures": result["n_structures"], - "chemsys": params.chemsys, - "output_dir": str(data_dir), - "structure_files": result["structure_files"], - "pickle_file": result["pickle_file"], - } - - -@mcp.tool( - name="plot_xanes", - description="Generate normalized XANES plots for completed FDMNES calculations.", -) -def plot_xanes(runs_dir: str): - """Generate XANES plots for all completed runs in a directory. - - Parameters - ---------- - runs_dir : str - Path to the ``fdmnes_batch_runs`` directory containing ``run_*`` - subdirectories with FDMNES outputs. - """ - from chemgraph.tools.xanes_tools import ( - plot_xanes_results, - _get_data_dir, - ) - - runs_path = Path(runs_dir) - if not runs_path.is_dir(): - raise ValueError(f"'{runs_dir}' is not a valid directory.") - - data_dir = _get_data_dir() - result = plot_xanes_results(data_dir, runs_path) - return { - "status": "success", - "n_plots": result["n_plots"], - "n_failed": result["n_failed"], - "plot_files": result["plot_files"], - "failed": result["failed"], - } - - -if __name__ == "__main__": - run_mcp_server(mcp, default_port=9007) diff --git a/src/chemgraph/mcp/xanes_mcp_parsl.py b/src/chemgraph/mcp/xanes_mcp_parsl.py index 55f2ef43..0ec794c1 100644 --- a/src/chemgraph/mcp/xanes_mcp_parsl.py +++ b/src/chemgraph/mcp/xanes_mcp_parsl.py @@ -1,6 +1,5 @@ import asyncio import json -import logging import os from pathlib import Path @@ -9,7 +8,7 @@ import parsl from parsl import bash_app -from chemgraph.mcp.server_utils import run_mcp_server +from chemgraph.mcp.server_utils import load_parsl_config, run_mcp_server from chemgraph.schemas.xanes_schema import ( xanes_input_schema, xanes_input_schema_ensemble, @@ -36,35 +35,6 @@ def run_fdmnes_parsl_app( return f'cd "{run_dir}" && "{fdmnes_exe}"' -def load_parsl_config(system_name: str): - """Dynamically import and return a Parsl config for the given HPC system. - - Parameters - ---------- - system_name : str - Target system name. Supported: ``polaris``, ``aurora``. - """ - system_name = system_name.lower() - run_dir = os.getcwd() - - logging.info("Initializing Parsl for system: %s", system_name) - - if system_name == "polaris": - from chemgraph.hpc_configs.polaris_parsl import get_polaris_config - - return get_polaris_config(run_dir=run_dir) - - elif system_name == "aurora": - from chemgraph.hpc_configs.aurora_parsl import get_aurora_config - - return get_aurora_config(run_dir=run_dir) - - else: - raise ValueError( - f"Unknown system specified: '{system_name}'. Supported: polaris, aurora" - ) - - # Load Parsl config at module level (same pattern as graspa_mcp_parsl.py) target_system = os.getenv("COMPUTE_SYSTEM", "polaris") parsl.load(load_parsl_config(target_system)) @@ -96,8 +66,19 @@ def load_parsl_config(system_name: str): description="Run a single XANES/FDMNES calculation for one input structure.", ) def run_xanes_single(params: xanes_input_schema): - """Run a single FDMNES calculation using the core engine.""" - from chemgraph.tools.xanes_tools import run_xanes_core + """Run a single FDMNES calculation using the core engine. + + Parameters + ---------- + params : xanes_input_schema + Input parameters for one XANES/FDMNES calculation. + + Returns + ------- + dict + XANES calculation result. + """ + from chemgraph.tools.xanes_core import run_xanes_core return run_xanes_core(params) @@ -122,7 +103,7 @@ async def run_xanes_ensemble(params: xanes_input_schema_ensemble): """ from ase.io import read as ase_read - from chemgraph.tools.xanes_tools import ( + from chemgraph.tools.xanes_core import ( write_fdmnes_input, extract_conv, ) @@ -194,6 +175,20 @@ async def run_xanes_ensemble(params: xanes_input_schema_ensemble): pending_tasks.append((task_meta, fut)) async def wait_for_task(meta, parsl_future): + """Await a Parsl FDMNES task and collect result metadata. + + Parameters + ---------- + meta : dict + Metadata describing the submitted structure/run. + parsl_future : concurrent.futures.Future + Parsl future returned by the submitted app. + + Returns + ------- + dict + Success or failure metadata for the task. + """ try: await asyncio.wrap_future(parsl_future) conv_data = extract_conv(meta["run_dir"]) @@ -235,8 +230,19 @@ async def wait_for_task(meta, parsl_future): description="Fetch optimized structures from Materials Project.", ) def fetch_mp_structures(params: mp_query_schema): - """Fetch structures from Materials Project and save as CIF files and pickle database.""" - from chemgraph.tools.xanes_tools import ( + """Fetch structures from Materials Project and save local artifacts. + + Parameters + ---------- + params : mp_query_schema + Materials Project query parameters. + + Returns + ------- + dict + Fetch summary including output directory and number of structures. + """ + from chemgraph.tools.xanes_core import ( fetch_materials_project_data, _get_data_dir, ) @@ -266,7 +272,7 @@ def plot_xanes(runs_dir: str): Path to the ``fdmnes_batch_runs`` directory containing ``run_*`` subdirectories with FDMNES outputs. """ - from chemgraph.tools.xanes_tools import ( + from chemgraph.tools.xanes_core import ( plot_xanes_results, _get_data_dir, ) diff --git a/src/chemgraph/memory/store.py b/src/chemgraph/memory/store.py index 0ecd1583..3d804897 100644 --- a/src/chemgraph/memory/store.py +++ b/src/chemgraph/memory/store.py @@ -60,6 +60,13 @@ class SessionStore: """ def __init__(self, db_path: Optional[str] = None): + """Initialize the SQLite session store. + + Parameters + ---------- + db_path : str, optional + SQLite database path. Defaults to ``~/.chemgraph/sessions.db``. + """ self.db_path = db_path or DEFAULT_DB_PATH os.makedirs(os.path.dirname(self.db_path), exist_ok=True) self._init_db() @@ -420,6 +427,16 @@ def _resolve_session_id(self, session_id: str) -> Optional[str]: Allows users to type just the first few characters of a UUID. Returns None if no match or ambiguous. + + Parameters + ---------- + session_id : str + Full session ID or prefix. + + Returns + ------- + str or None + Resolved full session ID, or ``None``. """ with self._connect() as conn: # Try exact match first diff --git a/src/chemgraph/models/openai.py b/src/chemgraph/models/openai.py index f904da67..e48fdb27 100644 --- a/src/chemgraph/models/openai.py +++ b/src/chemgraph/models/openai.py @@ -67,6 +67,18 @@ def _normalize_argo_model(model_name: str, base_url: str) -> str: names via ``ARGO_MODEL_MAP`` (e.g. ``argo:gpt-4o`` -> ``gpt4o``). * Other endpoints (ArgoProxy, custom): strip the ``argo:`` prefix and send the remainder as-is (e.g. ``argo:gpt-4o`` -> ``gpt-4o``). + + Parameters + ---------- + model_name : str + Requested model identifier. + base_url : str + Endpoint base URL used to choose normalization behavior. + + Returns + ------- + str + Endpoint-specific model name. """ if not model_name.startswith("argo:"): return model_name diff --git a/src/chemgraph/prompt/multi_agent_prompt.py b/src/chemgraph/prompt/multi_agent_prompt.py index 1251b6d4..6aed9cfb 100644 --- a/src/chemgraph/prompt/multi_agent_prompt.py +++ b/src/chemgraph/prompt/multi_agent_prompt.py @@ -20,10 +20,20 @@ 3. Each subtask must be independent — no task should depend on the result of another. 4. Include all relevant simulation parameters from the user's input (temperature, pressure, calculator, etc.) in each task prompt. +**PHASE 1b: Ask Human for Clarification** +- **Trigger:** The user query is missing critical information needed to generate tasks, or is ambiguous. +- **Action:** Set `next_step` to `"ask_human"` and provide a `clarification` string with a clear question. +- **When to ask:** + 1. Required simulation parameters are missing (e.g., no calculator specified, no temperature for thermochemistry, no molecule identified). + 2. The query is vague or could be interpreted in multiple ways (e.g., "calculate properties of water" — which properties?). + 3. An executor task failed and you need the user to decide how to proceed (e.g., retry with different parameters, use a different method, or skip). +- **Never guess or assume defaults** for critical parameters — always ask the human when in doubt. + **PHASE 2: Review Results (Subsequent invocations)** -- **Trigger:** You see executor results in the conversation history. +- **Trigger:** You see executor results or human clarification responses in the conversation history. - **Action:** Examine the results. Then either: - Set `next_step` to `"executor_subgraph"` with new `tasks` if more computation is needed (e.g., a task failed and should be retried, or intermediate results require follow-up calculations). + - Set `next_step` to `"ask_human"` with a `clarification` if you need further guidance from the user. - Set `next_step` to `"FINISH"` if all required data has been gathered. When finishing, include a comprehensive summary of all results in your `thought_process`, combining and aggregating values as needed to answer the user's original query. This summary is the final answer. **PHASE 2a: Handling Failed Tasks** @@ -50,6 +60,13 @@ ] } +When asking the human for clarification: +{ + "thought_process": "", + "next_step": "ask_human", + "clarification": "" +} + When finishing (all data gathered, final answer ready): { "thought_process": "", diff --git a/src/chemgraph/prompt/single_agent_prompt.py b/src/chemgraph/prompt/single_agent_prompt.py index 1028a56a..cac68e20 100644 --- a/src/chemgraph/prompt/single_agent_prompt.py +++ b/src/chemgraph/prompt/single_agent_prompt.py @@ -2,7 +2,16 @@ from chemgraph.schemas.agent_response import ResponseFormatter -single_agent_prompt = """You are an expert in computational chemistry, using advanced tools to solve complex problems. +_ASK_HUMAN_PROMPT_BLOCK = """\ +7. **Use the `ask_human` tool** in the following situations instead of guessing or failing silently: + - Required inputs are missing or ambiguous (e.g., molecule name, calculator type, temperature, pressure, basis set, or simulation method is not specified). + - You need confirmation before running a computationally expensive simulation (e.g., geometry optimization with a high-level calculator, large vibrational analysis). + - A previous tool call failed and you need the user to decide how to proceed (e.g., retry with different parameters, use a different calculator, or skip the step). + - The query is vague or could be interpreted in multiple ways. + Never crash or give up—always ask the human for help when you are uncertain. +""" + +single_agent_prompt = f"""You are an expert in computational chemistry, using advanced tools to solve complex problems. Instructions: 1. Extract all relevant inputs from the user's query, such as SMILES strings, molecule names, methods, software, properties, and conditions. @@ -11,7 +20,29 @@ 4. Review previous tool outputs. If they indicate failure, retry the tool with adjusted inputs if possible. 5. Use available simulation data directly. If data is missing, clearly state that a tool call is required. 6. If no tool call is needed, respond using factual domain knowledge. -""" +{_ASK_HUMAN_PROMPT_BLOCK}""" + + +def get_single_agent_prompt(human_supervised: bool = True) -> str: + """Return the single-agent system prompt. + + When *human_supervised* is ``False`` the ``ask_human`` instruction + block (item 7) is removed so the LLM is not instructed to call a + tool that is unavailable. + + Parameters + ---------- + human_supervised : bool, optional + Whether the ``ask_human`` tool will be available, by default True. + + Returns + ------- + str + The system prompt string. + """ + if human_supervised: + return single_agent_prompt + return single_agent_prompt.replace(_ASK_HUMAN_PROMPT_BLOCK, "") _response_schema_json = json.dumps(ResponseFormatter.model_json_schema(), indent=2) diff --git a/src/chemgraph/schemas/ase_input.py b/src/chemgraph/schemas/ase_input.py index fbe72076..ec1d151d 100644 --- a/src/chemgraph/schemas/ase_input.py +++ b/src/chemgraph/schemas/ase_input.py @@ -1,43 +1,114 @@ +import importlib.util import json +import os +import shutil from pydantic import BaseModel, Field, model_validator, field_validator from typing import Union, Optional, Any, List, Type from chemgraph.schemas.atomsdata import AtomsData -try: - from chemgraph.schemas.calculators.tblite_calc import TBLiteCalc -except ImportError: - TBLiteCalc = None from chemgraph.schemas.calculators.emt_calc import EMTCalc from chemgraph.schemas.calculators.nwchem_calc import NWChemCalc from chemgraph.schemas.calculators.orca_calc import OrcaCalc +from chemgraph.schemas.calculators.fairchem_calc import FAIRChemCalc +from chemgraph.schemas.calculators.mace_calc import MaceCalc +from chemgraph.schemas.calculators.tblite_calc import TBLiteCalc +from chemgraph.schemas.calculators.aimnet2_calc import AIMNET2Calc -try: - from chemgraph.schemas.calculators.aimnet2_calc import AIMNET2Calc -except ImportError: - AIMNET2Calc = None +# Gate optional calculators on whether their engine package is installed. +# Schema classes are always importable (internal to ChemGraph), so we must +# probe the underlying engine with importlib.util.find_spec(). +# find_spec() can raise ModuleNotFoundError for sub-packages when the parent +# package is missing, so we guard with try/except. + +def _engine_available(module_name: str) -> bool: + """Return whether a Python calculator engine module is importable. + + Parameters + ---------- + module_name : str + Module name passed to ``importlib.util.find_spec``. + + Returns + ------- + bool + ``True`` when the module can be found. + """ + try: + return importlib.util.find_spec(module_name) is not None + except (ImportError, ModuleNotFoundError): + return False + + +def _command_available(command_name: str, env_var_name: str) -> bool: + """Return whether a calculator command is configured or on ``PATH``. + + Parameters + ---------- + command_name : str + Executable name to locate. + env_var_name : str + Environment variable that can provide the command. + + Returns + ------- + bool + ``True`` when the command is configured or discoverable. + """ + return bool(os.environ.get(env_var_name)) or shutil.which(command_name) is not None -# Attempt to import optional calculators -try: - from chemgraph.schemas.calculators.fairchem_calc import FAIRChemCalc -except ImportError: +if not _engine_available("fairchem.core"): FAIRChemCalc = None -try: - from chemgraph.schemas.calculators.mace_calc import MaceCalc -except ImportError: +if not _engine_available("mace"): MaceCalc = None +if not _engine_available("tblite"): + TBLiteCalc = None + +if not _engine_available("aimnet2calc"): + AIMNET2Calc = None + +if not _command_available("nwchem", "ASE_NWCHEM_COMMAND"): + NWChemCalc = None + +if not _command_available("orca", "ASE_ORCA_COMMAND"): + OrcaCalc = None + + +_CALCULATOR_ALIASES = { + "xtb": "tbli", + "gfn1xtb": "tbli", + "gfn2xtb": "tbli", +} + + +def _calculator_key(name: str) -> str: + """Return a normalized calculator lookup key. + + Parameters + ---------- + name : str + Calculator name or alias. + + Returns + ------- + str + Four-character normalized key used for calculator matching. + """ + normalized = "".join(ch for ch in name.lower() if ch.isalnum()) + return _CALCULATOR_ALIASES.get(normalized, normalized[:4]) + # Define all possible calculator classes _all_calculator_classes: List[Optional[Type[BaseModel]]] = [ FAIRChemCalc, MaceCalc, - NWChemCalc, + AIMNET2Calc, TBLiteCalc, - OrcaCalc, EMTCalc, - AIMNET2Calc, + NWChemCalc, + OrcaCalc, ] # Filter out unavailable calculators @@ -48,15 +119,42 @@ # Create a union for type hinting CalculatorUnion = Union[tuple(available_calculator_classes)] -# Determine default calculator and names string -if FAIRChemCalc: - default_calculator = FAIRChemCalc -elif MaceCalc: - default_calculator = MaceCalc -else: - default_calculator = NWChemCalc +_calculator_name_items = [calc.__name__ for calc in available_calculator_classes] +_calculator_name_items = [ + "TBLiteCalc (aliases: xTB, GFN1-xTB, GFN2-xTB)" + if name == "TBLiteCalc" + else name + for name in _calculator_name_items +] +_calculator_names = ", ".join(_calculator_name_items) + +# Determine default calculator using only calculators detected as available. +default_calculator = available_calculator_classes[0] + -_calculator_names = ", ".join([calc.__name__ for calc in available_calculator_classes]) +def get_available_calculator_names() -> List[str]: + """Return calculator class names detected as available in this environment.""" + return [calc.__name__ for calc in available_calculator_classes] + + +def get_default_calculator_name() -> str: + """Return the default calculator class name selected for this environment.""" + return default_calculator.__name__ + + +def get_calculator_selection_context() -> str: + """Return prompt text describing available calculators and default choice.""" + return ( + "\n\nCalculator availability detected during ChemGraph initialization:\n" + f"- Available ASE calculators: {_calculator_names}.\n" + f"- Default calculator when the user does not specify one: " + f"{default_calculator.__name__}.\n" + "- When calling run_ase, choose only from the available calculators above. " + "If the user requests an unavailable calculator, choose the default " + "available calculator when that substitution is appropriate; otherwise " + "ask for clarification or explain that the requested calculator is not " + "available." + ) class ASEInputSchema(BaseModel): @@ -134,6 +232,18 @@ class ASEInputSchema(BaseModel): @model_validator(mode="before") @classmethod def _validate_calculator_type(cls, data: Any): + """Validate and coerce the calculator payload. + + Parameters + ---------- + data : Any + Raw ASE input payload before Pydantic validation. + + Returns + ------- + Any + Payload with calculator converted to an available calculator model. + """ if not isinstance(data, dict): return data @@ -142,32 +252,36 @@ def _validate_calculator_type(cls, data: Any): calc = default_calculator() data["calculator"] = calc - available_calcs = [c.__name__[:4].lower() for c in available_calculator_classes] + available_calcs = { + _calculator_key(c.__name__.removesuffix("Calc")): c + for c in available_calculator_classes + } + available_calc_names = [c.__name__ for c in available_calculator_classes] if isinstance(calc, dict): calc_name = calc.get("calculator_type") if not calc_name: raise ValueError("Calculator dictionary must have a 'calculator_type' key.") - if calc_name[:4].lower() not in available_calcs: + calc_key = _calculator_key(calc_name) + if calc_key not in available_calcs: raise ValueError( f"Calculator {calc_name} is not an allowed or available calculator. " - f"Available calculators are: {available_calcs}" + f"Available calculators are: {available_calc_names}" ) - for c in available_calculator_classes: - if c.__name__[:4].lower() == calc_name[:4].lower(): - init_args = calc.copy() - init_args.pop("calculator_type", None) - data["calculator"] = c(**init_args) - return data + init_args = calc.copy() + init_args.pop("calculator_type", None) + data["calculator"] = available_calcs[calc_key](**init_args) + return data elif hasattr(calc, "__class__"): calc_type_name = calc.__class__.__name__ - if calc_type_name[:4].lower() not in available_calcs: + calc_key = _calculator_key(calc_type_name.removesuffix("Calc")) + if calc_key not in available_calcs: raise ValueError( f"Calculator {calc_type_name} is not an allowed or available calculator. " - f"Available calculators are: {available_calcs}" + f"Available calculators are: {available_calc_names}" ) return data @@ -223,7 +337,18 @@ class ASEOutputSchema(BaseModel): @field_validator("vibrational_frequencies", "ir_data", "thermochemistry", mode="before") @classmethod def _coerce_json_string_to_dict(cls, v: Any) -> dict: - """Accept dict-like payloads serialized as JSON strings.""" + """Accept dict-like payloads serialized as JSON strings. + + Parameters + ---------- + v : Any + Raw field value. + + Returns + ------- + dict + Parsed dictionary or empty dictionary. + """ if v is None: return {} if isinstance(v, dict): @@ -242,7 +367,18 @@ def _coerce_json_string_to_dict(cls, v: Any) -> dict: @field_validator("error", mode="before") @classmethod def _coerce_error_to_string(cls, v: Any) -> str: - """Allow null/non-string error fields from intermediate tool payloads.""" + """Allow null/non-string error fields from intermediate tool payloads. + + Parameters + ---------- + v : Any + Raw error value. + + Returns + ------- + str + Normalized error string. + """ if v is None: return "" return v if isinstance(v, str) else str(v) diff --git a/src/chemgraph/schemas/calculators/__init__.py b/src/chemgraph/schemas/calculators/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/chemgraph/schemas/calculators/fairchem_calc.py b/src/chemgraph/schemas/calculators/fairchem_calc.py index 323b14ea..68cf6f49 100644 --- a/src/chemgraph/schemas/calculators/fairchem_calc.py +++ b/src/chemgraph/schemas/calculators/fairchem_calc.py @@ -1,6 +1,6 @@ -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, model_validator -from typing import Optional, Dict, Any +from typing import Any, Optional, Dict import torch import logging @@ -22,8 +22,11 @@ class FAIRChemCalc(BaseModel): Must match available tasks in the model. seed : int, optional Seed for model reproducibility. Default is 42. - spin : int, optional - Spin multiplicity. Default is 1. + multiplicity : int, optional + Spin multiplicity (2S+1) of the system. Default is 1 (singlet). + UMA/OMOL reads this from ``atoms.info["spin"]``; the schema field is named + ``multiplicity`` for consistency with other calculators (TBLite, ORCA). + The deprecated alias ``spin=`` is still accepted as input. charge : int, optional System charge. Default is 0. model_name: str @@ -41,7 +44,14 @@ class FAIRChemCalc(BaseModel): description="Prediction task. Options are 'omol', 'omat', 'oc20', 'odac', or 'omc", ) seed: int = Field(default=42, description="Random seed for inference reproducibility.") - spin: Optional[int] = Field(default=1, description="Total spin multiplicity of the system.") + multiplicity: Optional[int] = Field( + default=1, + description=( + "Spin multiplicity (2S+1) of the system. Default 1 (singlet). " + "Passed to UMA via atoms.info['spin']." + ), + ge=1, + ) charge: Optional[int] = Field(default=0, description="Total system charge.") model_name: str = Field( default="uma-s-1p1", description="Model names. Options are 'uma-s-1p1' and 'uma-m-1'" @@ -54,6 +64,28 @@ class FAIRChemCalc(BaseModel): default="default", description="Settings for inference. Can be 'default' or 'turbo'" ) + @model_validator(mode="before") + @classmethod + def _accept_spin_alias(cls, data: Any) -> Any: + """Accept deprecated ``spin`` input as ``multiplicity``. + + Parameters + ---------- + data : Any + Raw calculator payload before Pydantic validation. + + Returns + ------- + Any + Payload with ``spin`` converted when applicable. + """ + if isinstance(data, dict) and "spin" in data and "multiplicity" not in data: + logging.warning( + "FAIRChemCalc: field 'spin' is deprecated; use 'multiplicity' instead." + ) + data["multiplicity"] = data.pop("spin") + return data + def get_calculator(self) -> Any: """Return a configured FAIRChemCalculator. @@ -83,8 +115,16 @@ def get_calculator(self) -> Any: ) def get_atoms_properties(self) -> Dict[str, Optional[int]]: - """Return atom-level info keys to inject into atoms.info.""" + """Return atom-level info keys to inject into atoms.info. + + UMA/OMOL reads spin multiplicity from ``atoms.info["spin"]``; we keep + that key name here even though our schema field is ``multiplicity``. + """ return { - "spin": self.spin, + "spin": self.multiplicity, "charge": self.charge, } + + def get_multiplicity(self) -> Optional[int]: + """Return spin multiplicity (2S+1) for thermochemistry.""" + return self.multiplicity diff --git a/src/chemgraph/schemas/calculators/nwchem_calc.py b/src/chemgraph/schemas/calculators/nwchem_calc.py index 9c314457..426d516c 100644 --- a/src/chemgraph/schemas/calculators/nwchem_calc.py +++ b/src/chemgraph/schemas/calculators/nwchem_calc.py @@ -33,6 +33,14 @@ class NWChemCalc(BaseModel): command : str, optional Command to execute NWChem (e.g., 'nwchem PREFIX.nwi > PREFIX.nwo'), by default None + charge : int, optional + Total charge of the system, by default None + multiplicity : int, optional + Spin multiplicity (2S+1) of the system, by default None. + For molecular theories ('dft', 'scf', 'mp2', 'ccsd', 'tce', 'tddft') this is + injected into the theory block as ``mult``. For 'scf', NWChem expects + ``nopen`` (number of unpaired electrons); set ``scf={'nopen': N}`` manually + if you need finer control. """ calculator_type: str = Field( @@ -61,6 +69,18 @@ class NWChemCalc(BaseModel): default=None, description="Command to execute NWChem (e.g., 'nwchem PREFIX.nwi > PREFIX.nwo').", ) + charge: Optional[int] = Field( + default=None, description="Total charge of the system." + ) + multiplicity: Optional[int] = Field( + default=None, + description=( + "Spin multiplicity (2S+1). Injected into the theory block as 'mult' " + "for dft/mp2/ccsd/tce/tddft; for 'scf' theory NWChem expects 'nopen' " + "(unpaired electrons) which is not auto-set here." + ), + ge=1, + ) def get_calculator(self): """Get an ASE-compatible NWChem calculator instance. @@ -80,7 +100,7 @@ def get_calculator(self): "Invalid calculator_type. The only valid option is 'nwchem'." ) - return NWChem( + kwargs = dict( theory=self.theory, xc=self.xc, basis=self.basis, @@ -88,3 +108,18 @@ def get_calculator(self): directory=self.directory, command=self.command, ) + + # NWChem accepts charge/multiplicity inside the theory-specific block. + block: Dict[str, Union[int, str]] = {} + if self.charge is not None: + block["charge"] = self.charge + if self.multiplicity is not None and self.theory != "scf": + block["mult"] = self.multiplicity + if block and self.theory in {"dft", "mp2", "ccsd", "tce", "tddft"}: + kwargs[self.theory] = block + + return NWChem(**kwargs) + + def get_multiplicity(self) -> Optional[int]: + """Return spin multiplicity (2S+1) for thermochemistry.""" + return self.multiplicity diff --git a/src/chemgraph/schemas/calculators/orca_calc.py b/src/chemgraph/schemas/calculators/orca_calc.py index 2d7f024b..962ff3e3 100644 --- a/src/chemgraph/schemas/calculators/orca_calc.py +++ b/src/chemgraph/schemas/calculators/orca_calc.py @@ -121,3 +121,7 @@ def get_calculator(self): directory=self.directory, profile=profile, ) + + def get_multiplicity(self) -> Optional[int]: + """Return spin multiplicity (2S+1) for thermochemistry.""" + return self.multiplicity diff --git a/src/chemgraph/schemas/calculators/psi4_calc.py b/src/chemgraph/schemas/calculators/psi4_calc.py index 3c76f65d..c061f13a 100644 --- a/src/chemgraph/schemas/calculators/psi4_calc.py +++ b/src/chemgraph/schemas/calculators/psi4_calc.py @@ -1,3 +1,5 @@ +from typing import Optional + from pydantic import BaseModel, Field @@ -28,6 +30,10 @@ class Psi4Calc(BaseModel): 'cd' (Cholesky Decomposition), by default 'pk' maxiter : int, optional Maximum number of SCF iterations, by default 50 + charge : int, optional + Total charge of the system, by default 0 + multiplicity : int, optional + Spin multiplicity (2S+1) of the system, by default 1 (singlet) """ calculator_type: str = Field( @@ -58,6 +64,10 @@ class Psi4Calc(BaseModel): maxiter: int = Field( default=50, description="Maximum number of SCF iterations. Default is 50." ) + charge: int = Field(default=0, description="Total charge of the system.") + multiplicity: int = Field( + default=1, description="Spin multiplicity (2S+1) of the system.", ge=1 + ) def get_calculator(self) -> dict: """Get a dictionary of PSI4 calculation parameters. @@ -77,5 +87,11 @@ def get_calculator(self) -> dict: "reference": self.reference, "scf_type": self.scf_type, "maxiter": self.maxiter, + "charge": self.charge, + "multiplicity": self.multiplicity, } return params + + def get_multiplicity(self) -> Optional[int]: + """Return spin multiplicity (2S+1) for thermochemistry.""" + return self.multiplicity diff --git a/src/chemgraph/schemas/calculators/tblite_calc.py b/src/chemgraph/schemas/calculators/tblite_calc.py index e8b2ed01..547c8679 100644 --- a/src/chemgraph/schemas/calculators/tblite_calc.py +++ b/src/chemgraph/schemas/calculators/tblite_calc.py @@ -110,3 +110,7 @@ def get_calculator(self): cache_api=self.cache_api, verbosity=self.verbosity, ) + + def get_multiplicity(self) -> Optional[int]: + """Return spin multiplicity (2S+1) for thermochemistry.""" + return self.multiplicity diff --git a/src/chemgraph/schemas/mace_parsl_schema.py b/src/chemgraph/schemas/mace_parsl_schema.py new file mode 100644 index 00000000..e04ddba6 --- /dev/null +++ b/src/chemgraph/schemas/mace_parsl_schema.py @@ -0,0 +1,149 @@ +"""Pydantic schemas for MACE Parsl simulations. + +These schemas define the API contract for MACE single and ensemble +calculations dispatched via Parsl. +""" + +from __future__ import annotations + +from pydantic import BaseModel, Field + + +class mace_input_schema(BaseModel): + input_structure_file: str = Field( + description="Path to the input coordinate file (e.g., CIF, XYZ, POSCAR) containing the atomic structure for the simulation." + ) + output_result_file: str = Field( + default="output.json", + description="Path to a JSON file where simulation results will be saved.", + ) + driver: str = Field( + default=None, + description="Specifies the type of simulation to run. Options: 'energy' for single-point energy calculations, 'opt' for geometry optimization, 'vib' for vibrational frequency analysis, and 'thermo' for thermochemical properties (including enthalpy, entropy, and Gibbs free energy).", + ) + model: str = Field( + default="medium-mpa-0", + description="Path to the model. Default is medium-mpa-0." + "Options are 'small', 'medium', 'large', 'small-0b', 'medium-0b', 'small-0b2', 'medium-0b2','large-0b2', 'medium-0b3', 'medium-mpa-0', 'medium-omat-0', 'mace-matpes-pbe-0', 'mace-matpes-r2scan-0'", + ) + device: str = Field( + default="cpu", + description="Device to run the MACE calculation on. Options are cpu, cuda or xpu.", + ) + temperature: float = Field( + default=298.15, + description="Temperature for thermo property calculations in Kelvin (K).", + ) + pressure: float = Field( + default=101325.0, + description="Pressure for thermo property calculations in Pascal (Pa).", + ) + fmax: float = Field( + default=0.01, + description="The convergence criterion for forces (in eV/Å). Optimization stops when all force components are smaller than this value.", + ) + steps: int = Field( + default=1000, + description="Maximum number of optimization steps. The optimization will terminate if this number is reached, even if forces haven't converged to fmax.", + ) + optimizer: str = Field( + default="lbfgs", + description="The optimization algorithm used for geometry optimization. Options are 'bfgs', 'lbfgs', 'gpmin', 'fire', 'mdmin'", + ) + + +class mace_input_schema_ensemble(BaseModel): + input_structure_directory: str = Field( + description="Path to a folder of input structures containing the atomic structure for the simulations." + ) + output_result_file: str = Field( + default="output.json", + description="Path to a JSON file where simulation results will be saved.", + ) + driver: str = Field( + default=None, + description="Specifies the type of simulation to run. Options: 'energy' for single-point energy calculations, 'opt' for geometry optimization, 'vib' for vibrational frequency analysis, and 'thermo' for thermochemical properties (including enthalpy, entropy, and Gibbs free energy).", + ) + model: str = Field( + default="medium-mpa-0", + description="Path to the model. Default is medium-mpa-0." + "Options are 'small', 'medium', 'large', 'small-0b', 'medium-0b', 'small-0b2', 'medium-0b2','large-0b2', 'medium-0b3', 'medium-mpa-0', 'medium-omat-0', 'mace-matpes-pbe-0', 'mace-matpes-r2scan-0'", + ) + device: str = Field( + default="cpu", + description="Device to run the MACE calculation on. Options are cpu, cuda or xpu.", + ) + temperature: float = Field( + default=298.15, + description="Temperature for thermo property calculations in Kelvin (K).", + ) + pressure: float = Field( + default=101325.0, + description="Pressure for thermo property calculations in Pascal (Pa).", + ) + fmax: float = Field( + default=0.01, + description="The convergence criterion for forces (in eV/Å). Optimization stops when all force components are smaller than this value.", + ) + steps: int = Field( + default=1000, + description="Maximum number of optimization steps. The optimization will terminate if this number is reached, even if forces haven't converged to fmax.", + ) + optimizer: str = Field( + default="lbfgs", + description="The optimization algorithm used for geometry optimization. Options are 'bfgs', 'lbfgs', 'gpmin', 'fire', 'mdmin'", + ) + + +class mace_output_schema(BaseModel): + final_structure_file: str = Field( + description="Path to the final coordinate file (e.g., CIF, XYZ, POSCAR) containing the atomic structure for the simulation." + ) + output_result_file: str = Field( + description="Path to a JSON file where simulation results is saved.", + ) + model: str = Field( + default=None, description="Path to the model. Default is medium-mpa-0." + ) + device: str = Field( + default="cpu", + description="Device to run the MACE calculation on. Options are cpu, cuda or xpu.", + ) + temperature: float = Field( + default=298.15, + description="Temperature for thermo property calculations in Kelvin (K).", + ) + pressure: float = Field( + default=101325.0, + description="Pressure for thermo property calculations in Pascal (Pa).", + ) + fmax: float = Field( + default=0.01, + description="The convergence criterion for forces (in eV/Å). Optimization stops when all force components are smaller than this value.", + ) + steps: int = Field( + default=1000, + description="Maximum number of optimization steps. The optimization will terminate if this number is reached, even if forces haven't converged to fmax.", + ) + energy: float = Field( + description="The electronic energy of the system in eV", + ) + success: bool = Field( + description="Status of the simulation", + ) + vibrational_frequencies: dict = Field( + default={}, + description="Vibrational frequencies (in cm-1) and energies (in eV).", + ) + thermochemistry: dict = Field( + default={}, + description="Thermochemistry data in eV.", + ) + error: str = Field( + default="", + description="Error captured during the simulation", + ) + wall_time: float = Field( + default=None, + description="Total wall time (in seconds) taken to complete the simulation.", + ) diff --git a/src/chemgraph/schemas/multi_agent_response.py b/src/chemgraph/schemas/multi_agent_response.py index 68993699..af21a8cb 100644 --- a/src/chemgraph/schemas/multi_agent_response.py +++ b/src/chemgraph/schemas/multi_agent_response.py @@ -30,31 +30,53 @@ class PlannerResponse(BaseModel): Response model from the Planner agent. The planner acts as a router: it decides whether to dispatch tasks - to executor subgraphs (``executor_subgraph``) or to finish - (``FINISH``) when all work is done. + to executor subgraphs (``executor_subgraph``), ask the human user + for clarification (``ask_human``), or finish (``FINISH``) when all + work is done. Attributes: thought_process (str): The planner's reasoning for the current decision. - next_step (str): The next node to activate — either ``"executor_subgraph"`` - to fan-out tasks or ``"FINISH"`` to end the workflow. + next_step (str): The next node to activate — ``"executor_subgraph"`` + to fan-out tasks, ``"ask_human"`` to request human input, or + ``"FINISH"`` to end the workflow. tasks (list[WorkerTask] | None): Tasks to assign when routing to executors. + clarification (str | None): Question to ask the human when + ``next_step`` is ``"ask_human"``. """ thought_process: str = Field( description="Your reasoning for the current decision." ) - next_step: Literal["executor_subgraph", "FINISH"] = Field( + next_step: Literal["executor_subgraph", "ask_human", "FINISH"] = Field( description="The next node to activate in the workflow." ) tasks: list[WorkerTask] = Field( default=None, description="List of tasks to assign to executor subgraphs.", ) + clarification: Optional[str] = Field( + default=None, + description=( + "Question to present to the human user when next_step is 'ask_human'. " + "Must be provided when next_step is 'ask_human'." + ), + ) @model_validator(mode="before") @classmethod def normalize_planner_payload(cls, data: Any) -> Any: - """Accept common planner variants and coerce into PlannerResponse shape.""" + """Accept common planner variants and coerce into PlannerResponse shape. + + Parameters + ---------- + data : Any + Raw planner payload before Pydantic validation. + + Returns + ------- + Any + Normalized payload compatible with ``PlannerResponse``. + """ if isinstance(data, list): return { "thought_process": "Delegating parsed tasks to executors.", @@ -73,6 +95,13 @@ def normalize_planner_payload(cls, data: Any) -> Any: normalized["thought_process"] = ( "Delegating parsed tasks to executors." ) + # Accept legacy "question" key as clarification + if ( + normalized.get("next_step") == "ask_human" + and "clarification" not in normalized + and "question" in normalized + ): + normalized["clarification"] = normalized.pop("question") return normalized return data diff --git a/src/chemgraph/state/graspa_state.py b/src/chemgraph/state/graspa_state.py index 9ba4d285..1c85600c 100644 --- a/src/chemgraph/state/graspa_state.py +++ b/src/chemgraph/state/graspa_state.py @@ -5,7 +5,20 @@ def merge_dicts(a: dict, b: dict) -> dict: - """Reducer to merge dictionaries (for worker logs).""" + """Reducer to merge dictionaries for worker logs. + + Parameters + ---------- + a : dict + Existing accumulated dictionary. + b : dict + New dictionary update. + + Returns + ------- + dict + Merged dictionary where values from ``b`` override ``a``. + """ return {**a, **b} @@ -73,7 +86,18 @@ class PlannerResponse(BaseModel): @model_validator(mode="before") @classmethod def normalize_planner_payload(cls, data: Any) -> Any: - """Accept common planner variants and coerce into full PlannerResponse shape.""" + """Accept common planner variants and coerce into PlannerResponse shape. + + Parameters + ---------- + data : Any + Raw planner payload before Pydantic validation. + + Returns + ------- + Any + Normalized payload compatible with ``PlannerResponse``. + """ if isinstance(data, list): return { "thought_process": "Delegating parsed tasks to executors.", diff --git a/src/chemgraph/state/multi_agent_state.py b/src/chemgraph/state/multi_agent_state.py index 0b983dea..febd5a4a 100644 --- a/src/chemgraph/state/multi_agent_state.py +++ b/src/chemgraph/state/multi_agent_state.py @@ -1,11 +1,24 @@ import operator -from typing import TypedDict, Annotated, Any, Literal +from typing import TypedDict, Annotated, Any, Literal, Optional from langgraph.graph import add_messages def merge_dicts(a: dict, b: dict) -> dict: - """Reducer that merges dictionaries (used for executor logs).""" + """Reducer that merges dictionaries for executor logs. + + Parameters + ---------- + a : dict + Existing accumulated dictionary. + b : dict + New dictionary update. + + Returns + ------- + dict + Merged dictionary where values from ``b`` override ``a``. + """ return {**a, **b} @@ -44,12 +57,16 @@ class PlannerState(TypedDict): failures across iterations, enabling the planner to make informed retry decisions. Each entry is a dict with keys: ``task_index``, ``error``, ``retry_count``, and ``executor_id``. + + ``clarification`` holds the question text when the planner routes + to ``ask_human`` to request human input before proceeding. """ messages: Annotated[list, add_messages] - next_step: Literal["executor_subgraph", "FINISH"] + next_step: Literal["executor_subgraph", "ask_human", "FINISH"] tasks: list[dict[str, Any]] executor_results: Annotated[list, operator.add] executor_logs: Annotated[dict[str, list], merge_dicts] planner_iterations: int failed_tasks: Annotated[list, operator.add] + clarification: Optional[str] diff --git a/src/chemgraph/tools/ase_core.py b/src/chemgraph/tools/ase_core.py new file mode 100644 index 00000000..4e3dc915 --- /dev/null +++ b/src/chemgraph/tools/ase_core.py @@ -0,0 +1,697 @@ +"""Core simulation functions — the single source of truth. + +Every callable here is a plain Python function (no LangChain ``@tool``, +no MCP ``@mcp.tool``, no Parsl ``@python_app``). Framework-specific +wrappers in ``ase_tools.py``, ``mcp_tools.py``, and ``parsl_tools.py`` +simply delegate to these functions. +""" + +from __future__ import annotations + +import glob +import json +import os +import shutil +import tempfile +import time +from pathlib import Path +from typing import List, Optional + +import numpy as np + +from chemgraph.schemas.atomsdata import AtomsData +from chemgraph.schemas.ase_input import ASEInputSchema, ASEOutputSchema + + +# --------------------------------------------------------------------------- +# Path helpers +# --------------------------------------------------------------------------- + +def _resolve_path(path: str) -> str: + """Resolve a path relative to ``CHEMGRAPH_LOG_DIR`` when appropriate. + + Parameters + ---------- + path : str + Absolute or relative file path. + + Returns + ------- + str + Resolved path. + """ + log_dir = os.environ.get("CHEMGRAPH_LOG_DIR") + if log_dir and not os.path.isabs(path): + os.makedirs(log_dir, exist_ok=True) + return os.path.join(log_dir, path) + return path + + +# --------------------------------------------------------------------------- +# AtomsData <-> ASE Atoms conversions +# --------------------------------------------------------------------------- + +def atoms_to_atomsdata(atoms) -> AtomsData: + """Convert an ASE ``Atoms`` object to :class:`AtomsData`. + + Parameters + ---------- + atoms : ase.Atoms + ASE Atoms object. + + Returns + ------- + AtomsData + """ + return AtomsData( + numbers=atoms.numbers.tolist(), + positions=atoms.positions.tolist(), + cell=atoms.cell.tolist(), + pbc=atoms.pbc.tolist(), + ) + + +def atomsdata_to_atoms(atomsdata: AtomsData): + """Convert :class:`AtomsData` to an ASE ``Atoms`` object. + + Parameters + ---------- + atomsdata : AtomsData + + Returns + ------- + ase.Atoms + """ + from ase import Atoms + + return Atoms( + numbers=atomsdata.numbers, + positions=atomsdata.positions, + cell=atomsdata.cell, + pbc=atomsdata.pbc, + ) + + +# --------------------------------------------------------------------------- +# Molecular property helpers +# --------------------------------------------------------------------------- + +def is_linear_molecule(atomsdata: AtomsData, tol: float = 1e-3) -> bool: + """Determine whether a molecule is linear. + + Parameters + ---------- + atomsdata : AtomsData + Molecular structure. + tol : float, optional + Tolerance for the second singular value ratio, by default 1e-3. + + Returns + ------- + bool + ``True`` if the molecule is linear. + """ + coords = np.array(atomsdata.positions) + centered = coords - np.mean(coords, axis=0) + _, s, _ = np.linalg.svd(centered) + if s[0] == 0: + return False # degenerate — all atoms at one point + return (s[1] / s[0]) < tol + + +def get_symmetry_number(atomsdata: AtomsData) -> int: + """Return the rotational symmetry number using Pymatgen. + + Parameters + ---------- + atomsdata : AtomsData + + Returns + ------- + int + """ + from pymatgen.symmetry.analyzer import PointGroupAnalyzer + from ase import Atoms + from pymatgen.io.ase import AseAtomsAdaptor + + atoms = Atoms( + numbers=atomsdata.numbers, + positions=atomsdata.positions, + cell=atomsdata.cell, + pbc=atomsdata.pbc, + ) + aaa = AseAtomsAdaptor() + molecule = aaa.get_molecule(atoms) + pga = PointGroupAnalyzer(molecule) + return pga.get_rotational_symmetry_number() + + +# --------------------------------------------------------------------------- +# Calculator loading +# --------------------------------------------------------------------------- + +def load_calculator(calculator: dict) -> tuple[object, dict, object]: + """Instantiate an ASE calculator from a config dictionary. + + Parameters + ---------- + calculator : dict + Must contain a ``"calculator_type"`` key. + + Returns + ------- + tuple[object, dict, object] + ``(ase_calculator, extra_info, calc_schema_instance)`` + + Raises + ------ + ValueError + If the calculator type is unsupported. + """ + calc_type = calculator["calculator_type"].lower() + + if "emt" in calc_type: + from chemgraph.schemas.calculators.emt_calc import EMTCalc + calc = EMTCalc(**calculator) + elif "tblite" in calc_type or "xtb" in calc_type: + from chemgraph.schemas.calculators.tblite_calc import TBLiteCalc + calc = TBLiteCalc(**calculator) + elif "orca" in calc_type: + from chemgraph.schemas.calculators.orca_calc import OrcaCalc + calc = OrcaCalc(**calculator) + elif "nwchem" in calc_type: + from chemgraph.schemas.calculators.nwchem_calc import NWChemCalc + calc = NWChemCalc(**calculator) + elif "fairchem" in calc_type: + from chemgraph.schemas.calculators.fairchem_calc import FAIRChemCalc + calc = FAIRChemCalc(**calculator) + elif "mace" in calc_type: + from chemgraph.schemas.calculators.mace_calc import MaceCalc + calc = MaceCalc(**calculator) + elif "aimnet2" in calc_type: + from chemgraph.schemas.calculators.aimnet2_calc import AIMNET2Calc + calc = AIMNET2Calc(**calculator) + else: + raise ValueError( + f"Unsupported calculator: {calculator}. " + "Available calculators are EMT, TBLite (GFN2-xTB, GFN1-xTB), " + "Orca, NWChem, FAIRChem, MACE, or AIMNET2." + ) + + extra_info: dict = {} + if hasattr(calc, "get_atoms_properties"): + extra_info = calc.get_atoms_properties() + + return calc.get_calculator(), extra_info, calc + + +# --------------------------------------------------------------------------- +# Misc helpers (kept for backward compat / UI) +# --------------------------------------------------------------------------- + +def extract_ase_atoms_from_tool_result(tool_result: dict): + """Extract ``(atomic_numbers, positions)`` from a tool-result dict. + + Returns ``(None, None)`` if extraction fails. + + Parameters + ---------- + tool_result : dict + Tool result that may contain atom numbers and positions. + + Returns + ------- + tuple + ``(atomic_numbers, positions)`` or ``(None, None)``. + """ + for keyset in ({"numbers", "positions"}, {"atomic_numbers", "positions"}): + if keyset.issubset(tool_result.keys()): + return tool_result[keyset.pop()], tool_result["positions"] + + if "atoms" in tool_result: + atoms_data = tool_result["atoms"] + if {"numbers", "positions"}.issubset(atoms_data): + return atoms_data["numbers"], atoms_data["positions"] + + return None, None + + +def create_ase_atoms(atomic_numbers, positions): + """Create an ASE ``Atoms`` object from atomic numbers and positions. + + Parameters + ---------- + atomic_numbers : sequence + Atomic numbers for each atom. + positions : sequence + Cartesian coordinates for each atom. + + Returns + ------- + ase.Atoms or None + Constructed atoms object, or ``None`` if construction fails. + """ + from ase import Atoms + + try: + return Atoms(numbers=atomic_numbers, positions=positions) + except Exception as e: + print(f"Error creating ASE Atoms object: {e}") + return None + + +def create_xyz_string(atomic_numbers, positions) -> Optional[str]: + """Create an XYZ-format string from atomic numbers and positions. + + Parameters + ---------- + atomic_numbers : sequence + Atomic numbers for each atom. + positions : sequence + Cartesian coordinates for each atom. + + Returns + ------- + str or None + XYZ-format structure text, or ``None`` if conversion fails. + """ + from ase import Atoms + + try: + atoms = Atoms(numbers=atomic_numbers, positions=positions) + xyz_lines = [str(len(atoms)), "Generated by ChemGraph"] + for symbol, pos in zip(atoms.get_chemical_symbols(), atoms.positions): + xyz_lines.append( + f"{symbol:2s} {pos[0]:12.6f} {pos[1]:12.6f} {pos[2]:12.6f}" + ) + return "\n".join(xyz_lines) + except Exception as e: + print(f"Error creating XYZ string: {e}") + return None + + +# --------------------------------------------------------------------------- +# Unified ASE simulation core +# --------------------------------------------------------------------------- + +def run_ase_core(params: ASEInputSchema) -> dict: + """Run an ASE simulation — the single implementation for all call methods. + + This function implements energy, dipole, optimization, vibrational, + thermochemistry, and IR calculations. Framework-specific wrappers + (LangChain ``@tool``, MCP ``@mcp.tool``, Parsl) delegate here. + + Parameters + ---------- + params : ASEInputSchema + Fully validated simulation input. + + Returns + ------- + dict + Minimal result payload (status, message, key numbers). + """ + from ase.io import read + from ase.optimize import BFGS, LBFGS, GPMin, FIRE, MDMin + + # ---- unpack params ---- + try: + calculator = params.calculator.model_dump() + except Exception as e: + return { + "status": "failure", + "error_type": "ValidationError", + "message": f"Missing calculator parameter for the simulation. Raised exception: {e}", + } + + start_time = time.time() + + input_structure_file = params.input_structure_file + output_results_file = _resolve_path(params.output_results_file) + optimizer = params.optimizer + fmax = params.fmax + steps = params.steps + driver = params.driver + temperature = params.temperature + pressure = params.pressure + + # ---- input validation ---- + if not os.path.isfile(input_structure_file): + return { + "status": "failure", + "error_type": "FileNotFoundError", + "message": f"Input structure file {input_structure_file} does not exist.", + } + + if not output_results_file.endswith(".json"): + return { + "status": "failure", + "error_type": "ValueError", + "message": f"Output results file must end with '.json', got: {params.output_results_file}", + } + + calc, system_info, calc_model = load_calculator(calculator) + + if calc is None: + return { + "status": "failure", + "error_type": "ValueError", + "message": ( + f"Unsupported calculator: {calculator}. Available calculators are " + "MACE (mace_mp, mace_off, mace_anicc), EMT, TBLite (GFN2-xTB, GFN1-xTB), NWChem and Orca" + ), + } + + try: + atoms = read(input_structure_file) + except Exception as e: + return { + "status": "failure", + "error_type": type(e).__name__, + "message": f"Cannot read {input_structure_file} using ASE. Exception from ASE: {e}", + } + + atoms.info.update(system_info) + atoms.calc = calc + + # ------------------------------------------------------------------ + # Driver: energy / dipole (single-point, no optimization) + # ------------------------------------------------------------------ + if driver in ("energy", "dipole"): + energy = atoms.get_potential_energy() + final_structure = atoms_to_atomsdata(atoms) + + dipole: List[Optional[float]] = [None, None, None] + if driver == "dipole": + try: + dipole = [round(x, 4) for x in atoms.get_dipole_moment()] + except Exception: + pass + + end_time = time.time() + wall_time = end_time - start_time + + simulation_output = ASEOutputSchema( + input_structure_file=input_structure_file, + converged=True, + final_structure=final_structure, + simulation_input=params, + success=True, + dipole_value=dipole, + single_point_energy=energy, + wall_time=wall_time, + ) + with open(output_results_file, "w", encoding="utf-8") as wf: + wf.write(simulation_output.model_dump_json(indent=4)) + + if driver == "energy": + return { + "status": "success", + "message": f"Simulation completed. Results saved to {os.path.abspath(output_results_file)}", + "single_point_energy": energy, + "unit": "eV", + } + else: # dipole + return { + "status": "success", + "message": f"Simulation completed. Results saved to {os.path.abspath(output_results_file)}", + "dipole_moment": dipole, + "dipole_unit": "e * Angstrom", + } + + # ------------------------------------------------------------------ + # Drivers that require optimization: opt / vib / thermo / ir + # ------------------------------------------------------------------ + OPTIMIZERS = { + "bfgs": BFGS, + "lbfgs": LBFGS, + "gpmin": GPMin, + "fire": FIRE, + "mdmin": MDMin, + } + try: + optimizer_class = OPTIMIZERS.get(optimizer.lower()) + if optimizer_class is None: + raise ValueError(f"Unsupported optimizer: {optimizer}") + + if len(atoms) > 1: + dyn = optimizer_class(atoms) + converged = dyn.run(fmax=fmax, steps=steps) + else: + converged = True + + single_point_energy = float(atoms.get_potential_energy()) + final_structure = AtomsData( + numbers=atoms.numbers, + positions=atoms.positions, + cell=atoms.cell, + pbc=atoms.pbc, + ) + thermo_data: dict = {} + vib_data: dict = {} + ir_data: dict = {} + + # -------------------------------------------------------------- + # Vibrational / thermo / IR analysis + # -------------------------------------------------------------- + if driver in {"vib", "thermo", "ir"}: + from ase.vibrations import Vibrations + from ase import units + + ir_plot_path: Optional[str] = None + mol_stem = ( + Path(input_structure_file).stem if input_structure_file else "mol" + ) + + with tempfile.TemporaryDirectory( + prefix=f"chemgraph_vib_{mol_stem}_" + ) as tmpdir: + vib_name = os.path.join(tmpdir, "vib") + vib = Vibrations(atoms, name=vib_name) + vib.clean() + vib.run() + + vib_data = { + "energies": [], + "energy_unit": "meV", + "frequencies": [], + "frequency_unit": "cm-1", + } + + energies = vib.get_energies() + + for _idx, e in enumerate(energies): + is_imag = abs(e.imag) > 1e-8 + e_val = e.imag if is_imag else e.real + energy_meV = 1e3 * e_val + freq_cm1 = e_val / units.invcm + suffix = "i" if is_imag else "" + vib_data["energies"].append(f"{energy_meV}{suffix}") + vib_data["frequencies"].append(f"{freq_cm1}{suffix}") + + # Write frequencies CSV + freq_file_path = _resolve_path(f"frequencies_{mol_stem}.csv") + freq_file = Path(freq_file_path) + if freq_file.exists(): + freq_file.unlink() + with freq_file.open("w", encoding="utf-8") as f: + for i, freq in enumerate(vib_data["frequencies"], start=0): + f.write(f"{mol_stem}_vib.{i}.traj,{freq}\n") + + # Write normal-mode .traj files, then copy out of tmpdir + for i in range(len(energies)): + vib.write_mode(n=i, kT=units.kB * 300, nimages=30) + + traj_dest_dir = _resolve_path("") + if traj_dest_dir: + os.makedirs(traj_dest_dir, exist_ok=True) + for traj_file in glob.glob(os.path.join(tmpdir, "vib.*.traj")): + dest_name = f"{mol_stem}_{Path(traj_file).name}" + dest_path = ( + os.path.join(traj_dest_dir, dest_name) + if traj_dest_dir + else dest_name + ) + shutil.copy2(traj_file, dest_path) + + # ---- IR ---- + if driver == "ir": + from ase.vibrations import Infrared + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + ir_data["spectrum_frequencies"] = [] + ir_data["spectrum_frequencies_units"] = "cm-1" + ir_data["spectrum_intensities"] = [] + ir_data["spectrum_intensities_units"] = "D/Å^2 amu^-1" + + ir_name = os.path.join(tmpdir, "ir") + ir = Infrared(atoms, name=ir_name) + ir.clean() + ir.run() + + IR_SPECTRUM_START = 500 + IR_SPECTRUM_END = 4000 + freq_intensity = ir.get_spectrum( + start=IR_SPECTRUM_START, end=IR_SPECTRUM_END + ) + fig, ax = plt.subplots() + ax.plot(freq_intensity[0], freq_intensity[1]) + ax.set_xlabel("Frequency (cm⁻¹)") + ax.set_ylabel("Intensity (a.u.)") + ax.set_title("Infrared Spectrum") + ax.grid(True) + ir_plot_path = _resolve_path(f"ir_spectrum_{mol_stem}.png") + fig.savefig(ir_plot_path, format="png", dpi=300) + plt.close(fig) + + ir_data["IR Plot"] = f"Saved to {os.path.abspath(ir_plot_path)}" + ir_data["Normal mode data"] = ( + f"Normal modes saved as individual .traj files with prefix {mol_stem}_" + ) + + # ---- Thermochemistry ---- + if driver == "thermo": + if len(atoms) == 1: + thermo_data = { + "enthalpy": single_point_energy, + "entropy": 0.0, + "gibbs_free_energy": single_point_energy, + "unit": "eV", + } + else: + from ase.thermochemistry import IdealGasThermo + + linear = is_linear_molecule(final_structure) + geometry = "linear" if linear else "nonlinear" + symmetrynumber = get_symmetry_number(final_structure) + + # IdealGasThermo expects total spin S; calculators expose + # multiplicity (2S+1) via get_multiplicity() when supported. + multiplicity = ( + getattr(calc_model, "get_multiplicity", lambda: None)() + or 1 + ) + spin_S = (multiplicity - 1) / 2.0 + + thermo = IdealGasThermo( + vib_energies=energies, + potentialenergy=single_point_energy, + atoms=atoms, + geometry=geometry, + symmetrynumber=symmetrynumber, + spin=spin_S, + ) + thermo_data = { + "enthalpy": float( + thermo.get_enthalpy(temperature=temperature) + ), + "entropy": float( + thermo.get_entropy( + temperature=temperature, pressure=pressure + ) + ), + "gibbs_free_energy": float( + thermo.get_gibbs_energy( + temperature=temperature, pressure=pressure + ) + ), + "unit": "eV", + } + + # ---- serialise full output ---- + end_time = time.time() + wall_time = end_time - start_time + + simulation_output = ASEOutputSchema( + input_structure_file=input_structure_file, + converged=converged, + final_structure=final_structure, + simulation_input=params, + vibrational_frequencies=vib_data, + thermochemistry=thermo_data, + success=True, + ir_data=ir_data, + single_point_energy=single_point_energy, + wall_time=wall_time, + ) + with open(output_results_file, "w", encoding="utf-8") as wf: + wf.write(simulation_output.model_dump_json(indent=4)) + + # ---- minimal return payload ---- + abs_output = os.path.abspath(output_results_file) + if driver == "opt": + return { + "status": "success", + "message": f"Simulation completed. Results saved to {abs_output}", + "single_point_energy": single_point_energy, + "unit": "eV", + } + elif driver == "vib": + return { + "status": "success", + "result": {"vibrational_frequencies": vib_data}, + "message": ( + "Vibrational analysis completed; frequencies returned. " + f"Full results (structure, vibrations and metadata) saved to {abs_output}." + ), + } + elif driver == "thermo": + return { + "status": "success", + "result": {"thermochemistry": thermo_data}, + "message": ( + "Thermochemistry computed and returned. " + f"Full results (structure, vibrations, thermochemistry and metadata) saved to {abs_output}" + ), + } + elif driver == "ir": + return { + "status": "success", + "result": {"vibrational_frequencies": vib_data}, + "message": ( + "Infrared computed and returned. " + f"Full results (structure, vibrations, thermochemistry and metadata) saved to {abs_output}. " + f"IR plot saved to {os.path.abspath(ir_plot_path) if ir_plot_path else 'N/A'}. " + "Normal modes saved as individual .traj files" + ), + } + + except Exception as e: + return { + "status": "failure", + "error_type": type(e).__name__, + "message": str(e), + } + + +# --------------------------------------------------------------------------- +# JSON result loader +# --------------------------------------------------------------------------- + + +def extract_output_json_core(json_file: str) -> dict: + """Load simulation results from a JSON file produced by ``run_ase_core``. + + Parameters + ---------- + json_file : str + Path to the JSON file containing ASE simulation results. + + Returns + ------- + dict + Parsed results from the JSON file as a Python dictionary. + + Raises + ------ + FileNotFoundError + If the specified file does not exist. + json.JSONDecodeError + If the file is not valid JSON. + """ + with open(json_file, "r", encoding="utf-8") as f: + data = json.load(f) + return data diff --git a/src/chemgraph/tools/ase_tools.py b/src/chemgraph/tools/ase_tools.py index dc02663a..ff4650a3 100644 --- a/src/chemgraph/tools/ase_tools.py +++ b/src/chemgraph/tools/ase_tools.py @@ -1,114 +1,44 @@ -from pathlib import Path +"""LangChain ``@tool`` wrappers over :mod:`chemgraph.tools.ase_core`. + +Every public function here is a thin decorator that delegates to the +corresponding plain-Python implementation in ``ase_core.py``. +""" + +from __future__ import annotations + import os -import time -import json -import numpy as np from typing import Any, Dict from langchain_core.tools import tool + from chemgraph.schemas.atomsdata import AtomsData from chemgraph.schemas.ase_input import ASEInputSchema from chemgraph.schemas.calculators.mace_calc import _mace_lock -from chemgraph.tools.mcp_helper import _resolve_path +from chemgraph.tools.ase_core import ( + _resolve_path, + atoms_to_atomsdata, + extract_output_json_core, + run_ase_core, + is_linear_molecule as _is_linear_molecule, + get_symmetry_number as _get_symmetry_number, +) @tool def extract_output_json(json_file: str) -> Dict[str, Any]: - """ - Load simulation results from a JSON file produced by run_ase. + """Load simulation results from a JSON file produced by run_ase. Parameters ---------- json_file : str - Path to the JSON file containing ASE simulation results. - - Returns - ------- - Dict[str, Any] - Parsed results from the JSON file as a Python dictionary. - - Raises - ------ - FileNotFoundError - If the specified file does not exist. - json.JSONDecodeError - If the file is not valid JSON. - """ - with open(json_file, "r", encoding="utf-8") as f: - data = json.load(f) - return data - - -def extract_ase_atoms_from_tool_result(tool_result: dict): - """Extract ASE atoms data from tool result dictionary. - - Parameters - ---------- - tool_result : dict - Dictionary containing tool result data - - Returns - ------- - tuple - (atomic_numbers, positions) or (None, None) if extraction fails - """ - for keyset in ( - {"numbers", "positions"}, - {"atomic_numbers", "positions"}, - ): - if keyset.issubset(tool_result.keys()): - return tool_result[keyset.pop()], tool_result["positions"] - - if "atoms" in tool_result: - atoms_data = tool_result["atoms"] - if {"numbers", "positions"}.issubset(atoms_data): - return atoms_data["numbers"], atoms_data["positions"] - - return None, None - - -def atoms_to_atomsdata(atoms): - """Convert ASE Atoms object to AtomsData. - - Parameters - ---------- - atoms : ase.Atoms - ASE Atoms object - - Returns - ------- - AtomsData - ChemGraph AtomsData object - """ - return AtomsData( - numbers=atoms.numbers.tolist(), - positions=atoms.positions.tolist(), - cell=atoms.cell.tolist(), - pbc=atoms.pbc.tolist(), - ) - - -def atomsdata_to_atoms(atomsdata: AtomsData): - """Convert AtomsData to ASE Atoms object. - - Parameters - ---------- - atomsdata : AtomsData - ChemGraph AtomsData object + Path to the JSON output file. Returns ------- - ase.Atoms - ASE Atoms object + dict[str, Any] + Parsed simulation results. """ - from ase import Atoms - - return Atoms( - numbers=atomsdata.numbers, - positions=atomsdata.positions, - cell=atomsdata.cell, - pbc=atomsdata.pbc, - ) + return extract_output_json_core(json_file) @tool @@ -136,14 +66,7 @@ def file_to_atomsdata(fname: str) -> AtomsData: try: atoms = read(fname) - # Create AtomsData object from ASE Atoms object - atoms_data = AtomsData( - numbers=atoms.numbers.tolist(), - positions=atoms.positions.tolist(), - cell=atoms.cell.tolist(), - pbc=atoms.pbc.tolist(), - ) - return atoms_data + return atoms_to_atomsdata(atoms) except FileNotFoundError: raise FileNotFoundError(f"File not found: {fname}") except Exception as e: @@ -202,23 +125,7 @@ def get_symmetry_number(atomsdata: AtomsData) -> int: int Rotational symmetry number of the molecule """ - from pymatgen.symmetry.analyzer import PointGroupAnalyzer - from ase import Atoms - from pymatgen.io.ase import AseAtomsAdaptor - - atoms = Atoms( - numbers=atomsdata.numbers, - positions=atomsdata.positions, - cell=atomsdata.cell, - pbc=atomsdata.pbc, - ) - - aaa = AseAtomsAdaptor() - molecule = aaa.get_molecule(atoms) - pga = PointGroupAnalyzer(molecule) - symmetrynumber = pga.get_rotational_symmetry_number() - - return symmetrynumber + return _get_symmetry_number(atomsdata) @tool @@ -237,80 +144,7 @@ def is_linear_molecule(atomsdata: AtomsData, tol=1e-3) -> bool: bool True if the molecule is linear, False otherwise """ - coords = np.array(atomsdata.positions) - # Center the coordinates. - centered = coords - np.mean(coords, axis=0) - # Singular value decomposition. - U, s, Vt = np.linalg.svd(centered) - # For a linear molecule, only one singular value is significantly nonzero. - if s[0] == 0: - return False # degenerate case (all atoms at one point) - return (s[1] / s[0]) < tol - - -def load_calculator(calculator: dict) -> tuple[object, dict, dict]: - """Load an ASE calculator based on the provided configuration. - - Parameters - ---------- - calculator : dict - Dictionary containing calculator configuration parameters - - Returns - ------- - object - ASE calculator instance - - Raises - ------ - ValueError - If the calculator type is not supported - """ - calc_type = calculator["calculator_type"].lower() - - if "emt" in calc_type: - from chemgraph.schemas.calculators.emt_calc import EMTCalc - - calc = EMTCalc(**calculator) - elif "tblite" in calc_type: - from chemgraph.schemas.calculators.tblite_calc import TBLiteCalc - - calc = TBLiteCalc(**calculator) - elif "orca" in calc_type: - from chemgraph.schemas.calculators.orca_calc import OrcaCalc - - calc = OrcaCalc(**calculator) - - elif "nwchem" in calc_type: - from chemgraph.schemas.calculators.nwchem_calc import NWChemCalc - - calc = NWChemCalc(**calculator) - - elif "fairchem" in calc_type: - from chemgraph.schemas.calculators.fairchem_calc import FAIRChemCalc - - calc = FAIRChemCalc(**calculator) - - elif "mace" in calc_type: - from chemgraph.schemas.calculators.mace_calc import MaceCalc - - calc = MaceCalc(**calculator) - - elif "aimnet2" in calc_type: - from chemgraph.schemas.calculators.aimnet2_calc import AIMNET2Calc - - calc = AIMNET2Calc(**calculator) - - else: - raise ValueError( - f"Unsupported calculator: {calculator}. Available calculators are EMT, TBLite (GFN2-xTB, GFN1-xTB), Orca and FAIRChem or MACE or AIMNET2." - ) - # Extract additional args like spin/charge if the model defines it - extra_info = {} - if hasattr(calc, "get_atoms_properties"): - extra_info = calc.get_atoms_properties() - - return calc.get_calculator(), extra_info, calc + return _is_linear_molecule(atomsdata, tol) @tool @@ -335,404 +169,5 @@ def run_ase(params: ASEInputSchema) -> dict: calc_type = params.calculator.calculator_type.lower() if "mace" in calc_type: with _mace_lock: - return _run_ase_impl(params) - return _run_ase_impl(params) - - -def _run_ase_impl(params: ASEInputSchema): - """Core implementation of run_ase, separated to allow lock-guarded dispatch.""" - from ase.io import read - from ase.optimize import BFGS, LBFGS, GPMin, FIRE, MDMin - - try: - calculator = params.calculator.model_dump() - except Exception as e: - return f"Missing calculator parameter for the simulation. Raised exception: {str(e)}" - - # Calculate wall time. - start_time = time.time() - - input_structure_file = params.input_structure_file - output_results_file = _resolve_path(params.output_results_file) - optimizer = params.optimizer - fmax = params.fmax - steps = params.steps - driver = params.driver - temperature = params.temperature - pressure = params.pressure - - # # Validate that the input structure file exists - if not os.path.isfile(input_structure_file): - return f"Input structure file {input_structure_file} does not exist." - - # Validate the output results file (if provided) - if not output_results_file.endswith(".json"): - return f"Output results file must end with '.json', got: {params.output_results_file}" - - calc, system_info, calc_model = load_calculator(calculator) - - if calc is None: - return f"Unsupported calculator: {calculator}. Available calculators are MACE (mace_mp, mace_off, mace_anicc), EMT, TBLite (GFN2-xTB, GFN1-xTB), NWChem and Orca" - - try: - atoms = read(input_structure_file) - except Exception as e: - return f"Cannot read {input_structure_file} using ASE. Exception from ASE: {e}" - - atoms.info.update(system_info) - atoms.calc = calc - - if driver == "energy" or driver == "dipole": - energy = atoms.get_potential_energy() - final_structure = atoms_to_atomsdata(atoms=atoms) - - dipole = [None, None, None] - if driver == "dipole": - # Catch exception if calculator doesn't have get_dipole_moment() - try: - dipole = [round(x, 4) for x in atoms.get_dipole_moment()] - except Exception: - pass - - end_time = time.time() - wall_time = end_time - start_time - simulation_output = { - "input_structure_file": input_structure_file, - "converged": True, - "final_structure": final_structure.model_dump(), - "simulation_input": params.model_dump(), - "success": True, - "dipole_value": dipole, - "single_point_energy": energy, - "energy_unit": "eV", - "wall_time": wall_time, - } - with open(output_results_file, "w", encoding="utf-8") as wf: - json.dump(simulation_output, wf, indent=4, default=str) - - if driver == "energy": - return { - "status": "success", - "message": f"Simulation completed. Results saved to {os.path.abspath(output_results_file)}", - "single_point_energy": energy, - "unit": "eV", - } - elif driver == "dipole": - return { - "status": "success", - "message": f"Simulation completed. Results saved to {os.path.abspath(output_results_file)}", - "dipole_moment": dipole, - "dipole_unit": "e * angstrom", - } - - OPTIMIZERS = { - "bfgs": BFGS, - "lbfgs": LBFGS, - "gpmin": GPMin, - "fire": FIRE, - "mdmin": MDMin, - } - try: - optimizer_class = OPTIMIZERS.get(optimizer.lower()) - if optimizer_class is None: - raise ValueError(f"Unsupported optimizer: {optimizer_class}") - - # Do optimization only if number of atoms > 1 to avoid error. - if len(atoms) > 1: - dyn = optimizer_class(atoms) - converged = dyn.run(fmax=fmax, steps=steps) - else: - converged = True - - single_point_energy = float(atoms.get_potential_energy()) - final_structure = AtomsData( - numbers=atoms.numbers, - positions=atoms.positions, - cell=atoms.cell, - pbc=atoms.pbc, - ) - thermo_data = {} - vib_data = {} - ir_data = {} - - if driver in {"vib", "thermo", "ir"}: - from ase.vibrations import Vibrations - from ase import units - import tempfile - import shutil - import glob - - ir_plot_path = None # Will be set inside tmpdir block if driver == "ir" - # Use a temporary directory to isolate parallel vibration runs. - # ASE's Vibrations class writes cache files (vib/cache.*.json) and - # trajectory files (vib.*.traj) using the `name` parameter. Without - # isolation, parallel calls for different molecules write to the same - # files, causing shape-mismatch errors and corrupted thermochemistry. - mol_stem = ( - Path(input_structure_file).stem if input_structure_file else "mol" - ) - - with tempfile.TemporaryDirectory( - prefix=f"chemgraph_vib_{mol_stem}_" - ) as tmpdir: - vib_name = os.path.join(tmpdir, "vib") - vib = Vibrations(atoms, name=vib_name) - - vib.clean() - vib.run() - - vib_data = { - "energies": [], - "energy_unit": "meV", - "frequencies": [], - "frequency_unit": "cm-1", - } - - energies = vib.get_energies() - linear = is_linear_molecule.invoke({"atomsdata": final_structure}) - - for idx, e in enumerate(energies): - is_imag = abs(e.imag) > 1e-8 - e_val = e.imag if is_imag else e.real - energy_meV = 1e3 * e_val - freq_cm1 = e_val / units.invcm - suffix = "i" if is_imag else "" - vib_data["energies"].append(f"{energy_meV}{suffix}") - vib_data["frequencies"].append(f"{freq_cm1}{suffix}") - - # Write frequencies.csv to the resolved output directory - freq_file_path = _resolve_path(f"frequencies_{mol_stem}.csv") - freq_file = Path(freq_file_path) - if freq_file.exists(): - freq_file.unlink() - - with freq_file.open("w") as f: - for i, freq in enumerate(vib_data["frequencies"], start=0): - f.write(f"{mol_stem}_vib.{i}.traj,{freq}\n") - - # Write normal modes .traj files inside tmpdir, then copy out - for i in range(len(energies)): - vib.write_mode(n=i, kT=units.kB * 300, nimages=30) - - # Copy .traj files to the resolved output directory with molecule prefix - traj_dest_dir = _resolve_path("") - if traj_dest_dir: - os.makedirs(traj_dest_dir, exist_ok=True) - for traj_file in glob.glob(os.path.join(tmpdir, "vib.*.traj")): - dest_name = f"{mol_stem}_{Path(traj_file).name}" - dest_path = ( - os.path.join(traj_dest_dir, dest_name) - if traj_dest_dir - else dest_name - ) - shutil.copy2(traj_file, dest_path) - - if driver == "ir": - from ase.vibrations import Infrared - import matplotlib.pyplot as plt - - ir_data["spectrum_frequencies"] = [] - ir_data["spectrum_frequencies_units"] = "cm-1" - - ir_data["spectrum_intensities"] = [] - ir_data["spectrum_intensities_units"] = "D/Å^2 amu^-1" - - ir_name = os.path.join(tmpdir, "ir") - ir = Infrared(atoms, name=ir_name) - ir.clean() - ir.run() - - IR_SPECTRUM_START = 500 # Start of IR spectrum range - IR_SPECTRUM_END = 4000 # End of IR spectrum range - freq_intensity = ir.get_spectrum( - start=IR_SPECTRUM_START, end=IR_SPECTRUM_END - ) - # Generate IR spectrum plot - fig, ax = plt.subplots() - ax.plot(freq_intensity[0], freq_intensity[1]) - ax.set_xlabel("Frequency (cm⁻¹)") - ax.set_ylabel("Intensity (a.u.)") - ax.set_title("Infrared Spectrum") - ax.grid(True) - ir_plot_path = _resolve_path(f"ir_spectrum_{mol_stem}.png") - fig.savefig(ir_plot_path, format="png", dpi=300) - plt.close(fig) - - ir_data["IR Plot"] = f"Saved to {os.path.abspath(ir_plot_path)}" - ir_data["Normal mode data"] = ( - f"Normal modes saved as individual .traj files with prefix {mol_stem}_" - ) - - if driver == "thermo": - # Approximation for a single atom system. - if len(atoms) == 1: - thermo_data = { - "enthalpy": single_point_energy, - "entropy": 0.0, - "gibbs_free_energy": single_point_energy, - "unit": "eV", - } - else: - from ase.thermochemistry import IdealGasThermo - - linear = is_linear_molecule.invoke( - {"atomsdata": final_structure} - ) - geometry = "linear" if linear else "nonlinear" - symmetrynumber = get_symmetry_number.invoke( - {"atomsdata": final_structure} - ) - - thermo = IdealGasThermo( - vib_energies=energies, - potentialenergy=single_point_energy, - atoms=atoms, - geometry=geometry, - symmetrynumber=symmetrynumber, - spin=0, # Only support spin=0 - ) - thermo_data = { - "enthalpy": float( - thermo.get_enthalpy(temperature=temperature) - ), - "entropy": float( - thermo.get_entropy( - temperature=temperature, pressure=pressure - ) - ), - "gibbs_free_energy": float( - thermo.get_gibbs_energy( - temperature=temperature, pressure=pressure - ) - ), - "unit": "eV", - } - - end_time = time.time() - wall_time = end_time - start_time - simulation_output = { - "input_structure_file": input_structure_file, - "converged": converged, - "final_structure": final_structure.model_dump(), - "simulation_input": params.model_dump(), - "vibrational_frequencies": vib_data, - "thermochemistry": thermo_data, - "success": True, - "ir_data": ir_data, - "single_point_energy": single_point_energy, - "energy_unit": "eV", - "wall_time": wall_time, - } - - with open(output_results_file, "w", encoding="utf-8") as wf: - json.dump(simulation_output, wf, indent=4, default=str) - - # Return message based on driver. Keep the return output minimal. - if driver == "opt": - return { - "status": "success", - "message": f"Simulation completed. Results saved to {os.path.abspath(output_results_file)}", - "single_point_energy": single_point_energy, # small payload for LLMs - "unit": "eV", - } - elif driver == "vib": - return { - "status": "success", - "result": { - "vibrational_frequencies": vib_data, - }, # small payload for LLMs - "message": ( - "Vibrational analysis completed; frequencies returned. " - f"Full results (structure, vibrations and metadata) saved to {os.path.abspath(output_results_file)}." - ), - } - elif driver == "thermo": - return { - "status": "success", - "result": {"thermochemistry": thermo_data}, # small payload for LLMs - "message": ( - "Thermochemistry computed and returned. " - f"Full results (structure, vibrations, thermochemistry and metadata) saved to {os.path.abspath(output_results_file)}" - ), - } - elif driver == "ir": - return { - "status": "success", - "result": { - "vibrational_frequencies": vib_data - }, # small payload for LLMs - "message": ( - "Infrared computed and returned. " - f"Full results (structure, vibrations, thermochemistry and metadata) saved to {os.path.abspath(output_results_file)}. " - f"IR plot saved to {os.path.abspath(ir_plot_path) if ir_plot_path else 'N/A'}. Normal modes saved as individual .traj files" - ), - } - - except Exception as e: - return { - "status": "failure", - "error_type": type(e).__name__, - "message": str(e), - } - - -def create_ase_atoms(atomic_numbers, positions): - """Create an ASE Atoms object from atomic numbers and positions. - - Parameters - ---------- - atomic_numbers : list or array - List of atomic numbers - positions : list or array - List of atomic positions (3D coordinates) - - Returns - ------- - ase.Atoms - ASE Atoms object - """ - from ase import Atoms - - try: - atoms = Atoms(numbers=atomic_numbers, positions=positions) - return atoms - except Exception as e: - print(f"Error creating ASE Atoms object: {e}") - return None - - -def create_xyz_string(atomic_numbers, positions): - """Create an XYZ format string from atomic numbers and positions. - - Parameters - ---------- - atomic_numbers : list or array - List of atomic numbers - positions : list or array - List of atomic positions (3D coordinates) - - Returns - ------- - str - XYZ format string - """ - from ase import Atoms - - try: - atoms = Atoms(numbers=atomic_numbers, positions=positions) - - # Create XYZ string manually - xyz_lines = [str(len(atoms))] - xyz_lines.append("Generated by ChemGraph") - - for i, (symbol, pos) in enumerate( - zip(atoms.get_chemical_symbols(), atoms.positions) - ): - xyz_lines.append( - f"{symbol:2s} {pos[0]:12.6f} {pos[1]:12.6f} {pos[2]:12.6f}" - ) - - return "\n".join(xyz_lines) - except Exception as e: - print(f"Error creating XYZ string: {e}") - return None + return run_ase_core(params) + return run_ase_core(params) diff --git a/src/chemgraph/tools/cheminformatics_core.py b/src/chemgraph/tools/cheminformatics_core.py new file mode 100644 index 00000000..0ffe13b3 --- /dev/null +++ b/src/chemgraph/tools/cheminformatics_core.py @@ -0,0 +1,187 @@ +"""Pure-Python cheminformatics helpers (no LangChain / MCP decorators). + +Provides a single implementation for PubChem lookups and RDKit +SMILES-to-3D conversion, used by both the LangChain ``@tool`` wrappers +in :mod:`cheminformatics_tools` and the MCP wrappers in +:mod:`chemgraph.mcp.mcp_tools`. +""" + +from __future__ import annotations + +import os +from typing import Literal + +import pubchempy as pcp + +from chemgraph.schemas.atomsdata import AtomsData +from chemgraph.tools.ase_core import _resolve_path + + +# --------------------------------------------------------------------------- +# SMILES → 3D coordinates (single implementation) +# --------------------------------------------------------------------------- + + +def smiles_to_3d( + smiles: str, seed: int = 2025 +) -> tuple[list[int], list[list[float]]]: + """Convert a SMILES string to 3D coordinates via RDKit. + + Parameters + ---------- + smiles : str + SMILES string representation of the molecule. + seed : int, optional + Random seed for reproducible 3D embedding, by default 2025. + + Returns + ------- + tuple[list[int], list[list[float]]] + ``(atomic_numbers, positions)`` where *positions* is a list of + ``[x, y, z]`` lists in Angstroms. + + Raises + ------ + ValueError + If the SMILES string is invalid or 3D generation/optimization fails. + """ + from rdkit import Chem + from rdkit.Chem import AllChem + + mol = Chem.MolFromSmiles(smiles) + if mol is None: + raise ValueError("Invalid SMILES string.") + + mol = Chem.AddHs(mol) + if AllChem.EmbedMolecule(mol, randomSeed=seed) != 0: + raise ValueError("Failed to generate 3D coordinates.") + if AllChem.UFFOptimizeMolecule(mol) != 0: + raise ValueError("Failed to optimize 3D geometry.") + + conf = mol.GetConformer() + numbers = [atom.GetAtomicNum() for atom in mol.GetAtoms()] + positions = [list(conf.GetAtomPosition(i)) for i in range(mol.GetNumAtoms())] + return numbers, positions + + +# --------------------------------------------------------------------------- +# PubChem name → SMILES +# --------------------------------------------------------------------------- + + +def molecule_name_to_smiles_core(name: str) -> str: + """Resolve a molecule name to its canonical SMILES via PubChem. + + Parameters + ---------- + name : str + Common or IUPAC molecule name. + + Returns + ------- + str + Canonical SMILES string. + + Raises + ------ + ValueError + If no PubChem match is found or the returned SMILES is empty. + """ + if not name or not str(name).strip(): + raise ValueError("Parameter 'name' must be a non-empty string.") + + comps = pcp.get_compounds(str(name).strip(), "name") + if not comps: + raise ValueError(f"No PubChem compound found for name: {name!r}") + + smiles = comps[0].canonical_smiles + if not smiles: + raise ValueError(f"PubChem returned an empty SMILES for {name!r}.") + return smiles + + +# --------------------------------------------------------------------------- +# SMILES → coordinate file +# --------------------------------------------------------------------------- + + +def smiles_to_coordinate_file_core( + smiles: str, + output_file: str = "molecule.xyz", + seed: int = 2025, + fmt: Literal["xyz"] = "xyz", +) -> dict: + """Convert a SMILES string to a coordinate file on disk. + + Parameters + ---------- + smiles : str + SMILES string representation of the molecule. + output_file : str, optional + Path to save the output coordinate file. + seed : int, optional + Random seed for RDKit 3D structure generation, by default 2025. + fmt : {"xyz"}, optional + Output format. Only ``"xyz"`` is supported currently. + + Returns + ------- + dict + ``{"ok": True, "artifact": "coordinate_file", "path": ..., + "smiles": ..., "natoms": ...}`` + + Raises + ------ + ValueError + If the SMILES string is invalid or 3D generation fails. + """ + from ase import Atoms + from ase.io import write as ase_write + + numbers, positions = smiles_to_3d(smiles, seed=seed) + atoms = Atoms(numbers=numbers, positions=positions) + + final_output_file = _resolve_path(output_file) + ase_write(final_output_file, atoms) + + return { + "ok": True, + "artifact": "coordinate_file", + "path": os.path.abspath(final_output_file), + "smiles": smiles, + "natoms": len(numbers), + } + + +# --------------------------------------------------------------------------- +# SMILES → AtomsData +# --------------------------------------------------------------------------- + + +def smiles_to_atomsdata_core(smiles: str, seed: int = 2025) -> AtomsData: + """Convert a SMILES string to an :class:`~chemgraph.schemas.atomsdata.AtomsData`. + + Parameters + ---------- + smiles : str + SMILES string representation of the molecule. + seed : int, optional + Random seed for RDKit 3D structure generation, by default 2025. + + Returns + ------- + AtomsData + Structure with no periodic boundary conditions. + + Raises + ------ + ValueError + If the SMILES string is invalid or 3D generation fails. + """ + numbers, positions = smiles_to_3d(smiles, seed=seed) + return AtomsData( + numbers=numbers, + positions=positions, + cell=[[0, 0, 0], [0, 0, 0], [0, 0, 0]], + pbc=[False, False, False], + ) diff --git a/src/chemgraph/tools/cheminformatics_tools.py b/src/chemgraph/tools/cheminformatics_tools.py index 32e6c12e..857ec4bc 100644 --- a/src/chemgraph/tools/cheminformatics_tools.py +++ b/src/chemgraph/tools/cheminformatics_tools.py @@ -1,13 +1,21 @@ -import os +"""LangChain ``@tool`` wrappers for cheminformatics functions. + +Each tool delegates to the pure-Python implementation in +:mod:`chemgraph.tools.cheminformatics_core`. +""" + +from __future__ import annotations + from typing import Literal -import pubchempy from langchain_core.tools import tool -from ase.io import write as ase_write -from ase import Atoms from chemgraph.schemas.atomsdata import AtomsData -from chemgraph.tools.mcp_helper import _resolve_path +from chemgraph.tools.cheminformatics_core import ( + molecule_name_to_smiles_core, + smiles_to_atomsdata_core, + smiles_to_coordinate_file_core, +) @tool @@ -23,13 +31,8 @@ def molecule_name_to_smiles(name: str) -> dict: ------- dict A JSON-serializable dict with the resolved SMILES. - - Raises - ------ - IndexError - If the molecule name is not found in PubChem. """ - smiles = pubchempy.get_compounds(str(name), "name")[0].connectivity_smiles + smiles = molecule_name_to_smiles_core(name) return {"name": str(name), "smiles": smiles} @@ -48,39 +51,8 @@ def smiles_to_atomsdata(smiles: str, randomSeed: int = 2025) -> AtomsData: ------- AtomsData AtomsData object containing the molecular structure. - - Raises - ------ - ValueError - If the SMILES string is invalid or if 3D structure generation fails. """ - from rdkit import Chem - from rdkit.Chem import AllChem - - # Generate the molecule object - mol = Chem.MolFromSmiles(smiles) - if mol is None: - raise ValueError("Invalid SMILES string.") - - # Add hydrogens and optimize 3D structure - mol = Chem.AddHs(mol) - if AllChem.EmbedMolecule(mol, randomSeed=randomSeed) != 0: - raise ValueError("Failed to generate 3D coordinates.") - if AllChem.UFFOptimizeMolecule(mol) != 0: - raise ValueError("Failed to optimize 3D geometry.") - # Extract atomic information - conf = mol.GetConformer() - numbers = [atom.GetAtomicNum() for atom in mol.GetAtoms()] - positions = [list(conf.GetAtomPosition(i)) for i in range(mol.GetNumAtoms())] - - # Create AtomsData object - atoms_data = AtomsData( - numbers=numbers, - positions=positions, - cell=[[0, 0, 0], [0, 0, 0], [0, 0, 0]], - pbc=[False, False, False], # No periodic boundary conditions - ) - return atoms_data + return smiles_to_atomsdata_core(smiles, seed=randomSeed) @tool @@ -106,47 +78,8 @@ def smiles_to_coordinate_file( Returns ------- str - A single-line JSON string LLMs can parse, e.g. - {"ok": true, "artifact": "coordinate_file", "format": "xyz", "path": "...", "smiles": "...", "natoms": 12} - - Raises - ------ - ValueError - If the SMILES string is invalid or if 3D structure generation fails. + A single-line JSON string LLMs can parse. """ - from rdkit import Chem - from rdkit.Chem import AllChem - - # Generate the molecule object - mol = Chem.MolFromSmiles(smiles) - if mol is None: - raise ValueError("Invalid SMILES string.") - - # Add hydrogens and optimize 3D structure - mol = Chem.AddHs(mol) - if AllChem.EmbedMolecule(mol, randomSeed=randomSeed) != 0: - raise ValueError("Failed to generate 3D coordinates.") - if AllChem.UFFOptimizeMolecule(mol) != 0: - raise ValueError("Failed to optimize 3D geometry.") - # Extract atomic information - conf = mol.GetConformer() - numbers = [atom.GetAtomicNum() for atom in mol.GetAtoms()] - positions = [list(conf.GetAtomPosition(i)) for i in range(mol.GetNumAtoms())] - - # Create Atoms object - atoms = Atoms(numbers=numbers, positions=positions) - - final_output_file = _resolve_path(output_file) - ase_write( - final_output_file, - atoms, + return smiles_to_coordinate_file_core( + smiles, output_file=output_file, seed=randomSeed, fmt=fmt ) - - # Return dict for LLM/tool chaining - return { - "ok": True, - "artifact": "coordinate_file", - "path": os.path.abspath(final_output_file), - "smiles": smiles, - "natoms": len(numbers), - } diff --git a/src/chemgraph/tools/generic_tools.py b/src/chemgraph/tools/generic_tools.py index fa20feb4..ec93ef5e 100644 --- a/src/chemgraph/tools/generic_tools.py +++ b/src/chemgraph/tools/generic_tools.py @@ -1,9 +1,12 @@ +import io import math import numexpr +import traceback +from contextlib import redirect_stderr, redirect_stdout from langchain_core.tools import Tool from langchain_core.tools import tool -from langchain_experimental.utilities import PythonREPL +from langgraph.types import interrupt @tool @@ -61,6 +64,70 @@ def calculator(expression: str) -> str: return f"Error evaluating expression: {e!s}" +@tool +def ask_human(question: str) -> str: + """Ask the human user for clarification, confirmation, or additional details. + + Use this tool when: + - Required inputs are missing or ambiguous (e.g., molecule name, calculator + type, temperature, pressure, or simulation method). + - You need confirmation before running a computationally expensive simulation + (e.g., geometry optimization, vibrational analysis, thermochemistry). + - A previous tool call failed and you need the user to decide how to proceed + (e.g., retry with different parameters, skip the step, or abort). + + The graph execution will pause until the human responds. The human's + answer is returned as a string. + + Parameters + ---------- + question : str + The question or request to present to the human user. + + Returns + ------- + str + The human's response. + """ + response = interrupt({"question": question}) + if isinstance(response, dict): + return response.get("answer", response.get("response", str(response))) + return str(response) + + +class PythonREPL: + """Small persistent Python REPL used by the python_repl tool.""" + + def __init__(self): + """Initialize an empty persistent global namespace.""" + self.globals = {} + + def run(self, command: str) -> str: + """Execute Python code in the persistent REPL namespace. + + Parameters + ---------- + command : str + Python code to execute. + + Returns + ------- + str + Captured stdout/stderr and traceback text, if any. + """ + cleaned_command = command.strip() + if not cleaned_command: + return "" + + output = io.StringIO() + try: + with redirect_stdout(output), redirect_stderr(output): + exec(cleaned_command, self.globals, self.globals) + except Exception: + return output.getvalue() + traceback.format_exc() + return output.getvalue() + + python_repl = PythonREPL() repl_tool = Tool( name="python_repl", diff --git a/src/chemgraph/tools/graspa_core.py b/src/chemgraph/tools/graspa_core.py new file mode 100644 index 00000000..0294b75f --- /dev/null +++ b/src/chemgraph/tools/graspa_core.py @@ -0,0 +1,349 @@ +"""Pure-Python gRASPA simulation helpers (no LangChain / MCP decorators). + +Contains the core workflow functions for running gRASPA-SYCL +simulations, parsing output, and mock simulations for testing. +Used by the LangChain ``@tool`` wrapper in :mod:`graspa_tools` and the +MCP/Parsl wrappers in :mod:`chemgraph.mcp.graspa_mcp_parsl`. +""" + +from __future__ import annotations + +import glob +import os +import random +import shutil +import subprocess +import time +from pathlib import Path + +import ase +import numpy as np +from ase.io import read as ase_read + +from chemgraph.schemas.graspa_schema import graspa_input_schema + +# Template directory for gRASPA-SYCL input files +_file_dir = Path(__file__).parent / "files" / "template_graspa_sycl" + +# gRASPA-SYCL command +graspa_cmd = ( + "export OMP_NUM_THREADS=1; " + "export ZE_FLAT_DEVICE_HIERARCHY=FLAT; " + "/lus/flare/projects/IQC/thang/soft/gRASPA/graspa-sycl/bin/sycl.out" +) + + +# --------------------------------------------------------------------------- +# Output parsing +# --------------------------------------------------------------------------- + + +def _read_graspa_sycl_output( + output_path: str, + adsorbate: str = "H2O", + cifname: str = None, + output_fname: str = "raspa.log", + temperature: float = None, + pressure: float = None, +) -> dict: + """Parse gRASPA output and return uptake results. + + Parameters + ---------- + output_path : str + Directory containing the gRASPA output files. + adsorbate : str + Name of the adsorbate molecule. + cifname : str, optional + Stem name of the CIF file (without extension). + output_fname : str + Name of the gRASPA log file. + temperature : float, optional + Simulation temperature in Kelvin. + pressure : float, optional + Simulation pressure in Pascal. + + Returns + ------- + dict + Parsed results including uptake, status, and CIF path. + """ + result = { + "status": "failure", + "uptake_in_mol_kg": 0, + "adsorbate": adsorbate, + "temperature_in_K": None, + "pressure_in_Pa": None, + "cif_path": None, + } + + target_file = Path(output_path) / Path(output_fname).name + + # --- Resolve CIF Path --- + if cifname is None: + cif_list = glob.glob(os.path.join(output_path, "*.cif")) + if len(cif_list) != 1: + cifpath = None + else: + cifpath = os.path.abspath(cif_list[0]) + else: + cifpath = os.path.abspath(os.path.join(output_path, f"{cifname}.cif")) + + result["cif_path"] = cifpath + + # --- Check Log Existence --- + if not os.path.exists(target_file): + return result + + # --- Parse Log --- + unitcell_line = None + uptake_line = None + + with open(target_file, "r") as rf: + for line in rf: + if "UnitCells" in line: + unitcell_line = line.strip() + elif "Overall: Average:" in line: + uptake_line = line.strip() + + if unitcell_line is None or uptake_line is None: + return result + + try: + if cifpath is None: + raise ValueError(f"Could not resolve CIF path in {output_path}") + + uptake_total_molecule = float(uptake_line.split()[2][:-1]) + unitcell = unitcell_line.split()[4:] + unitcell = [int(float(i)) for i in unitcell] + + atoms = ase_read(cifpath) + framework_mass = ( + sum(atoms.get_masses()) * unitcell[0] * unitcell[1] * unitcell[2] + ) + + uptake_mol_kg = round((uptake_total_molecule / framework_mass) * 1000, 2) + result["uptake_in_mol_kg"] = float(uptake_mol_kg) + result["status"] = "success" + result["temperature_in_K"] = temperature + result["pressure_in_Pa"] = pressure + except Exception as e: + print(f"Error parsing results in {output_path}: {e}") + result["status"] = "failure" + + return result + + +# --------------------------------------------------------------------------- +# Mock simulation (for testing) +# --------------------------------------------------------------------------- + + +def mock_graspa(params: graspa_input_schema) -> dict: + """Return mock gRASPA results for testing without the SYCL runtime. + + Parameters + ---------- + params : graspa_input_schema + Input parameters (only ``adsorbates`` is used to determine output shape). + + Returns + ------- + dict + Simulated uptake results. + """ + + def rand_uptake( + low: float, high: float, ndigits: int = 3, min_positive: float | None = None + ) -> float: + """Generate a rounded mock uptake value. + + Parameters + ---------- + low : float + Lower bound for the random value. + high : float + Upper bound for the random value. + ndigits : int, optional + Number of decimal places to round to. + min_positive : float, optional + Replacement value when rounding produces zero. + + Returns + ------- + float + Mock uptake value. + """ + value = random.uniform(low, high) + value = round(value, ndigits) + if min_positive is not None and value == 0.0: + value = min_positive + return value + + time.sleep(random.uniform(20, 40)) + n_ads = len(params.adsorbates) + + if n_ads == 1: + uptake_co2 = rand_uptake(0, 2, ndigits=3) + return {"co2_uptake_mol_per_kg": uptake_co2} + + elif n_ads == 2: + uptake_co2 = rand_uptake(0, 2, ndigits=3) + uptake_n2 = rand_uptake(0, 0.5, ndigits=3, min_positive=1e-3) + try: + selectivity = uptake_co2 / uptake_n2 + except Exception: + selectivity = 1e4 + return { + "co2_uptake_mol_per_kg": uptake_co2, + "n2_uptake_mol_per_kg": uptake_n2, + "co2_n2_selectivity": round(selectivity, 2), + } + + elif n_ads == 3: + uptake_co2 = rand_uptake(0, 2, ndigits=3) + uptake_n2 = rand_uptake(0, 0.5, ndigits=3, min_positive=1e-3) + uptake_h2o = rand_uptake(0, 5, ndigits=3) + try: + selectivity = uptake_co2 / uptake_n2 + except Exception: + selectivity = 1e4 + return { + "co2_uptake_mol_per_kg": uptake_co2, + "n2_uptake_mol_per_kg": uptake_n2, + "h2o_uptake_mol_per_kg": uptake_h2o, + "co2_n2_selectivity": round(selectivity, 2), + } + + else: + raise ValueError("Only supports 1-3 adsorbates only.") + + +# --------------------------------------------------------------------------- +# Core simulation runner +# --------------------------------------------------------------------------- + + +def run_graspa_core(params: graspa_input_schema) -> dict: + """Run a single gRASPA calculation using specified input parameters. + + Parameters + ---------- + params : graspa_input_schema + Input parameters for the gRASPA calculation. + + Returns + ------- + dict + Parsed simulation results including uptake and status. + """ + + def _calculate_cell_size( + atoms: ase.Atoms, cutoff: float = 12.8 + ) -> list[int]: + """Calculate unit-cell replication for GCMC with the given cutoff. + + Parameters + ---------- + atoms : ase.Atoms + Unit-cell structure. + cutoff : float, optional + Minimum replicated cell length in angstrom. + + Returns + ------- + list[int] + Replication factors along the three lattice vectors. + """ + unit_cell = atoms.cell[:] + a = unit_cell[0] + b = unit_cell[1] + c = unit_cell[2] + + wa = np.divide( + np.linalg.norm(np.dot(np.cross(b, c), a)), + np.linalg.norm(np.cross(b, c)), + ) + wb = np.divide( + np.linalg.norm(np.dot(np.cross(c, a), b)), + np.linalg.norm(np.cross(c, a)), + ) + wc = np.divide( + np.linalg.norm(np.dot(np.cross(a, b), c)), + np.linalg.norm(np.cross(a, b)), + ) + + uc_x = int(np.ceil(cutoff / (0.5 * wa))) + uc_y = int(np.ceil(cutoff / (0.5 * wb))) + uc_z = int(np.ceil(cutoff / (0.5 * wc))) + + return [uc_x, uc_y, uc_z] + + cif_path = Path(params.input_structure_file).resolve() + if not cif_path.exists(): + raise FileNotFoundError(f"CIF file does not exist: {cif_path}") + + base_dir = cif_path.parent + + cifname = cif_path.stem + temperature = params.temperature + pressure = params.pressure + adsorbate = params.adsorbate + n_cycle = params.n_cycles + + folder_name = f"{cifname}--{adsorbate}-{temperature}-{pressure:g}" + sim_dir = base_dir / folder_name + sim_dir.mkdir(parents=True, exist_ok=True) + + for item in _file_dir.iterdir(): + dest = sim_dir / item.name + if item.is_dir(): + if dest.exists(): + shutil.rmtree(dest) + shutil.copytree(item, dest) + else: + shutil.copy2(item, sim_dir) + + # Copy the specific CIF file + shutil.copy2(cif_path, sim_dir / f"{cifname}.cif") + + atoms = ase_read(cif_path) + [uc_x, uc_y, uc_z] = _calculate_cell_size(atoms) + + input_file = sim_dir / "simulation.input" + temp_file = sim_dir / "simulation.input.tmp" + + with open(input_file, "r") as f_in, open(temp_file, "w") as f_out: + for line in f_in: + if "NCYCLE" in line: + line = line.replace("NCYCLE", str(n_cycle)) + if "ADSORBATE" in line: + line = line.replace("ADSORBATE", adsorbate) + if "TEMPERATURE" in line: + line = line.replace("TEMPERATURE", str(temperature)) + if "PRESSURE" in line: + line = line.replace("PRESSURE", str(pressure)) + if "UC_X UC_Y UC_Z" in line: + line = line.replace("UC_X UC_Y UC_Z", f"{uc_x} {uc_y} {uc_z}") + if "CUTOFF" in line: + line = line.replace("CUTOFF", str(12.8)) + if "CIFFILE" in line: + line = line.replace("CIFFILE", cifname) + f_out.write(line) + + shutil.move(temp_file, input_file) + output_filename = Path(params.output_result_file).name + with ( + open(os.path.join(sim_dir, output_filename), "w") as fp, + open(os.path.join(sim_dir, "raspa.err"), "w") as fe, + ): + subprocess.run(graspa_cmd, cwd=sim_dir, stdout=fp, stderr=fe, shell=True) + + return _read_graspa_sycl_output( + output_path=str(sim_dir), + adsorbate=adsorbate, + cifname=cifname, + output_fname=params.output_result_file, + temperature=temperature, + pressure=pressure, + ) diff --git a/src/chemgraph/tools/graspa_tools.py b/src/chemgraph/tools/graspa_tools.py index c0a41b1c..332c33e1 100644 --- a/src/chemgraph/tools/graspa_tools.py +++ b/src/chemgraph/tools/graspa_tools.py @@ -1,306 +1,57 @@ -import subprocess -import os -from pathlib import Path -import shutil -import random -import time -import numpy as np -import glob +"""LangChain ``@tool`` wrapper for gRASPA simulations. -import ase -from ase.io import read as ase_read +Delegates to the pure-Python implementation in +:mod:`chemgraph.tools.graspa_core`. +""" + +from __future__ import annotations from langchain_core.tools import tool -from chemgraph.schemas.graspa_schema import ( - graspa_input_schema, +from chemgraph.schemas.graspa_schema import graspa_input_schema +from chemgraph.tools.graspa_core import ( + # Re-export core helpers so existing ``from graspa_tools import ...`` + # statements (e.g. in graspa_mcp_parsl.py) continue to work. + _read_graspa_sycl_output, + mock_graspa, + run_graspa_core, ) -# LangGraph gRASPA tools. -_file_dir = Path(__file__).parent / "files" / "template_graspa_sycl" - -# gRASPA-SYCL command -graspa_cmd = "export OMP_NUM_THREADS=1; export ZE_FLAT_DEVICE_HIERARCHY=FLAT; /lus/flare/projects/IQC/thang/soft/gRASPA/graspa-sycl/bin/sycl.out" - - -def _read_graspa_sycl_output( - output_path: str, - adsorbate: str = "H2O", - cifname: str = None, - output_fname: str = "raspa.log", - temperature: float = None, - pressure: float = None, -): - """ - Parses gRASPA output and includes the full path to the CIF file used. - """ - result = { - "status": "failure", - "uptake_in_mol_kg": 0, - "adsorbate": adsorbate, - "temperature_in_K": None, - "pressure_in_Pa": None, - "cif_path": None, - } - - target_file = Path(output_path) / Path(output_fname).name - - # --- Resolve CIF Path --- - # We resolve this early so we can return it even if the log parsing fails - if cifname is None: - cif_list = glob.glob(os.path.join(output_path, "*.cif")) - if len(cif_list) != 1: - # If explicit name not provided and we can't auto-detect unique CIF, we can't resolve path - cifpath = None - else: - cifpath = os.path.abspath(cif_list[0]) - else: - # Construct absolute path based on output_path - cifpath = os.path.abspath(os.path.join(output_path, f"{cifname}.cif")) - - result["cif_path"] = cifpath - - # --- Check Log Existence --- - if not os.path.exists(target_file): - return result - - # --- Parse Log --- - unitcell_line = None - uptake_line = None - - with open(target_file, "r") as rf: - for line in rf: - if "UnitCells" in line: - unitcell_line = line.strip() - elif "Overall: Average:" in line: - uptake_line = line.strip() - - if unitcell_line is None or uptake_line is None: - return result - - try: - if cifpath is None: - raise ValueError(f"Could not resolve CIF path in {output_path}") - - uptake_total_molecule = float(uptake_line.split()[2][:-1]) - # Parse UnitCells (robust to whitespace) - unitcell = unitcell_line.split()[4:] - unitcell = [int(float(i)) for i in unitcell] - - atoms = ase_read(cifpath) - framework_mass = ( - sum(atoms.get_masses()) * unitcell[0] * unitcell[1] * unitcell[2] - ) - - uptake_mol_kg = round((uptake_total_molecule / framework_mass) * 1000, 2) - result["uptake_in_mol_kg"] = float(uptake_mol_kg) - # result["error_in_mol_kg"] = float(error_mol_kg) - result["status"] = "success" - result["temperature_in_K"] = temperature - result["pressure_in_Pa"] = pressure - except Exception as e: - print(f"Error parsing results in {output_path}: {e}") - result["status"] = "failure" - - return result - - -def mock_graspa(params: graspa_input_schema) -> dict: - def rand_uptake( - low: float, high: float, ndigits: int = 3, min_positive: float | None = None - ) -> float: - """Random uptake with rounding and optional minimum positive value.""" - value = random.uniform(low, high) - value = round(value, ndigits) - if min_positive is not None and value == 0.0: - value = min_positive - return value - - time.sleep(random.uniform(20, 40)) - n_ads = len(params.adsorbates) - - if n_ads == 1: - uptake_co2 = rand_uptake(0, 2, ndigits=3) - return { - "co2_uptake_mol_per_kg": uptake_co2, - } - - elif n_ads == 2: - uptake_co2 = rand_uptake(0, 2, ndigits=3) - # prevent rounded value from becoming exactly zero - uptake_n2 = rand_uptake(0, 0.5, ndigits=3, min_positive=1e-3) - - try: - selectivity = uptake_co2 / uptake_n2 - except Exception: - selectivity = 1e4 - - return { - "co2_uptake_mol_per_kg": uptake_co2, - "n2_uptake_mol_per_kg": uptake_n2, - "co2_n2_selectivity": round(selectivity, 2), - } - - elif n_ads == 3: - uptake_co2 = rand_uptake(0, 2, ndigits=3) - uptake_n2 = rand_uptake(0, 0.5, ndigits=3, min_positive=1e-3) - uptake_h2o = rand_uptake(0, 5, ndigits=3) - - try: - selectivity = uptake_co2 / uptake_n2 - except Exception: - selectivity = 1e4 +__all__ = [ + "_read_graspa_sycl_output", + "mock_graspa", + "run_graspa_core", + "run_graspa", +] - return { - "co2_uptake_mol_per_kg": uptake_co2, - "n2_uptake_mol_per_kg": uptake_n2, - "h2o_uptake_mol_per_kg": uptake_h2o, - "co2_n2_selectivity": round(selectivity, 2), - } - - else: - raise ValueError("Only supports 1–3 adsorbates only.") +@tool +def run_graspa(graspa_input: graspa_input_schema): + """Run a gRASPA simulation using the core engine and return the uptakes. -def run_graspa_core(params: graspa_input_schema): - """Run a single gRASPA calculations using specified input parameters. + This tool acts as a wrapper for the agentic workflow. Parameters ---------- - params : graspa_input_schema - Input parameters for the gRASPA calculation - """ - - def _calculate_cell_size( - atoms: ase.Atoms, cutoff: float = 12.8 - ) -> list[int, int, int]: - """Method to calculate Unitcells (for periodic boundary condition) for GCMC - - Args: - atoms (ase.Atoms): ASE atom object - cutoff (float, optional): Cutoff in Angstrom. Defaults to 12.8. - - Returns: - list[int, int, int]: Unit cell in x, y and z - """ - unit_cell = atoms.cell[:] - # Unit cell vectors - a = unit_cell[0] - b = unit_cell[1] - c = unit_cell[2] - # minimum distances between unit cell faces - wa = np.divide( - np.linalg.norm(np.dot(np.cross(b, c), a)), - np.linalg.norm(np.cross(b, c)), - ) - wb = np.divide( - np.linalg.norm(np.dot(np.cross(c, a), b)), - np.linalg.norm(np.cross(c, a)), - ) - wc = np.divide( - np.linalg.norm(np.dot(np.cross(a, b), c)), - np.linalg.norm(np.cross(a, b)), - ) - - uc_x = int(np.ceil(cutoff / (0.5 * wa))) - uc_y = int(np.ceil(cutoff / (0.5 * wb))) - uc_z = int(np.ceil(cutoff / (0.5 * wc))) - - return [uc_x, uc_y, uc_z] - - cif_path = Path(params.input_structure_file).resolve() - if not cif_path.exists(): - raise FileNotFoundError(f"CIF file does not exist: {cif_path}") - - base_dir = cif_path.parent - - cifname = cif_path.stem - temperature = params.temperature - pressure = params.pressure - adsorbate = params.adsorbate - n_cycle = params.n_cycles + graspa_input : graspa_input_schema + Legacy gRASPA tool input. - folder_name = f"{cifname}--{adsorbate}-{temperature}-{pressure:g}" - sim_dir = base_dir / folder_name - sim_dir.mkdir(parents=True, exist_ok=True) - - for item in _file_dir.iterdir(): - dest = sim_dir / item.name - if item.is_dir(): - if dest.exists(): - shutil.rmtree(dest) - shutil.copytree(item, dest) - else: - shutil.copy2(item, sim_dir) - - # Copy the specific CIF file - shutil.copy2(cif_path, sim_dir / f"{cifname}.cif") - - atoms = ase_read(cif_path) - [uc_x, uc_y, uc_z] = _calculate_cell_size(atoms) - - input_file = sim_dir / "simulation.input" - temp_file = sim_dir / "simulation.input.tmp" - - with open(input_file, "r") as f_in, open(temp_file, "w") as f_out: - for line in f_in: - if "NCYCLE" in line: - line = line.replace("NCYCLE", str(n_cycle)) - if "ADSORBATE" in line: - line = line.replace("ADSORBATE", adsorbate) - if "TEMPERATURE" in line: - line = line.replace("TEMPERATURE", str(temperature)) - if "PRESSURE" in line: - line = line.replace("PRESSURE", str(pressure)) - if "UC_X UC_Y UC_Z" in line: - line = line.replace("UC_X UC_Y UC_Z", f"{uc_x} {uc_y} {uc_z}") - if "CUTOFF" in line: - line = line.replace( - "CUTOFF", str(12.8) - ) # Default or params.cutoff if added - if "CIFFILE" in line: - line = line.replace("CIFFILE", cifname) - f_out.write(line) - - shutil.move(temp_file, input_file) - output_filename = Path(params.output_result_file).name - with ( - open(os.path.join(sim_dir, output_filename), "w") as fp, - open(os.path.join(sim_dir, "raspa.err"), "w") as fe, - ): - subprocess.run(graspa_cmd, cwd=sim_dir, stdout=fp, stderr=fe, shell=True) - - return _read_graspa_sycl_output( - output_path=str(sim_dir), - adsorbate=adsorbate, - cifname=cifname, - output_fname=params.output_result_file, - temperature=temperature, - pressure=pressure, - ) - - -@tool -def run_graspa(graspa_input: graspa_input_schema): - """ - Run a gRASPA simulation using the core engine and return the uptakes. - This tool acts as a wrapper for the agentic workflow. + Returns + ------- + float + Uptake in mol/kg from the core gRASPA result. """ - # Map GRASPAInputSchema fields to the internal schema expected by run_graspa_core params = graspa_input_schema( input_structure_file=graspa_input.cif_path, adsorbate=graspa_input.adsorbate, temperature=graspa_input.temperature, pressure=graspa_input.pressure, n_cycles=graspa_input.n_cycle, - output_result_file="raspa.log" + output_result_file="raspa.log", ) - # Execute core logic result = run_graspa_core(params) - # Return the parsed metrics - # Note: run_graspa_core returns a dict; we extract what the tool usually expects if result["status"] == "success": return result["uptake_in_mol_kg"] else: diff --git a/src/chemgraph/tools/mcp_helper.py b/src/chemgraph/tools/mcp_helper.py deleted file mode 100644 index b858bc26..00000000 --- a/src/chemgraph/tools/mcp_helper.py +++ /dev/null @@ -1,158 +0,0 @@ -import os -from chemgraph.schemas.atomsdata import AtomsData - - -def _resolve_path(path: str) -> str: - """If CHEMGRAPH_LOG_DIR is set and path is relative, prepend it.""" - log_dir = os.environ.get("CHEMGRAPH_LOG_DIR") - if log_dir and not os.path.isabs(path): - # Create directory if it doesn't exist (race condition safe-ish) - os.makedirs(log_dir, exist_ok=True) - return os.path.join(log_dir, path) - return path - - -def load_calculator(calculator: dict) -> tuple[object, dict, dict]: - """Load an ASE calculator based on the provided configuration. - - Parameters - ---------- - calculator : dict - Dictionary containing calculator configuration parameters - - Returns - ------- - object - ASE calculator instance - - Raises - ------ - ValueError - If the calculator type is not supported - """ - calc_type = calculator["calculator_type"].lower() - - if "emt" in calc_type: - from chemgraph.schemas.calculators.emt_calc import EMTCalc - - calc = EMTCalc(**calculator) - elif "tblite" in calc_type: - from chemgraph.schemas.calculators.tblite_calc import TBLiteCalc - - calc = TBLiteCalc(**calculator) - elif "orca" in calc_type: - from chemgraph.schemas.calculators.orca_calc import OrcaCalc - - calc = OrcaCalc(**calculator) - - elif "nwchem" in calc_type: - from chemgraph.schemas.calculators.nwchem_calc import NWChemCalc - - calc = NWChemCalc(**calculator) - - elif "fairchem" in calc_type: - from chemgraph.schemas.calculators.fairchem_calc import FAIRChemCalc - - calc = FAIRChemCalc(**calculator) - - elif "mace" in calc_type: - from chemgraph.schemas.calculators.mace_calc import MaceCalc - - calc = MaceCalc(**calculator) - - elif "aimnet2" in calc_type: - from chemgraph.schemas.calculators.aimnet2_calc import AIMNET2Calc - - calc = AIMNET2Calc(**calculator) - - else: - raise ValueError( - f"Unsupported calculator: {calculator}. Available calculators are EMT, TBLite (GFN2-xTB, GFN1-xTB), Orca and FAIRChem or MACE or AIMNET2." - ) - # Extract additional args like spin/charge if the model defines it - extra_info = {} - if hasattr(calc, "get_atoms_properties"): - extra_info = calc.get_atoms_properties() - - return calc.get_calculator(), extra_info, calc - - -def atoms_to_atomsdata(atoms): - """Convert ASE Atoms object to AtomsData. - - Parameters - ---------- - atoms : ase.Atoms - ASE Atoms object - - Returns - ------- - AtomsData - ChemGraph AtomsData object - """ - return AtomsData( - numbers=atoms.numbers.tolist(), - positions=atoms.positions.tolist(), - cell=atoms.cell.tolist(), - pbc=atoms.pbc.tolist(), - ) - - -def is_linear_molecule(atomsdata: AtomsData, tol=1e-3) -> bool: - """Determine if a molecule is linear or not. - - Parameters - ---------- - atomsdata : AtomsData - AtomsData object containing the molecular structure - tol : float, optional - Tolerance to check for linear molecule, by default 1e-3 - - Returns - ------- - bool - True if the molecule is linear, False otherwise - """ - import numpy as np - - coords = np.array(atomsdata.positions) - # Center the coordinates. - centered = coords - np.mean(coords, axis=0) - # Singular value decomposition. - _, s, _ = np.linalg.svd(centered) - # For a linear molecule, only one singular value is significantly nonzero. - if s[0] == 0: - return False # degenerate case (all atoms at one point) - return (s[1] / s[0]) < tol - - -def get_symmetry_number(atomsdata: AtomsData) -> int: - """Get the rotational symmetry number of a molecule using Pymatgen. - - Parameters - ---------- - atomsdata : AtomsData - AtomsData object containing the molecular structure - - Returns - ------- - int - Rotational symmetry number of the molecule - """ - from pymatgen.symmetry.analyzer import PointGroupAnalyzer - from ase import Atoms - from pymatgen.io.ase import AseAtomsAdaptor - - atoms = Atoms( - numbers=atomsdata.numbers, - positions=atomsdata.positions, - cell=atomsdata.cell, - pbc=atomsdata.pbc, - ) - - aaa = AseAtomsAdaptor() - molecule = aaa.get_molecule(atoms) - pga = PointGroupAnalyzer(molecule) - symmetrynumber = pga.get_rotational_symmetry_number() - - return symmetrynumber diff --git a/src/chemgraph/tools/parsl_tools.py b/src/chemgraph/tools/parsl_tools.py index db0eb9b0..908ac29c 100644 --- a/src/chemgraph/tools/parsl_tools.py +++ b/src/chemgraph/tools/parsl_tools.py @@ -1,406 +1,79 @@ -from pathlib import Path -import time -import glob -import os +"""Parsl-oriented MACE helpers. -from pydantic import BaseModel, Field +``run_mace_core`` converts the MACE-specific input schema to +:class:`ASEInputSchema` and delegates to :func:`chemgraph.tools.ase_core.run_ase_core`. +""" +from __future__ import annotations -from chemgraph.tools.mcp_helper import ( - get_symmetry_number, - is_linear_molecule, - atoms_to_atomsdata, +from chemgraph.tools.ase_core import run_ase_core +from chemgraph.schemas.ase_input import ASEInputSchema +from chemgraph.schemas.mace_parsl_schema import ( + mace_input_schema, + mace_input_schema_ensemble, + mace_output_schema, ) +# Re-export schemas so existing ``from chemgraph.tools.parsl_tools import …`` +# statements continue to work. +__all__ = [ + "mace_input_schema", + "mace_input_schema_ensemble", + "mace_output_schema", + "run_mace_core", +] -class mace_input_schema(BaseModel): - input_structure_file: str = Field( - description="Path to the input coordinate file (e.g., CIF, XYZ, POSCAR) containing the atomic structure for the simulation." - ) - output_result_file: str = Field( - default="output.json", - description="Path to a JSON file where simulation results will be saved.", - ) - driver: str = Field( - default=None, - description="Specifies the type of simulation to run. Options: 'energy' for single-point energy calculations, 'opt' for geometry optimization, 'vib' for vibrational frequency analysis, and 'thermo' for thermochemical properties (including enthalpy, entropy, and Gibbs free energy).", - ) - model: str = Field( - default="medium-mpa-0", - description="Path to the model. Default is medium-mpa-0." - "Options are 'small', 'medium', 'large', 'small-0b', 'medium-0b', 'small-0b2', 'medium-0b2','large-0b2', 'medium-0b3', 'medium-mpa-0', 'medium-omat-0', 'mace-matpes-pbe-0', 'mace-matpes-r2scan-0'", - ) - device: str = Field( - default="cpu", - description="Device to run the MACE calculation on. Options are cpu, cuda or xpu.", - ) - temperature: float = Field( - default=298.15, - description="Temperature for thermo property calculations in Kelvin (K).", - ) - pressure: float = Field( - default=101325.0, - description="Pressure for thermo property calculations in Pascal (Pa).", - ) - fmax: float = Field( - default=0.01, - description="The convergence criterion for forces (in eV/Å). Optimization stops when all force components are smaller than this value.", - ) - steps: int = Field( - default=1000, - description="Maximum number of optimization steps. The optimization will terminate if this number is reached, even if forces haven't converged to fmax.", - ) - optimizer: str = Field( - default="lbfgs", - description="The optimization algorithm used for geometry optimization. Options are 'bfgs', 'lbfgs', 'gpmin', 'fire', 'mdmin'", - ) +# --------------------------------------------------------------------------- +# Core execution — delegates to the unified implementation +# --------------------------------------------------------------------------- -class mace_input_schema_ensemble(BaseModel): - input_structure_directory: str = Field( - description="Path to a folder of input structures containing the atomic structure for the simulations." - ) - output_result_file: str = Field( - default="output.json", - description="Path to a JSON file where simulation results will be saved.", - ) - driver: str = Field( - default=None, - description="Specifies the type of simulation to run. Options: 'energy' for single-point energy calculations, 'opt' for geometry optimization, 'vib' for vibrational frequency analysis, and 'thermo' for thermochemical properties (including enthalpy, entropy, and Gibbs free energy).", - ) - model: str = Field( - default="medium-mpa-0", - description="Path to the model. Default is medium-mpa-0." - "Options are 'small', 'medium', 'large', 'small-0b', 'medium-0b', 'small-0b2', 'medium-0b2','large-0b2', 'medium-0b3', 'medium-mpa-0', 'medium-omat-0', 'mace-matpes-pbe-0', 'mace-matpes-r2scan-0'", - ) - device: str = Field( - default="cpu", - description="Device to run the MACE calculation on. Options are cpu, cuda or xpu.", - ) - temperature: float = Field( - default=298.15, - description="Temperature for thermo property calculations in Kelvin (K).", - ) - pressure: float = Field( - default=101325.0, - description="Pressure for thermo property calculations in Pascal (Pa).", - ) - fmax: float = Field( - default=0.01, - description="The convergence criterion for forces (in eV/Å). Optimization stops when all force components are smaller than this value.", - ) - steps: int = Field( - default=1000, - description="Maximum number of optimization steps. The optimization will terminate if this number is reached, even if forces haven't converged to fmax.", - ) - optimizer: str = Field( - default="lbfgs", - description="The optimization algorithm used for geometry optimization. Options are 'bfgs', 'lbfgs', 'gpmin', 'fire', 'mdmin'", - ) - -class mace_output_schema(BaseModel): - final_structure_file: str = Field( - description="Path to the final coordinate file (e.g., CIF, XYZ, POSCAR) containing the atomic structure for the simulation." - ) - output_result_file: str = Field( - description="Path to a JSON file where simulation results is saved.", - ) - model: str = Field( - default=None, description="Path to the model. Default is medium-mpa-0." - ) - device: str = Field( - default="cpu", - description="Device to run the MACE calculation on. Options are cpu, cuda or xpu.", - ) - temperature: float = Field( - default=298.15, - description="Temperature for thermo property calculations in Kelvin (K).", - ) - pressure: float = Field( - default=101325.0, - description="Pressure for thermo property calculations in Pascal (Pa).", - ) - fmax: float = Field( - default=0.01, - description="The convergence criterion for forces (in eV/Å). Optimization stops when all force components are smaller than this value.", - ) - steps: int = Field( - default=1000, - description="Maximum number of optimization steps. The optimization will terminate if this number is reached, even if forces haven't converged to fmax.", - ) - energy: float = Field( - description="The electronic energy of the system in eV", - ) - success: bool = Field( - description="Status of the simulation", - ) - vibrational_frequencies: dict = Field( - default={}, - description="Vibrational frequencies (in cm-1) and energies (in eV).", - ) - thermochemistry: dict = Field( - default={}, - description="Thermochemistry data in eV.", - ) - error: str = Field( - default="", - description="Error captured during the simulation", - ) - wall_time: float = Field( - default=None, - description="Total wall time (in seconds) taken to complete the simulation.", - ) - - -def run_mace_core(mace_input_schema: mace_input_schema): - """Run a single MACE calculations using specified input parameters. +def _mace_input_to_ase_input(params: mace_input_schema) -> ASEInputSchema: + """Convert a MACE-specific input schema to a generic ASEInputSchema. Parameters ---------- - mace_input_schema : mace_input_schema - Input parameters for the MACE calculation - """ - from ase.io import read - from ase.io import write as ase_write - from ase.optimize import BFGS, LBFGS, GPMin, FIRE, MDMin - - from mace.calculators import mace_mp - - # Input validations - if not os.path.isfile(mace_input_schema.input_structure_file): - err = f"Input structure file {mace_input_schema.input_structure_file} does not exist." - raise ValueError(err) - # Validate the output results file (if provided) - if not mace_input_schema.output_result_file.endswith(".json"): - err = f"Output results file must end with '.json', got: {mace_input_schema.output_result_file}" - raise ValueError(err) - - # Validate the input structure with ASE io - try: - atoms = read(mace_input_schema.input_structure_file) - except Exception as e: - err = f"Cannot read {mace_input_schema.input_structure_file} using ASE. Exception from ASE: {e}" - raise ValueError(err) - - # Start time - start_time = time.time() - - calc = mace_mp(model=mace_input_schema.model, device=mace_input_schema.device) - atoms.calc = calc - - # Create parent directory for output file - output_path = Path(mace_input_schema.output_result_file) - output_path.parent.mkdir(parents=True, exist_ok=True) - - if mace_input_schema.driver == "energy": - try: - energy = atoms.get_potential_energy() - except Exception as e: - err = f"Error encountered with the simulation. Exception: {e}" - raise ValueError(err) - - final_structure_file = os.path.abspath(mace_input_schema.input_structure_file) - - end_time = time.time() - wall_time = end_time - start_time - simulation_output = mace_output_schema( - final_structure_file=final_structure_file, - success=True, - energy=energy, - wall_time=wall_time, - output_result_file=os.path.abspath(mace_input_schema.output_result_file), - model=mace_input_schema.model, - driver=mace_input_schema.driver, - device=mace_input_schema.device, - ) - with open(mace_input_schema.output_result_file, "w") as wf: - wf.write(simulation_output.model_dump_json(indent=4)) - return { - "status": "success", - "message": f"Simulation completed. Results saved to {os.path.abspath(mace_input_schema.output_result_file)}", - "single_point_energy": energy, - "unit": "eV", - } - else: - OPTIMIZERS = { - "bfgs": BFGS, - "lbfgs": LBFGS, - "gpmin": GPMin, - "fire": FIRE, - "mdmin": MDMin, - } - try: - optimizer_class = OPTIMIZERS.get(mace_input_schema.optimizer.lower()) - if optimizer_class is None: - raise ValueError(f"Unsupported optimizer: {optimizer_class}") - - # Do optimization only if number of atoms > 1 to avoid error. - if len(atoms) > 1: - dyn = optimizer_class(atoms) - dyn.run( - fmax=mace_input_schema.fmax, - steps=mace_input_schema.steps, - ) - - # Get the single-point energy of the optimized structure - opt_energy = float(atoms.get_potential_energy()) - - # Write the optimized structure - input_path = mace_input_schema.input_structure_file - root, ext = os.path.splitext(input_path) - opt_path = root + "_opt" + ext - ase_write(opt_path, atoms) - final_structure_file = os.path.abspath(opt_path) - - # Initiate thermo and vibrational data - thermo_data = {} - vib_data = {} + params : mace_input_schema + MACE-specific tool input. - if mace_input_schema.driver in {"vib", "thermo"}: - from ase.vibrations import Vibrations - from ase import units - - vib = Vibrations(atoms) - - vib.clean() - vib.run() - - vib_data = { - "energies": [], - "energy_unit": "meV", - "frequencies": [], - "frequency_unit": "cm-1", - } - - energies = vib.get_energies() - - for idx, e in enumerate(energies): - is_imag = abs(e.imag) > 1e-8 - e_val = e.imag if is_imag else e.real - energy_meV = 1e3 * e_val - freq_cm1 = e_val / units.invcm - suffix = "i" if is_imag else "" - vib_data["energies"].append(f"{energy_meV}{suffix}") - vib_data["frequencies"].append(f"{freq_cm1}{suffix}") - - # Remove existing frequencies.txt and .traj files - for traj_file in glob.glob("*.traj"): - os.remove(traj_file) - - # Write frequencies into frequencies.txt - freq_file = Path("frequencies.csv") - if freq_file.exists(): - freq_file.unlink() - - with freq_file.open("w") as f: - for i, freq in enumerate(vib_data["frequencies"], start=0): - f.write(f"vib.{i}.traj,{freq}\n") - - # Write normal modes .traj files - for i in range(len(energies)): - vib.write_mode(n=None, kT=units.kB * 300, nimages=30) - - if mace_input_schema.driver == "thermo": - # Approximation for a single atom system. - if len(atoms) == 1: - thermo_data = { - "enthalpy": opt_energy, - "entropy": 0.0, - "gibbs_free_energy": opt_energy, - "unit": "eV", - } - else: - from ase.thermochemistry import IdealGasThermo - - final_structure = atoms_to_atomsdata(atoms) - linear = is_linear_molecule(final_structure) - geometry = "linear" if linear else "nonlinear" - symmetrynumber = get_symmetry_number(final_structure) + Returns + ------- + ASEInputSchema + Equivalent generic ASE input. + """ + return ASEInputSchema( + input_structure_file=params.input_structure_file, + output_results_file=params.output_result_file, + driver=params.driver, + optimizer=params.optimizer, + calculator={ + "calculator_type": "mace_mp", + "model": params.model, + "device": params.device, + }, + fmax=params.fmax, + steps=params.steps, + temperature=params.temperature, + pressure=params.pressure, + ) - thermo = IdealGasThermo( - vib_energies=energies, - potentialenergy=opt_energy, - atoms=atoms, - geometry=geometry, - symmetrynumber=symmetrynumber, - spin=0, - ) - thermo_data = { - "enthalpy": float( - thermo.get_enthalpy( - temperature=mace_input_schema.temperature - ) - ), - "entropy": float( - thermo.get_entropy( - temperature=mace_input_schema.temperature, - pressure=mace_input_schema.pressure, - ) - ), - "gibbs_free_energy": float( - thermo.get_gibbs_energy( - temperature=mace_input_schema.temperature, - pressure=mace_input_schema.pressure, - ) - ), - "unit": "eV", - } - end_time = time.time() - wall_time = end_time - start_time +def run_mace_core(params: mace_input_schema) -> dict: + """Run a single MACE calculation. - simulation_output = mace_output_schema( - final_structure_file=final_structure_file, - success=True, - energy=opt_energy, - wall_time=wall_time, - output_result_file=os.path.abspath( - mace_input_schema.output_result_file - ), - model=mace_input_schema.model, - driver=mace_input_schema.driver, - device=mace_input_schema.device, - vibrational_frequencies=vib_data, - thermochemistry=thermo_data, - ) - with open(mace_input_schema.output_result_file, "w") as wf: - wf.write(simulation_output.model_dump_json(indent=4)) + Converts the MACE-specific schema to :class:`ASEInputSchema` and + delegates to :func:`chemgraph.tools.ase_core.run_ase_core`. - # Return message based on driver. Keep the return output minimal. - if mace_input_schema.driver == "opt": - return { - "status": "success", - "message": f"Simulation completed. Results saved to {mace_input_schema.output_result_file}", - "single_point_energy": opt_energy, # small payload for LLMs - "unit": "eV", - } - elif mace_input_schema.driver == "vib": - return { - "status": "success", - "result": { - "vibrational_frequencies": vib_data, - }, # small payload for LLMs - "message": ( - "Vibrational analysis completed; frequencies returned. " - f"Full results (structure, vibrations and metadata) saved to {mace_input_schema.output_result_file}." - ), - } - elif mace_input_schema.driver == "thermo": - return { - "status": "success", - "result": {"thermochemistry": thermo_data}, - "message": ( - "Thermochemistry computed and returned. " - f"Full results (structure, vibrations, thermochemistry and metadata) saved to {mace_input_schema.output_result_file}" - ), - } - except Exception as e: - err = f"ASE simulation gave an exception:{e}" - return { - "status": "failure", - "error_type": type(e).__name__, - "message": str(e), - } + Parameters + ---------- + params : mace_input_schema + MACE-specific input parameters. - return True + Returns + ------- + dict + Simulation result payload. + """ + ase_params = _mace_input_to_ase_input(params) + return run_ase_core(ase_params) diff --git a/src/chemgraph/tools/rag_tools.py b/src/chemgraph/tools/rag_tools.py index b22fe665..eb6c9b04 100644 --- a/src/chemgraph/tools/rag_tools.py +++ b/src/chemgraph/tools/rag_tools.py @@ -119,6 +119,16 @@ def _get_embeddings(provider: str = "openai"): Supports OpenAI-compatible custom endpoints via OPENAI_BASE_URL. Falls back to HuggingFace if OpenAI embeddings are unavailable. + + Parameters + ---------- + provider : str, optional + Preferred embedding provider. + + Returns + ------- + Embeddings + LangChain-compatible embeddings object. """ if provider == "openai": try: diff --git a/src/chemgraph/tools/report_tools.py b/src/chemgraph/tools/report_tools.py index b531452f..9761c1e6 100644 --- a/src/chemgraph/tools/report_tools.py +++ b/src/chemgraph/tools/report_tools.py @@ -4,7 +4,7 @@ from typing import Optional from langchain_core.tools import tool -from ase.data import chemical_symbols +from ase.data import chemical_symbols as _chemical_symbols from chemgraph.schemas.ase_input import ASEOutputSchema from chemgraph.tools.ase_tools import is_linear_molecule @@ -392,7 +392,7 @@ def generate_html( for num, pos in zip( ase_output.final_structure.numbers, ase_output.final_structure.positions ): - element = chemical_symbols[num] if num < len(chemical_symbols) else f"X{num}" + element = _chemical_symbols[num] if num < len(_chemical_symbols) else f"X{num}" x, y, z = pos xyz_lines.append(f"{element} {x:.6f} {y:.6f} {z:.6f}") @@ -449,7 +449,7 @@ def add_additional_info_to_html(html_content: str, ase_output: ASEOutputSchema) for num, pos in zip( ase_output.final_structure.numbers, ase_output.final_structure.positions ): - element = chemical_symbols[num] if num < len(chemical_symbols) else f"X{num}" + element = _chemical_symbols[num] if num < len(_chemical_symbols) else f"X{num}" x, y, z = pos xyz_lines.append(f"{element} {x:.6f} {y:.6f} {z:.6f}") diff --git a/src/chemgraph/tools/xanes_core.py b/src/chemgraph/tools/xanes_core.py new file mode 100644 index 00000000..9edf1ede --- /dev/null +++ b/src/chemgraph/tools/xanes_core.py @@ -0,0 +1,561 @@ +"""Pure-Python XANES/FDMNES helpers (no LangChain / MCP decorators). + +Contains all core workflow functions for FDMNES input generation, +execution, result parsing, Materials Project data fetching, and plotting. +Used by the LangChain ``@tool`` wrappers in :mod:`xanes_tools` and the +MCP wrappers in :mod:`chemgraph.mcp.xanes_mcp_parsl`. +""" + +from __future__ import annotations + +import logging +import os +import pickle +import shutil +import subprocess +from pathlib import Path +from typing import List, Optional + +import numpy as np +from ase import Atoms +from ase.io import read as ase_read, write as ase_write + +from chemgraph.schemas.xanes_schema import xanes_input_schema, mp_query_schema + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Helper Functions +# --------------------------------------------------------------------------- + + +def write_fdmnes_input( + ase_atoms: Atoms, + z_absorber: int = None, + input_file_dir: Path = None, + radius: float = 6.0, + magnetism: bool = False, +): + """Write FDMNES input files (fdmfile.txt and fdmnes_in.txt) for a structure. + + Parameters + ---------- + ase_atoms : ase.Atoms + Atomic structure to compute XANES for. + z_absorber : int, optional + Atomic number of the X-ray absorbing atom. + Defaults to the heaviest element in the structure. + input_file_dir : Path, optional + Directory to write input files into. Defaults to cwd. + radius : float + Cluster radius in Angstrom. Default 6.0. + magnetism : bool + Enable magnetic contributions. Default False. + """ + if not isinstance(ase_atoms, Atoms): + raise TypeError("ase_atoms must be an ase.Atoms object") + + atomic_numbers = ase_atoms.get_atomic_numbers() + if z_absorber is None: + z_absorber = int(atomic_numbers.max()) + + if input_file_dir is None: + input_file_dir = Path.cwd() + + with open(input_file_dir / "fdmfile.txt", "w") as f: + f.write("1\n") + f.write("fdmnes_in.txt\n") + + with open(input_file_dir / "fdmnes_in.txt", "w") as f: + f.write("Filout\n") + f.write(f"{input_file_dir.name}\n\n") + + # Energy mesh + f.write("Range\n") + f.write("-55. 1.0 -10. 0.01 5. 0.1 150.\n\n") + + # Cluster radius + f.write("Radius\n") + f.write(f"{radius}\n\n") + + # Absorbing atom + f.write("Z_absorber\n") + f.write(f"{z_absorber}\n\n") + + # Magnetic contributions + if magnetism: + f.write("Magnetism\n\n") + + f.write("Green\n") + f.write("Density_all\n") + f.write("Quadrupole\n") + f.write("Spherical\n") + f.write("SCF\n\n") + + if all(ase_atoms.pbc): + f.write("Crystal\n") + f.write(" ".join(map(str, ase_atoms.cell.cellpar())) + "\n") + positions = np.round(ase_atoms.get_scaled_positions(), 6) + else: + f.write("Molecule\n") + cell_length = abs(ase_atoms.get_positions().max()) + abs( + ase_atoms.get_positions().min() + ) + f.write(f"{cell_length} {cell_length} {cell_length} 90 90 90\n") + positions = np.round(ase_atoms.get_positions(), 6) + + for i, position in enumerate(positions): + f.write(f"{atomic_numbers[i]} " + " ".join(map(str, position)) + "\n") + + f.write("\n") + f.write("Convolution\n") + f.write("End") + + +def get_normalized_xanes( + conv_file: Path | str, + pre_edge_width: float = 20.0, + post_edge_width: float = 50.0, + calc_E0: bool = False, +) -> tuple[np.ndarray, np.ndarray]: + """Normalize a XANES spectrum from an FDMNES convolution output file. + + Parameters + ---------- + conv_file : Path or str + Path to the FDMNES ``*_conv.txt`` output file. + pre_edge_width : float + Width of the pre-edge region in eV for baseline fitting. + post_edge_width : float + Width of the post-edge region in eV for step normalization. + calc_E0 : bool + If True, determine the edge energy E0 from the maximum of dmu/dE. + Otherwise E0 is assumed to be 0 (the FDMNES convention). + + Returns + ------- + normalized : np.ndarray + (N, 2) array of [energy, normalized_mu]. + raw : np.ndarray + (N, 2) array of [energy, raw_mu] as read from the file. + """ + energy_xas = np.loadtxt(conv_file, skiprows=1) + + E = energy_xas[:, 0].astype(float) + mu = energy_xas[:, 1].astype(float) + + if calc_E0: + dmu_dE = np.gradient(mu, E) + E0 = E[np.argmax(dmu_dE)] + else: + E0 = 0 + + pre_mask = E <= (E0 - pre_edge_width) + post_mask = E >= (E0 + post_edge_width) + + m_pre, b_pre = np.polyfit(E[pre_mask], mu[pre_mask], 1) + m_post, b_post = np.polyfit(E[post_mask], mu[post_mask], 1) + + pre_line = m_pre * E + b_pre + mu_corr = mu - pre_line + + step = (m_post * E0 + b_post) - (m_pre * E0 + b_pre) + mu_norm = mu_corr / step + + return np.column_stack([E, mu_norm]), energy_xas + + +def extract_conv(fdmnes_output_dir: Path | str) -> dict: + """Extract all convolution output files from an FDMNES run directory. + + Parameters + ---------- + fdmnes_output_dir : Path or str + Directory containing FDMNES output files. + + Returns + ------- + dict + Mapping of index to (N, 2) arrays of [energy, mu]. + """ + if not isinstance(fdmnes_output_dir, Path): + fdmnes_output_dir = Path(fdmnes_output_dir) + + energy_xas = {} + for i, conv_file in enumerate(fdmnes_output_dir.glob("*conv.txt")): + energy_xas[i] = np.loadtxt(conv_file, skiprows=1) + + return energy_xas + + +# --------------------------------------------------------------------------- +# Data directory helper +# --------------------------------------------------------------------------- + + +def _get_data_dir() -> Path: + """Return the working data directory for XANES workflows.""" + cwd = Path.cwd() + if "PBS_O_WORKDIR" in os.environ: + cwd = Path(os.environ["PBS_O_WORKDIR"]) + + data_dir = cwd / "xanes_data" + if not data_dir.exists(): + data_dir.mkdir(parents=True) + return data_dir + + +# --------------------------------------------------------------------------- +# Core Workflow Functions +# --------------------------------------------------------------------------- + + +def run_xanes_core(params: xanes_input_schema) -> dict: + """Run a single XANES/FDMNES calculation for one structure. + + This is the core function analogous to ``run_graspa_core``. It: + 1. Reads the input structure file via ASE. + 2. Creates FDMNES input files via ``write_fdmnes_input``. + 3. Runs FDMNES via subprocess. + 4. Parses the convolution output if available. + + Parameters + ---------- + params : xanes_input_schema + Input parameters for the FDMNES calculation. + + Returns + ------- + dict + Result dictionary with keys: status, output_dir, conv_data (if success), + error (if failure). + """ + fdmnes_exe = os.environ.get("FDMNES_EXE") + if not fdmnes_exe: + raise ValueError( + "FDMNES_EXE environment variable is not set. " + "Set it to the path of the FDMNES executable." + ) + + input_path = Path(params.input_structure_file).resolve() + if not input_path.exists(): + raise FileNotFoundError(f"Input structure file not found: {input_path}") + + atoms = ase_read(str(input_path)) + + # Determine output directory + if params.output_dir is not None: + run_dir = Path(params.output_dir).resolve() + else: + run_dir = input_path.parent / f"fdmnes_{input_path.stem}" + run_dir.mkdir(parents=True, exist_ok=True) + + # Write FDMNES input files + write_fdmnes_input( + ase_atoms=atoms, + z_absorber=params.z_absorber, + input_file_dir=run_dir, + radius=params.radius, + magnetism=params.magnetism, + ) + + # Save the atoms object alongside the inputs for provenance + formula = atoms.get_chemical_formula() + z_abs = params.z_absorber or int(atoms.get_atomic_numbers().max()) + mp_id = atoms.info.get("MP-id", "local") + pkl_filename = f"Z{z_abs}_{mp_id}_{formula}.pkl" + with open(run_dir / pkl_filename, "wb") as f: + pickle.dump(atoms, f) + + # Run FDMNES + logger.info("Running FDMNES in %s", run_dir) + with ( + open(run_dir / "fdmnes_stdout.txt", "w") as fp_out, + open(run_dir / "fdmnes_stderr.txt", "w") as fp_err, + ): + proc = subprocess.run( + fdmnes_exe, + cwd=str(run_dir), + stdout=fp_out, + stderr=fp_err, + shell=True, + ) + + if proc.returncode != 0: + logger.error( + "FDMNES failed with return code %d in %s", proc.returncode, run_dir + ) + return { + "status": "failure", + "output_dir": str(run_dir), + "error": f"FDMNES exited with return code {proc.returncode}", + } + + # Parse results + conv_data = extract_conv(run_dir) + if not conv_data: + logger.warning("No convolution output found in %s", run_dir) + return { + "status": "failure", + "output_dir": str(run_dir), + "error": "No *conv.txt output files found after FDMNES execution.", + } + + logger.info("FDMNES completed successfully in %s", run_dir) + return { + "status": "success", + "output_dir": str(run_dir), + "n_conv_files": len(conv_data), + } + + +def fetch_materials_project_data( + params: mp_query_schema, + db_path: Path, +) -> dict: + """Fetch optimized structures from Materials Project. + + Parameters + ---------- + params : mp_query_schema + Query parameters including chemical formulas and API key. + db_path : Path + Directory to save the fetched structures. + + Returns + ------- + dict + atoms_list : list[Atoms] -- fetched ASE Atoms objects + structure_files : list[str] -- absolute paths to saved CIF files + pickle_file : str -- absolute path to atoms_db.pkl + n_structures : int -- number of structures fetched + """ + from mp_api.client import MPRester + from pymatgen.io.ase import AseAtomsAdaptor + + api_key = params.mp_api_key or os.environ.get("MP_API_KEY") + if not api_key: + raise ValueError( + "No Materials Project API key provided. " + "Pass it via mp_api_key or set the MP_API_KEY environment variable." + ) + + logger.info("Fetching data from Materials Project for: %s", params.chemsys) + atoms_list = [] + + with MPRester(api_key) as mpr: + doc_list = mpr.materials.summary.search( + fields=["material_id", "structure"], + energy_above_hull=(0, params.energy_above_hull), + formula=params.chemsys, + deprecated=False, + ) + + for doc in doc_list: + ase_atoms = AseAtomsAdaptor.get_atoms(doc.structure) + ase_atoms.info.update({"MP-id": str(doc.material_id)}) + atoms_list.append(ase_atoms) + + if not db_path.exists(): + db_path.mkdir(parents=True) + + # Save pickle database + pkl_path = db_path / "atoms_db.pkl" + with open(pkl_path, "wb") as f: + pickle.dump(atoms_list, f) + + # Save individual CIF files + structure_files = [] + for atoms in atoms_list: + mp_id = atoms.info.get("MP-id", "unknown") + formula = atoms.get_chemical_formula() + cif_path = db_path / f"{mp_id}_{formula}.cif" + ase_write(str(cif_path), atoms) + structure_files.append(str(cif_path)) + + logger.info( + "Saved %d structures (%s) and pickle database to %s", + len(atoms_list), + [Path(f).name for f in structure_files], + db_path, + ) + + return { + "atoms_list": atoms_list, + "structure_files": structure_files, + "pickle_file": str(pkl_path), + "n_structures": len(atoms_list), + } + + +def create_fdmnes_inputs( + root_dir: Path, + atoms_list: Optional[List[Atoms]] = None, + z_absorber: Optional[int] = None, + radius: float = 6.0, + magnetism: bool = False, +) -> Path: + """Create FDMNES input files for a batch of structures. + + Parameters + ---------- + root_dir : Path + Root directory for the batch. A ``fdmnes_batch_runs`` subdirectory + will be created containing per-structure run directories. + atoms_list : list[ase.Atoms], optional + Structures to process. If None, loads from ``root_dir/atoms_db.pkl``. + z_absorber : int, optional + Atomic number of the absorbing atom. Defaults to heaviest per structure. + radius : float + Cluster radius in Angstrom. + magnetism : bool + Enable magnetic contributions. + + Returns + ------- + Path + Path to the ``fdmnes_batch_runs`` directory. + """ + logger.info("Creating FDMNES inputs in %s", root_dir) + runs_dir = root_dir / "fdmnes_batch_runs" + + start_idx = 0 + if runs_dir.exists(): + for subdir in runs_dir.iterdir(): + try: + start_idx = max(start_idx, int(subdir.name.split("_")[-1])) + except ValueError: + continue + last_run = runs_dir / f"run_{start_idx}" + if last_run.exists(): + shutil.rmtree(last_run) + else: + runs_dir.mkdir(parents=True) + + if atoms_list is None: + db_path = root_dir / "atoms_db.pkl" + if not db_path.exists(): + raise FileNotFoundError(f"No atoms provided and {db_path} not found.") + with open(db_path, "rb") as f: + atoms_list = pickle.load(f) + + for i, atoms in enumerate(atoms_list, start=start_idx): + curr_run_dir = runs_dir / f"run_{i}" + curr_run_dir.mkdir(parents=True, exist_ok=True) + + current_z = ( + z_absorber + if z_absorber is not None + else int(max(atoms.get_atomic_numbers())) + ) + write_fdmnes_input( + ase_atoms=atoms, + input_file_dir=curr_run_dir, + z_absorber=current_z, + radius=radius, + magnetism=magnetism, + ) + + mp_id = atoms.info.get("MP-id", "local") + formula = atoms.get_chemical_formula() + pkl_filename = f"Z{current_z}_{mp_id}_{formula}.pkl" + with open(curr_run_dir / pkl_filename, "wb") as f: + pickle.dump(atoms, f) + + return runs_dir + + +def expand_database_results(root_dir: Path, runs_dir: Path) -> None: + """Expand the atoms database with XANES convolution results. + + For each completed run directory, loads the pickled Atoms object, + attaches the FDMNES convolution data to ``atoms.info``, and saves + all expanded structures to ``root_dir/atoms_db_expanded.pkl``. + + Parameters + ---------- + root_dir : Path + Root directory where the expanded database will be saved. + runs_dir : Path + Directory containing ``run_*`` subdirectories with FDMNES outputs. + """ + logger.info("Expanding database with XANES results...") + expanded_atoms_list = [] + + for sub_dir in sorted(runs_dir.glob("run_*")): + atoms_pkl_files = list(sub_dir.glob("*.pkl")) + if not atoms_pkl_files: + continue + + with open(atoms_pkl_files[0], "rb") as f: + ase_atoms = pickle.load(f) + + conv_data = extract_conv(fdmnes_output_dir=sub_dir) + ase_atoms.info.update({"FDMNES-xanes": conv_data}) + expanded_atoms_list.append(ase_atoms) + + with open(root_dir / "atoms_db_expanded.pkl", "wb") as f: + pickle.dump(expanded_atoms_list, f) + + logger.info( + "Saved %d expanded structures to %s", + len(expanded_atoms_list), + root_dir / "atoms_db_expanded.pkl", + ) + + +def plot_xanes_results(root_dir: Path, runs_dir: Path) -> dict: + """Generate normalized XANES plots for completed FDMNES calculations. + + For each run directory containing a ``*_conv.txt`` file, produces + a ``xanes_plot.png`` with the normalized absorption spectrum. + + Parameters + ---------- + root_dir : Path + Root data directory (unused currently, reserved for summary plots). + runs_dir : Path + Directory containing ``run_*`` subdirectories with FDMNES outputs. + + Returns + ------- + dict + plot_files : list[str] -- absolute paths to generated plot images + n_plots : int -- number of plots successfully generated + n_failed : int -- number of runs that failed to plot + failed : list[str] -- names of run directories that failed + """ + import matplotlib.pyplot as plt + + logger.info("Plotting XANES results from %s", runs_dir) + + plot_files = [] + failed = [] + + for sub_dir in sorted(runs_dir.glob("run_*")): + conv_file = next(sub_dir.glob("*_conv.txt"), None) + if conv_file: + try: + norm_energy, _raw = get_normalized_xanes(conv_file) + plot_path = sub_dir / "xanes_plot.png" + plt.figure() + plt.plot(norm_energy[:, 0], norm_energy[:, 1], label=sub_dir.name) + plt.xlabel("Energy [eV]") + plt.ylabel("Normalized Absorption") + plt.title(f"XANES for {sub_dir.name}") + plt.legend() + plt.savefig(plot_path, dpi=150) + plt.close() + plot_files.append(str(plot_path)) + logger.info("Plotted %s", sub_dir.name) + except Exception as e: + logger.error("Failed to plot %s: %s", sub_dir.name, e) + failed.append(sub_dir.name) + + return { + "plot_files": plot_files, + "n_plots": len(plot_files), + "n_failed": len(failed), + "failed": failed, + } diff --git a/src/chemgraph/tools/xanes_tools.py b/src/chemgraph/tools/xanes_tools.py index 9e78043b..c4aac946 100644 --- a/src/chemgraph/tools/xanes_tools.py +++ b/src/chemgraph/tools/xanes_tools.py @@ -1,562 +1,63 @@ -import logging -import os -import pickle -import subprocess -import shutil -from pathlib import Path -from typing import List, Optional - -import numpy as np -from ase import Atoms -from ase.io import read as ase_read, write as ase_write -from langchain_core.tools import tool - -from chemgraph.schemas.xanes_schema import xanes_input_schema, mp_query_schema - -logger = logging.getLogger(__name__) - -# ----------------------------------------------------------------------------- -# Helper Functions -# ----------------------------------------------------------------------------- - - -def write_fdmnes_input( - ase_atoms: Atoms, - z_absorber: int = None, - input_file_dir: Path = None, - radius: float = 6.0, - magnetism: bool = False, -): - """Write FDMNES input files (fdmfile.txt and fdmnes_in.txt) for a structure. - - Parameters - ---------- - ase_atoms : ase.Atoms - Atomic structure to compute XANES for. - z_absorber : int, optional - Atomic number of the X-ray absorbing atom. - Defaults to the heaviest element in the structure. - input_file_dir : Path, optional - Directory to write input files into. Defaults to cwd. - radius : float - Cluster radius in Angstrom. Default 6.0. - magnetism : bool - Enable magnetic contributions. Default False. - """ - if not isinstance(ase_atoms, Atoms): - raise TypeError("ase_atoms must be an ase.Atoms object") - - atomic_numbers = ase_atoms.get_atomic_numbers() - if z_absorber is None: - z_absorber = int(atomic_numbers.max()) - - if input_file_dir is None: - input_file_dir = Path.cwd() - - with open(input_file_dir / "fdmfile.txt", "w") as f: - f.write("1\n") - f.write("fdmnes_in.txt\n") - - with open(input_file_dir / "fdmnes_in.txt", "w") as f: - f.write("Filout\n") - f.write(f"{input_file_dir.name}\n\n") - - # Energy mesh - f.write("Range\n") - f.write("-55. 1.0 -10. 0.01 5. 0.1 150.\n\n") - - # Cluster radius - f.write("Radius\n") - f.write(f"{radius}\n\n") - - # Absorbing atom - f.write("Z_absorber\n") - f.write(f"{z_absorber}\n\n") - - # Magnetic contributions - if magnetism: - f.write("Magnetism\n\n") - - f.write("Green\n") - f.write("Density_all\n") - f.write("Quadrupole\n") - f.write("Spherical\n") - f.write("SCF\n\n") - - if all(ase_atoms.pbc): - f.write("Crystal\n") - f.write(" ".join(map(str, ase_atoms.cell.cellpar())) + "\n") - positions = np.round(ase_atoms.get_scaled_positions(), 6) - else: - f.write("Molecule\n") - cell_length = abs(ase_atoms.get_positions().max()) + abs( - ase_atoms.get_positions().min() - ) - f.write(f"{cell_length} {cell_length} {cell_length} 90 90 90\n") - positions = np.round(ase_atoms.get_positions(), 6) - - for i, position in enumerate(positions): - f.write(f"{atomic_numbers[i]} " + " ".join(map(str, position)) + "\n") - - f.write("\n") - f.write("Convolution\n") - f.write("End") - - -def get_normalized_xanes( - conv_file: Path | str, - pre_edge_width: float = 20.0, - post_edge_width: float = 50.0, - calc_E0: bool = False, -) -> tuple[np.ndarray, np.ndarray]: - """Normalize a XANES spectrum from an FDMNES convolution output file. - - Parameters - ---------- - conv_file : Path or str - Path to the FDMNES ``*_conv.txt`` output file. - pre_edge_width : float - Width of the pre-edge region in eV for baseline fitting. - post_edge_width : float - Width of the post-edge region in eV for step normalization. - calc_E0 : bool - If True, determine the edge energy E0 from the maximum of dmu/dE. - Otherwise E0 is assumed to be 0 (the FDMNES convention). - - Returns - ------- - normalized : np.ndarray - (N, 2) array of [energy, normalized_mu]. - raw : np.ndarray - (N, 2) array of [energy, raw_mu] as read from the file. - """ - energy_xas = np.loadtxt(conv_file, skiprows=1) - - E = energy_xas[:, 0].astype(float) - mu = energy_xas[:, 1].astype(float) - - if calc_E0: - dmu_dE = np.gradient(mu, E) - E0 = E[np.argmax(dmu_dE)] - else: - E0 = 0 - - pre_mask = E <= (E0 - pre_edge_width) - post_mask = E >= (E0 + post_edge_width) - - m_pre, b_pre = np.polyfit(E[pre_mask], mu[pre_mask], 1) - m_post, b_post = np.polyfit(E[post_mask], mu[post_mask], 1) - - pre_line = m_pre * E + b_pre - mu_corr = mu - pre_line - - step = (m_post * E0 + b_post) - (m_pre * E0 + b_pre) - mu_norm = mu_corr / step +"""LangChain ``@tool`` wrappers for XANES/FDMNES functions. - return np.column_stack([E, mu_norm]), energy_xas +Each tool delegates to the pure-Python implementation in +:mod:`chemgraph.tools.xanes_core`. +""" +from __future__ import annotations -def extract_conv(fdmnes_output_dir: Path | str) -> dict: - """Extract all convolution output files from an FDMNES run directory. - - Parameters - ---------- - fdmnes_output_dir : Path or str - Directory containing FDMNES output files. - - Returns - ------- - dict - Mapping of index to (N, 2) arrays of [energy, mu]. - """ - if not isinstance(fdmnes_output_dir, Path): - fdmnes_output_dir = Path(fdmnes_output_dir) - - energy_xas = {} - for i, conv_file in enumerate(fdmnes_output_dir.glob("*conv.txt")): - energy_xas[i] = np.loadtxt(conv_file, skiprows=1) - - return energy_xas +from pathlib import Path +from langchain_core.tools import tool -# ----------------------------------------------------------------------------- -# Core Workflow Functions -# ----------------------------------------------------------------------------- +from chemgraph.schemas.xanes_schema import xanes_input_schema, mp_query_schema +from chemgraph.tools.xanes_core import ( + # Re-export core helpers so existing ``from xanes_tools import ...`` + # statements in MCP servers continue to work during the transition. + write_fdmnes_input, + get_normalized_xanes, + extract_conv, + _get_data_dir, + run_xanes_core, + fetch_materials_project_data, + create_fdmnes_inputs, + expand_database_results, + plot_xanes_results, +) + +# Make re-exports explicit for linters. +__all__ = [ + "write_fdmnes_input", + "get_normalized_xanes", + "extract_conv", + "_get_data_dir", + "run_xanes_core", + "fetch_materials_project_data", + "create_fdmnes_inputs", + "expand_database_results", + "plot_xanes_results", + "run_xanes", + "fetch_xanes_data", + "plot_xanes_data", +] -def run_xanes_core(params: xanes_input_schema) -> dict: - """Run a single XANES/FDMNES calculation for one structure. +@tool +def run_xanes(params: xanes_input_schema) -> str: + """Run a single XANES/FDMNES calculation for one structure file. - This is the core function analogous to ``run_graspa_core``. It: - 1. Reads the input structure file via ASE. - 2. Creates FDMNES input files via ``write_fdmnes_input``. - 3. Runs FDMNES via subprocess. - 4. Parses the convolution output if available. + This tool reads the structure, generates FDMNES input files, runs FDMNES, + and returns the result status. Requires the FDMNES_EXE environment variable. Parameters ---------- params : xanes_input_schema - Input parameters for the FDMNES calculation. + Input parameters for the XANES/FDMNES calculation. Returns ------- - dict - Result dictionary with keys: status, output_dir, conv_data (if success), - error (if failure). - """ - fdmnes_exe = os.environ.get("FDMNES_EXE") - if not fdmnes_exe: - raise ValueError( - "FDMNES_EXE environment variable is not set. " - "Set it to the path of the FDMNES executable." - ) - - input_path = Path(params.input_structure_file).resolve() - if not input_path.exists(): - raise FileNotFoundError(f"Input structure file not found: {input_path}") - - atoms = ase_read(str(input_path)) - - # Determine output directory - if params.output_dir is not None: - run_dir = Path(params.output_dir).resolve() - else: - run_dir = input_path.parent / f"fdmnes_{input_path.stem}" - run_dir.mkdir(parents=True, exist_ok=True) - - # Write FDMNES input files - write_fdmnes_input( - ase_atoms=atoms, - z_absorber=params.z_absorber, - input_file_dir=run_dir, - radius=params.radius, - magnetism=params.magnetism, - ) - - # Save the atoms object alongside the inputs for provenance - formula = atoms.get_chemical_formula() - z_abs = params.z_absorber or int(atoms.get_atomic_numbers().max()) - mp_id = atoms.info.get("MP-id", "local") - pkl_filename = f"Z{z_abs}_{mp_id}_{formula}.pkl" - with open(run_dir / pkl_filename, "wb") as f: - pickle.dump(atoms, f) - - # Run FDMNES - logger.info("Running FDMNES in %s", run_dir) - with ( - open(run_dir / "fdmnes_stdout.txt", "w") as fp_out, - open(run_dir / "fdmnes_stderr.txt", "w") as fp_err, - ): - proc = subprocess.run( - fdmnes_exe, - cwd=str(run_dir), - stdout=fp_out, - stderr=fp_err, - shell=True, - ) - - if proc.returncode != 0: - logger.error( - "FDMNES failed with return code %d in %s", proc.returncode, run_dir - ) - return { - "status": "failure", - "output_dir": str(run_dir), - "error": f"FDMNES exited with return code {proc.returncode}", - } - - # Parse results - conv_data = extract_conv(run_dir) - if not conv_data: - logger.warning("No convolution output found in %s", run_dir) - return { - "status": "failure", - "output_dir": str(run_dir), - "error": "No *conv.txt output files found after FDMNES execution.", - } - - logger.info("FDMNES completed successfully in %s", run_dir) - return { - "status": "success", - "output_dir": str(run_dir), - "n_conv_files": len(conv_data), - } - - -def fetch_materials_project_data( - params: mp_query_schema, - db_path: Path, -) -> dict: - """Fetch optimized structures from Materials Project. - - Parameters - ---------- - params : mp_query_schema - Query parameters including chemical formulas and API key. - db_path : Path - Directory to save the fetched structures. - - Returns - ------- - dict - atoms_list : list[Atoms] — fetched ASE Atoms objects - structure_files : list[str] — absolute paths to saved CIF files - pickle_file : str — absolute path to atoms_db.pkl - n_structures : int — number of structures fetched - """ - from mp_api.client import MPRester - from pymatgen.io.ase import AseAtomsAdaptor - - api_key = params.mp_api_key or os.environ.get("MP_API_KEY") - if not api_key: - raise ValueError( - "No Materials Project API key provided. " - "Pass it via mp_api_key or set the MP_API_KEY environment variable." - ) - - logger.info("Fetching data from Materials Project for: %s", params.chemsys) - atoms_list = [] - - with MPRester(api_key) as mpr: - doc_list = mpr.materials.summary.search( - fields=["material_id", "structure"], - energy_above_hull=(0, params.energy_above_hull), - formula=params.chemsys, - deprecated=False, - ) - - for doc in doc_list: - ase_atoms = AseAtomsAdaptor.get_atoms(doc.structure) - ase_atoms.info.update({"MP-id": str(doc.material_id)}) - atoms_list.append(ase_atoms) - - if not db_path.exists(): - db_path.mkdir(parents=True) - - # Save pickle database - pkl_path = db_path / "atoms_db.pkl" - with open(pkl_path, "wb") as f: - pickle.dump(atoms_list, f) - - # Save individual CIF files - structure_files = [] - for atoms in atoms_list: - mp_id = atoms.info.get("MP-id", "unknown") - formula = atoms.get_chemical_formula() - cif_path = db_path / f"{mp_id}_{formula}.cif" - ase_write(str(cif_path), atoms) - structure_files.append(str(cif_path)) - - logger.info( - "Saved %d structures (%s) and pickle database to %s", - len(atoms_list), - [Path(f).name for f in structure_files], - db_path, - ) - - return { - "atoms_list": atoms_list, - "structure_files": structure_files, - "pickle_file": str(pkl_path), - "n_structures": len(atoms_list), - } - - -def create_fdmnes_inputs( - root_dir: Path, - atoms_list: Optional[List[Atoms]] = None, - z_absorber: Optional[int] = None, - radius: float = 6.0, - magnetism: bool = False, -) -> Path: - """Create FDMNES input files for a batch of structures. - - Parameters - ---------- - root_dir : Path - Root directory for the batch. A ``fdmnes_batch_runs`` subdirectory - will be created containing per-structure run directories. - atoms_list : list[ase.Atoms], optional - Structures to process. If None, loads from ``root_dir/atoms_db.pkl``. - z_absorber : int, optional - Atomic number of the absorbing atom. Defaults to heaviest per structure. - radius : float - Cluster radius in Angstrom. - magnetism : bool - Enable magnetic contributions. - - Returns - ------- - Path - Path to the ``fdmnes_batch_runs`` directory. - """ - logger.info("Creating FDMNES inputs in %s", root_dir) - runs_dir = root_dir / "fdmnes_batch_runs" - - start_idx = 0 - if runs_dir.exists(): - for subdir in runs_dir.iterdir(): - try: - start_idx = max(start_idx, int(subdir.name.split("_")[-1])) - except ValueError: - continue - last_run = runs_dir / f"run_{start_idx}" - if last_run.exists(): - shutil.rmtree(last_run) - else: - runs_dir.mkdir(parents=True) - - if atoms_list is None: - db_path = root_dir / "atoms_db.pkl" - if not db_path.exists(): - raise FileNotFoundError(f"No atoms provided and {db_path} not found.") - with open(db_path, "rb") as f: - atoms_list = pickle.load(f) - - for i, atoms in enumerate(atoms_list, start=start_idx): - curr_run_dir = runs_dir / f"run_{i}" - curr_run_dir.mkdir(parents=True, exist_ok=True) - - current_z = ( - z_absorber - if z_absorber is not None - else int(max(atoms.get_atomic_numbers())) - ) - write_fdmnes_input( - ase_atoms=atoms, - input_file_dir=curr_run_dir, - z_absorber=current_z, - radius=radius, - magnetism=magnetism, - ) - - mp_id = atoms.info.get("MP-id", "local") - formula = atoms.get_chemical_formula() - pkl_filename = f"Z{current_z}_{mp_id}_{formula}.pkl" - with open(curr_run_dir / pkl_filename, "wb") as f: - pickle.dump(atoms, f) - - return runs_dir - - -def expand_database_results(root_dir: Path, runs_dir: Path) -> None: - """Expand the atoms database with XANES convolution results. - - For each completed run directory, loads the pickled Atoms object, - attaches the FDMNES convolution data to ``atoms.info``, and saves - all expanded structures to ``root_dir/atoms_db_expanded.pkl``. - - Parameters - ---------- - root_dir : Path - Root directory where the expanded database will be saved. - runs_dir : Path - Directory containing ``run_*`` subdirectories with FDMNES outputs. - """ - logger.info("Expanding database with XANES results...") - expanded_atoms_list = [] - - for sub_dir in sorted(runs_dir.glob("run_*")): - atoms_pkl_files = list(sub_dir.glob("*.pkl")) - if not atoms_pkl_files: - continue - - with open(atoms_pkl_files[0], "rb") as f: - ase_atoms = pickle.load(f) - - conv_data = extract_conv(fdmnes_output_dir=sub_dir) - ase_atoms.info.update({"FDMNES-xanes": conv_data}) - expanded_atoms_list.append(ase_atoms) - - with open(root_dir / "atoms_db_expanded.pkl", "wb") as f: - pickle.dump(expanded_atoms_list, f) - - logger.info( - "Saved %d expanded structures to %s", - len(expanded_atoms_list), - root_dir / "atoms_db_expanded.pkl", - ) - - -def plot_xanes_results(root_dir: Path, runs_dir: Path) -> dict: - """Generate normalized XANES plots for completed FDMNES calculations. - - For each run directory containing a ``*_conv.txt`` file, produces - a ``xanes_plot.png`` with the normalized absorption spectrum. - - Parameters - ---------- - root_dir : Path - Root data directory (unused currently, reserved for summary plots). - runs_dir : Path - Directory containing ``run_*`` subdirectories with FDMNES outputs. - - Returns - ------- - dict - plot_files : list[str] — absolute paths to generated plot images - n_plots : int — number of plots successfully generated - n_failed : int — number of runs that failed to plot - failed : list[str] — names of run directories that failed - """ - import matplotlib.pyplot as plt - - logger.info("Plotting XANES results from %s", runs_dir) - - plot_files = [] - failed = [] - - for sub_dir in sorted(runs_dir.glob("run_*")): - conv_file = next(sub_dir.glob("*_conv.txt"), None) - if conv_file: - try: - norm_energy, _raw = get_normalized_xanes(conv_file) - plot_path = sub_dir / "xanes_plot.png" - plt.figure() - plt.plot(norm_energy[:, 0], norm_energy[:, 1], label=sub_dir.name) - plt.xlabel("Energy [eV]") - plt.ylabel("Normalized Absorption") - plt.title(f"XANES for {sub_dir.name}") - plt.legend() - plt.savefig(plot_path, dpi=150) - plt.close() - plot_files.append(str(plot_path)) - logger.info("Plotted %s", sub_dir.name) - except Exception as e: - logger.error("Failed to plot %s: %s", sub_dir.name, e) - failed.append(sub_dir.name) - - return { - "plot_files": plot_files, - "n_plots": len(plot_files), - "n_failed": len(failed), - "failed": failed, - } - - -# ----------------------------------------------------------------------------- -# Data directory helper -# ----------------------------------------------------------------------------- - - -def _get_data_dir() -> Path: - """Return the working data directory for XANES workflows.""" - cwd = Path.cwd() - if "PBS_O_WORKDIR" in os.environ: - cwd = Path(os.environ["PBS_O_WORKDIR"]) - - data_dir = cwd / "xanes_data" - if not data_dir.exists(): - data_dir.mkdir(parents=True) - return data_dir - - -@tool -def run_xanes(params: xanes_input_schema) -> str: - """Run a single XANES/FDMNES calculation for one structure file. - - This tool reads the structure, generates FDMNES input files, runs FDMNES, - and returns the result status. Requires the FDMNES_EXE environment variable. + str + Human-readable completion summary. """ result = run_xanes_core(params) if result["status"] == "success": @@ -578,6 +79,16 @@ def fetch_xanes_data(params: mp_query_schema) -> str: Requires a Materials Project API key via the mp_api_key parameter or the MP_API_KEY environment variable. + + Parameters + ---------- + params : mp_query_schema + Materials Project query parameters. + + Returns + ------- + str + Human-readable fetch summary. """ data_dir = _get_data_dir() result = fetch_materials_project_data(params, data_dir) diff --git a/src/chemgraph/utils/async_utils.py b/src/chemgraph/utils/async_utils.py index 5b438fe4..bdf0da2f 100644 --- a/src/chemgraph/utils/async_utils.py +++ b/src/chemgraph/utils/async_utils.py @@ -13,6 +13,16 @@ def run_async_callable(fn: Callable[..., Any]) -> Any: If no event loop is running, uses ``asyncio.run`` directly. Otherwise, spawns a daemon thread so that the call does not conflict with an already-running loop (e.g. inside Streamlit). + + Parameters + ---------- + fn : Callable[..., Any] + Zero-argument callable that returns an awaitable. + + Returns + ------- + Any + Result of the awaited callable. """ try: asyncio.get_running_loop() @@ -23,6 +33,7 @@ def run_async_callable(fn: Callable[..., Any]) -> Any: error_container: dict[str, Exception] = {} def runner() -> None: + """Run the awaitable in a background event loop.""" try: result_container["value"] = asyncio.run(fn()) except Exception as exc: diff --git a/src/chemgraph/utils/config_utils.py b/src/chemgraph/utils/config_utils.py index 7abb5334..8d5706a1 100644 --- a/src/chemgraph/utils/config_utils.py +++ b/src/chemgraph/utils/config_utils.py @@ -19,7 +19,18 @@ def flatten_config(config: Dict[str, Any]) -> Dict[str, Any]: - """Flatten nested TOML-like config into top-level keys used by the CLI.""" + """Flatten nested TOML-like config into top-level keys used by the CLI. + + Parameters + ---------- + config : dict[str, Any] + Nested configuration dictionary. + + Returns + ------- + dict[str, Any] + Flattened configuration with section names included in keys. + """ flattened: Dict[str, Any] = {} if "general" in config: @@ -50,7 +61,18 @@ def flatten_config(config: Dict[str, Any]) -> Dict[str, Any]: def normalize_openai_base_url(base_url: Optional[str]) -> Optional[str]: - """Normalize Argo-style URLs to OpenAI-compatible /v1 URLs.""" + """Normalize Argo-style URLs to OpenAI-compatible /v1 URLs. + + Parameters + ---------- + base_url : str, optional + Provider base URL. + + Returns + ------- + str or None + Normalized URL, or ``None`` when no URL was provided. + """ if not base_url: return base_url if ( @@ -67,7 +89,20 @@ def normalize_openai_base_url(base_url: Optional[str]) -> Optional[str]: def get_base_url_for_model_from_nested_config( model_name: str, config: Dict[str, Any] ) -> Optional[str]: - """Resolve provider base URL using nested config structure.""" + """Resolve provider base URL using nested config structure. + + Parameters + ---------- + model_name : str + Model identifier. + config : dict[str, Any] + Nested configuration dictionary. + + Returns + ------- + str or None + Matching provider base URL, or ``None`` when not configured. + """ api = config.get("api", {}) if model_name in supported_argo_models: @@ -90,7 +125,20 @@ def get_base_url_for_model_from_nested_config( def get_base_url_for_model_from_flat_config( model_name: str, config: Dict[str, Any] ) -> Optional[str]: - """Resolve provider base URL using flattened config keys.""" + """Resolve provider base URL using flattened config keys. + + Parameters + ---------- + model_name : str + Model identifier. + config : dict[str, Any] + Flattened configuration dictionary. + + Returns + ------- + str or None + Matching provider base URL, or ``None`` when not configured. + """ if model_name in supported_argo_models: return normalize_openai_base_url( config.get("api_openai_base_url") or ARGO_DEFAULT_BASE_URL @@ -113,6 +161,16 @@ def get_model_options_for_nested_config(config: Dict[str, Any]) -> list[str]: Always show all curated models so users can switch providers from the UI. If Argo endpoint is configured, prioritize Argo model IDs at the top. + + Parameters + ---------- + config : dict[str, Any] + Nested configuration dictionary. + + Returns + ------- + list[str] + Model identifiers for UI selection. """ base_url = config.get("api", {}).get("openai", {}).get("base_url") if base_url and "argoapi" in base_url: @@ -122,7 +180,18 @@ def get_model_options_for_nested_config(config: Dict[str, Any]) -> list[str]: def get_argo_user_from_nested_config(config: Dict[str, Any]) -> Optional[str]: - """Resolve Argo user from nested config (if configured).""" + """Resolve Argo user from nested config. + + Parameters + ---------- + config : dict[str, Any] + Nested configuration dictionary. + + Returns + ------- + str or None + Configured Argo username, or ``None``. + """ value = config.get("api", {}).get("openai", {}).get("argo_user") if isinstance(value, str): value = value.strip() @@ -130,7 +199,18 @@ def get_argo_user_from_nested_config(config: Dict[str, Any]) -> Optional[str]: def get_argo_user_from_flat_config(config: Dict[str, Any]) -> Optional[str]: - """Resolve Argo user from flattened config (if configured).""" + """Resolve Argo user from flattened config. + + Parameters + ---------- + config : dict[str, Any] + Flattened configuration dictionary. + + Returns + ------- + str or None + Configured Argo username, or ``None``. + """ value = config.get("api_openai_argo_user") if isinstance(value, str): value = value.strip() diff --git a/src/chemgraph/utils/get_workflow_from_llm.py b/src/chemgraph/utils/get_workflow_from_llm.py index 51837657..3376464e 100644 --- a/src/chemgraph/utils/get_workflow_from_llm.py +++ b/src/chemgraph/utils/get_workflow_from_llm.py @@ -85,6 +85,13 @@ def get_workflow_from_state(state) -> dict: workflow_dict = {"tool_calls": []} def recurse(obj): + """Collect AI tool calls from nested workflow state objects. + + Parameters + ---------- + obj : Any + Dictionary, list, or scalar value from the serialized workflow. + """ if isinstance(obj, dict): # Extract tool_calls if it's an AI message if obj.get("type") == "ai": diff --git a/src/chemgraph/utils/parsing.py b/src/chemgraph/utils/parsing.py index 228cee42..ce9d8083 100644 --- a/src/chemgraph/utils/parsing.py +++ b/src/chemgraph/utils/parsing.py @@ -18,6 +18,16 @@ def extract_json_block(text: str) -> str | None: Handles markdown-fenced blocks (```json ... ```) and bare JSON objects. Returns the extracted string or *None* if nothing looks like JSON. + + Parameters + ---------- + text : str + Text that may contain a JSON object. + + Returns + ------- + str or None + Extracted JSON text, or ``None`` when no object is found. """ # Try markdown-fenced JSON first m = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL) @@ -40,6 +50,11 @@ def parse_response_formatter( fields ``None``) so the pipeline never breaks -- the raw text is still available in the agent's message history. + Parameters + ---------- + raw_text : str + Raw LLM output to parse. + Returns ------- tuple[ResponseFormatter, str | None] diff --git a/src/chemgraph/utils/tool_call_eval.py b/src/chemgraph/utils/tool_call_eval.py index e8cd0abb..21639204 100644 --- a/src/chemgraph/utils/tool_call_eval.py +++ b/src/chemgraph/utils/tool_call_eval.py @@ -1,7 +1,7 @@ """Module for quick LLM evaluations""" from deepdiff import DeepDiff -from chemgraph.models.ase_input import ASEInputSchema +from chemgraph.schemas.ase_input import ASEInputSchema def remove_ignored_fields(obj, ignored_keys=("cell", "pbc")): @@ -27,9 +27,21 @@ def remove_ignored_fields(obj, ignored_keys=("cell", "pbc")): def apply_defaults(args: dict, schema: dict) -> dict: - """ - Recursively fills missing fields with default values based on a JSON-like schema. + """Fill missing fields with default values from a JSON-like schema. + Handles nested objects and anyOf/default combinations. + + Parameters + ---------- + args : dict + Tool-call arguments to augment. + schema : dict + JSON schema containing properties and defaults. + + Returns + ------- + dict + Arguments with applicable default values filled in. """ if not isinstance(args, dict): return args # Only process dicts @@ -65,7 +77,18 @@ def apply_defaults(args: dict, schema: dict) -> dict: def lowercase_dict(obj): - """Recursively lowercases string keys and string values in a dict/list structure.""" + """Recursively lowercase string keys and string values. + + Parameters + ---------- + obj : Any + Dictionary, list, string, or scalar value to normalize. + + Returns + ------- + Any + Normalized object with lowercased string keys/values. + """ if isinstance(obj, dict): return {(k.lower() if isinstance(k, str) else k): lowercase_dict(v) for k, v in obj.items()} elif isinstance(obj, list): diff --git a/src/ui/_pages/configuration.py b/src/ui/_pages/configuration.py index 96c1b5f5..ebf7094e 100644 --- a/src/ui/_pages/configuration.py +++ b/src/ui/_pages/configuration.py @@ -28,12 +28,36 @@ def normalize_workflow_name(value: str) -> str: + """Normalize workflow aliases to internal workflow names. + + Parameters + ---------- + value : str + Workflow name or alias from configuration/UI state. + + Returns + ------- + str + Canonical workflow name. + """ if not value: return value return WORKFLOW_ALIASES.get(value, value) def get_model_options(config: Dict[str, Any]) -> list: + """Return model options for the configuration UI. + + Parameters + ---------- + config : dict[str, Any] + Nested UI configuration dictionary. + + Returns + ------- + list + Model names shown in the model selector. + """ from chemgraph.utils.config_utils import get_model_options_for_nested_config return get_model_options_for_nested_config(config) @@ -91,6 +115,13 @@ def render() -> None: def _render_general_settings(config: dict) -> None: + """Render and update general configuration widgets. + + Parameters + ---------- + config : dict + Mutable draft configuration dictionary. + """ st.subheader("General Settings") col1, col2 = st.columns(2) @@ -152,6 +183,12 @@ def _render_general_settings(config: dict) -> None: value=config["general"]["report"], key="config_report", ) + config["general"]["human_supervised"] = st.checkbox( + "Human Supervised", + value=config["general"].get("human_supervised", False), + key="config_human_supervised", + help="Enable the ask_human tool so the agent can pause and request human input.", + ) config["general"]["verbose"] = st.checkbox( "Verbose Output", value=config["general"]["verbose"], @@ -246,6 +283,13 @@ def _render_general_settings(config: dict) -> None: def _render_api_settings(config: dict) -> None: + """Render and update API configuration widgets. + + Parameters + ---------- + config : dict + Mutable draft configuration dictionary. + """ st.subheader("API Settings") st.markdown("**API Keys (Session Only)**") @@ -422,6 +466,13 @@ def _render_api_settings(config: dict) -> None: def _render_raw_toml(config: dict) -> None: + """Render raw TOML editor for the draft configuration. + + Parameters + ---------- + config : dict + Mutable draft configuration dictionary. + """ st.subheader("Raw TOML Configuration") st.markdown( """ @@ -455,6 +506,13 @@ def _render_raw_toml(config: dict) -> None: def _render_action_buttons(config: dict) -> None: + """Render save/reload/reset/download configuration actions. + + Parameters + ---------- + config : dict + Mutable draft configuration dictionary. + """ st.markdown("---") col1, col2, col3, col4 = st.columns(4) @@ -495,6 +553,13 @@ def _render_action_buttons(config: dict) -> None: def _render_config_summary(config: dict) -> None: + """Render a compact summary of the draft configuration. + + Parameters + ---------- + config : dict + Draft configuration dictionary. + """ with st.expander("\U0001f4ca Configuration Summary", expanded=False): st.write("**Current Configuration:**") st.write(f"- Model: {config['general']['model']}") diff --git a/src/ui/_pages/main_interface.py b/src/ui/_pages/main_interface.py index 0d53b0da..6346f41f 100644 --- a/src/ui/_pages/main_interface.py +++ b/src/ui/_pages/main_interface.py @@ -1,8 +1,13 @@ """Main chat interface page for ChemGraph.""" -import html as html_mod +import asyncio import logging import os +import pprint +import queue +import threading +import uuid +from datetime import datetime from pathlib import Path from typing import Any, Dict, Optional @@ -10,23 +15,25 @@ import streamlit as st from ase.io import read as ase_read +from chemgraph.agent.llm_agent import HumanInputRequired from chemgraph.memory.store import SessionStore from chemgraph.models.supported_models import supported_argo_models +from chemgraph.schemas.ase_input import ( + get_available_calculator_names, + get_default_calculator_name, +) from chemgraph.utils.config_utils import ( get_argo_user_from_nested_config, get_base_url_for_model_from_nested_config, - get_model_options_for_nested_config, ) -from ui.agent_manager import initialize_agent, run_async_callable +from ui.agent_manager import initialize_agent from ui.branding import LOGO_IMAGES, first_existing_asset from ui.config import load_config from ui.endpoint import check_local_model_endpoint from ui.file_utils import ( - changed_recently, extract_log_dir_from_messages, find_latest_xyz_file_in_dir, - resolve_output_path, ) from ui.message_utils import ( extract_messages_from_result, @@ -37,6 +44,7 @@ has_structure_signal, is_infrared_requested, normalize_message_content, + split_markdown_latex_blocks, strip_viewer_from_report_html, ) from ui.session_utils import ( @@ -63,11 +71,76 @@ def _get_base_url_for_model(model_name: str, config: Dict[str, Any]) -> Optional[str]: + """Resolve the configured base URL for a model. + + Parameters + ---------- + model_name : str + Selected model identifier. + config : dict[str, Any] + Nested UI configuration. + + Returns + ------- + str or None + Provider base URL, or ``None`` when not configured. + """ return get_base_url_for_model_from_nested_config(model_name, config) -def _get_model_options(config: Dict[str, Any]) -> list: - return get_model_options_for_nested_config(config) +def _initial_ui_log_root() -> str: + """Return the root directory for per-chat UI artifacts.""" + env_log_dir = os.environ.get("CHEMGRAPH_LOG_DIR") + if env_log_dir: + path = Path(env_log_dir).expanduser() + if path.name.startswith(("session_", "ui_session_")): + path = path.parent + return str(path.resolve()) + return str((Path.cwd() / "cg_logs").resolve()) + + +def _ensure_chat_log_dir() -> str: + """Create and activate a log directory owned by the current chat.""" + if not st.session_state.get("ui_log_root"): + st.session_state.ui_log_root = _initial_ui_log_root() + + chat_log_dir = st.session_state.get("current_chat_log_dir") + if not chat_log_dir: + timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + suffix = str(uuid.uuid4())[:8] + chat_log_dir = str( + Path(st.session_state.ui_log_root) / f"ui_session_{timestamp}_{suffix}" + ) + st.session_state.current_chat_log_dir = chat_log_dir + + os.makedirs(chat_log_dir, exist_ok=True) + os.environ["CHEMGRAPH_LOG_DIR"] = chat_log_dir + return chat_log_dir + + +def _resolve_structured_output_for_model( + model_name: str, structured_output: bool +) -> tuple[bool, Optional[str]]: + """Disable structured output for Argo models, including quick overrides. + + Parameters + ---------- + model_name : str + Selected model identifier. + structured_output : bool + Requested structured-output setting. + + Returns + ------- + tuple[bool, str | None] + Effective structured-output setting and optional warning message. + """ + if model_name in supported_argo_models and structured_output: + return ( + False, + "Structured output is disabled for Argo models to avoid JSON parsing errors.", + ) + return structured_output, None # --------------------------------------------------------------------------- @@ -85,15 +158,9 @@ def render() -> None: selected_output = config["general"]["output"] structured_output = config["general"]["structured"] generate_report = config["general"]["report"] + human_supervised = config["general"].get("human_supervised", False) thread_id = config["general"]["thread"] - # Argo models: disable structured output - if selected_model in supported_argo_models and structured_output: - structured_output = False - st.session_state.ui_notice = ( - "Structured output is disabled for Argo models to avoid JSON parsing errors." - ) - # ----- Header ----- logo_image = first_existing_asset(LOGO_IMAGES) if logo_image: @@ -101,17 +168,23 @@ def render() -> None: else: st.title("\U0001f9ea ChemGraph") - st.markdown( - """ + st.markdown(""" ChemGraph enables you to perform various **computational chemistry** tasks with natural-language queries using AI agents. - """ - ) + """) + + # ----- Calculator availability sidebar ----- + _render_available_calculators_sidebar() + _render_chat_controls() - # ----- Quick settings sidebar ----- - selected_model, thread_id = _render_quick_settings( - config, selected_model, thread_id + structured_output, ui_notice = _resolve_structured_output_for_model( + selected_model, structured_output ) + st.session_state.ui_notice = ui_notice + st.session_state.active_model = selected_model + st.session_state.active_workflow = selected_workflow + if ui_notice: + st.info(ui_notice) selected_base_url = _get_base_url_for_model(selected_model, config) endpoint_status = check_local_model_endpoint(selected_base_url) @@ -126,9 +199,7 @@ def render() -> None: st.rerun() # ----- Agent status sidebar ----- - _render_agent_status( - selected_model, selected_workflow, thread_id, endpoint_status - ) + _render_agent_status(selected_model, selected_workflow, thread_id, endpoint_status) # ----- Auto-initialize agent ----- _auto_initialize_agent( @@ -138,22 +209,41 @@ def render() -> None: structured_output, selected_output, generate_report, + human_supervised, selected_base_url, ) # ----- Conversation history ----- _render_conversation_history(thread_id) - # ----- Query input ----- - query = _render_query_input(config, selected_model) + # ----- Pending interrupt display ----- + _render_pending_interrupt() + + # ----- Example queries ----- + _render_example_queries(config, selected_model) - # ----- Submit ----- - _handle_query_submission( - query, thread_id, endpoint_status, selected_base_url + # ----- Chat input (handles both normal queries and interrupt responses) ----- + is_interrupt = st.session_state.pending_human_question is not None + prompt = st.chat_input( + ( + "Type your response..." + if is_interrupt + else "Ask a computational chemistry question..." + ), ) - # ----- Footer ----- - _render_footer() + # Check for example query submitted via button click + example_query = st.session_state.pop("_pending_example_query", None) + if example_query: + prompt = example_query + + if prompt: + if is_interrupt: + _handle_human_response(prompt, thread_id) + else: + _handle_query_submission( + prompt, thread_id, endpoint_status, selected_base_url + ) # --------------------------------------------------------------------------- @@ -161,40 +251,79 @@ def render() -> None: # --------------------------------------------------------------------------- -def _render_quick_settings( - config: dict, selected_model: str, thread_id: int -) -> tuple[str, int]: - with st.sidebar.expander("\U0001f527 Quick Settings"): - st.write("Override settings for this session:") - - if st.checkbox("Override Model"): - model_options = _get_model_options(config) - selected_model = st.selectbox( - "Select Model", - model_options, - index=( - model_options.index(selected_model) - if selected_model in model_options - else 0 - ), - ) - quick_custom_model = st.text_input( - "Custom model ID (optional)", - value="", - key="quick_custom_model", - help="If set, this overrides the selected model for this session.", - ).strip() - if quick_custom_model: - selected_model = quick_custom_model - - if st.checkbox("Override Thread ID"): - thread_id = st.number_input( - "Thread ID", min_value=1, max_value=1000, value=thread_id - ) +def _render_markdown_with_math(text: str) -> None: + """Render Markdown text, sending display math blocks through ``st.latex``. - st.info("\U0001f4a1 To make permanent changes, use the Configuration page.") + Parameters + ---------- + text : str + Markdown text that may contain display math blocks. + """ + for block_type, content in split_markdown_latex_blocks(text): + if block_type == "latex": + st.latex(_prepare_latex_block(content)) + else: + st.markdown(content) - return selected_model, thread_id + +def _prepare_latex_block(content: str) -> str: + """Clean display math for Streamlit's KaTeX renderer. + + Parameters + ---------- + content : str + Raw LaTeX block content. + + Returns + ------- + str + KaTeX-compatible display math. + """ + lines = [line.strip() for line in content.splitlines() if line.strip()] + if not lines: + return "" + if len(lines) == 1 or r"\begin{" in content: + return " ".join(lines) + return "\\begin{aligned}\n" + (r" \\" + "\n").join(lines) + "\n\\end{aligned}" + + +def _format_calculator_label(calculator_name: str) -> str: + """Format calculator class names for display. + + Parameters + ---------- + calculator_name : str + Calculator class name or label. + + Returns + ------- + str + Human-readable calculator label. + """ + label = calculator_name.removesuffix("Calc") + if label == "TBLite": + return "TBLite (xTB, GFN1-xTB, GFN2-xTB)" + return label + + +def _render_available_calculators_sidebar() -> None: + """Render the calculators detected during ChemGraph initialization.""" + available = get_available_calculator_names() + default = get_default_calculator_name() + + with st.sidebar.expander("\U0001f9ee Available Calculators", expanded=True): + st.caption("Detected during ChemGraph initialization.") + + for calculator_name in available: + label = _format_calculator_label(calculator_name) + if calculator_name == default: + st.success(f"{label} (default)") + else: + st.markdown(f"- {label}") + + st.caption( + "The agent uses this list when choosing calculators for ASE simulations." + ) def _start_new_chat() -> None: @@ -206,6 +335,19 @@ def _start_new_chat() -> None: st.session_state.last_run_error = None st.session_state.last_run_result = None st.session_state.last_run_query = None + st.session_state.pop("_pending_example_query", None) + st.session_state.agent = None + st.session_state.last_config = None + st.session_state.current_chat_log_dir = None + os.environ.pop("CHEMGRAPH_LOG_DIR", None) + _clear_interrupt_state() + + +def _render_chat_controls() -> None: + """Render chat-level actions that must be available even without memory.""" + if st.sidebar.button("New Chat", key="new_chat_btn", use_container_width=True): + _start_new_chat() + st.rerun() def _render_session_sidebar() -> None: @@ -215,11 +357,6 @@ def _render_session_sidebar() -> None: return with st.sidebar.expander("\U0001f4c2 Sessions", expanded=False): - # New Chat button - if st.button("\u2795 New Chat", key="new_chat_btn", use_container_width=True): - _start_new_chat() - st.rerun() - # Show current session info current_sid = st.session_state.get("current_session_id") if current_sid: @@ -270,12 +407,20 @@ def _render_session_sidebar() -> None: if current_sid == s.session_id: _start_new_chat() except Exception as exc: - logger.warning("Failed to delete session %s: %s", s.session_id, exc) + logger.warning( + "Failed to delete session %s: %s", s.session_id, exc + ) st.rerun() def _load_session(session_id: str) -> None: - """Load a stored session into the active conversation.""" + """Load a stored session into the active conversation. + + Parameters + ---------- + session_id : str + Session ID or prefix selected in the sidebar. + """ store: Optional[SessionStore] = st.session_state.get("session_store") if store is None: return @@ -293,20 +438,44 @@ def _load_session(session_id: str) -> None: st.session_state.last_run_error = None st.session_state.last_run_result = None st.session_state.last_run_query = None + st.session_state.current_chat_log_dir = session.log_dir + if session.log_dir: + os.environ["CHEMGRAPH_LOG_DIR"] = session.log_dir + + +def _active_session_metadata() -> tuple[str, str]: + """Return model/workflow metadata matching the active UI run.""" + config = st.session_state.config + model = ( + st.session_state.get("pending_interrupt_model") + or st.session_state.get("active_model") + or config["general"]["model"] + ) + workflow = ( + st.session_state.get("pending_interrupt_workflow") + or st.session_state.get("active_workflow") + or config["general"]["workflow"] + ) + return model, normalize_workflow_name(workflow) def _save_exchange_to_store(query: str, result: Any) -> None: """Persist a single query/result exchange to the SessionStore. Creates the session DB row on the first call, then appends messages. + + Parameters + ---------- + query : str + User query text. + result : Any + Agent result to persist as session messages. """ store: Optional[SessionStore] = st.session_state.get("session_store") if store is None: return - config = st.session_state.config - model = config["general"]["model"] - workflow = normalize_workflow_name(config["general"]["workflow"]) + model, workflow = _active_session_metadata() try: # Create the session row on the first exchange @@ -319,6 +488,7 @@ def _save_exchange_to_store(query: str, result: Any) -> None: model_name=model, workflow_type=workflow, title=title, + log_dir=st.session_state.get("current_chat_log_dir"), ) st.session_state.session_created = True @@ -326,9 +496,7 @@ def _save_exchange_to_store(query: str, result: Any) -> None: entry = {"query": query, "result": result} messages = conversation_entry_to_messages(entry) if messages: - store.save_messages( - st.session_state.current_session_id, messages - ) + store.save_messages(st.session_state.current_session_id, messages) except Exception as exc: # Best-effort persistence -- don't break the UI. logger.warning("Failed to save exchange to session store: %s", exc) @@ -340,7 +508,20 @@ def _render_agent_status( thread_id: int, endpoint_status: dict, ) -> None: - st.sidebar.header("\U0001f171\U0001f172 Agent Status") + """Render sidebar status for the active agent. + + Parameters + ---------- + selected_model : str + Selected model name. + selected_workflow : str + Selected workflow name. + thread_id : int + Current LangGraph thread ID. + endpoint_status : dict + Local endpoint status dictionary. + """ + st.sidebar.header("Agent Status") if st.session_state.agent: st.sidebar.success("\u2705 Agents Ready") @@ -354,11 +535,15 @@ def _render_agent_status( st.sidebar.caption(f"LLM endpoint: {endpoint_status['message']}") else: st.sidebar.error(f"LLM endpoint issue: {endpoint_status['message']}") + if st.session_state.pending_human_question is not None: + st.sidebar.warning("Waiting for your input...") if st.session_state.last_run_error: st.sidebar.error("Last run error (see verbose info).") if st.sidebar.button("\U0001f504 Refresh Agents"): st.session_state.agent = None + # Checkpoint is lost on re-init, so clear interrupt state + _clear_interrupt_state() st.rerun() else: st.sidebar.error("\u274c Agents Not Ready") @@ -381,92 +566,153 @@ def _auto_initialize_agent( structured_output: bool, selected_output: str, generate_report: bool, + human_supervised: bool, selected_base_url: Optional[str], ) -> None: + """Initialize or refresh the cached Streamlit agent when config changes. + + Parameters + ---------- + config : dict + Nested UI configuration. + selected_model : str + Selected model name. + selected_workflow : str + Selected workflow name. + structured_output : bool + Effective structured-output setting. + selected_output : str + Agent return mode. + generate_report : bool + Whether report generation is enabled. + human_supervised : bool + Whether human-supervision tools are enabled. + selected_base_url : str, optional + Model endpoint URL. + """ current_config = ( selected_model, selected_workflow, structured_output, selected_output, generate_report, + human_supervised, config["general"]["recursion_limit"], selected_base_url, get_argo_user_from_nested_config(config), + st.session_state.get("current_chat_log_dir"), ) - if ( - st.session_state.agent is None - or st.session_state.last_config != current_config - ): + if st.session_state.agent is None or st.session_state.last_config != current_config: with st.spinner("\U0001f680 Initializing ChemGraph agents..."): - st.session_state.agent = initialize_agent( + chat_log_dir = _ensure_chat_log_dir() + agent = initialize_agent( selected_model, selected_workflow, structured_output, selected_output, generate_report, + human_supervised, config["general"]["recursion_limit"], selected_base_url, get_argo_user_from_nested_config(config), + log_dir=chat_log_dir, ) - st.session_state.last_config = current_config + st.session_state.agent = agent + if agent is not None: + st.session_state.last_config = ( + selected_model, + selected_workflow, + structured_output, + selected_output, + generate_report, + human_supervised, + config["general"]["recursion_limit"], + selected_base_url, + get_argo_user_from_nested_config(config), + chat_log_dir, + ) + else: + st.session_state.last_config = None def _render_conversation_history(thread_id: int) -> None: + """Render all saved conversation exchanges. + + Parameters + ---------- + thread_id : int + Current LangGraph thread ID. + """ if not st.session_state.conversation_history: return - st.subheader("\U0001f5e8\ufe0f Conversation History") - for idx, entry in enumerate(st.session_state.conversation_history, 1): _render_single_exchange(idx, entry, thread_id) - st.markdown("---") def _render_single_exchange(idx: int, entry: dict, thread_id: int) -> None: - """Render one user-query / agent-response exchange.""" - # User bubble - st.markdown( - f""" -
- \U0001f464 You:
{html_mod.escape(entry["query"])} -
""", - unsafe_allow_html=True, - ) + """Render one user-query / agent-response exchange. + + Parameters + ---------- + idx : int + One-based exchange index. + entry : dict + Conversation-history entry. + thread_id : int + Current LangGraph thread ID. + """ + # User message + with st.chat_message("user"): + st.markdown(entry["query"]) + + # Interrupt exchanges (if any occurred during this query) + for exch in entry.get("interrupt_exchanges", []): + with st.chat_message("assistant"): + _render_markdown_with_math(exch["question"]) + with st.chat_message("user"): + st.markdown(exch["answer"]) messages = extract_messages_from_result(entry["result"]) # Find final AI response final_answer = _extract_final_answer(messages) - # Display the AI response - if final_answer: - st.markdown( - f""" -
- \U0001f171\U0001f172 ChemGraph:
{html_mod.escape(final_answer).replace(chr(10), "
")} -
""", - unsafe_allow_html=True, - ) + # Display the AI response with visualizations + with st.chat_message("assistant"): + if final_answer: + _render_markdown_with_math(final_answer) - # Structure visualisation - html_filename = find_html_filename(messages) - _render_structure_section(idx, messages, final_answer, entry, html_filename) + # Structure visualisation + html_filename = find_html_filename(messages) + _render_structure_section(idx, messages, final_answer, entry, html_filename) - # HTML report - if html_filename: - _render_html_report(html_filename, messages) + # HTML report + if html_filename: + _render_html_report(idx, html_filename, messages, entry) - # IR spectrum - if is_infrared_requested(messages): - _render_ir_spectrum(idx) + # IR spectrum + if is_infrared_requested(messages): + _render_ir_spectrum(idx, messages, entry) - # Debug expander - _render_verbose_info(idx, messages, entry) + # Debug expander + _render_verbose_info(idx, messages, entry) def _extract_final_answer(messages: list) -> str: - """Walk messages in reverse to find the last non-JSON AI message.""" + """Walk messages in reverse to find the last non-JSON AI message. + + Parameters + ---------- + messages : list + Message-like objects or dictionaries. + + Returns + ------- + str + Final displayable answer text, or an empty string. + """ final_answer = "" for message in reversed(messages): if hasattr(message, "content") and hasattr(message, "type"): @@ -490,9 +736,7 @@ def _extract_final_answer(messages: list) -> str: final_answer = content break elif hasattr(message, "content"): - content = normalize_message_content( - getattr(message, "content", "") - ).strip() + content = normalize_message_content(getattr(message, "content", "")).strip() if content and not ( content.startswith("{") and content.endswith("}") @@ -510,6 +754,21 @@ def _render_structure_section( entry: dict, html_filename: Optional[str], ) -> None: + """Render molecular structure artifacts for an exchange. + + Parameters + ---------- + idx : int + One-based exchange index. + messages : list + Message-like objects from the exchange. + final_answer : str + Final assistant answer text. + entry : dict + Conversation-history entry. + html_filename : str, optional + HTML report path/filename, if detected. + """ structure = find_structure_in_messages(messages) if structure: display_molecular_structure( @@ -527,7 +786,7 @@ def _render_structure_section( ) elif not html_filename: if has_structure_signal(messages, entry.get("query", ""), final_answer): - log_dir = extract_log_dir_from_messages(messages) + log_dir = _artifact_log_dir(messages, entry) if log_dir and os.path.isdir(log_dir): latest_xyz = find_latest_xyz_file_in_dir(log_dir) if latest_xyz: @@ -542,10 +801,28 @@ def _render_structure_section( st.warning(f"Failed to load XYZ structure: {exc}") -def _render_html_report(html_filename: str, messages: list) -> None: +def _render_html_report( + idx: int, html_filename: str, messages: list, entry: dict +) -> None: + """Render an HTML report expander and download button. + + Parameters + ---------- + idx : int + One-based exchange index. + html_filename : str + HTML report path or filename. + messages : list + Message-like objects from the exchange. + entry : dict + Conversation-history entry. + """ with st.expander("\U0001f4ca Report", expanded=False): try: - resolved_html = resolve_output_path(html_filename) + resolved_html = _resolve_artifact_path( + html_filename, + _artifact_log_dir(messages, entry), + ) with open(resolved_html, "r", encoding="utf-8") as f: html_content = f.read() @@ -554,10 +831,17 @@ def _render_html_report(html_filename: str, messages: list) -> None: display_molecular_structure( report_structure["atomic_numbers"], report_structure["positions"], - title="Molecular Structure", + title=f"Molecular Structure (Report {idx})", ) cleaned_html = strip_viewer_from_report_html(html_content) + st.download_button( + "Download HTML Report", + data=html_content, + file_name=Path(resolved_html).name, + mime="text/html", + key=f"download_report_{idx}", + ) st.components.v1.html(cleaned_html, height=600, scrolling=True) except FileNotFoundError: st.warning(f"HTML file '{html_filename}' not found") @@ -565,104 +849,208 @@ def _render_html_report(html_filename: str, messages: list) -> None: st.error(f"Error displaying HTML: {e}") -def _render_ir_spectrum(idx: int) -> None: - if changed_recently(): - with st.expander("\U0001f50d IR Spectrum", expanded=True): - col1, col2 = st.columns(2, border=True) - - with col1: - ir_path = resolve_output_path("ir_spectrum.png") - if os.path.exists(ir_path): - st.image(ir_path) - else: - st.warning("IR spectrum plot not found.") - - with col2: - freq_path = resolve_output_path("frequencies.csv") - if not os.path.exists(freq_path): - st.warning("Frequencies file not found.") - else: - df = pd.read_csv( - freq_path, - index_col=False, - names=["filename", "frequency"], - ).iloc[6:] - - if not df.empty: - st.write("**Select a frequency to visualize:**") - freq_options = { - f"{float(row['frequency'].strip('i')):.2f} cm\u207b\u00b9": i - for i, row in df.iterrows() - } - selected_freq = st.selectbox( - "Frequency", - list(freq_options.keys()), - index=0, - key=f"ir_frequency_select_{idx}", - ) - traj_file = df.loc[freq_options[selected_freq]]["filename"] - traj_path = resolve_output_path(traj_file) - if not os.path.exists(traj_path): - st.warning( - f"Trajectory file '{traj_file}' not found." - ) - elif not STMOL_AVAILABLE: - st.info( - "3D viewer not available; install stmol to animate trajectories." - ) - else: - import stmol - from ase.io.trajectory import Trajectory - - traj = Trajectory(traj_path) - view = visualize_trajectory(traj) - view.zoomTo() - stmol.showmol(view, height=400, width=500) - else: - st.warning("No vibrational frequencies found.") - else: +def _artifact_log_dir(messages: list, entry: dict) -> Optional[str]: + """Return the log directory tied to a specific conversation entry. + + Parameters + ---------- + messages : list + Message-like objects from the exchange. + entry : dict + Conversation-history entry. + + Returns + ------- + str or None + Artifact/log directory, if found. + """ + entry_log_dir = entry.get("log_dir") + if entry_log_dir: + return entry_log_dir + return extract_log_dir_from_messages(messages) + + +def _latest_artifact_path(directory: Optional[str], pattern: str) -> Optional[str]: + """Return the newest shallow match for an output artifact pattern. + + Parameters + ---------- + directory : str, optional + Directory to search. + pattern : str + Glob pattern to match. + + Returns + ------- + str or None + Newest matching file path, or ``None``. + """ + if not directory or not os.path.isdir(directory): + return None + + candidates: list[Path] = [] + try: + candidates.extend(path for path in Path(directory).glob(pattern) if path.is_file()) + except OSError: + return None + + if not candidates: + return None + return str(max(candidates, key=lambda path: path.stat().st_mtime)) + + +def _resolve_artifact_path(filename: str, directory: Optional[str]) -> str: + """Resolve an artifact path relative to its run directory when known. + + Parameters + ---------- + filename : str + Absolute or relative artifact path. + directory : str, optional + Run artifact directory. + + Returns + ------- + str + Resolved artifact path. + """ + if os.path.isabs(filename): + return filename + if directory: + return str(Path(directory) / filename) + return filename + + +def _render_ir_spectrum(idx: int, messages: list, entry: dict) -> None: + """Render IR spectrum plot, frequency table, and trajectory viewer. + + Parameters + ---------- + idx : int + One-based exchange index. + messages : list + Message-like objects from the exchange. + entry : dict + Conversation-history entry. + """ + log_dir = _artifact_log_dir(messages, entry) + ir_path = _latest_artifact_path(log_dir, "ir_spectrum*.png") + freq_path = _latest_artifact_path(log_dir, "frequencies*.csv") + + if not ir_path and not freq_path: st.warning("IR spectrum not found.") + return + + with st.expander("\U0001f50d IR Spectrum", expanded=True): + col1, col2 = st.columns(2, border=True) + + with col1: + if ir_path and os.path.exists(ir_path): + st.image(ir_path) + else: + st.warning("IR spectrum plot not found.") + + with col2: + if not freq_path or not os.path.exists(freq_path): + st.warning("Frequencies file not found.") + return + + df = pd.read_csv( + freq_path, + index_col=False, + names=["filename", "frequency"], + ) + modes = df.iloc[6:] if len(df) > 6 else df + + if modes.empty: + st.warning("No vibrational frequencies found.") + return + + st.write("**Select a frequency to visualize:**") + freq_options = {} + for mode_idx, row in modes.iterrows(): + freq_text = str(row["frequency"]).strip() + suffix = "i" if freq_text.endswith("i") else "" + try: + freq_value = float(freq_text.rstrip("i")) + label = f"Mode {mode_idx}: {freq_value:.2f}{suffix} cm\u207b\u00b9" + except ValueError: + label = f"Mode {mode_idx}: {freq_text} cm\u207b\u00b9" + freq_options[label] = mode_idx + + selected_freq = st.selectbox( + "Frequency", + list(freq_options.keys()), + index=0, + key=f"ir_frequency_select_{idx}", + ) + traj_file = str(modes.loc[freq_options[selected_freq]]["filename"]) + traj_path = _resolve_artifact_path(traj_file, log_dir) + if not os.path.exists(traj_path): + st.warning(f"Trajectory file '{traj_file}' not found.") + elif not STMOL_AVAILABLE: + st.info("3D viewer not available; install stmol to animate trajectories.") + else: + import stmol + from ase.io.trajectory import Trajectory + + traj = Trajectory(traj_path) + view = visualize_trajectory(traj) + view.zoomTo() + stmol.showmol(view, height=400, width=500) def _render_verbose_info(idx: int, messages: list, entry: dict) -> None: + """Render raw result/debug information for an exchange. + + Parameters + ---------- + idx : int + One-based exchange index. + messages : list + Message-like objects from the exchange. + entry : dict + Conversation-history entry. + """ structure = find_structure_in_messages(messages) with st.expander(f"\U0001f50d Verbose Info (Query {idx})", expanded=False): st.write(f"**Number of messages:** {len(messages)}") st.write(f"**Structure found:** {'Yes' if structure else 'No'}") + raw_result = entry.get("result") if st.session_state.last_run_query == entry.get("query"): if st.session_state.last_run_error: st.write("**Last run error:**") st.code(str(st.session_state.last_run_error)) if st.session_state.last_run_result is not None: - st.write("**Raw result (repr):**") - st.code(repr(st.session_state.last_run_result)) - - for i, msg in enumerate(messages): - if hasattr(msg, "type"): - msg_type = msg.type - content = normalize_message_content(msg.content) - elif isinstance(msg, dict): - msg_type = msg.get("type", "unknown") - content = normalize_message_content(msg.get("content", "")) - else: - msg_type = type(msg).__name__ - content = normalize_message_content( - getattr(msg, "content", str(msg)) - ) - content_preview = ( - (content[:100] + "...") if len(content) > 100 else content - ) - st.write(f" **Message {i+1}:** `{msg_type}` - {content_preview}") + raw_result = st.session_state.last_run_result + + st.write("**Raw result:**") + st.code(pprint.pformat(raw_result, width=1, compact=False), language="text") -def _render_query_input(config: dict, selected_model: str) -> str: - with st.expander("\U0001f4a1 Example Queries"): +def _render_example_queries(config: dict, selected_model: str) -> None: + """Show example queries that the user can click to submit directly. + + Parameters + ---------- + config : dict + Nested UI configuration. + selected_model : str + Selected model name. + """ + # Hide after the first message or during an interrupt + if ( + st.session_state.conversation_history + or st.session_state.pending_human_question is not None + ): + return + + with st.expander("Example Queries", expanded=False): st.markdown("**Based on your current configuration:**") st.markdown(f"- Model: {selected_model}") st.markdown( f"- Default Calculator: {config['chemistry']['calculators']['default']}" ) - st.markdown("- Temperature: 0.0 (optimized for tool calling)") examples = [ "What is the SMILES string for caffeine?", @@ -672,34 +1060,247 @@ def _render_query_input(config: dict, selected_model: str) -> str: ] for ex in examples: if st.button(ex, key=f"ex_{ex}"): - st.session_state.query_input = ex + st.session_state._pending_example_query = ex st.rerun() - if "query_input" not in st.session_state: - st.session_state.query_input = "" - query = st.text_area( - "Enter your computational chemistry query:", - value=st.session_state.query_input, - height=100, - key="query_text_area", - ) +def _render_pending_interrupt() -> None: + """Show the agent's pending question and any prior interrupt exchanges.""" + question = st.session_state.pending_human_question + if question is None: + return - if query != st.session_state.query_input: - st.session_state.query_input = query + # Show the original user query that triggered the interrupt + original_query = st.session_state.pending_interrupt_query + if original_query: + with st.chat_message("user"): + _render_markdown_with_math(original_query) + + # Show any prior interrupt exchanges in this chain + for exch in st.session_state.interrupt_exchanges: + with st.chat_message("assistant"): + _render_markdown_with_math(exch["question"]) + with st.chat_message("user"): + _render_markdown_with_math(exch["answer"]) + + # Show the current pending question + with st.chat_message("assistant"): + st.info("The agent needs your input to continue.", icon="\u2753") + _render_markdown_with_math(question) + + # Cancel button + if st.button("Cancel", key="cancel_interrupt"): + _clear_interrupt_state() + st.rerun() - col_send, col_clear, col_refresh = st.columns([2, 1, 1]) - st.session_state._send_clicked = col_send.button( - "\U0001f680 Send", type="primary", use_container_width=True - ) - if col_clear.button("\U0001f5d1\ufe0f Clear Chat", use_container_width=True): - _start_new_chat() - st.rerun() - if col_refresh.button("\U0001f504 Refresh", use_container_width=True): - st.rerun() +def _clear_interrupt_state() -> None: + """Clear all interrupt-related session state.""" + st.session_state.pending_human_question = None + st.session_state.pending_interrupt_config = None + st.session_state.pending_interrupt_query = None + st.session_state.pending_interrupt_thread_id = None + st.session_state.pending_interrupt_prev_msg_count = 0 + st.session_state.pending_interrupt_model = None + st.session_state.pending_interrupt_workflow = None + st.session_state.pending_interrupt_log_dir = None + st.session_state.interrupt_count = 0 + st.session_state.interrupt_exchanges = [] + + +def _classify_message(msg): + """Classify a LangGraph message for UI display. + + Parameters + ---------- + msg : Any + LangGraph/LangChain message to classify. + + Returns + ------- + tuple or None + ``("tool_call", [tool_names])``, ``("tool_result", tool_name)``, or + ``None`` when not relevant for display. + """ + tool_calls = getattr(msg, "tool_calls", None) + if tool_calls: + names = [tc.get("name", "unknown") for tc in tool_calls if isinstance(tc, dict)] + if names: + return ("tool_call", names) + if getattr(msg, "type", None) == "tool": + name = getattr(msg, "name", None) + if name: + return ("tool_result", name) + return None + + +def _stream_workflow(stream_input, config, agent, msg_queue): + """Run the agent workflow in a background thread, pushing events to a queue. + + Events pushed: + ("tool_call", [tool_names]) — agent is calling tool(s) + ("tool_result", tool_name) — a tool finished + ("interrupt", question_str) + ("done", last_state) + ("error", exception) + + Parameters + ---------- + stream_input : dict or Command + Initial workflow input or resume command. + config : dict + LangGraph run configuration. + agent : ChemGraph + Active ChemGraph agent. + msg_queue : queue.Queue + Queue receiving stream events for the UI thread. + """ + from langgraph.errors import GraphInterrupt - return query + async def _run(): + """Stream the workflow and enqueue UI events.""" + prev_msgs: list = [] + last_st = None + interrupt_val = None + + try: + async for s in agent.workflow.astream( + stream_input, stream_mode="values", config=config + ): + if "__interrupt__" in s: + int_data = s["__interrupt__"] + if isinstance(int_data, (list, tuple)) and int_data: + interrupt_val = int_data[0].value + elif hasattr(int_data, "value"): + interrupt_val = int_data.value + else: + interrupt_val = {"question": "The workflow needs your input."} + + if "messages" in s and s["messages"] != prev_msgs: + new_message = s["messages"][-1] + classified = _classify_message(new_message) + if classified: + msg_queue.put(classified) + prev_msgs = s["messages"] + last_st = s + except GraphInterrupt as gi: + interrupts = gi.args[0] if gi.args else [] + if interrupts: + interrupt_val = interrupts[0].value + else: + interrupt_val = {"question": "The workflow needs your input."} + + # Check checkpoint for pending interrupts + if interrupt_val is None: + try: + snapshot = agent.workflow.get_state(config) + if snapshot and snapshot.tasks: + for t in snapshot.tasks: + t_interrupts = getattr(t, "interrupts", None) + if t_interrupts: + interrupt_val = t_interrupts[0].value + break + except Exception: + pass + + if interrupt_val is not None: + if isinstance(interrupt_val, dict): + q = interrupt_val.get( + "question", + interrupt_val.get("message", str(interrupt_val)), + ) + else: + q = str(interrupt_val) + msg_queue.put(("interrupt", q)) + else: + msg_queue.put(("done", last_st)) + + try: + asyncio.run(_run()) + except HumanInputRequired as hir: + msg_queue.put(("interrupt", hir.question)) + except Exception as exc: + msg_queue.put(("error", exc)) + + +def _poll_and_display(msg_queue, status_container, placeholder, thread): + """Poll the message queue and render a compact tool-call log. + + Uses a single ``st.empty()`` placeholder to re-render the full list + each time, so completed tools get a checkmark and only the active + tool shows a spinner indicator. + + Returns: + ("done", last_state) | ("interrupt", question) | ("error", exception) + + Parameters + ---------- + msg_queue : queue.Queue + Queue receiving stream events. + status_container : DeltaGenerator + Streamlit status container. + placeholder : DeltaGenerator + Placeholder used for the tool-call log. + thread : threading.Thread + Background stream thread. + + Returns + ------- + tuple + ``("done", state)``, ``("interrupt", question)``, or + ``("error", exception)``. + """ + completed: list[str] = [] # tools that finished + active: list[str] = [] # tools currently running + + def _render(): + """Render the current tool-call status list.""" + lines = [] + for name in completed: + lines.append(f"- :green[**{name}**] :white_check_mark:") + for name in active: + lines.append(f"- **{name}** :hourglass_flowing_sand:") + placeholder.markdown("\n".join(lines) if lines else "") + + while True: + try: + event_type, event_data = msg_queue.get(timeout=0.1) + except queue.Empty: + if not thread.is_alive(): + try: + event_type, event_data = msg_queue.get_nowait() + except queue.Empty: + return ("error", RuntimeError("Stream ended without result.")) + else: + continue + + if event_type == "tool_call": + # Mark previously active tools as completed + completed.extend(active) + active.clear() + active.extend(event_data) + label = ", ".join(event_data) + status_container.update(label=f"Running {label}", state="running") + _render() + elif event_type == "tool_result": + # Move this specific tool from active to completed + if event_data in active: + active.remove(event_data) + if event_data not in completed: + completed.append(event_data) + if active: + status_container.update( + label=f"Running {', '.join(active)}", state="running" + ) + else: + status_container.update(label="Thinking...", state="running") + _render() + elif event_type in ("done", "interrupt", "error"): + # Final render — mark everything as completed + completed.extend(active) + active.clear() + _render() + return (event_type, event_data) def _handle_query_submission( @@ -708,9 +1309,19 @@ def _handle_query_submission( endpoint_status: dict, selected_base_url: Optional[str], ) -> None: - if not st.session_state.get("_send_clicked", False): - return - + """Handle a submitted user query and stream the workflow response. + + Parameters + ---------- + query : str + User query text. + thread_id : int + Current LangGraph thread ID. + endpoint_status : dict + Local endpoint status dictionary. + selected_base_url : str, optional + Model endpoint URL used in error messages. + """ if not endpoint_status["ok"]: msg = ( f"Cannot reach local model endpoint `{selected_base_url}`. " @@ -718,56 +1329,221 @@ def _handle_query_submission( ) st.session_state.last_run_error = RuntimeError(msg) st.error(msg) - elif not st.session_state.agent: - st.error("\u274c Agent not ready. Please check configuration and try again.") - if st.button("\U0001f504 Try Again"): - st.rerun() - elif not query.strip(): - st.warning("Please enter a question.") - else: - with st.spinner("ChemGraph agents working...", show_time=True): + return + if not st.session_state.agent: + st.error("Agent not ready. Please check configuration and try again.") + return + if not query.strip(): + return + + trimmed_query = query.strip() + agent = st.session_state.agent + cfg = {"configurable": {"thread_id": str(thread_id)}} + cfg["recursion_limit"] = agent.recursion_limit + st.session_state.last_run_query = trimmed_query + st.session_state.last_run_error = None + st.session_state.last_run_result = None + + # Agent setup (mirroring agent.run() preamble) + if agent.log_dir: + os.environ["CHEMGRAPH_LOG_DIR"] = agent.log_dir + try: + agent._ensure_session(trimmed_query) + except Exception: + pass + + # Snapshot message count before streaming so we can isolate new messages + prev_msg_count = 0 + try: + snapshot = agent.workflow.get_state(cfg) + if snapshot and snapshot.values: + prev_msg_count = len(snapshot.values.get("messages", [])) + except Exception: + pass + + # Show the user's message immediately + with st.chat_message("user"): + st.markdown(trimmed_query) + + # Stream agent response with live tool-call display + with st.chat_message("assistant"): + msg_q: queue.Queue = queue.Queue() + inputs = {"messages": trimmed_query} + + stream_thread = threading.Thread( + target=_stream_workflow, + args=(inputs, cfg, agent, msg_q), + daemon=True, + ) + + status = st.status("Thinking...", expanded=True) + with status: + tool_log = st.empty() + stream_thread.start() + event_type, event_data = _poll_and_display( + msg_q, status, tool_log, stream_thread + ) + stream_thread.join(timeout=5) + + if event_type == "done": + status.update(label="Complete", state="complete", expanded=False) + last_state = event_data + if last_state is None: + st.error("Workflow produced no output.") + return + + # Only keep messages from this query (not prior thread history) + all_msgs = last_state.get("messages", []) + new_msgs = all_msgs[prev_msg_count:] + result = {"messages": new_msgs} + + # Save messages to persistent session store (best-effort) try: - cfg = {"configurable": {"thread_id": thread_id}} - st.session_state.last_run_query = query.strip() - st.session_state.last_run_error = None - st.session_state.last_run_result = None - # Capture references eagerly so the lambda never touches - # st.session_state from the background thread (thread safety). - agent = st.session_state.agent - trimmed_query = query.strip() - result = run_async_callable( - lambda: agent.run(trimmed_query, config=cfg) - ) - st.session_state.last_run_result = result - st.session_state.conversation_history.append( - { - "query": query.strip(), - "result": result, - "thread_id": thread_id, - } - ) - # Persist the exchange to the session store - _save_exchange_to_store(query.strip(), result) + agent._save_messages_to_store(last_state, trimmed_query) + except Exception: + pass + + st.session_state.last_run_result = result + st.session_state.conversation_history.append( + { + "query": trimmed_query, + "result": result, + "thread_id": thread_id, + "log_dir": agent.log_dir, + } + ) + _save_exchange_to_store(trimmed_query, result) + st.session_state.query_input = "" + st.rerun() - st.session_state.query_input = "" - st.success("\u2705 Done!") - st.rerun() - except Exception as exc: - st.session_state.last_run_error = exc - st.error(f"Processing error: {exc}") + elif event_type == "interrupt": + status.update(label="Waiting for input", state="complete", expanded=False) + cfg_for_resume = dict(cfg) + st.session_state.pending_human_question = event_data + st.session_state.pending_interrupt_config = cfg_for_resume + st.session_state.pending_interrupt_query = trimmed_query + st.session_state.pending_interrupt_thread_id = thread_id + st.session_state.pending_interrupt_prev_msg_count = prev_msg_count + st.session_state.pending_interrupt_model = st.session_state.get( + "active_model" + ) + st.session_state.pending_interrupt_workflow = st.session_state.get( + "active_workflow" + ) + st.session_state.pending_interrupt_log_dir = agent.log_dir + st.session_state.interrupt_count = 1 + st.session_state.interrupt_exchanges = [] + st.rerun() + else: # error + status.update(label="Error", state="error", expanded=False) + st.session_state.last_run_error = event_data + st.error(f"Processing error: {event_data}") -def _render_footer() -> None: - st.markdown("---") - st.markdown( - """ - ### Quick Help - **Main Features:** Molecular optimization, vibrational frequencies, SMILES \u2194 structure conversions, 3D visualization +def _handle_human_response(answer: str, thread_id: int) -> None: + """Resume the agent workflow with the human's answer. - \U0001f4d6 For detailed information, documentation, and links to research papers, visit the **About ChemGraph** page. + Parameters + ---------- + answer : str + Human response to the pending interrupt question. + thread_id : int + Current LangGraph thread ID. """ + from langgraph.types import Command + + agent = st.session_state.agent + resume_config = st.session_state.pending_interrupt_config + original_query = st.session_state.pending_interrupt_query + current_question = st.session_state.pending_human_question + interrupt_count = st.session_state.interrupt_count + + if agent is None or resume_config is None: + st.error("Agent was re-initialized. Please submit your query again.") + _clear_interrupt_state() + return + if agent.log_dir: + os.environ["CHEMGRAPH_LOG_DIR"] = agent.log_dir + + MAX_INTERRUPTS = 10 + + # Record this exchange + st.session_state.interrupt_exchanges.append( + {"question": current_question, "answer": answer} ) - if st.session_state.ui_notice: - st.info(st.session_state.ui_notice) + # Show the user's reply immediately + with st.chat_message("user"): + st.markdown(answer) + + # Stream resumed agent response + with st.chat_message("assistant"): + msg_q: queue.Queue = queue.Queue() + resume_cmd = Command(resume=answer) + + stream_thread = threading.Thread( + target=_stream_workflow, + args=(resume_cmd, resume_config, agent, msg_q), + daemon=True, + ) + + status = st.status("Processing your response...", expanded=True) + with status: + tool_log = st.empty() + stream_thread.start() + event_type, event_data = _poll_and_display( + msg_q, status, tool_log, stream_thread + ) + stream_thread.join(timeout=5) + + if event_type == "done": + status.update(label="Complete", state="complete", expanded=False) + result_state = event_data + + if result_state is None: + st.error("Resume produced no output.") + _clear_interrupt_state() + return + + # Only keep messages from this query (not prior thread history) + prev_msg_count = st.session_state.get("pending_interrupt_prev_msg_count", 0) + all_msgs = result_state.get("messages", []) + new_msgs = all_msgs[prev_msg_count:] + final_result = {"messages": new_msgs} + + exchanges = list(st.session_state.interrupt_exchanges) + st.session_state.last_run_result = final_result + st.session_state.conversation_history.append( + { + "query": original_query, + "result": final_result, + "thread_id": thread_id, + "log_dir": st.session_state.get("pending_interrupt_log_dir") + or agent.log_dir, + "interrupt_exchanges": exchanges, + } + ) + _save_exchange_to_store(original_query, final_result) + st.session_state.query_input = "" + _clear_interrupt_state() + st.rerun() + + elif event_type == "interrupt": + status.update(label="Waiting for input", state="complete", expanded=False) + new_count = interrupt_count + 1 + if new_count > MAX_INTERRUPTS: + st.error( + "Agent exceeded maximum number of follow-up questions. Aborting." + ) + _clear_interrupt_state() + return + st.session_state.pending_human_question = event_data + st.session_state.interrupt_count = new_count + st.rerun() + + else: # error + status.update(label="Error", state="error", expanded=False) + st.session_state.last_run_error = event_data + st.error(f"Error during resume: {event_data}") + _clear_interrupt_state() diff --git a/src/ui/agent_manager.py b/src/ui/agent_manager.py index 358ed048..d1963493 100644 --- a/src/ui/agent_manager.py +++ b/src/ui/agent_manager.py @@ -11,9 +11,11 @@ def initialize_agent( structured_output: bool, return_option: str, generate_report: bool, + human_supervised: bool, recursion_limit: int, base_url: Optional[str], argo_user: Optional[str], + log_dir: Optional[str] = None, ): """Create a :class:`ChemGraph` agent instance. @@ -22,6 +24,34 @@ def initialize_agent( ``st.session_state.agent`` and ``st.session_state.last_config``. Using the decorator caused failed initialisations (``None``) to be permanently cached with no way to retry. + + Parameters + ---------- + model_name : str + LLM model identifier. + workflow_type : str + ChemGraph workflow name. + structured_output : bool + Whether structured final output is requested. + return_option : str + Agent return mode. + generate_report : bool + Whether report generation is enabled. + human_supervised : bool + Whether human-supervision tools are enabled. + recursion_limit : int + LangGraph recursion limit. + base_url : str, optional + Custom model endpoint URL. + argo_user : str, optional + Argo username for Argo-hosted models. + log_dir : str, optional + Directory for ChemGraph run logs. + + Returns + ------- + ChemGraph or None + Initialized agent, or ``None`` if initialization fails. """ try: from chemgraph.agent.llm_agent import ChemGraph @@ -35,6 +65,8 @@ def initialize_agent( generate_report=generate_report, return_option=return_option, recursion_limit=recursion_limit, + human_supervised=human_supervised, + log_dir=log_dir, ) except Exception as exc: st.error(f"Failed to initialize agent: {exc}") @@ -42,7 +74,18 @@ def initialize_agent( def run_async_callable(fn): - """Run an async callable and return its result in a sync context.""" + """Run an async callable and return its result in a sync context. + + Parameters + ---------- + fn : Callable + Zero-argument callable returning an awaitable. + + Returns + ------- + Any + Result produced by the awaited callable. + """ from chemgraph.utils.async_utils import run_async_callable as _impl return _impl(fn) diff --git a/src/ui/branding.py b/src/ui/branding.py index bf3707ab..70711257 100644 --- a/src/ui/branding.py +++ b/src/ui/branding.py @@ -19,7 +19,18 @@ def first_existing_asset(paths: tuple[Path, ...]) -> str | None: - """Return the first available Streamlit-compatible asset path.""" + """Return the first available Streamlit-compatible asset path. + + Parameters + ---------- + paths : tuple[pathlib.Path, ...] + Candidate asset paths in priority order. + + Returns + ------- + str or None + First existing asset path as a string, or ``None``. + """ for path in paths: if path.exists(): return str(path) diff --git a/src/ui/config.py b/src/ui/config.py index bbdfa807..99e3fe9d 100644 --- a/src/ui/config.py +++ b/src/ui/config.py @@ -9,7 +9,18 @@ def load_config(config_path: str = "config.toml") -> Dict[str, Any]: - """Load configuration from TOML file.""" + """Load configuration from a TOML file. + + Parameters + ---------- + config_path : str, optional + Path to the TOML configuration file. + + Returns + ------- + dict[str, Any] + Nested configuration dictionary with defaults filled in. + """ try: if os.path.exists(config_path): with open(config_path, "r") as f: @@ -47,7 +58,20 @@ def load_config(config_path: str = "config.toml") -> Dict[str, Any]: def save_config(config: Dict[str, Any], config_path: str = "config.toml") -> bool: - """Save configuration to TOML file.""" + """Save configuration to a TOML file. + + Parameters + ---------- + config : dict[str, Any] + Nested configuration dictionary to write. + config_path : str, optional + Destination TOML file path. + + Returns + ------- + bool + ``True`` if the file was written successfully. + """ try: with open(config_path, "w") as f: toml.dump(config, f) @@ -68,6 +92,7 @@ def get_default_config() -> Dict[str, Any]: "report": False, "thread": 1, "recursion_limit": 20, + "human_supervised": False, "verbose": False, }, "api": { @@ -102,5 +127,16 @@ def get_default_config() -> Dict[str, Any]: def flatten_config(config: Dict[str, Any]) -> Dict[str, Any]: - """Flatten nested configuration for easier access.""" + """Flatten nested configuration for easier access. + + Parameters + ---------- + config : dict[str, Any] + Nested configuration dictionary. + + Returns + ------- + dict[str, Any] + Flattened configuration dictionary. + """ return _flatten_config(config) diff --git a/src/ui/endpoint.py b/src/ui/endpoint.py index e906a31f..c0d941c2 100644 --- a/src/ui/endpoint.py +++ b/src/ui/endpoint.py @@ -9,13 +9,36 @@ def _is_local_address(hostname: str) -> bool: + """Return whether a hostname points to the local machine. + + Parameters + ---------- + hostname : str + Hostname parsed from a URL. + + Returns + ------- + bool + ``True`` for localhost-style addresses. + """ host = (hostname or "").strip().lower() return host in {"localhost", "127.0.0.1", "0.0.0.0", "::1"} @st.cache_data(ttl=10) def check_local_model_endpoint(base_url: Optional[str]) -> Dict[str, Any]: - """Quick reachability check for local OpenAI-compatible endpoints.""" + """Quick reachability check for local OpenAI-compatible endpoints. + + Parameters + ---------- + base_url : str, optional + Base URL to probe. + + Returns + ------- + dict[str, Any] + Status dictionary with ``ok`` and ``message`` keys. + """ if not base_url: return {"ok": True, "message": "No base URL configured."} diff --git a/src/ui/file_utils.py b/src/ui/file_utils.py index f1acee4b..addd1bdc 100644 --- a/src/ui/file_utils.py +++ b/src/ui/file_utils.py @@ -12,7 +12,18 @@ def resolve_output_path(path: str) -> str: - """Resolve output paths relative to CHEMGRAPH_LOG_DIR when set.""" + """Resolve output paths relative to CHEMGRAPH_LOG_DIR when set. + + Parameters + ---------- + path : str + Absolute or relative output path. + + Returns + ------- + str + Resolved output path. + """ if not path: return path if os.path.isabs(path): @@ -24,7 +35,20 @@ def resolve_output_path(path: str) -> str: def changed_recently(path: str = "ir_spectrum.png", window_seconds: int = 300) -> bool: - """Return True if *path* exists and was modified within *window_seconds*.""" + """Return True if a file was modified within a recent time window. + + Parameters + ---------- + path : str, optional + File path to inspect. + window_seconds : int, optional + Recency window in seconds. + + Returns + ------- + bool + ``True`` when the file exists and is recent. + """ p = Path(resolve_output_path(path)) if not p.exists(): return False @@ -59,7 +83,18 @@ def find_latest_xyz_file() -> Optional[str]: def find_latest_xyz_file_in_dir(directory: str) -> Optional[str]: - """Find the most recently modified ``.xyz`` file under *directory*.""" + """Find the most recently modified ``.xyz`` file under a directory. + + Parameters + ---------- + directory : str + Directory to search recursively. + + Returns + ------- + str or None + Latest XYZ file path, or ``None`` when none is found. + """ if not directory or not os.path.isdir(directory): return None latest_path: Optional[str] = None @@ -76,7 +111,18 @@ def find_latest_xyz_file_in_dir(directory: str) -> Optional[str]: def extract_log_dir_from_messages(messages: Any) -> Optional[str]: - """Extract a directory path from message content that references an output file.""" + """Extract a log directory from messages that reference output files. + + Parameters + ---------- + messages : Any + Message object, dictionary, list, or text to scan. + + Returns + ------- + str or None + Parent directory of a referenced output file, or ``None``. + """ if not messages: return None patterns = [ @@ -87,6 +133,18 @@ def extract_log_dir_from_messages(messages: Any) -> Optional[str]: ] def _scan_value(value: Any) -> Optional[str]: + """Recursively scan a value for absolute output-file references. + + Parameters + ---------- + value : Any + Message content, mapping, list, or scalar to scan. + + Returns + ------- + str or None + Parent directory of a referenced file, or ``None``. + """ if isinstance(value, str): for pattern in patterns: match = re.search(pattern, value) diff --git a/src/ui/message_utils.py b/src/ui/message_utils.py index 3ff1b627..6a590aac 100644 --- a/src/ui/message_utils.py +++ b/src/ui/message_utils.py @@ -7,7 +7,7 @@ import ast import json import re -from typing import Any, Optional +from typing import Any, Literal, Optional from ase.data import chemical_symbols @@ -18,7 +18,18 @@ def normalize_message_content(content: Any) -> str: - """Convert varying message content payloads (str/list/dict) into plain text.""" + """Convert varying message content payloads into plain text. + + Parameters + ---------- + content : Any + Message content as text, list blocks, dictionary, or scalar. + + Returns + ------- + str + Normalized plain-text content. + """ if content is None: return "" if isinstance(content, str): @@ -45,13 +56,269 @@ def normalize_message_content(content: Any) -> str: return str(content) +def normalize_latex_delimiters(text: str) -> str: + """Convert common LLM math delimiters into Streamlit-renderable Markdown. + + Parameters + ---------- + text : str + Markdown text that may contain LLM-style math delimiters. + + Returns + ------- + str + Text with display and inline math delimiters normalized. + """ + if not text: + return "" + + text = _convert_square_bracket_math(text) + + return _convert_parenthetical_math_outside_display(text) + + +def _is_latex_square_delimiter(text: str, index: int) -> bool: + """Return whether a bracket belongs to ``\\left``/``\\right``. + + Parameters + ---------- + text : str + Source text. + index : int + Bracket index to inspect. + + Returns + ------- + bool + ``True`` when the bracket is part of a LaTeX delimiter command. + """ + prefix = text[max(0, index - 6) : index] + return prefix.endswith(r"\left") or prefix.endswith(r"\right") + + +def _find_square_math_close(text: str, start: int) -> int | None: + """Find the matching closing square bracket for display math. + + Parameters + ---------- + text : str + Source text. + start : int + Index of the opening square bracket. + + Returns + ------- + int or None + Closing bracket index, or ``None`` if unmatched. + """ + depth = 1 + index = start + 1 + while index < len(text): + if text.startswith(r"\left[", index): + index += len(r"\left[") + continue + if text.startswith(r"\right]", index): + index += len(r"\right]") + continue + + char = text[index] + if char == "[": + depth += 1 + elif char == "]": + depth -= 1 + if depth == 0: + return index + index += 1 + return None + + +def _convert_square_bracket_math(text: str) -> str: + """Convert LLM-style ``[ TeX ]`` display math while preserving links. + + Parameters + ---------- + text : str + Source Markdown text. + + Returns + ------- + str + Text with display math converted to ``$$`` blocks. + """ + chunks: list[str] = [] + index = 0 + while index < len(text): + char = text[index] + if char != "[" or _is_latex_square_delimiter(text, index): + chunks.append(char) + index += 1 + continue + + close_index = _find_square_math_close(text, index) + if close_index is None: + chunks.append(char) + index += 1 + continue + + body = text[index + 1 : close_index].strip() + has_latex_command = bool(re.search(r"\\[A-Za-z]+", body)) + is_markdown_link = close_index + 1 < len(text) and text[close_index + 1] == "(" + if body and has_latex_command and not is_markdown_link: + chunks.append(f"$$\n{body}\n$$") + else: + chunks.append(text[index : close_index + 1]) + index = close_index + 1 + + return "".join(chunks) + + +def _find_parenthesis_close(text: str, start: int) -> int | None: + """Find the matching closing parenthesis. + + Parameters + ---------- + text : str + Source text. + start : int + Index of the opening parenthesis. + + Returns + ------- + int or None + Closing parenthesis index, or ``None`` if unmatched. + """ + depth = 1 + index = start + 1 + while index < len(text): + char = text[index] + if char == "(": + depth += 1 + elif char == ")": + depth -= 1 + if depth == 0: + return index + index += 1 + return None + + +def _convert_parenthetical_inline_math(text: str) -> str: + """Convert parenthesized math-like text to inline math. + + Parameters + ---------- + text : str + Text outside display-math blocks. + + Returns + ------- + str + Text with math-like parenthetical content wrapped in ``$``. + """ + chunks: list[str] = [] + index = 0 + while index < len(text): + char = text[index] + if char != "(": + chunks.append(char) + index += 1 + continue + + close_index = _find_parenthesis_close(text, index) + if close_index is None: + chunks.append(char) + index += 1 + continue + + body = text[index + 1 : close_index].strip() + has_latex = bool(re.search(r"\\[A-Za-z]+|\\\s|[_^{}]", body)) + if body and has_latex: + chunks.append(f"${body}$") + else: + chunks.append(text[index : close_index + 1]) + index = close_index + 1 + + return "".join(chunks) + + +def _convert_parenthetical_math_outside_display(text: str) -> str: + """Convert parenthetical math outside display-math blocks. + + Parameters + ---------- + text : str + Markdown text that may contain display-math blocks. + + Returns + ------- + str + Text with inline parenthetical math normalized. + """ + chunks: list[str] = [] + last_end = 0 + for match in re.finditer(r"(?s)\$\$.*?\$\$", text): + chunks.append(_convert_parenthetical_inline_math(text[last_end : match.start()])) + chunks.append(match.group(0)) + last_end = match.end() + chunks.append(_convert_parenthetical_inline_math(text[last_end:])) + return "".join(chunks) + + +def split_markdown_latex_blocks( + text: str, +) -> list[tuple[Literal["markdown", "latex"], str]]: + """Split text into Markdown and display-LaTeX blocks. + + Parameters + ---------- + text : str + Markdown text that may contain display math blocks. + + Returns + ------- + list[tuple[Literal["markdown", "latex"], str]] + Ordered render blocks for Streamlit. + """ + normalized = normalize_latex_delimiters(text) + if not normalized: + return [] + + parts: list[tuple[Literal["markdown", "latex"], str]] = [] + last_end = 0 + for match in re.finditer(r"(?s)\$\$\s*(.*?)\s*\$\$", normalized): + markdown = normalized[last_end : match.start()].strip() + if markdown: + parts.append(("markdown", markdown)) + + latex = match.group(1).strip() + if latex: + parts.append(("latex", latex)) + last_end = match.end() + + trailing = normalized[last_end:].strip() + if trailing: + parts.append(("markdown", trailing)) + + return parts + + # --------------------------------------------------------------------------- # Message extraction # --------------------------------------------------------------------------- def extract_messages_from_result(result: Any) -> list: - """Extract messages from a result object, handling different formats.""" + """Extract messages from a result object, handling different formats. + + Parameters + ---------- + result : Any + Agent result, state dictionary, message list, or scalar result. + + Returns + ------- + list + Extracted message-like objects. + """ if isinstance(result, list): return result elif isinstance(result, dict) and "messages" in result: @@ -76,7 +343,18 @@ def extract_messages_from_result(result: Any) -> list: def extract_molecular_structure(message_content: str) -> Optional[dict]: - """Return ``{atomic_numbers, positions}`` if structure data is embedded.""" + """Return embedded molecular structure data if present. + + Parameters + ---------- + message_content : str + Message text that may contain JSON or plain-text structure data. + + Returns + ------- + dict or None + Dictionary with ``atomic_numbers`` and ``positions``, or ``None``. + """ if not message_content: return None @@ -150,8 +428,19 @@ def extract_molecular_structure(message_content: str) -> Optional[dict]: def find_structure_in_messages(messages: list) -> Optional[dict]: - """Look through all messages to find structure data.""" - for message in messages: + """Look through messages in reverse to find the latest structure data. + + Parameters + ---------- + messages : list + Message-like objects or dictionaries to scan. + + Returns + ------- + dict or None + Latest embedded structure dictionary, or ``None``. + """ + for message in reversed(messages): if hasattr(message, "content") or isinstance(message, dict): raw_content = ( getattr(message, "content", "") @@ -168,7 +457,22 @@ def find_structure_in_messages(messages: list) -> Optional[dict]: def has_structure_signal( messages: list, query_text: str = "", final_answer: str = "" ) -> bool: - """Return True when the interaction appears to include structure artifacts.""" + """Return True when an interaction appears to include structure artifacts. + + Parameters + ---------- + messages : list + Message-like objects or dictionaries to inspect. + query_text : str, optional + Original user query text. + final_answer : str, optional + Final assistant answer text. + + Returns + ------- + bool + ``True`` when structure-related tools, artifacts, or text are found. + """ structure_tools = { "smiles_to_coordinate_file", "run_ase", @@ -219,6 +523,16 @@ def find_html_filename(messages: list) -> Optional[str]: """Scan *messages* in reverse for the first ``*.html`` reference. Returns the matched substring (path or bare filename) or ``None``. + + Parameters + ---------- + messages : list + Message-like objects, dictionaries, or strings to scan. + + Returns + ------- + str or None + First HTML path/filename found from the end of the message list. """ pattern = r"[\w./-]+\.html\b" @@ -243,7 +557,18 @@ def find_html_filename(messages: list) -> Optional[str]: def extract_xyz_from_report_html(html_content: str) -> Optional[dict]: - """Decode base64-encoded XYZ data from an HTML report's ``atob()`` call.""" + """Decode base64-encoded XYZ data from an HTML report. + + Parameters + ---------- + html_content : str + HTML report content containing an ``atob()`` XYZ payload. + + Returns + ------- + dict or None + Structure dictionary with atomic numbers and positions, or ``None``. + """ import base64 as _b64 match = re.search(r'atob\(["\']([A-Za-z0-9+/=]+)["\']\)', html_content) @@ -295,6 +620,16 @@ def strip_viewer_from_report_html(html_content: str) -> str: Strips the ``
`` element, the NGL ``