November 5, 2024

Building a Python React Agent Class: A Step-by-Step Guide

Roy Pasternak | CTO at Neradot
Building a Python React Agent Class: A Step-by-Step Guide

In this blog post, we’ll walk through the creation of a React agent class in Python. After going over our previous post (ReAct: Merging Reasoning and Action to Elevate AI Task Solving) you should already be familiar with the React flow. Here is a flow diagram to quickly recall how React handles the reason-action loop:

Now, let’s jump to the coding part. We’ll start with a minimal setup and gradually add functionality until we have a fully functional agent. Each step is highlighted with code snippets that focus on the differences from the previous step, providing a clear understanding of the incremental changes.

NOTE: If you’d like to see the full code directly, you can check out the complete implementation on react-simple
React flow diagram of reasoning and action
Step 0: Define the class interface

Here is the React class structure we are going to create together:

class React:
    def __init__(self, tools: list[callable], max_steps: int = 10, verbose: bool = False):
        # Initializes tools, settings, and internal state variables, such as `max_steps` and `verbose`.
        ...
        
    def reason(self) -> tuple[str, str, str, bool]:
        # Generates reasoning steps and determines the next action.
        ...
        
    def act(self, action: str, action_input: dict | str) -> str:
        # Executes actions by calling tool functions based on the current reasoning.
        ...
        
    def start(self, question: str):
        # Begins a new session with a given question.
        ...
        
    def next(self):
        # Advances the agent by processing reasoning, acting, and handling completion.
        ...
        
    def finish(self):
        # Ends the session, storing the final answer and completion status.
        ...

    def add_step(self, step: dict):
        # Logs each step of the agent’s process, such as thoughts, actions, and observations.
        ...
        
    def build_messages(self) -> list[dict]:
        # Constructs prompt messages for the reasoning process, integrating available tools and instructions.
        ...

  • __init__: Initializes the agent with tools, step limit (max_steps), and verbosity option (verbose).
  • reason: Produces reasoning steps, decides the next action, and determines if the agent should continue.
  • act: Executes a specified action using a tool with provided inputs, returning the result (observation).
  • start: Begins a session by setting up the agent with a question.
  • next: Advances the agent by running one reasoning/action cycle and checking for completion.
  • finish: Ends the session, saving the final answer and marking it complete.
  • add_step: Logs each step of the process, including thoughts, actions, and observations.
  • build_messages: Creates prompt messages that integrate tools and instructions for reasoning.

Step 1: Minimal Setup
import time

# Minimal setup with initialization and start method
class React:
    def __init__(self, verbose: bool = False):
        self.start_time: float | None = None
        self.is_started: bool = False
        self.intermediate_steps: list[dict] = []
        self.verbose = verbose

    def add_intermediate_step(self, intermediate_step: dict):
        self.intermediate_steps.append(intermediate_step)
        if self.verbose:
            print(intermediate_step)

    def start(self, question: str):
        if self.verbose:
          print(f"Starting agent with:")
        self.start_time = time.time()
        self.is_started = True
        self.intermediate_steps = []
        self.add_intermediate_step({"question": question})

# Testing the initial start function
react_agent = React(verbose=True)
react_agent.start("What is 2 + 2?")

In this initial step, we set up the basic structure for the React class. This minimal version includes an initialization method, where we define the attributes to track the agent’s state, and a simple start method to begin an interaction. Here, we also introduce the intermediate_steps list, which will store each phase of the agent’s process, such as questions, thoughts, actions, and observations. This list acts as a log for the interaction, making it possible for the agent to reference prior steps and maintain a coherent flow of actions and reasoning as we expand its functionality.

Output:
Starting agent with:
{'question': 'What is 2 + 2?'}

At this point, we’ve created a basic class that can start an interaction by saving the question as an intermediate step.

Step 2: Adding Reasoning with reason() and build_messages() Methods
import json
from utils import parse_json, get_completion
from prompts import (
    REACT_TOOLS_DESCRIPTION, 
    REACT_VALID_ACTIONS, 
    REACT_JSON_FORMAT, 
    REACT_PROCESS_FORMAT, 
    REACT_INTERMEDIATE_STEPS, 
    REACT_ADDITIONAL_INSTRUCTIONS
)


