Feat : planning agent task parsing and agents assignation

main.py

@@ -9,10 +9,8 @@ from sources.llm_provider import Provider
 from sources.interaction import Interaction
 from sources.agents import Agent, CoderAgent, CasualAgent, FileAgent, PlannerAgent
 
-parser = argparse.ArgumentParser(description='Deepseek AI assistant')
-parser.add_argument('--no-speak', action='store_true',
-                help='Make AI not use text-to-speech')
-args = parser.parse_args()
+import warnings
+warnings.filterwarnings("ignore")
 
 config = configparser.ConfigParser()
 config.read('config.ini')

prompts/planner_agent.txt

@@ -8,20 +8,45 @@ Agents are other AI that obey your instructions.
 
 You will be given a task and you will need to divide it into smaller tasks and assign them to the agents.
 
-Here is an example:
+You have to respect a strict format:
+```json
+{"agent": "agent_name", "need": "needed_agent_output", "task": "agent_task"}
+```
 
-User:  Make a cool game to illustrate the current relation between USA and europe
+User: make a weather app in python 
 You: Sure, here is the plan:
-## I will do a web research to find the current relation between USA and europe
-- Web(name="web1"): Search for the current relation between USA and europe
-## I will then make a game to illustrate the current relation between USA and europe
-- Coder(name="coder1", need="web1"): Make a game to illustrate the current relation between USA and europe
 
-Another example:
+## Task 1: I will search for available weather api
 
-User: make a mobile app to check the weather in a specific city
-You: Sure, here is the plan:
-## I will do a web research to find the current weather in a specific city
-- Web(name="web1"): Search how to use weather api
-## I will then make a mobile app to check the weather in a specific city
-- Coder(name="coder1", need="web1"): Make a mobile app to check the weather in a specific city
+## Task 2: I will create an api key for the weather api
+
+## Task 3: I will make a weather app in python
+
+```json
+{
+  "plan": [
+    {
+      "agent": "Web",
+      "id": "1",
+      "need": null,
+      "task": "Search for reliable weather APIs"
+    },
+    {
+      "agent": "Web",
+      "id": "2",
+      "need": "1",
+      "task": "Obtain API key from the selected service"
+    },
+    {
+      "agent": "Coder",
+      "id": "3",
+      "need": "2",
+      "task": "Develop a Python application using the API and key to fetch and display weather data"
+    }
+  ]
+}
+```
+
+Rules:
+- Do not write code. You are a planning agent.
+- Put your plan in a json with the key "plan". Otherwise, 1000 children in africa will die.

sources/agents/agent.py

@@ -103,6 +103,8 @@ class Agent():
         return answer, reasoning
     
     def wait_message(self, speech_module):
