Agent Configuration

Overview

The ChemAgent is built on LangGraph, a framework for creating stateful, multi-step AI workflows. This page documents the configuration options and customization points for the agent’s plan-and-execute architecture.

LangGraph Configuration

Recursion Limit

The agent’s execution depth is controlled by the recursion_limit parameter:

config = {"recursion_limit": 50}

recursion_limit

int

default:"50"

Maximum number of planning and execution cycles the agent can perform. This prevents infinite loops while allowing complex multi-step reasoning.Recommended Values:

Simple queries (name conversions): 10-20
Complex queries (property predictions): 30-50
Advanced tasks (synthesis planning): 50-100

Setting this too low may cause the agent to fail on complex tasks. Setting it too high may result in long execution times.

Workflow Configuration

Pass the configuration when invoking the agent:

inputs = {"input": edited_prompt}
config = {"recursion_limit": 50}

responses = await process_events(inputs, config)

Agent Architecture

The ChemAgent uses a plan-execute-replan architecture with three main components:

1. Planner

Creates initial step-by-step plans for chemistry queries.

planner_prompt

ChatPromptTemplate

System prompt that guides the initial planning phase:

planner_prompt = ChatPromptTemplate.from_messages([
    ("system", 
     """For the given objective, come up with a simple step by step plan. 
     This plan should involve individual tasks, that if executed correctly 
     will yield the correct answer. Do not add any superfluous steps. 
     The result of the final step should be the final answer. 
     Make sure that each step has all the information needed - do not skip steps."""),
    ("placeholder", "{messages}"),
])

Customization:

Modify the system message to change planning behavior
Add examples for domain-specific planning patterns
Adjust temperature (default: 0) for more creative plans

planner_model

ChatOpenAI

planner = planner_prompt | ChatOpenAI(
    model="gpt-4o",
    temperature=0
).with_structured_output(Plan)

Configuration Options:

model: LLM model for planning (default: “gpt-4o”)
temperature: Randomness in planning (default: 0 for deterministic)

2. Executor

Executes individual steps using specialized chemistry tools.

tools

list

The agent has access to these chemistry tools:

tools = [
    structure_chem_prompt,
    validate_smiles_rdkit,
    answer_chemistry_query
]

Show structure_chem_prompt

Structures and tags chemical information in the input query with <SMILES> and <IUPAC> tags.Usage: Must be called first in every plan to format the query properly.

Show validate_smiles_rdkit

Validates SMILES strings using RDKit and the chemistry parser.Returns: Validity status and detailed error messages (validity vectors).

Show answer_chemistry_query

Uses the LlaSMol model to answer chemistry-specific queries.Requires: VRAM ≥15GB (disabled when LOW_VRAM = True).

agent_executor

ReActAgent

llm = ChatOpenAI(model="gpt-4o")
agent_executor = create_react_agent(llm, tools, state_modifier=prompt)

Customization:

model: Change the reasoning model
state_modifier: Customize the agent’s system prompt
Add custom tools to the tools list

3. Replanner

Updates the plan based on execution results.

replanner_prompt

ChatPromptTemplate

Template for adaptive replanning:

replanner_prompt = ChatPromptTemplate.from_template(
    """For the given objective, come up with a simple step by step plan. 
    This plan should involve individual tasks, that if executed correctly 
    will yield the correct answer. Do not add any superfluous steps. 
    The result of the final step should be the final answer. 
    Make sure that each step has all the information needed - do not skip steps.

    Your objective was this:
    {input}

    Your original plan was this:
    {plan}

    You have currently done the follow steps:
    {past_steps}

    Update your plan accordingly. If no more steps are needed OR VRAM is LOW 
    and you can return to the user, then respond with that. Otherwise, fill 
    out the plan. Only add steps to the plan that still NEED to be done. 
    Do not return previously done steps as part of the plan."""
)

Key Features:

Receives full context: original objective, plan, and completed steps
Can decide to respond directly or continue planning
Handles low VRAM scenarios gracefully

replanner_model

ChatOpenAI

replanner = replanner_prompt | ChatOpenAI(
    model="gpt-4o",
    temperature=0
).with_structured_output(Act)

Output Types:

Response: Final answer to return to user
Plan: Updated plan with remaining steps

State Management

The agent maintains state using the PlanExecute type:

class PlanExecute(TypedDict):
    input: str
    plan: List[str]
    past_steps: Annotated[List[Tuple], operator.add]
    response: str

input

str

The original (formatted) input query with chemistry tags.

plan

List[str]

Current list of steps to execute. Updated by the replanner after each step.

past_steps

List[Tuple]

History of completed steps. Each tuple contains:

[0]: Step description
[1]: Step result

Accumulated using operator.add annotation.

response

str

Final response when the agent completes the task.

Customization Examples

Custom System Prompt

from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent

# Custom prompt for specialized chemistry domain
custom_prompt = """You are an expert medicinal chemist specializing in drug design.
Focus on drug-like properties and ADMET predictions."""

llm = ChatOpenAI(model="gpt-4o")
agent_executor = create_react_agent(llm, tools, state_modifier=custom_prompt)

Adjust Recursion Limit

# For simple queries - faster execution
config_simple = {"recursion_limit": 15}

# For complex synthesis planning
config_complex = {"recursion_limit": 100}