class React: 

	...
	
	# Adding the reasoning logic
	def reason(self) -> tuple[str, str, str, bool]:
	    messages = self.build_messages()
	    completion_response = get_completion(messages, model="gpt-4o", temperature=0, max_tokens=256, stop=["</reasoning>"])
	    completion = completion_response.choices[0].message.content
	    parsed_completion = parse_json(completion)
	    thought = parsed_completion["thought"]
	    action = parsed_completion["action"]
	    action_input = parsed_completion["action_input"]
	    is_final_answer = action == "Final Answer"
	    return thought, action, action_input, is_final_answer
	
	# Building the prompt messages from templates
	def build_messages(self) -> list[dict]:
      question = self.intermediate_steps[0]["question"]
      intermediate_steps=json.dumps(self.intermediate_steps[1:])
      system_prompt_message = \
              REACT_TOOLS_DESCRIPTION.format(tools_description="") + \
              REACT_VALID_ACTIONS.format(tools_names="") + \
              REACT_JSON_FORMAT + \
              REACT_PROCESS_FORMAT + \
              REACT_ADDITIONAL_INSTRUCTIONS + \
              REACT_INTERMEDIATE_STEPS.format(question=question, intermediate_steps=intermediate_steps)
        messages = [{ "role": "system", "content": system_prompt_message }]
        return messages
	 

Step 3: Driving Agent Progress with next() method
class React: 

	...
	
  @property
  def steps_count(self):
      return int(len(self.intermediate_steps) - 1 )
      
  def next(self):
      if not self.is_started:
          raise ValueError("React was not started")
      
      if self.verbose:
          print(f"Step {self.steps_count}")
      thought, action, action_input, is_final_answer = self.reason() 
      self.add_intermediate_step({"thought": thought, "action": action, "action_input": action_input})

react_agent = React(verbose=True)
react_agent.start("What is 2 + 2?")
react_agent.next()

With the reasoning structure in place, we now add the next method to control the agent's progression through steps. The next method initiates a reasoning process by calling reason, captures the generated thought, action, and action_input, and adds them as a new entry in intermediate_steps. we also introduce a steps_count property. This property calculates the number of reasoning of intermediate_steps (excluding the initial question).

Output:
Starting agent with:
{'question': 'What is 2 + 2?'}
Step 0
{'thought': 'The question asks for the sum of 2 + 2, which is a basic arithmetic operation that I can solve without any tools.', 'action': 'Final Answer', 'action_input': 4}

Here, the agent now has the capability to reason, generating an action and thought about the question. But, we are still missing the tool usage!

Step 4: Introducing Tools
class React: 
	...
	
# Adding tools functionality
	def __init__(self, tools: list[callable], verbose: bool = False):
			...
	    self.tools_dict: dict[str, callable] = {tool.__name__: tool for tool in tools}
	    self.tools_description: str = "\n".join([f"{tool_name}: {tool.__doc__}" for tool_name, tool in self.tools_dict.items()])
	    self.tools_names: list[str] = list(self.tools_dict.keys())
      if self.verbose:
          print("Initialized agent with tools:")
          print(f"{self.tools_description}")
          print()
          
def build_messages(self) -> list[dict]:
        ...
        system_prompt_message = \
                REACT_TOOLS_DESCRIPTION.format(tools_description=self.tools_description) + \
                REACT_VALID_ACTIONS.format(tools_names=self.tools_names) + \
        ...
def calculator(expression: str) -> float:
  """Evaluates a mathematical expression and returns the result"""
  return eval(expression)
react_agent = React(tools=[calculator], verbose=True)
react_agent.start("What is 2 + 2?")
react_agent.next()
react_agent.next()

In this step, we enhance the agent's abilities by adding a tool system. Tools are specific functions that the agent can call to perform certain actions, such as calculations or data retrieval. Each tool is passed to the React class upon initialization, allowing the agent to select and use them as needed.

Inside the __init__ method, we create a dictionary, tools_dict, that maps each tool’s name to its function. We also generate tools_description, a formatted string that provides a description of each tool’s purpose (using the tool’s docstring), and tools_names, a list of tool names.

Additionally, we update the build_messages method to integrate tools_description and tools_names directly into the system prompt, allowing the agent to reference the tools available during its reasoning process.

Expected Output:
Initialized agent with tools:
calculator: Evaluates a mathematical expression and returns the result

Starting agent with:
{'question': 'What is 2 + 2?'}
Step 0
{'thought': 'To solve the problem of finding what 2 + 2 equals, I should perform a simple arithmetic addition calculation.', 'action': 'calculator', 'action_input': '2 + 2'}
Step 1
{'thought': 'I should calculate the sum of 2 and 2 to find the answer.', 'action': 'calculator', 'action_input': '2 + 2'}

