Simple ReAct Agent using Hugging Face Transformers - No Agent Framework, No paid APIs

from transformers import AutoTokenizer, AutoModelForCausalLM, MistralForCausalLM, StoppingCriteria, StoppingCriteriaList
import json
import re
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
print("Using device:", device)

model_name = "Qwen/Qwen2.5-3B-Instruct"

tokenizer = AutoTokenizer.from_pretrained(model_name)
# the tokenizer is built within the model, so you always use .from_pretrained(model_name)

dtype = torch.float16 if device == "cuda" else torch.float32
# use dtype, to load model weights in FP16 instead of FP32 for less powerful processors,halves memory, often faster.

# this may take a few minutes to load the model depending on your runtime
# you may ignore the accelerate warning

if device == "cuda":
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        dtype=dtype,
        device_map="auto",
        low_cpu_mem_usage=True
    )
else:
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        dtype=dtype,
        low_cpu_mem_usage=True,
    ).to(device)

import math

# define operations calculator will be allowed to execute
ALLOWED_NAMES = {
    name: value
    for name, value in math.__dict__.items()
    if not name.startswith("__")
}

list(ALLOWED_NAMES.keys())[:10]

len(ALLOWED_NAMES)

def calculator(expression: str) -> str:
    try:
        result = eval(
            expression,
            {"__builtins__": {}},  #removes python builtin functions from execution
            ALLOWED_NAMES          #specifies the only operations allowed
        )
        return f"The result is {result}"
    except Exception as e:
        return f"Error in calculation: {e}"

tools = [{
    "type": "function",
    "function": {
        "name": "calculator",
        "description": "Evaluates a mathematical expression and returns the result. Supports basic arithmetic operations like addition (+), subtraction (-), multiplication (*), division (/), and exponentiation (**).",
        "parameters": {
            "type": "object",
            "properties": {
                "expression": {
                    "type": "string",
                    "description": "The mathematical expression to evaluate, e.g., '2 + 2' or '(10 * 5) / 2'"
                }
            },
            "required": ["expression"],
            "additionalProperties": False
        }
    }
}]

# Define a global tool registry
TOOL_REGISTRY = {
    "calculator": calculator
}

class ReActAgent:
    def __init__(self, model, tokenizer, tools):
        self.model = model
        self.tokenizer = tokenizer
        self.tools = tools
        self.tool_names = [tool["function"]["name"] for tool in tools]

        # Automatically map tool names to functions from registry
        self.tool_functions = {
            name: TOOL_REGISTRY[name]
            for name in self.tool_names
            if name in TOOL_REGISTRY
        }

available_ops = ", ".join(sorted(ALLOWED_NAMES.keys()))
available_ops

def format_prompt(self, question, max_iterations=10):
    """Construct the exact prompt template with tool descriptions"""
    import json

    tool_names_str = ", ".join(self.tool_names)
    tools_json = json.dumps(self.tools, indent=2)
    available_ops = ", ".join(sorted(ALLOWED_NAMES.keys()))

    system_prompt = f"""You are a ReAct agent capable of using tools to answer questions.
You will think through each problem step-by-step, use tools as necessary, and provide accurate answers.

You have access to the following tools:

{tools_json}

You must always use the tools for evaluating mathematical operations. If needed, you may break down a problem into multiple
tool calls to evaluate the final answer.
When a problem requires multiple steps (multiple tool calls), do the following:
    1. make a plan
    2. at each step, review the plan and make sure you are on track
    3. execute all the steps before you answer the question.

IMPORTANT CALCULATOR CONSTRAINTS:
- The calculator can ONLY use these operations: {available_ops}
- The calculator uses Python's math module - use functions like sqrt(x), pow(x,y), sin(x), etc.
- Expression examples: "sqrt(144)", "pow(2, 3)", "sin(pi/2)"
- For rounding: If a result has more than 2 decimal places, round it using this expression: "floor(result * 100 + 0.5) / 100"
  Do not use round(), you can't execute Python builtins
  Example: To round 3.14159 to 2 decimals, use "floor(3.14159 * 100 + 0.5) / 100" which gives 3.14

Use the tools by specifying a json blob with an 'action' key (tool name) and an 'action_input' key (the tool input, matching the parameters schema above).

Valid actions: {tool_names_str}

The $JSON_BLOB must only contain a SINGLE action and must be formatted as markdown. Do NOT return a list of multiple actions.

Example:
Action:
```json
{{
    "action": "calculator",
    "action_input": "5+2"
}}
```

ALWAYS use the following format:

Question: the input question you must answer
Thought: you should always think about what action to take. Only one action at a time.
Action:
```json
{{JSON_BLOB}}
```
Observation: the result of the action

This Thought/Action/Observation cycle can repeat up to {max_iterations} times. Take several steps as needed, but use your iterations wisely.

You must always end your output with the following format:
Thought: I now know the final answer
Final Answer: the final answer to the original input question

Now begin! Reminder to ALWAYS use the exact characters 'Final Answer:' when you provide a definitive answer.

Question: {question}"""

    return system_prompt