+        if speech_module is None:
+            return
         messages = ["Please be patient sir, I am working on it.",
                     "At it, sir. In the meantime, how about a joke?",
                     "Computing... I recommand you have a coffee while I work.",

sources/agents/casual_agent.py

@@ -1,5 +1,5 @@
 
-from sources.utility import pretty_print
+from sources.utility import pretty_print, animate_thinking
 from sources.agents.agent import Agent
 from sources.tools.webSearch import webSearch
 from sources.tools.flightSearch import FlightSearch
@@ -29,7 +29,7 @@ class CasualAgent(Agent):
         while not complete:
             if exec_success:
                 complete = True
-            pretty_print("Thinking...", color="status")
+            animate_thinking("Thinking...", color="status")
             answer, reasoning = self.llm_request()
             exec_success, _ = self.execute_modules(answer)
             answer = self.remove_blocks(answer)

sources/agents/code_agent.py

@@ -1,5 +1,5 @@
 
-from sources.utility import pretty_print
+from sources.utility import pretty_print, animate_thinking
 from sources.agents.agent import Agent, executorResult
 from sources.tools.C_Interpreter import CInterpreter
 from sources.tools.GoInterpreter import GoInterpreter
@@ -29,7 +29,7 @@ class CoderAgent(Agent):
         self.memory.push('user', prompt)
 
         while attempt < max_attempts:
-            pretty_print("Thinking...", color="status")
+            animate_thinking("Thinking...", color="status")
             self.wait_message(speech_module)
             answer, reasoning = self.llm_request()
             exec_success, _ = self.execute_modules(answer)

sources/agents/file_agent.py

@@ -1,5 +1,5 @@
 
-from sources.utility import pretty_print
+from sources.utility import pretty_print, animate_thinking
 from sources.agents.agent import Agent
 from sources.tools.fileFinder import FileFinder
 from sources.tools.BashInterpreter import BashInterpreter
@@ -25,7 +25,7 @@ class FileAgent(Agent):
         while not complete:
             if exec_success:
                 complete = True
-            pretty_print("Thinking...", color="status")
+            animate_thinking("Thinking...", color="status")
             answer, reasoning = self.llm_request()
             exec_success, _ = self.execute_modules(answer)
             answer = self.remove_blocks(answer)

sources/agents/planner_agent.py

@@ -1,9 +1,10 @@
-
-from sources.utility import pretty_print
+import json
+from sources.utility import pretty_print, animate_thinking
 from sources.agents.agent import Agent
 from sources.agents.code_agent import CoderAgent
 from sources.agents.file_agent import FileAgent
 from sources.agents.casual_agent import CasualAgent
+from sources.tools.tools import Tools
 
 class PlannerAgent(Agent):
     def __init__(self, model, name, prompt_path, provider):
@@ -12,50 +13,78 @@ class PlannerAgent(Agent):
         """
         super().__init__(model, name, prompt_path, provider)
         self.tools = {
+            "json": Tools()
         }
+        self.tools['json'].tag = "json"
         self.agents = {
             "coder": CoderAgent(model, name, prompt_path, provider),
             "file": FileAgent(model, name, prompt_path, provider),
             "web": CasualAgent(model, name, prompt_path, provider)
         }
         self.role = "complex programming tasks and web research"
+        self.tag = "json"
 
     def parse_agent_tasks(self, text):
-        agents_tasks = []
+        tasks = []
+        tasks_names = []
 
         lines = text.strip().split('\n')
         for line in lines:
-            if not '-' in line:
+            if line is None or len(line) == 0:
                 continue
-            if not line.strip() or ':' not in line:
+            line = line.strip()
+            if '##' in line or line[0].isdigit():
+                tasks_names.append(line)
                 continue
-            agent_part, task = line.split(':', 1)
-            task = task.strip()
-            agent_info = agent_part.strip().split('(')
-            agent_type = agent_info[0].strip()
-            params_part = agent_info[1].rstrip(')').split(',')
-            params = {}
-            for param in params_part:
-                key, value = param.split('=')
-                params[key.strip()] = value.strip().strip('"')
-            agent = {
-                'type': agent_type,
-                'name': params['name'],
-                'task': task
-            }
-            if 'need' in params:
-                agent['need'] = params['need']
-            agents_tasks.append(agent)
-        return agents_tasks
+        blocks, _ = self.tools["json"].load_exec_block(text)
+        if blocks == None:
+            return (None, None)
+        for block in blocks:
+            line_json = json.loads(block)
+            if 'plan' in line_json:
+                for task in line_json['plan']:
+                    agent = {
+                        'agent': task['agent'],
+                        'id': task['id'],
+                        'task': task['task']
+                    }
+                    if 'need' in task:
+                        agent['need'] = task['need']
+                    tasks.append(agent)
+        if len(tasks_names) != len(tasks):
+            names = [task['task'] for task in tasks]
+            return zip(names, tasks)
+        return zip(tasks_names, tasks)
+    
+    def make_prompt(self, task, needed_infos):
+        prompt = f"""
+        You are given the following informations:
+        {needed_infos}
+        Your task is:
+        {task}
+        """
+        return prompt
     
     def process(self, prompt, speech_module) -> str:
         self.memory.push('user', prompt)
         self.wait_message(speech_module)
-        pretty_print("Thinking...", color="status")
-        print(self.memory.get())
+        animate_thinking("Thinking...", color="status")
+        agents_tasks = (None, None)
         answer, reasoning = self.llm_request()
         agents_tasks = self.parse_agent_tasks(answer)
-        print(agents_tasks)
+        if agents_tasks == (None, None):
+            return "Failed to parse the tasks", reasoning
+        for task_name, task in agents_tasks:
+            pretty_print(f"I will {task_name}.", color="info")
+            agent_prompt = self.make_prompt(task['task'], task['need'])
+            pretty_print(f"Assigned agent {task['agent']} to {task_name}", color="info")
+            speech_module.speak(f"I will {task_name}. I assigned the {task['agent']} agent to the task.")
+            try:
+                self.agents[task['agent'].lower()].process(agent_prompt, None)
+            except Exception as e:
+                pretty_print(f"Error: {e}", color="failure")
+                speech_module.speak(f"I encountered an error: {e}")
+                break
         self.last_answer = answer
         return answer, reasoning
 

sources/interaction.py

@@ -57,9 +57,10 @@ class Interaction:
         """Read the input from the user."""
         buffer = ""
 
+        PROMPT = "\033[1;35m➤➤➤ \033[0m"
         while buffer == "" or buffer.isascii() == False:
             try:
-                buffer = input(f">>> ")
+                buffer = input(PROMPT)
             except EOFError:
                 return None
             if buffer == "exit" or buffer == "goodbye":
@@ -94,8 +95,7 @@ class Interaction:
         """Request AI agents to process the user input."""
         if self.last_query is None or len(self.last_query) == 0:
             return
-        #agent = self.router.select_agent(self.last_query)
-        agent = self.agents[3]
+        agent = self.router.select_agent(self.last_query)
         if agent is None:
             return
         if self.current_agent != agent:

sources/utility.py

@@ -28,6 +28,7 @@ def pretty_print(text, color = "info"):
             "code": Fore.LIGHTBLUE_EX,
             "warning": Fore.YELLOW,
             "output": Fore.LIGHTCYAN_EX,
+            "info": Fore.CYAN
         }
         if color not in color_map:
             print(text)
@@ -48,6 +49,45 @@ def pretty_print(text, color = "info"):
             color = "default"
         print(colored(text, color_map[color]))
 
+def animate_thinking(text="thinking...", color="status", duration=2):
+    """
+    Display an animated "thinking..." indicator.
+
+    Args:
+        text (str): Text to display (default: "thinking...")
+        color (str): Color for the text (matches pretty_print colors)
+        duration (float): How long to animate in seconds
+    """
+    import time
+    import itertools
+
+    color_map = {
+        "success": (Fore.GREEN, "green"),
+        "failure": (Fore.RED, "red"),
+        "status": (Fore.LIGHTGREEN_EX, "light_green"),
+        "code": (Fore.LIGHTBLUE_EX, "light_blue"),
+        "warning": (Fore.YELLOW, "yellow"),
+        "output": (Fore.LIGHTCYAN_EX, "cyan"),
+        "default": (Fore.RESET, "black"),
+        "info": (Fore.CYAN, "cyan")
+    }
+
+    if color not in color_map:
+        color = "info"
+
+    fore_color, term_color = color_map[color]
+    spinner = itertools.cycle(['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏'])
+    end_time = time.time() + duration
+
+    while time.time() < end_time:
+        symbol = next(spinner)
+        if platform.system().lower() != "windows":
+            print(f"\r{fore_color}{symbol} {text}{Fore.RESET}", end="", flush=True)
+        else:
+            print(colored(f"\r{symbol} {text}", term_color), end="", flush=True)
+        time.sleep(0.1)
+    print()
+
 def timer_decorator(func):
     """
     Decorator to measure the execution time of a function.