Chemgraph

What is ChemGraph

ChemGraph is a Python framework (package name chemgraphagent) built at Argonne National Laboratory that automates computational chemistry workflows using LLMs. It connects natural language queries to molecular simulations via an agent architecture built on LangGraph/LangChain, ASE (Atomic Simulation Environment), RDKit, and MCP (Model Context Protocol) servers.

Key capabilities: molecule lookup (PubChem), 3D structure generation (RDKit), geometry optimization, vibrational analysis, thermochemistry, IR spectra, and HPC-scale ensemble simulations via Parsl.

Project layout

ChemGraph/
  src/
    chemgraph/              # Core package
      agent/                # Main ChemGraph agent class (llm_agent.py)
      graphs/               # LangGraph workflow definitions (9 workflows)
      tools/                # LangChain tool implementations
      mcp/                  # FastMCP server implementations
      models/               # LLM provider integrations (OpenAI, Anthropic, Gemini, Groq, Ollama, ALCF, Argo)
      prompt/               # System prompt templates per model/workflow
      schemas/              # Pydantic data models (AtomsData, ASEInput/Output, calculators)
      memory/               # Session memory (SQLite-backed persistence, schemas)
      state/                # LangGraph state definitions
      hpc_configs/          # Parsl configs for ALCF Polaris/Aurora
      utils/                # Config, logging, evaluation utilities
    ui/                     # Streamlit web app (app.py) and Rich CLI (cli.py)
  tests/                    # pytest test suite (20+ files)
  scripts/                  # MCP examples, Parsl examples
  notebooks/                # Jupyter demo notebooks
  docs/                     # MkDocs documentation source
  config.toml               # Default runtime configuration
  pyproject.toml            # Package metadata and dependencies
  docker-compose.yml        # Multi-profile Docker (jupyter, streamlit, mcp, cli)

Architecture overview

Agent entry point

src/chemgraph/agent/llm_agent.py contains the ChemGraph class. This is the central orchestrator:

• Selects and loads LLM models from any supported provider
• Dispatches to the correct workflow graph
• Runs async execution via LangGraph's astream
• Handles state serialization and logging

Workflows (graphs/)

Each file defines a LangGraph StateGraph. The 9 workflows are:

Tools (tools/)

LangChain @tool-decorated functions. Key files:

• ase_tools.py -- run_ase (energy/opt/vib/thermo), save_atomsdata_to_file, file_to_atomsdata
• cheminformatics_tools.py -- molecule_name_to_smiles, smiles_to_coordinate_file, smiles_to_atomsdata
• generic_tools.py -- calculator (safe math eval), Python REPL
• report_tools.py -- generate_html (interactive HTML reports with NGL 3D viewer)
• graspa_tools.py -- gRASPA simulation tools
• architector_tools.py -- Metal complex tools
• pormake_tools.py -- MOF topology/structure tools
• parsl_tools.py -- MACE with Parsl for HPC parallel execution

MCP servers (mcp/)

FastMCP-based servers. Each exposes chemistry tools over stdio or HTTP:

• mcp_tools.py -- General chemistry MCP server (name-to-SMILES, structure gen, ASE simulations, file I/O). Port 9003.
• mace_mcp_parsl.py -- MACE ML potential with Parsl HPC. Port 9004.
• graspa_mcp_parsl.py -- gRASPA simulation with Parsl HPC. Port 9005.
• data_analysis_mcp.py -- Data analysis (CIF splitting, JSONL aggregation, isotherm plotting). Port 9006.
• server_utils.py -- Shared startup utility (run_mcp_server), handles stdio vs streamable_http transport, logging to stderr.

Schemas (schemas/)

Pydantic models for data validation:

• atomsdata.py -- AtomsData (numbers, positions, cell, pbc)
• ase_input.py -- ASEInputSchema / ASEOutputSchema
• agent_response.py -- ResponseFormatter, VibrationalFrequency, IRSpectrum, etc.
• calculators/ -- One schema per calculator: mace_calc.py, emt_calc.py, tblite_calc.py, nwchem_calc.py, orca_calc.py, psi4_calc.py, fairchem_calc.py, mopac_calc.py, aimnet2_calc.py

Memory (memory/)

SQLite-backed session persistence:

• store.py -- SessionStore class: CRUD for sessions, context building for resume, prefix-based session ID lookup. Database at ~/.chemgraph/sessions.db.
• schemas.py -- SessionMessage (role, content, tool_name, timestamp), Session (full record with messages), SessionSummary (lightweight listing model)

State (state/)

LangGraph state definitions:

• state.py -- State (messages + remaining_steps), MultiAgentState
• multi_agent_state.py -- ManagerWorkerState for Planner/Executor/Aggregator
• graspa_state.py, mof_state.py -- Domain-specific states

How to add a new LangChain tool

