Chemgraph

What is ChemGraph

ChemGraph is a Python framework (package name chemgraph) 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 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), ask_human
• 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].connectivity_smiles  # .smiles keeps stereochemistry

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