Agent Workflow - ChemAgent

ChemAgent implements a dynamic plan-execute-replan pattern using LangGraph, allowing it to adapt its strategy based on validation results and handle complex multi-step chemistry queries.

Overview

The agent workflow consists of three primary nodes that form a continuous cycle:

Planner

Creates a step-by-step execution plan from the user’s query

Executor (Agent)

Executes the first step of the plan using available tools

Replanner

Evaluates results and either returns a response or creates a new plan

Workflow Graph

The LangGraph state machine is defined as follows:

from langgraph.graph import StateGraph, START, END

workflow = StateGraph(PlanExecute)

# Add nodes
workflow.add_node("planner", plan_step)
workflow.add_node("agent", execute_step)
workflow.add_node("replan", replan_step)

# Define edges
workflow.add_edge(START, "planner")
workflow.add_edge("planner", "agent")
workflow.add_edge("agent", "replan")
workflow.add_conditional_edges("replan", should_end, ["agent", END])

app = workflow.compile()

See plan_execute_agent/rdkit_agent.py:199 for the complete graph definition.

State Structure

The agent maintains state across the workflow:

from typing import Annotated, List, Tuple
from typing_extensions import TypedDict
import operator

class PlanExecute(TypedDict):
    input: str                                    # Original user query
    plan: List[str]                               # Current execution plan
    past_steps: Annotated[List[Tuple], operator.add]  # History of completed steps
    response: str                                 # Final response (when ready)

The past_steps field uses operator.add to automatically accumulate results across iterations. See plan_execute_agent/rdkit_agent.py:75 for state definition.

Node 1: Planner

Purpose

The planner analyzes the input query and generates a structured execution plan.

Implementation

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

See plan_execute_agent/rdkit_agent.py:172 for implementation.

Prompt Structure

The planner uses a specialized prompt to ensure chemistry-specific planning:

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}"),
])

See plan_execute_agent/rdkit_agent.py:95 for the complete prompt.

Structured Output

The planner returns a Pydantic model for type safety:

from pydantic import BaseModel, Field

class Plan(BaseModel):
    """Plan to follow in future"""
    steps: List[str] = Field(
        description="different steps to follow, should be in sorted order"
    )

Example Plan

For the query: "What is the SMILES for <IUPAC> aspirin </IUPAC>?" The planner generates:

{
  "steps": [
    "Call 'structure_chem_prompt' to tag chemical information",
    "Call 'answer_chemistry_query' to get SMILES representation",
    "Call 'validate_smiles_rdkit' to validate the SMILES output",
    "Return the validated SMILES to the user"
  ]
}

Node 2: Executor (Agent)

Purpose

Executes the first step of the current plan using the appropriate tool.

Implementation

async def execute_step(state: PlanExecute):
    plan = state["plan"]
    plan_str = "\n".join(f"{i+1}. {step}" for i, step in enumerate(plan))
    task = plan[0]
    task_formatted = f"""For the following plan:
{plan_str}

You are tasked with executing step {1}, {task}."""
    
    agent_response = await agent_executor.ainvoke(
        {"messages": [("user", task_formatted)]}
    )
    
    return {
        "past_steps": [(task, agent_response["messages"][-1].content)],
    }

See plan_execute_agent/rdkit_agent.py:158 for implementation.

Tool Selection

The executor (powered by GPT-4o) autonomously selects which tool to use based on the task description:

from plan_execute_agent.chem_tools import *

tools = [
    structure_chem_prompt,      # Step 1: Tag chemical information
    answer_chemistry_query,     # Step 2: Query LlaSMol
    validate_smiles_rdkit       # Step 3: Validate output
]

agent_executor = create_react_agent(llm, tools, state_modifier=prompt)

See plan_execute_agent/rdkit_agent.py:45 for tool registration.

Execution Example

Task: “Call ‘structure_chem_prompt’ to tag chemical information” Agent reasoning:

Thought: I need to use the structure_chem_prompt tool to tag the IUPAC name.
Action: structure_chem_prompt
Action Input: "What is the SMILES for aspirin?"
Observation: {"new_prompt": "What is the SMILES for <IUPAC> aspirin </IUPAC>?"}

Result added to past_steps:

("Call 'structure_chem_prompt' to tag chemical information", 
 "Successfully structured prompt: What is the SMILES for <IUPAC> aspirin </IUPAC>?")

Node 3: Replanner

Purpose

Evaluates completed steps and decides whether to:

End: Return final response to user
Continue: Generate new plan and execute next step

Implementation

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}

See plan_execute_agent/rdkit_agent.py:177 for implementation.

Decision Models

The replanner uses union types to represent decisions:

from typing import Union

class Response(BaseModel):
    """Response to user."""
    response: str

class Act(BaseModel):
    """Action to perform."""
    action: Union[Response, Plan] = Field(
        description="Action to perform. If you want to respond to user, use Response. "
        "If you need to further use tools to get the answer, use Plan."
    )

See plan_execute_agent/rdkit_agent.py:115 for decision model definitions.

Replanner Prompt