The agent now has a calculator tool it can reference to answer questions involving calculations.

However, while the agent can select the tool, it doesn’t yet actually call the tool function to produce results. This is why the output shows no progress based on the action taken.

Step 4: Taking Actions and Observations
class React:
		...
    @property
    def steps_count(self):
+        return int((len(self.intermediate_steps) - 1 ) / 2)

# Adding action and observation handling
def act(self, action: str, action_input: dict | str) -> str:
    tool_func = self.tools_dict[action]
    if isinstance(action_input, dict):
        tool_result = tool_func(**action_input)
    else:
        tool_result = tool_func(action_input)
    return tool_result

# Updating `next` method to include observation step
def next(self):
    ...
+    thought, action, action_input, is_final_answer = self.reason() 
+    self.add_intermediate_step({"thought": thought, "action": action, "action_input": action_input})
+    observation = self.act(action, action_input)
+    self.add_intermediate_step({"observation": observation})

react_agent = React(tools=[calculator], verbose=True)
react_agent.start("What is 2 + 2?")
react_agent.next()
react_agent.next()

Now, we add the act method, allowing the agent to execute actions based on parsed inputs from the reason method. This act method calls a relevant tool function and records the observation for each action taken.

Expected Output:
Initialized agent with tools:
calculator: Evaluates a mathematical expression and returns the result

Starting agent with:
{'question': 'What is 2 + 2?'}
Step 0
{'thought': 'To solve the problem, I need to add the numbers 2 and 2.', 'action': 'calculator', 'action_input': '2 + 2'}
{'observation': 4}
Step 1
{'thought': 'The calculation has been performed and the result is the addition of the two numbers.', 'action': 'Final Answer', 'action_input': 4}

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
...
KeyError: 'Final Answer'

The agent now has an act method to execute tool functions, and an observation step that records the results of each action. This allows the agent to see progress after calling a tool.

However, because the act method expects a tool name, the agent raises an error if the action is "Final Answer" rather than a tool. In the next step, we’ll handle this case to allow the agent to complete its reasoning when a final answer is reached.

Step 5: Adding a Finish Step
class React:
	...
	def __init__(self, tools: list[callable], verbose: bool = False):
			...
			self.not_finished: bool = True
	# Adding the finish method
  def finish(self):
      self.not_finished = False
      self.add_intermediate_step({"answer": self.last["action_input"]})

	# Updating `next` method to finalize when the final answer is reached
  def next(self):
    ...
    thought, action, action_input, is_final_answer = self.reason() 
    self.add_intermediate_step({"thought": thought, "action": action, "action_input": action_input})
    if is_final_answer:
        self.finish()
    else:
        observation = self.act(action, action_input)
        self.add_intermediate_step({"observation": observation})

react_agent = React(tools=[calculator], verbose=True)
react_agent.start("What is 2 + 2?")
while react_agent.not_finished:
    react_agent.next()

In this step, we add the finish method to mark the agent's session as complete. We add  additional attribute: not_finished, which becomes False when the agent has completed its task.

Additionally, we update the next method to handle the "Final Answer" action gracefully. If reason returns "Final Answer", next now calls finish rather than attempting to execute a tool, resolving the error from the previous step where the agent tried to run "Final Answer" as a tool name.

Output:
Initialized agent with tools:
calculator: Evaluates a mathematical expression and returns the result

Starting agent with:
{'question': 'What is 2 + 2?'}
Step 0
{'thought': 'I can calculate the sum of 2 + 2 using a basic mathematical operation.', 'action': 'calculator', 'action_input': '2 + 2'}
{'observation': 4}
Step 1
{'thought': 'The sum of 2 + 2 is calculated using a simple arithmetic operation, and the result was confirmed as 4.', 'action': 'Final Answer', 'action_input': 4}
{'answer': 4}

With the addition of the finish method and the updated next method, the agent now recognizes when it has reached a final answer and stops further actions, fully resolving the error from the previous step.

Final Touch: Adding Max Steps, Finish Reason, Duration, and Token Usage
class React: 
    def __init__(self, tools: list[callable], max_steps: int = 10, verbose: bool = False):
+       self.max_steps = max_steps
+       self.finish_reason: Literal["final answer", "max_steps_reached"] | None = None
+       self.completion_records: list[openai.types.completion.Completion] = []
  
    @property
