Agent Workflow

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()

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)

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

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

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

{
  "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"
  ]
}

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

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)

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>?"}

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

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}

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

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

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

{
  "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"
  ]
}

# 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"

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)

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

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

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

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

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

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