feat : stop button integration

sources/agents/agent.py

@@ -46,6 +46,7 @@ class Agent():
         self.last_answer = ""
         self.last_reasoning = ""
         self.status_message = "Haven't started yet"
+        self.stop = False
         self.verbose = verbose
         self.executor = ThreadPoolExecutor(max_workers=1)
     
@@ -119,6 +120,13 @@ class Agent():
         except Exception as e:
             raise e
     
+    def request_stop(self) -> None:
+        """
+        Request the agent to stop.
+        """
+        self.stop = True
+        self.status_message = "Stopped"
+    
     @abstractmethod
     def process(self, prompt, speech_module) -> str:
         """

sources/agents/browser_agent.py

@@ -181,6 +181,7 @@ class BrowserAgent(Agent):
         animate_thinking("Thinking...", color="status")
         self.memory.push('user', prompt)
         answer, reasoning = await self.llm_request()
+        self.last_reasoning = reasoning
         if show_reasoning:
             pretty_print(reasoning, color="failure")
         pretty_print(answer, color="output")
@@ -349,11 +350,13 @@ class BrowserAgent(Agent):
         self.show_search_results(search_result)
         prompt = self.make_newsearch_prompt(user_prompt, search_result)
         unvisited = [None]
-        while not complete and len(unvisited) > 0:
-
+        while not complete and len(unvisited) > 0 and not self.stop:
             self.memory.clear()
             unvisited = self.select_unvisited(search_result)
             answer, reasoning = await self.llm_decide(prompt, show_reasoning = False)
+            if self.stop:
+                pretty_print(f"Requested stop.", color="failure")
+                break
             if self.last_answer == answer:
                 prompt = self.stuck_prompt(user_prompt, unvisited)
                 continue
@@ -424,4 +427,4 @@ class BrowserAgent(Agent):
         return answer, reasoning
 
 if __name__ == "__main__":
-    pass
+    pass

sources/agents/code_agent.py

@@ -51,10 +51,12 @@ class CoderAgent(Agent):
         self.memory.push('user', prompt)
         clarify_trigger = "REQUEST_CLARIFICATION"
 
-        while attempt < max_attempts:
+        while attempt < max_attempts and not self.stop:
+            print("Stopped?", self.stop)
             animate_thinking("Thinking...", color="status")
             await self.wait_message(speech_module)
             answer, reasoning = await self.llm_request()
+            self.last_reasoning = reasoning
             if clarify_trigger in answer:
                 self.last_answer = answer
                 await asyncio.sleep(0)

sources/agents/file_agent.py

@@ -28,10 +28,11 @@ class FileAgent(Agent):
         exec_success = False
         prompt += f"\nYou must work in directory: {self.work_dir}"
         self.memory.push('user', prompt)
-        while exec_success is False:
+        while exec_success is False and not self.stop:
             await self.wait_message(speech_module)
             animate_thinking("Thinking...", color="status")
             answer, reasoning = await self.llm_request()
+            self.last_reasoning = reasoning
             exec_success, _ = self.execute_modules(answer)
             answer = self.remove_blocks(answer)
             self.last_answer = answer

sources/agents/planner_agent.py

@@ -83,11 +83,15 @@ class PlannerAgent(Agent):
                         self.logger.warning(f"Agent {task['agent']} does not exist.")
                         pretty_print(f"Agent {task['agent']} does not exist.", color="warning")
                         return []
-                    agent = {
-                        'agent': task['agent'],
-                        'id': task['id'],
-                        'task': task['task']
-                    }
+                    try:
+                        agent = {
+                            'agent': task['agent'],
+                            'id': task['id'],
+                            'task': task['task']
+                        }
+                    except:
+                        self.logger.warning("Missing field in json plan.")
+                        return []
                     self.logger.info(f"Created agent {task['agent']} with task: {task['task']}")
                     if 'need' in task:
                         self.logger.info(f"Agent {task['agent']} was given info:\n {task['need']}")
@@ -156,6 +160,7 @@ class PlannerAgent(Agent):
                 return []
             agents_tasks = self.parse_agent_tasks(answer)
             if agents_tasks == []:
+                self.show_plan(agents_tasks, answer)
                 prompt = f"Failed to parse the tasks. Please write down your task followed by a json plan within ```json. Do not ask for clarification.\n"
                 pretty_print("Failed to make plan. Retrying...", color="warning")
                 continue
@@ -178,7 +183,11 @@ class PlannerAgent(Agent):
         last_agent_work = agents_work_result[id]
         tool_success_str = "success" if success else "failure"
         pretty_print(f"Agent {id} work {tool_success_str}.", color="success" if success else "failure")
-        if int(id) == len(agents_tasks):
+        try:
+            id_int = int(id)
+        except Exception as e:
+            return agents_tasks
+        if id_int == len(agents_tasks):
             next_task = "No task follow, this was the last step. If it failed add a task to recover."
         else:
             next_task = f"Next task is: {agents_tasks[int(id)][0]}."
@@ -258,7 +267,7 @@ class PlannerAgent(Agent):
             return "Failed to parse the tasks.", ""
         i = 0
         steps = len(agents_tasks)
-        while i < steps:
+        while i < steps and not self.stop:
             task_name, task = agents_tasks[i][0], agents_tasks[i][1]
             self.status_message = "Starting agents..."
             pretty_print(f"I will {task_name}.", color="info")
@@ -272,6 +281,8 @@ class PlannerAgent(Agent):
                 answer, success = await self.start_agent_process(task, required_infos)
             except Exception as e:
                 raise e
+            if self.stop:
+                pretty_print(f"Requested stop.", color="failure")
             agents_work_result[task['id']] = answer
             agents_tasks = await self.update_plan(goal, agents_tasks, agents_work_result, task['id'], success)
             steps = len(agents_tasks)