Reduced duplicate code by adding 'itf_agent'

itf_agent.py

@@ -0,0 +1,99 @@
+from llama_index.core.workflow import Context
+
+
+import os
+from typing import List
+
+from args import LLMInterface
+from llm_factory import LLMFactory
+from llama_index.core.agent.workflow import AgentWorkflow
+
+
+class IAgent():
+    def __init__(self, temperature, max_tokens, sys_prompt_file, llm_itf: LLMInterface):
+        self.temperature, self.max_tokens = temperature, max_tokens
+        # Load the system prompt from a file
+        system_prompt_path = os.path.join(os.getcwd(), "system_prompts", sys_prompt_file)
+        self.system_prompt = ""
+        with open(system_prompt_path, "r") as file:
+            self.system_prompt = file.read().strip()
+        # Initialize the tool agents
+        self.tools = self.setup_tools()
+        self.slaves: List[IAgent] = self.setup_slaves()
+        # Define the LLM and agent
+        self.llm = LLMFactory.create(llm_itf, self.system_prompt, temperature, max_tokens)
+        self.agent = self._setup_agent()
+        self.ctx = Context(self.agent)
+    
+    def setup_tools(self) -> List:
+        """
+        Abstract method to set up the tools.
+        This method must be overridden by subclasses to define custom tools this agent can use.
+        """
+        raise NotImplementedError("Subclasses must implement the setup_tools method.")
+    
+    def setup_slaves(self) -> List:
+        """
+        Abstract method to set up the slave agents.
+        This method must be overridden by subclasses to define custom sub-agents this agent can use.
+        """
+        raise NotImplementedError("Subclasses must implement the setup_slaves method.")
+    
+    def _setup_agent(self) -> AgentWorkflow:
+        """
+        Initializes and returns an agent workflow based on the presence of tools and slaves.
+        If both `self.tools` and `self.slaves` are empty, it sets up a default agent using the provided language model (`self.llm`).
+        Otherwise, it creates an agent workflow using the combined list of tools and slaves with the language model.
+        Returns:
+            AgentWorkflow: An instance of the agent workflow configured with the appropriate tools and language model.
+        """
+        if not self.tools and not self.slaves:
+            return AgentWorkflow.setup_agent(llm=self.llm)
+        
+        # Create tools from slaves: each tool calls slave.query(question) asynchronously
+        slave_tools = []
+        for slave in self.slaves:
+            slave_tools.append(slave.query)
+
+        self.tools.extend(slave_tools)
+
+        return AgentWorkflow.from_tools_or_functions(
+        self.tools,
+        llm=self.llm
+        )
+
+    def get_system_prompt(self) -> str:
+        """
+        Retrieves the system prompt.
+
+        Returns:
+            str: The system prompt string.
+        """
+        return self.system_prompt
+
+    async def query(self, question: str) -> str:
+        """
+        Asynchronously queries the agent with a given question and returns the response.
+
+        Args:
+            question (str): The question to be sent to the agent.
+
+        Returns:
+            str: The response from the agent as a string.
+        """
+        response = await self.agent.run(question, ctx=self.ctx)
+        response = str(response)
+        return response
+
+    def clear_context(self):
+        """
+        Clears the current context of the agent, resetting any conversation history.
+        This is useful when starting a new conversation or when the context needs to be refreshed.
+        """
+        self.ctx = Context(self.agent)
+
+        if not self.slaves:
+            return
+
+        for slave in self.slaves:
+            slave.clear_context()

management.py

@@ -72,18 +72,18 @@ class Manager:
         self.solver = Solver(temperature, max_tokens)
         self.summarizer = Summarizer(temperature, max_tokens)
 
-    async def query(self, question: str, remember = True) -> str:
+    async def query(self, question: str, has_context = True) -> str:
         """
         Process a question using the manager agent and return a response.
         
         Args:
             question: The question or task to process
-            remember: Whether to maintain context between queries (default: True)
+            has_context: Whether to maintain context between queries (default: True)
             
         Returns:
             The agent's response as a string
         """
-        if remember:
+        if has_context:
             response = await self.agent.run(question, ctx=self.ctx)
         else:
             response = await self.agent.run(question)
@@ -117,7 +117,7 @@ class Manager:
         observation = ""
         if self.current_depth < self.max_depth:
             for task in tasks:
-                solution = await self.query(task, remember=False)
+                solution = await self.query(task, has_context=False)
                 response = f"For task:\n\n{task}\n\nThe following break up has been provided:\n\n{solution}\n\n"
                 observation += response
         elif try_solving:

solver.py

@@ -1,413 +1,131 @@
-from llama_index.core.agent.workflow import AgentWorkflow
-from llama_index.core.tools import FunctionTool
-from llama_index.core.workflow import Context
+from typing import List
 
-import asyncio
-import os
-
-from llm_factory import LLMFactory
+from itf_agent import IAgent
 from toolbox import Toolbox
 from args import Args
 
 
-class Summarizer:
+class Summarizer(IAgent):
     def __init__(self, temperature, max_tokens):
-        # Load the system prompt from a file
-        system_prompt_path = os.path.join(os.getcwd(), "system_prompts", "04_summarizer.txt")
-        self.system_prompt = ""
-        with open(system_prompt_path, "r") as file:
-            self.system_prompt = file.read().strip()
-        # Define the LLM and agent
-        llm = LLMFactory.create(Args.primary_llm_interface, self.system_prompt, temperature, max_tokens)
-        self.agent = AgentWorkflow.setup_agent(llm=llm)
-        self.ctx = Context(self.agent)
-    
-    def get_system_prompt(self) -> str:
-        """
-        Retrieves the system prompt.
-
-        Returns:
-            str: The system prompt string.
-        """
-        return self.system_prompt
+        super().__init__(temperature, max_tokens, "04_summarizer.txt", Args.primary_llm_interface)
 
-    async def query(self, question: str) -> str:
+    def setup_tools(self) -> List:
         """
-        Asynchronously queries the agent with a given question and returns the response.
-
-        Args:
-            question (str): The question to be sent to the agent.
-
-        Returns:
-            str: The response from the agent as a string.
+        Abstract method to set up the tools.
+        This method must be overridden by subclasses to define custom tools this agent can use.
         """
-        response = await self.agent.run(question, ctx=self.ctx)
-        response = str(response)
-        return response
+        return []
 
-    def clear_context(self):
+    def setup_slaves(self) -> List:
         """
-        Clears the current context of the agent, resetting any conversation history.
-        This is useful when starting a new conversation or when the context needs to be refreshed.
+        Abstract method to set up the slave agents.
+        This method must be overridden by subclasses to define custom sub-agents this agent can use.
         """
-        self.ctx = Context(self.agent)
+        return []
 
 
-class Researcher:
+class Researcher(IAgent):
     def __init__(self, temperature, max_tokens):
-        # Load the system prompt from a file
-        system_prompt_path = os.path.join(os.getcwd(), "system_prompts", "05_researcher.txt")
-        self.system_prompt = ""
-        with open(system_prompt_path, "r") as file:
-            self.system_prompt = file.read().strip()
-        # Define the LLM and agent
-        llm = LLMFactory.create(Args.primary_llm_interface, self.system_prompt, temperature, max_tokens)
-        self.agent = AgentWorkflow.from_tools_or_functions(
-            Toolbox.web_search.duck_duck_go_tools,
-            llm=llm
-        )
-        self.ctx = Context(self.agent)
-    
-    def get_system_prompt(self) -> str:
-        """
-        Retrieves the system prompt.
-
-        Returns:
-            str: The system prompt string.
-        """
-        return self.system_prompt
-
-    async def query(self, question: str) -> str:
-        """
-        Asynchronously queries the agent with a given question and returns the response.
-
-        Args:
-            question (str): The question to be sent to the agent.
+        super().__init__(temperature, max_tokens, "05_researcher.txt", Args.primary_llm_interface)
 
-        Returns:
-            str: The response from the agent as a string.
-        """
-        response = await self.agent.run(question, ctx=self.ctx)
-        response = str(response)
-        return response
+    def setup_tools(self) -> List:
+        return [Toolbox.web_search.duck_duck_go_tools]
 
-    def clear_context(self):
-        """
-        Clears the current context of the agent, resetting any conversation history.
-        This is useful when starting a new conversation or when the context needs to be refreshed.
-        """
-        self.ctx = Context(self.agent)
+    def setup_slaves(self) -> List:
+        return []
 
 
-class EncryptionExpert:
+class EncryptionExpert(IAgent):
     def __init__(self, temperature, max_tokens):