+   def is_max_steps_reached(self):
+       return self.steps_count >= self.max_steps
    
    @property
+   def duration(self):
+       if not self.start_time or not self.end_time:
+           return None
+       return round((self.end_time - self.start_time)*1000, 3)
    
    @property
+   def token_usage(self):
+       usage = {'prompt_tokens': 0, 'completion_tokens': 0, 'total_tokens': 0}
+       for res in self.completion_records:
+           usage['prompt_tokens'] += res.usage.prompt_tokens
+           usage['completion_tokens'] += res.usage.completion_tokens
+           usage['total_tokens'] += res.usage.total_tokens
+       return usage

    def next(self):
        ...
        else:
            observation = self.act(action, action_input)
            self.add_intermediate_step({"observation": observation})
+           if self.is_max_steps_reached:
+               self.finish()
            
    def reason(self) -> tuple[str, str, str, bool]:
        ...
+       self.completion_records.append(completion_response)
        ...
    
    def finish(self):
        self.not_finished = False
        if not self.is_max_steps_reached:
            self.finish_reason = "final answer"
            self.add_intermediate_step({"answer": self.last["action_input"]})
        else:
            self.finish_reason = "max_steps_reached"
        self.end_time = time.time()

react_agent = React(tools=[calculator], verbose=True)
react_agent.start("What is 2 + 2?")
while react_agent.not_finished:
    react_agent.next()

In this final step, we add few features to enhance the agent's performance and monitoring:

  1. Max Steps: The max_steps parameter limits the number of reasoning steps the agent can take. If this limit is reached, finish is called automatically to prevent looping.
  2. Finish Reason: The finish_reason attribute records why the agent stopped: either "final answer" when a solution is reached or "max_steps_reached" if the step limit is hit.
  3. Duration: The duration property calculates the total runtime from start to finish in milliseconds, giving insight into performance.
  4. Token Usage: The token_usage property sums up tokens used in all completions, tracking prompt and completion tokens for resource monitoring.

With these additions, the agent is now fully functional, self-limiting, and capable of providing detailed usage insights.

Output:
Initialized agent with tools:
calculator: Evaluates a mathematical expression and returns the result

Starting agent with:
{'question': 'What is 2 + 2?'}
Step 0
{'thought': 'To find the sum of 2 + 2, I will use the calculator tool to perform the arithmetic operation.', 'action': 'calculator', 'action_input': '2 + 2'}
{'observation': 4}
Step 1
{'thought': 'I have used the calculator tool to perform the addition of 2 + 2 and obtained the result.', 'action': 'Final Answer', 'action_input': 4}
{'answer': 4}
Finish reason: final answer
Duration: 2686.926 ms
Token Usage: {'prompt_tokens': 652, 'completion_tokens': 100, 'total_tokens': 752}

Next

In the next post, we’ll dive into practical methods for evaluating agent performance, focusing on metrics that are particularly relevant for the React agent, such as correctness, efficiency, and response time. We’ll explore techniques for assessing the accuracy of the agent’s reasoning and action choices, testing for correctness in both intermediate steps and final answers.

Full Prompt
The available tools are:
{tools_description}

Use the tool name as the "action" value in the JSON response.

Valid actions are:
- "Final Answer"
- Any of {tools_names}

Respond in the following $JSON_BLOB format:
<reasoning>
{
        "thought": // Explain your reasoning for the chosen action, consider previous and subsequent steps
        "action": // The name of the tool to use
        "action_input": // The input required for the tool
}
</reasoning>

After you select an action, you will receive an observation. Then you can select another action or provide a final answer.
The pattern looks like this:

<reasoning>
$JSON_BLOB
</reasoning>
<observation>
{"observation": // action result}
</observation>
<reasoning>
$JSON_BLOB
</reasoning>
<observation>
{"observation": // action result}
</observation>

... (repeated until you have enough observations to answer the question)

<reasoning>
{
        "thought": // Explain why you have enough information to provide a final answer,
        "action": "Final Answer",
        "action_input": // Your final answer to the question
}
</reasoning>

Instructions:
1. Do not use comments in your JSON answer;
2. ALWAYS respond with a valid json blob of a single action;
3. ALWAYS think before choosing an action;
4. Respond in a JSON blob no matter what.

Here is the user question: 
"{question}"

Here are the intermediate steps so far:
{intermediate_steps}