# Attach the method to the class
ReActAgent.format_prompt = format_prompt

#test
agent = ReActAgent(model, tokenizer, tools)
prompt = agent.format_prompt(question="What is the square root of 25?")
print(prompt)

class StopOnObservation(StoppingCriteria):
    def __init__(self, target_sequence, prompt, tokenizer):
        """
        Stop generation when target_sequence (e.g., 'Observation:') appears in generated text.

        Args:
            target_sequence: String to watch for (e.g., 'Observation:')
            prompt: The original prompt (to exclude it from checking)
            tokenizer: Tokenizer to decode tokens
        """
        self.target_sequence = target_sequence
        self.prompt = prompt
        self.prompt_length = len(tokenizer.encode(prompt, add_special_tokens=True))
        self.tokenizer = tokenizer

    def __call__(self, input_ids, scores, **kwargs):
        # Get only the tokens AFTER the prompt (the newly generated tokens)
        generated_token_ids = input_ids[0][self.prompt_length:]

        # Decode only the new tokens
        new_text = self.tokenizer.decode(generated_token_ids, skip_special_tokens=True)

        # Check if target sequence appears in the newly generated text only
        if self.target_sequence in new_text:
            return True

        return False

# Test
stop_criteria = StopOnObservation(
    target_sequence="Observation:",
    prompt=prompt,
    tokenizer=tokenizer
)
# Inspect the object
print("Target sequence:", stop_criteria.target_sequence)
print("Prompt length:", len(stop_criteria.prompt))
print("Tokenizer EOS token:", stop_criteria.tokenizer.eos_token)

def generate_response(self, prompt, stopping_criteria):
    """Generate a response from the model given a prompt and stopping criteria"""
    inputs = self.tokenizer(prompt, return_tensors="pt") #return_tensors="pt" returns the tokenized output as PyTorch tensors
    input_ids = inputs.input_ids.to(device) # move the tensor to the specified device (CPU or GPU)
    attention_mask = inputs.attention_mask.to(device) # move the attention mask tensor to the specified device

    outputs = self.model.generate(
        input_ids,
        attention_mask=attention_mask,
        max_new_tokens=750,
        temperature=0.8,
        pad_token_id=self.tokenizer.eos_token_id,# pad (fill shorter sequences) with end-of-string token
        repetition_penalty=1.1,
        do_sample=True, #sample the next token from the probability distribution
        stopping_criteria=stopping_criteria
    )

    # Decode the FULL output
    full_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Remove the prompt
    if full_text.startswith(prompt):
        generated_text = full_text[len(prompt):].strip()
    else:
        # Fallback: The prompt might not match exactly due to tokenization
        # In this case, warn and return what we can
        print("WARNING: Generated text doesn't start with prompt - returning full output")
        print(f"Prompt ends with: ...{prompt[-100:]}")
        print(f"Output starts with: {full_text[:100]}...")
        generated_text = full_text

    return generated_text

# Attach the method to the class
ReActAgent.generate_response = generate_response

# Test - this may take a few minutes on colab, depending on the model used
response = agent.generate_response(prompt, stopping_criteria=StoppingCriteriaList([stop_criteria]))
response

def extract_action(self, response):
    """Extract the JSON blob related to an action from the model's response."""

    # Try 1: Look for JSON with markdown code fence
    match = re.search(
        r'```json\s*(\{.*?\})\s*```',
        response,
        re.DOTALL
    )

    if match:
        json_str = match.group(1).strip()
        try:
            return json.loads(json_str)
        except json.JSONDecodeError as e:
            print(f"JSON parsing error: {e}")
            print(f"Malformed JSON: {json_str}")
            return None

    # Try 2: Look for JSON object with "action" and "action_input" keys (no markdown fence)
    match = re.search(
        r'\{[^{}]*"action"[^{}]*"action_input"[^{}]*\}',
        response,
        re.DOTALL
    )

    if match:
        json_str = match.group(0).strip()
        try:
            return json.loads(json_str)
        except json.JSONDecodeError as e:
            print(f"JSON parsing error: {e}")
            print(f"Malformed JSON: {json_str}")
            return None

    return None

# Attach the method to the class
ReActAgent.extract_action = extract_action

# Test
action = extract_action(agent, response)
action

