react_agent.py

from tool import Tool
from dotenv import load_dotenv
from groq import Groq
from message import Message
import re
import json
from colorama import Fore
from extraction_result import ExtractionResult

load_dotenv()

SYSTEM_PROMPT = Message("system", """
You are an AI model that is capable of function calls(one or more call) and you approach request in the following steps: Thought, Action, Observation
You will be provided with functions and its signatures under <tools></tools>
You may call function to assist user's request. During function call do not make any assumptions regarding function's name and it's properties. Refer to <tools> section to find
appropriate function and pay special attention to parameters so you have a better understanding of what types function expects as arguments to be passed.

If you decide that function call is required construct a JSON object(with double quotes) and place it within <tool_call></tool_call> Here's an example:
<tool_call>
{'name': 'function_name', 'arguments': {'arg1': 'val1', 'arg2': 'val2': 'arg2'}
</tool_call>

Here are the list of available tools that could be used:
<tools>
%s
</tools>

Here is an example request flow:
<question>What's the best restaurant for Khinkali in Tbilisi?</question>
<thought>User is willing to know which restaurant is considered best to eat khinkali in Tbilisi</thought>
<tool_call>{'name': 'get_restaurant_suggestion', 'parameters': {'city': 'tbilisi', 'food': 'Khinkali'}}</tool_call>

To break down: <question></question> Contains user's request.
<thought></thought> is your internal reasoning
<tool_call></tool_call> is constructed by you once you decide that tool_call is required and return this response.

If you return the <tool_call> to user, You will then be called with JSON that will be enclosed with <observation></observation> That contains the result of the previous tool_call.
Here's the example:
<observation>{'name': 'Pictogram', 'location': '31 Atoneli Street'}</observation>

Then you can respond to user with:
<response>To eat best khinkali in Tbilisi I suggest you to visit Restaurant Pictogram that is located on 31 Atoneli street</response>


Constraints to keep in mind:
If there is no tool available that could fulfil the user's request, proceed with regular flow. and enclose answer with <response></response>
DO NOT construct a <tool_call> if to be called tool does not exist within list of available tools under <tools></tools>
""")


class ReactAgent:
    def __init__(self, tools: list[Tool] | Tool, include_internal_logs=False) -> None:
        self.tools = tools if isinstance(tools, list) else [tools]
        self.tool_signatures = [t.signature for t in self.tools]
        self.tools_dict = {tool.name: tool for tool in self.tools}
        self.model = "llama-3.3-70b-versatile"
        self.agent = Groq()
        self.include_internal_logs = include_internal_logs

    def execute(self, message: Message) -> str:
        prompt_history: list[dict] = []
        SYSTEM_PROMPT.content = SYSTEM_PROMPT.content % "".join(self.tool_signatures)

        message.content = f'<question>{message.content}</question>'

        prompt_history.append(SYSTEM_PROMPT.to_dict())
        prompt_history.append(message.to_dict())

        while True:
            ai_response = self.agent.chat.completions.create(messages=prompt_history, model=self.model).choices[
                0].message.content
            prompt_history.append(Message("assistant", ai_response).to_dict())
            thought = self.extract_content_from_tag("thought", ai_response)
            tool_call = self.extract_content_from_tag("tool_call", ai_response)

            if self.include_internal_logs and thought.exists:
                print(Fore.RED + "Internal Thought process: " + thought.content)

            if tool_call.exists:
                if self.include_internal_logs:
                    print(Fore.CYAN + "Tool Call: " + tool_call.content)

                parameters = json.loads(tool_call.content)
                tool = self.tools_dict[parameters["name"]]
                result = tool(**parameters["arguments"])
                prompt_history.append(Message("user", f"<observation>{result}</observation>").to_dict())

            else:
                final_response = self.extract_content_from_tag("response", ai_response)
                break

        return final_response.content

    def extract_content_from_tag(self, tag_name: str, ai_response: str):
        re_match = re.search(f"<{tag_name}>(.*?)</{tag_name}>", ai_response, re.DOTALL)

        if re_match:
            content = re_match[1]
            return ExtractionResult(content, True)
        else:
            return ExtractionResult(None, False)