1. Create or edit a file in src/chemgraph/tools/
2. Define the function with the @tool decorator from langchain_core.tools
3. Use Pydantic schemas for structured input (see schemas/ase_input.py for the pattern)
4. Import the tool in the relevant graph file (src/chemgraph/graphs/) and add it to the tools list
5. Add tests in tests/

Example pattern from cheminformatics_tools.py:

from langchain_core.tools import tool

@tool
def molecule_name_to_smiles(name: str) -> str:
    """Convert a molecule name to SMILES using PubChem."""
    import pubchempy as pcp
    comps = pcp.get_compounds(name.strip(), "name")
    if not comps:
        raise ValueError(f"No PubChem compound found for: {name}")
    return comps[0].canonical_smiles

How to add a new MCP server tool

MCP tools are defined in src/chemgraph/mcp/ using FastMCP:

from mcp.server.fastmcp import FastMCP

mcp = FastMCP(name="My Server", instructions="...")

@mcp.tool(name="my_tool", description="What it does")
async def my_tool(param: str) -> dict:
    # implementation
    return {"status": "success", "result": ...}

if __name__ == "__main__":
    from chemgraph.mcp.server_utils import run_mcp_server
    run_mcp_server(mcp, default_port=9007)

The run_mcp_server utility handles:

• --transport stdio (default) for LangGraph/OpenCode MCP clients
• --transport streamable_http with --port and --host for HTTP access
• Logging to stderr (critical for stdio mode) and optional file logging via CHEMGRAPH_LOG_DIR

How to add a new calculator

1. Create a Pydantic schema in src/chemgraph/schemas/calculators/ (follow mace_calc.py pattern)
2. The schema must define calculator_type, implement get_calculator() returning an ASE calculator, and optionally get_atoms_properties()
3. Register it in src/chemgraph/tools/mcp_helper.py load_calculator() with a new elif branch
4. Add tests

How to add a new workflow

1. Create a new graph file in src/chemgraph/graphs/
2. Define a StateGraph with nodes, edges, and conditional routing
3. Register it in src/chemgraph/agent/llm_agent.py in the workflow dispatch logic
4. Add a prompt template in src/chemgraph/prompt/ if needed
5. Add a state class in src/chemgraph/state/ if the workflow needs custom state

Running tests

# Run all tests (excluding LLM-dependent tests)
pytest tests/

# Run with LLM tests
pytest tests/ --run-llm

# Run specific test file
pytest tests/test_mcp.py

# Run async tests
pytest tests/test_mcp.py -v

Test markers:

• @pytest.mark.llm -- requires LLM API access (skipped by default)
• @pytest.mark.asyncio -- async tests

Running MCP servers

# stdio mode (for LangGraph / OpenCode / Claude Desktop)
python -m chemgraph.mcp.mcp_tools

# HTTP mode
python -m chemgraph.mcp.mcp_tools --transport streamable_http --port 9003

# With log directory
CHEMGRAPH_LOG_DIR=/tmp/chemgraph_logs python -m chemgraph.mcp.mcp_tools

Running the CLI

# Single query
chemgraph --query "Calculate the energy of water using MACE"

# Interactive mode
chemgraph --interactive

# List supported models
chemgraph --list-models

# Session management
chemgraph --list-sessions
chemgraph --show-session a3b2
chemgraph --delete-session a3b2c1d4
chemgraph -q "Follow-up query" --resume a3b2

Running the Streamlit UI

streamlit run src/ui/app.py

Configuration

config.toml at the project root controls runtime settings:

• [general] -- model, workflow, recursion_limit, verbosity
• [chemistry.calculators] -- default calculator (mace_mp), fallback (emt)
• [chemistry.optimization] -- optimizer method, fmax, steps
• [api.*] -- LLM provider base URLs and timeouts

Coding conventions

• Python >= 3.10, formatted with Ruff (line-length 88)
• Pydantic for all data models and tool input schemas
• Async-first for MCP tool implementations
• All MCP server logging must go to stderr (stdout is reserved for stdio transport)
• Energies in eV, frequencies in cm^-1, distances in Angstroms
• Pre-commit hooks configured via .pre-commit-config.yaml (Ruff linter + formatter)

Docker

# Streamlit UI
docker compose --profile streamlit up

# MCP server
docker compose --profile mcp up

# Jupyter notebooks
docker compose --profile jupyter up

Key dependencies

• langgraph + langchain -- agent orchestration
• ase -- atomic simulation environment
• rdkit -- cheminformatics, 3D structure generation
• pubchempy -- PubChem molecule lookup
• mcp + fastmcp -- Model Context Protocol servers
• mace-torch -- MACE ML potentials
• pydantic -- data validation
• parsl -- HPC parallel execution (optional)
• streamlit + stmol -- web UI (optional)

When to use this skill

Use this skill when:

• Developing new features, tools, or workflows for ChemGraph
• Adding or modifying MCP servers
• Adding new calculator integrations
• Writing or debugging tests
• Understanding the project architecture
• Refactoring or extending existing code

Source: argonne-lcf/ChemGraph — distributed by TomeVault.