def interact(self, question, max_iterations=10):
    """Answer the question iteratively using the structured ReAct process"""

    print(f"Question: {question}\n")
    print("="*60)


    # Generate the initial prompt
    initial_prompt = self.format_prompt(question, max_iterations)
    prompt = initial_prompt
    print(f"\nInitial prompt contains {len(prompt)} characters")
    print(f"Max iterations: {max_iterations}\n")

    # The ReAct loop
    iteration = 0

    while iteration < max_iterations:
        iteration += 1

        print(f"\nITERATION {iteration}")
        print("-"*60)


        # Check if we should use stopping criteria
        # Only check the conversation history (not the initial system prompt)
        conversation_history = prompt[len(initial_prompt):]

        if "I now know the final answer" in conversation_history or "Final Answer:" in conversation_history:
            # Final answer mode - let the model complete naturally without stopping
            print("Final answer mode detected - disabling stop criteria")
            stopping_criteria_list = StoppingCriteriaList([])
        else:
            # Normal ReAct iteration - stop at "Observation:"
            stopping_criteria = StopOnObservation(
                target_sequence="Observation:",
                prompt=prompt,
                tokenizer=self.tokenizer
            )
            stopping_criteria_list = StoppingCriteriaList([stopping_criteria])

        print("Generating model response...")
        new_content = self.generate_response(
            prompt,
            stopping_criteria=stopping_criteria_list
        )

        print("\nModel Output:")
        print(new_content)

        # Check if we have a final answer
        if "Final Answer:" in new_content:
            final_answer_match = re.search(r"Final Answer:\s*(.+)", new_content, re.DOTALL)
            if final_answer_match:
                final_answer = final_answer_match.group(1).strip()
                print(f"\n{'='*60}")
                print("AGENT COMPLETED SUCCESSFULLY")
                print(f"{'='*60}\n")
                return final_answer
            else:
                # "Final Answer:" found but nothing after it
                print("Warning: 'Final Answer:' found but no answer provided")
                print("Continuing to next iteration...")
                continue

        # Check if we have an action to execute
        if "Action:" in new_content:
            print("\nExtracting and executing action...")
            action_json = self.extract_action(new_content)

            if action_json is None:
                print("Failed to parse action")
                return "I am unable to answer the question (failed to parse action)."

            tool_name = action_json.get("action")
            tool_input = action_json.get("action_input")

            print(f"   Tool: {tool_name}")
            print(f"   Input: {tool_input}")

            if tool_name not in self.tool_names:
                print(f"Unknown tool: {tool_name}")
                return f"I am unable to answer the question (unknown tool: {tool_name})."

            # Execute the tool
            if tool_name in self.tool_functions:
                tool_func = self.tool_functions[tool_name]

                # Handle both dict and string input formats
                # The model should provide a dict, but sometimes provides a string directly
                if isinstance(tool_input, dict):
                    # Correct format: {"expression": "..."}
                    result = tool_func(**tool_input)
                elif isinstance(tool_input, str):
                    # Fallback: model provided string directly
                    # Assume the tool takes 'expression' as the parameter name
                    result = tool_func(expression=tool_input)
                else:
                    result = f"Error: Unexpected tool_input type: {type(tool_input)}"

                print(f"\nTool Result: {result}")
            else:
                result = f"Error: Tool function '{tool_name}' not found"
                print(f"{result}")

            # MEMORY MECHANISM: Append the observation result to the prompt
            # This is how the agent maintains context across iterations.
            # The prompt grows to include all previous thoughts, actions, and observations,
            # allowing the model to see the full reasoning chain and build on previous results.
            prompt += f"\n{new_content} {result}\n"

            print(f"\nUpdated context (prompt now contains {len(prompt)} characters)")
            print("The model will see this history in the next iteration.\n")
        else:
            print("No action or final answer found")
            return "I am unable to answer the question (no action or final answer found)."

    print(f"\n{'='*60}")
    print("MAX ITERATIONS REACHED")
    print(f"{'='*60}\n")
    return "I am unable to answer the question (max iterations reached)."

# Attach the method to the class
ReActAgent.interact = interact

# Test 1
agent = ReActAgent(model, tokenizer, tools)
answer = agent.interact("What is the square root of 144?")
print(f"Final Answer: {answer}")

# Test 2
agent = ReActAgent(model, tokenizer, tools)
answer = agent.interact("What is (12 + 8) * 5 / 2?")
print(f"Final Answer: {answer}")

# Test 3
agent = ReActAgent(model, tokenizer, tools)
answer = agent.interact("What is the square root of (10 + 8) * 5 / 2?")
print(f"Final Answer: {answer}")