From 24cdf526d81133f127affc0424b5232429ae154d Mon Sep 17 00:00:00 2001
From: Andrew <15331990+ahuang11@users.noreply.github.com>
Date: Sat, 7 Dec 2024 04:19:43 -0800
Subject: [PATCH] Ability to break while loops (#821)

Co-authored-by: Philipp Rudiger <prudiger@anaconda.com>
---
 lumen/ai/agents.py      |  3 +++
 lumen/ai/controls.py    | 11 +++++++++++
 lumen/ai/coordinator.py |  4 ++++
 3 files changed, 18 insertions(+)

diff --git a/lumen/ai/agents.py b/lumen/ai/agents.py
index bbde0bbdb..f13db095c 100644
--- a/lumen/ai/agents.py
+++ b/lumen/ai/agents.py
@@ -218,6 +218,9 @@ async def respond(
         self.interface.send(source_controls, respond=False, user="SourceAgent")
         while not source_controls._add_button.clicks > 0:
             await asyncio.sleep(0.05)
+            if source_controls._cancel_button.clicks > 0:
+                self.interface.undo()
+                return None
         return source_controls
 
 
diff --git a/lumen/ai/controls.py b/lumen/ai/controls.py
index c39c8914b..40bca7e75 100644
--- a/lumen/ai/controls.py
+++ b/lumen/ai/controls.py
@@ -73,6 +73,8 @@ class SourceControls(Viewer):
 
     add = param.Event(doc="Use table(s)")
 
+    cancel = param.Event(doc="Cancel")
+
     memory = param.ClassSelector(class_=_Memory, default=None, doc="""
         Local memory which will be used to provide the agent context.
         If None the global memory will be used.""")
@@ -126,9 +128,18 @@ def __init__(self, **params):
             visible=False,
             button_type="success",
         )
+
+        self._cancel_button = Button.from_param(
+            self.param.cancel,
+            name="Cancel",
+            icon="circle-x",
+            visible=True,
+        )
+
         self.menu = Column(
             self._input_tabs if self.select_existing else self._input_tabs[0],
             self._add_button,
+            self._cancel_button,
             self.tables_tabs,
             sizing_mode="stretch_width",
         )
diff --git a/lumen/ai/coordinator.py b/lumen/ai/coordinator.py
index 8ec4e25a7..e0c3047c0 100644
--- a/lumen/ai/coordinator.py
+++ b/lumen/ai/coordinator.py
@@ -684,6 +684,7 @@ async def _compute_execution_graph(self, messages: list[Message], agents: dict[s
         unmet_dependencies = set()
         schemas = {}
         execution_graph = []
+        attempts = 0
         with self.interface.add_step(title="Planning how to solve user query...", user="Assistant") as istep:
             while not planned:
                 try:
@@ -699,8 +700,11 @@ async def _compute_execution_graph(self, messages: list[Message], agents: dict[s
                 execution_graph, unmet_dependencies = await self._resolve_plan(plan, agents, messages)
                 if unmet_dependencies:
                     istep.stream(f"The plan didn't account for {unmet_dependencies!r}", replace=True)
+                    attempts += 1
                 else:
                     planned = True
+                if attempts > 5:
+                    raise ValueError(f"Could not find a suitable plan for {messages[-1]['content']!r}")
             self._memory['plan'] = plan
             istep.stream('\n\nHere are the steps:\n\n')
             for i, step in enumerate(plan.steps):