The replanner receives full context of the conversation:

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.
    
    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."""
)

See plan_execute_agent/rdkit_agent.py:130 for the complete prompt.

Handling Validation Errors

The replanner is crucial for handling SMILES validation failures:

Error Detection

When validate_smiles_rdkit returns invalid SMILES:

{
  "valid": False,
  "error_message": "Unclosed Ring with Validity Vector: 111110,Invalid Character with Validity Vector: 111101"
}

See plan_execute_agent/chem_tools.py:196 for error recording.

Replanning Strategy

The replanner analyzes the error and creates a corrective plan: Example:

{
  "plan": [
    "Review the validation error indicating unclosed ring",
    "Call 'answer_chemistry_query' again with clearer instructions",
    "Call 'validate_smiles_rdkit' to verify the corrected SMILES"
  ]
}

This allows the agent to self-correct based on validation feedback.

Recursion Limit and Attempt Tracking

Preventing Infinite Loops

The system includes safeguards against infinite replanning:

# Configuration
config = {"recursion_limit": 50}

# Attempt tracking
replanning_attempts = 1

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

See plan_execute_agent/rdkit_agent.py:188 for the termination logic.

Recursion Error Handling

When the limit is exceeded, the system catches the error:

try:
    responses = await process_events(inputs, config)
    curr_completed = True
except GraphRecursionError:
    print("Run failed due to Graph Recursion Error")
    curr_completed = False
    replanning_attempts = 0
    return (None, curr_completed, replanning_attempts, None, "", formatted_input)

See plan_execute_agent/rdkit_agent.py:417 for error handling.

Tracking Metrics

The system logs attempt statistics for analysis:

result, completed, attempts, llasmol_response, llasmol_errors, formatted_input = 
    await process_input(input_prompt, image_path, use_rag=use_rag)

writer.writerow({
    "query": input_prompt,
    "attempts": attempts,
    "completed": completed,
    "errors_present": bool(llasmol_errors),
    "llasmol_errors": llasmol_errors.strip()[:500],
})

See plan_execute_agent/rdkit_agent.py:489 for logging implementation.

Complete Workflow Example

Let’s trace a full execution for: "What is the molecular formula of aspirin?"

Iteration 1

Planner Output:

{"plan": ["Structure the query", "Answer with LlaSMol", "Validate result", "Return answer"]}

Executor: Calls structure_chem_prompt("What is the molecular formula of aspirin?") Result: {"new_prompt": "What is the molecular formula of <IUPAC> aspirin </IUPAC>?"} Replanner Decision: Continue with remaining steps

Iteration 2

Replanner Output:

{"plan": ["Answer with LlaSMol", "Validate result", "Return answer"]}

Executor: Calls answer_chemistry_query("What is the molecular formula of <IUPAC> aspirin </IUPAC>?") Result: "<MOLFORMULA> C9H8O4 </MOLFORMULA>" Replanner Decision: Continue to validation

Iteration 3

Replanner Output:

{"plan": ["Validate result", "Return answer"]}

Executor: Validation not needed (no SMILES output) Replanner Decision: Return response

Iteration 4

Replanner Output:

{"response": "The molecular formula of aspirin is C9H8O4."}

Condition: should_end() returns END Final: Return to user

Low VRAM Mode

When LOW_VRAM=True in configuration:

if LOW_VRAM:
    raise RuntimeError(
        "answer_chemistry_query tool cannot be used with LOW_VRAM enabled."
    )

The agent will either:

Respond directly without using LlaSMol
Use only GPT-4o for general chemistry questions
Return early with available information

See plan_execute_agent/chem_tools.py:148 for LOW_VRAM handling.

Best Practices

Optimal Planning

Keep plans simple (3-4 steps maximum)
Ensure each step has clear success criteria
Include validation steps for SMILES outputs

Error Recovery

Use validity vectors to identify specific errors
Limit replanning attempts to avoid infinite loops
Provide clear error messages to the replanner

Performance

Enable RAG only when additional context is needed
Monitor replanning_attempts to identify problematic queries
Use LOW_VRAM mode for CPU-only deployments

Next Steps

Architecture

Learn about the system architecture

Chemistry Tags

Master the tag system for queries

​Overview

​Workflow Graph

​State Structure

​Node 1: Planner

​Purpose

​Implementation

​Prompt Structure

​Structured Output

​Example Plan

​Node 2: Executor (Agent)

​Purpose

​Implementation

​Tool Selection

​Execution Example

​Node 3: Replanner

​Purpose

​Implementation

​Decision Models

​Replanner Prompt

​Handling Validation Errors

​Error Detection

​Replanning Strategy

​Recursion Limit and Attempt Tracking

​Preventing Infinite Loops

​Recursion Error Handling

​Tracking Metrics

​Complete Workflow Example

​Iteration 1

​Iteration 2

​Iteration 3

​Iteration 4

​Low VRAM Mode

​Best Practices

​Next Steps

Architecture

Chemistry Tags

Overview

Workflow Graph

State Structure

Node 1: Planner

Purpose

Implementation

Prompt Structure

Structured Output

Example Plan

Node 2: Executor (Agent)

Purpose

Implementation

Tool Selection

Execution Example

Node 3: Replanner

Purpose

Implementation

Decision Models

Replanner Prompt

Handling Validation Errors

Error Detection

Replanning Strategy

Recursion Limit and Attempt Tracking

Preventing Infinite Loops

Recursion Error Handling

Tracking Metrics

Complete Workflow Example

Iteration 1

Iteration 2

Iteration 3

Iteration 4

Low VRAM Mode

Best Practices

Next Steps