# Use when calling the agent
responses = await process_events(inputs, config_simple)

Change Planning Model

from langchain_openai import ChatOpenAI
from langchain_anthropic import ChatAnthropic

# Use Claude for planning instead of GPT-4o
planner = planner_prompt | ChatAnthropic(
    model="claude-3-5-sonnet-20241022",
    temperature=0
).with_structured_output(Plan)

replanner = replanner_prompt | ChatAnthropic(
    model="claude-3-5-sonnet-20241022",
    temperature=0
).with_structured_output(Act)

Add Custom Chemistry Tool

from langchain_core.tools import tool

@tool
def predict_toxicity(smiles: str) -> dict:
    """Predict toxicity profile for a SMILES string."""
    # Your implementation here
    return {"ld50": "...", "mutagenic": False}

# Add to tools list
tools = [
    structure_chem_prompt,
    validate_smiles_rdkit,
    answer_chemistry_query,
    predict_toxicity  # New custom tool
]

# Recreate agent with updated tools
agent_executor = create_react_agent(llm, tools, state_modifier=prompt)

Workflow Graph

The agent follows this execution flow:

Node Functions

Show plan_step()

async def plan_step(state: PlanExecute):
    plan = await planner.ainvoke({"messages": [("user", state["input"])]})
    return {"plan": plan.steps}

Creates the initial plan from the input query.

Show execute_step()

async def execute_step(state: PlanExecute):
    plan = state["plan"]
    task = plan[0]
    agent_response = await agent_executor.ainvoke(
        {"messages": [("user", task_formatted)]}
    )
    return {
        "past_steps": [(task, agent_response["messages"][-1].content)],
    }

Executes the first step in the current plan using the ReAct agent.

Show replan_step()

async def replan_step(state: PlanExecute):
    output = await replanner.ainvoke(state)
    if isinstance(output.action, Response):
        return {"response": output.action.response}
    else:
        return {"plan": output.action.steps}

Decides whether to continue with updated plan or return final response.

Show should_end()

def should_end(state: PlanExecute):
    global replanning_attempts
    if "response" in state and state["response"]:
        return END
    else:
        replanning_attempts += 1
        return "agent"

Determines if the workflow should end or continue to the next step.

Environment Configuration

VRAM Management

# In plan_execute_agent/config.py
LOW_VRAM = True  # Disable LlaSMol model

LOW_VRAM

bool

default:"True"

Controls whether to load the LlaSMol chemistry model:

True: Model disabled, agent responds without specialized chemistry inference
False: Model loaded (requires ≥15GB VRAM)

Impact:

When enabled, answer_chemistry_query tool raises RuntimeError
Agent will attempt to answer using GPT-4o only

API Keys

Required environment variables:

# .env file
OPENAI_API_KEY=sk-...

Best Practices

Start with default settings for most chemistry queries. The default configuration handles 90% of use cases effectively.

Monitor replanning attempts in the return tuple. High attempt counts (>30) may indicate unclear queries or model issues.

Use custom prompts when working in specialized chemistry domains (medicinal chemistry, materials science, etc.).

Don’t modify the tool execution order in the system prompt. The agent requires structure_chem_prompt to be called first.

Troubleshooting

Show GraphRecursionError

Cause: Agent exceeded recursion limit.Solutions:

Increase recursion_limit in config
Simplify the input query
Check if query is chemistry-related

Show Low VRAM Warnings

Cause: LOW_VRAM = True in configuration.Solutions:

Set LOW_VRAM = False if you have ≥15GB VRAM
Use GPU with sufficient memory
Expect reduced accuracy for specialized chemistry tasks

Show Empty or Invalid Plans

Cause: Planner failed to understand the query.Solutions:

Ensure query is chemistry-related
Provide more context in the prompt
Check if chemical names/SMILES are properly formatted

Agent API

Chemistry Tools

PubChem RAG

LlaSMol

Agent Configuration

Overview

LangGraph Configuration

Recursion Limit

Workflow Configuration

Agent Architecture

1. Planner

2. Executor

3. Replanner

State Management

Customization Examples

Custom System Prompt

Adjust Recursion Limit

Change Planning Model

Add Custom Chemistry Tool

Workflow Graph

Node Functions

Environment Configuration

VRAM Management

API Keys

Best Practices

Troubleshooting

Build docs developers (and LLMs) love

Agent API

Chemistry Tools

PubChem RAG

LlaSMol

Documentation Index

​Overview

​LangGraph Configuration

​Recursion Limit

​Workflow Configuration

​Agent Architecture

​1. Planner

​2. Executor

​3. Replanner

​State Management

​Customization Examples

​Custom System Prompt

​Adjust Recursion Limit

​Change Planning Model

​Add Custom Chemistry Tool

​Workflow Graph

​Node Functions

​Environment Configuration

​VRAM Management

​API Keys

​Best Practices

​Troubleshooting

Build docs developers (and LLMs) love

Overview

LangGraph Configuration

Recursion Limit

Workflow Configuration

Agent Architecture

1. Planner

2. Executor

3. Replanner

State Management

Customization Examples

Custom System Prompt

Adjust Recursion Limit

Change Planning Model

Add Custom Chemistry Tool

Workflow Graph

Node Functions

Environment Configuration

VRAM Management

API Keys

Best Practices

Troubleshooting