-        # Load the system prompt from a file
-        system_prompt_path = os.path.join(os.getcwd(), "system_prompts", "06_encryption_expert.txt")
-        self.system_prompt = ""
-        with open(system_prompt_path, "r") as file:
-            self.system_prompt = file.read().strip()
-        # Define the LLM and agent
-        llm = LLMFactory.create(Args.primary_llm_interface, self.system_prompt, temperature, max_tokens)
-        self.agent = AgentWorkflow.from_tools_or_functions(
-            [
+        super().__init__(temperature, max_tokens, "06_encryption_expert.txt", Args.primary_llm_interface)
+
+    def setup_tools(self) -> List:
+        return [
                 Toolbox.encryption.base64_encode,
                 Toolbox.encryption.base64_decode,
                 Toolbox.encryption.caesar_cipher_encode,
                 Toolbox.encryption.caesar_cipher_decode,
                 Toolbox.encryption.reverse_string
                 # TODO: Add more encryption tools
-            ],
-            llm=llm
-        )
-        self.ctx = Context(self.agent)
-        # Initialize the tool agents
-        self.math_expert = MathExpert(temperature, max_tokens)
-        self.reasoner = Reasoner(temperature, max_tokens)
-    
-    def get_system_prompt(self) -> str:
-        """
-        Retrieves the system prompt.
-
-        Returns:
-            str: The system prompt string.
-        """
-        return self.system_prompt
-
-    async def query(self, question: str) -> str:
-        """
-        Asynchronously queries the agent with a given question and returns the response.
-
-        Args:
-            question (str): The question to be sent to the agent.
-
-        Returns:
-            str: The response from the agent as a string.
-        """
-        response = await self.agent.run(question, ctx=self.ctx)
-        response = str(response)
-        return response
+            ]
 
-    def clear_context(self):
-        """
-        Clears the current context of the agent, resetting any conversation history.
-        This is useful when starting a new conversation or when the context needs to be refreshed.
-        Also clears the context of any tool agents.
-        """
-        self.ctx = Context(self.agent)
-        # Clear context for tool agents
-        self.math_expert.clear_context()
-        self.reasoner.clear_context()
+    def setup_slaves(self) -> List:
+        return []
 
 
-class MathExpert:
+class MathExpert(IAgent):
     def __init__(self, temperature, max_tokens):
