import os from datetime import datetime from typing import List, Optional from dotenv import load_dotenv from loguru import logger from pydantic import BaseModel, Field from swarm_models import OpenAIChat from swarms import Agent from swarms.prompts.finance_agent_sys_prompt import ( FINANCIAL_AGENT_SYS_PROMPT, ) load_dotenv() # Get the OpenAI API key from the environment variable api_key = os.getenv("OPENAI_API_KEY") # Create an instance of the OpenAIChat class model = OpenAIChat( openai_api_key=api_key, model_name="gpt-4o-mini", temperature=0.1, max_tokens=2000, ) # Initialize the agent agent = Agent( agent_name="Financial-Analysis-Agent", system_prompt=FINANCIAL_AGENT_SYS_PROMPT, llm=model, max_loops=1, autosave=True, dashboard=False, verbose=True, dynamic_temperature_enabled=True, saved_state_path="finance_agent.json", user_name="swarms_corp", retry_attempts=1, context_length=200000, return_step_meta=False, # output_type="json", output_type=str, ) class ThoughtLog(BaseModel): """ Pydantic model to log each thought generated by the agent. """ thought: str timestamp: datetime = Field(default_factory=datetime.now) recursion_depth: int class MemoryLog(BaseModel): """ Pydantic model to log memory states during the agent's execution. """ thoughts: List[ThoughtLog] = [] final_result: Optional[str] = None completion_status: bool = False task: str class RecursiveAgent(Agent): """ An autonomous agent built on top of the Swarms Agent framework. Capable of recursively exploring tasks using a Tree of Thoughts mechanism. Attributes: - agent_name (str): The name of the agent. - system_prompt (str): The system prompt guiding the agent's behavior. - max_loops (int): The maximum depth for recursion in the Tree of Thoughts. - memory_limit (int): The maximum number of thought logs to store. - memory (MemoryLog): Pydantic model to store thoughts and logs. """ def __init__( self, agent_name: str, system_prompt: str, max_loops: int, memory_limit: int = 5, agent: Agent = agent, *args, **kwargs, ) -> None: """ Initialize the RecursiveAgent. :param agent_name: Name of the agent. :param system_prompt: The prompt guiding the agent's behavior. :param max_loops: The maximum number of recursive loops allowed. :param memory_limit: Maximum number of memory entries. :param kwargs: Additional arguments passed to the base Agent. """ super().__init__(agent_name=agent_name, **kwargs) self.system_prompt = system_prompt self.max_loops = max_loops self.memory = MemoryLog(task="") self.memory_limit = memory_limit # Max thoughts to store self.finished = False # Task completion flag self.agent = agent( agent_name=agent_name, system_prompt=system_prompt, max_loops=max_loops, ) logger.info( f"Initialized agent {self.agent_name} with recursion limit of {self.max_loops}" ) def add_to_memory( self, thought: str, recursion_depth: int ) -> None: """ Add a thought to the agent's memory using the Pydantic ThoughtLog model. :param thought: The thought generated by the agent. :param recursion_depth: The depth of the current recursion. """ if len(self.memory.thoughts) >= self.memory_limit: logger.debug( "Memory limit reached, discarding the oldest memory entry." ) self.memory.thoughts.pop(0) # Maintain memory size thought_log = ThoughtLog( thought=thought, recursion_depth=recursion_depth ) self.memory.thoughts.append(thought_log) logger.info( f"Added thought to memory at depth {recursion_depth}: {thought}" ) def check_if_finished(self, current_thought: str) -> bool: """ Check if the task is finished by evaluating the current thought. :param current_thought: The current thought or reasoning result. :return: True if task completion keywords are found, else False. """ # Define criteria for task completion based on keywords completion_criteria = [ "criteria met", "task completed", "done", "fully solved", ] if any( keyword in current_thought.lower() for keyword in completion_criteria ): self.finished = True self.memory.completion_status = True logger.info( f"Task completed with thought: {current_thought}" ) return self.finished def run_tree_of_thoughts( self, task: str, current_depth: int = 0 ) -> Optional[str]: """ Recursively explore thought branches based on the Tree of Thoughts mechanism. :param task: The task or query to be reasoned upon. :param current_depth: The current recursion depth. :return: The final solution or message indicating task completion or failure. """ logger.debug(f"Current recursion depth: {current_depth}") if current_depth >= self.max_loops: logger.warning( "Max recursion depth reached, task incomplete." ) return "Max recursion depth reached, task incomplete." # Generate multiple possible thoughts/branches using Swarms logic response = self.generate_thoughts(task) thoughts = self.extract_thoughts(response) self.memory.task = task # Log the task in memory # Store thoughts in memory for idx, thought in enumerate(thoughts): logger.info( f"Exploring thought {idx + 1}/{len(thoughts)}: {thought}" ) self.add_to_memory(thought, current_depth) if self.check_if_finished(thought): self.memory.final_result = ( thought # Log the final result ) return f"Task completed with thought: {thought}" # Recursive exploration result = self.run_tree_of_thoughts( thought, current_depth + 1 ) if self.finished: return result return "Exploration done but no valid solution found." def generate_thoughts(self, task: str) -> str: """ Generate thoughts for the task using the Swarms framework. :param task: The task or query to generate thoughts for. :return: A string representing multiple thought branches generated by Swarms logic. """ logger.debug(f"Generating thoughts for task: {task}") response = self.agent.run( task ) # Assuming Swarms uses an LLM for thought generation return response def extract_thoughts(self, response: str) -> List[str]: """ Extract individual thoughts/branches from the LLM's response. :param response: The response string containing multiple thoughts. :return: A list of extracted thoughts. """ logger.debug(f"Extracting thoughts from response: {response}") return [ thought.strip() for thought in response.split("\n") if thought ] def reflect(self) -> str: """ Reflect on the task and thoughts stored in memory, providing a summary of the process. The reflection will be generated by the LLM based on the stored thoughts. :return: Reflection output generated by the LLM. """ logger.debug("Running reflection on the task.") # Compile all thoughts into a prompt for reflection thoughts_for_reflection = "\n".join( [ f"Thought {i + 1}: {log.thought}" for i, log in enumerate(self.memory.thoughts) ] ) reflection_prompt = ( f"Reflect on the following task and thoughts:\n" f"Task: {self.memory.task}\n" f"Thoughts:\n{thoughts_for_reflection}\n" "What did we learn from this? How could this process be improved?" ) # Use the agent's LLM to generate a reflection based on the memory reflection_response = self.agent.run(reflection_prompt) self.memory.final_result = reflection_response logger.info(f"Reflection generated: {reflection_response}") return reflection_response # # Example usage of the RecursiveAgent # if __name__ == "__main__": # # Example initialization and running # agent_name = "Autonomous-Financial-Agent" # system_prompt = "You are a highly intelligent agent designed to handle financial queries efficiently." # max_loops = 1 # # Initialize the agent using Swarms # agent = RecursiveAgent( # agent_name=agent_name, # system_prompt=system_prompt, # max_loops=max_loops # ) # # Define the task for the agent # task = "How can I establish a ROTH IRA to buy stocks and get a tax break? What are the criteria?" # # Run the tree of thoughts mechanism # result = agent.run_tree_of_thoughts(task) # logger.info(f"Final result: {result}") # # Perform reflection # reflection = agent.reflect() # logger.info(f"Reflection: {reflection}")