-        # Load the system prompt from a file
-        system_prompt_path = os.path.join(os.getcwd(), "system_prompts", "07_math_expert.txt")
-        self.system_prompt = ""
-        with open(system_prompt_path, "r") as file:
-            self.system_prompt = file.read().strip()
-        # Define the LLM and agent
-        llm = LLMFactory.create(Args.primary_llm_interface, self.system_prompt, temperature, max_tokens)
-        self.agent = AgentWorkflow.from_tools_or_functions(
-            [
+        super().__init__(temperature, max_tokens, "07_math_expert.txt", Args.primary_llm_interface)
+
+    def setup_tools(self) -> List:
+        return [
             Toolbox.math.symbolic_calc,
             Toolbox.math.unit_converter,
-            ],
-            llm=llm
-        )
-        self.ctx = Context(self.agent)
-        # Initialize the tool agents
-        self.reasoner = Reasoner(temperature, max_tokens)
-    
-    def get_system_prompt(self) -> str:
-        """
-        Retrieves the system prompt.
-
-        Returns:
-            str: The system prompt string.
-        """
-        return self.system_prompt
-
-    async def query(self, question: str) -> str:
-        """
-        Asynchronously queries the agent with a given question and returns the response.
+            ]
 
-        Args:
-            question (str): The question to be sent to the agent.
+    def setup_slaves(self) -> List:
+        reasoner = Reasoner(self.temperature, self.max_tokens)
+        return [reasoner]
 
-        Returns:
-            str: The response from the agent as a string.
-        """
-        response = await self.agent.run(question, ctx=self.ctx)
-        response = str(response)
-        return response
-
-    def clear_context(self):
-        """
-        Clears the current context of the agent, resetting any conversation history.
-        This is useful when starting a new conversation or when the context needs to be refreshed.
-        Also clears the context of any tool agents.
-        """
-        self.ctx = Context(self.agent)
-        self.reasoner.clear_context()
 
-
-class Reasoner:
+class Reasoner(IAgent):
     def __init__(self, temperature, max_tokens):
-        # Load the system prompt from a file
-        system_prompt_path = os.path.join(os.getcwd(), "system_prompts", "08_reasoner.txt")
-        self.system_prompt = ""
-        with open(system_prompt_path, "r") as file:
-            self.system_prompt = file.read().strip()
-        # Define the LLM and agent
-        llm = LLMFactory.create(Args.primary_llm_interface, self.system_prompt, temperature, max_tokens)
-        self.agent = AgentWorkflow.setup_agent(llm=llm)
-        self.ctx = Context(self.agent)
-
-    async def query(self, question: str) -> str:
-        """
-        Asynchronously queries the agent with a given question and returns the response.
-
-        Args:
-            question (str): The question to be sent to the agent.
-
-        Returns:
-            str: The response from the agent as a string.
-        """
-        response = await self.agent.run(question, ctx=self.ctx)
-        response = str(response)
-        return response
+        super().__init__(temperature, max_tokens, "08_reasoner.txt", Args.primary_llm_interface)
 
-    def get_system_prompt(self) -> str:
-        """
-        Retrieves the system prompt.
+    def setup_tools(self) -> List:
+        return []
 
-        Returns:
-            str: The system prompt string.
-        """
-        return self.system_prompt
-    
-    def clear_context(self):
-        """
-        Clears the current context of the agent, resetting any conversation history.
-        This is useful when starting a new conversation or when the context needs to be refreshed.
-        """
-        self.ctx = Context(self.agent)
+    def setup_slaves(self) -> List:
+        return []
 
 
-class ImageHandler:
+class ImageHandler(IAgent):
     def __init__(self, temperature, max_tokens):
-        # Load the system prompt from a file
-        system_prompt_path = os.path.join(os.getcwd(), "system_prompts", "09_image_handler.txt")
-        self.system_prompt = ""
-        with open(system_prompt_path, "r") as file:
-            self.system_prompt = file.read().strip()
-        pass
-
-    def get_system_prompt(self) -> str:
-        """
-        Retrieves the system prompt.
+        super().__init__(temperature, max_tokens, "09_image_handler.txt", Args.vlm_interface)
 
-        Returns:
-            str: The system prompt string.
-        """
-        return self.system_prompt
+    def setup_tools(self) -> List:
+        return []
 
-    def clear_context(self):
-        """
-        Clears the current context of the agent, resetting any conversation history.
-        This is useful when starting a new conversation or when the context needs to be refreshed.
-        """
-        if hasattr(self, 'ctx') and hasattr(self, 'agent'):
-            self.ctx = Context(self.agent)
+    def setup_slaves(self) -> List:
+        return []
 
 
-class VideoHandler:
+class VideoHandler(IAgent):
     def __init__(self, temperature, max_tokens):
-        # Load the system prompt from a file
-        system_prompt_path = os.path.join(os.getcwd(), "system_prompts", "10_video_handler.txt")
-        self.system_prompt = ""
-        with open(system_prompt_path, "r") as file:
-            self.system_prompt = file.read().strip()
-        # No implementation yet
-        pass
-    
-    def get_system_prompt(self) -> str:
-        """
-        Retrieves the system prompt.
+        super().__init__(temperature, max_tokens, "10_video_handler.txt", Args.vlm_interface)
 
-        Returns:
-            str: The system prompt string.
-        """
-        return self.system_prompt
+    def setup_tools(self) -> List:
+        return []
 
-    def clear_context(self):
-        """
-        Clears the current context of the agent, resetting any conversation history.
-        This is useful when starting a new conversation or when the context needs to be refreshed.
-        """
-        if hasattr(self, 'ctx') and hasattr(self, 'agent'):
-            self.ctx = Context(self.agent)
+    def setup_slaves(self) -> List:
+        return []
 
 
-class Solver:
+class Solver(IAgent):
     def __init__(self, temperature, max_tokens):
-        # Load the system prompt from a file
-        system_prompt_path = os.path.join(os.getcwd(), "system_prompts", "03_solver.txt")
-        self.system_prompt = ""
-        with open(system_prompt_path, "r") as file:
-            self.system_prompt = file.read().strip()
-        # Define the LLM and agent
-        llm = LLMFactory.create(Args.primary_llm_interface, self.system_prompt, temperature, max_tokens)
-        self.agent = AgentWorkflow.from_tools_or_functions(
-            [
-            self.call_summarizer,
-            self.call_researcher,
-            self.call_encryption_expert,
-            self.call_math_expert,
-            self.call_reasoner,
-            self.call_image_handler,
-            self.call_video_handler
-            ],
-            llm=llm
-        )
-        self.ctx = Context(self.agent)
-        # Initialize the tool agents
-        self.summarizer = Summarizer(temperature, max_tokens)
-        self.researcher = Researcher(temperature, max_tokens)
-        self.encryption_expert = EncryptionExpert(temperature, max_tokens)
-        self.math_expert = MathExpert(temperature, max_tokens)
-        self.reasoner = Reasoner(temperature, max_tokens)
-        self.image_handler = ImageHandler(temperature, max_tokens)
-        self.video_handler = VideoHandler(temperature, max_tokens)
-    
-    def get_system_prompt(self) -> str:
-        """
-        Retrieves the system prompt.
-
-        Returns:
-            str: The system prompt string.
-        """
-        return self.system_prompt
-
-    async def query(self, question: str) -> str:
-        """
-        Asynchronously queries the agent with a given question and returns the response.
-
-        Args:
-            question (str): The question to be sent to the agent.
-
-        Returns:
-            str: The response from the agent as a string.
-        """
-        response = await self.agent.run(question, ctx=self.ctx)
-        response = str(response)
-        return response
-
-    def clear_context(self):
-        """
-        Clears the current context of the agent, resetting any conversation history.
-        This is useful when starting a new conversation or when the context needs to be refreshed.
-        Also clears the context of all tool agents.
-        """
-        self.ctx = Context(self.agent)
-        # Clear context for all tool agents
-        self.summarizer.clear_context()
-        self.researcher.clear_context()
-        self.encryption_expert.clear_context()
-        self.math_expert.clear_context()
-        self.reasoner.clear_context()
-        self.image_handler.clear_context()
-        self.video_handler.clear_context()
-
-    async def call_summarizer(self, question: str) -> str:
-        return await self.summarizer.query(question)
-
-    async def call_researcher(self, question: str) -> str:
-        return await self.researcher.query(question)
-
-    async def call_encryption_expert(self, question: str) -> str:
-        return await self.encryption_expert.query(question)
-
-    async def call_math_expert(self, question: str) -> str:
-        return await self.math_expert.query(question)
-
-    async def call_reasoner(self, question: str) -> str:
-        return await self.reasoner.query(question)
-
-    async def call_image_handler(self, question: str) -> str:
-        # ImageHandler may not have a query method yet, but following the pattern
-        if hasattr(self.image_handler, 'query'):
-            return await self.image_handler.query(question)
-        return "Image handling is not implemented yet."
-        # TODO
-
-    async def call_video_handler(self, question: str) -> str:
-        # VideoHandler may not have a query method yet, but following the pattern
-        if hasattr(self.video_handler, 'query'):
-            return await self.video_handler.query(question)
-        return "Video handling is not implemented yet."
-        # TODO
+        super().__init__(temperature, max_tokens, "03_solver.txt", Args.primary_llm_interface)
+
+    def setup_tools(self) -> List:
+        return []
+
+    def setup_slaves(self) -> List:
+        summarizer = Summarizer(self.temperature, self.max_tokens)
+        researcher = Researcher(self.temperature, self.max_tokens)
+        encryption_expert = EncryptionExpert(self.temperature, self.max_tokens)
+        math_expert = MathExpert(self.temperature, self.max_tokens)
+        reasoner = Reasoner(self.temperature, self.max_tokens)
+        image_handler = ImageHandler(self.temperature, self.max_tokens)
+        video_handler = VideoHandler(self.temperature, self.max_tokens)
+
+        return [
+            summarizer,
+            researcher,
+            encryption_expert,
+            math_expert,
+            reasoner,
+            image_handler,
+            video_handler
+            ]
 
 
 # if __name__ == "__main__":