error handling and step retry

7 files changed, 63 insertions, 201 deletions

diff --git a/continuedev/src/continuedev/core/agent.py b/continuedev/src/continuedev/core/agent.py
deleted file mode 100644
index 1996abb1..00000000
--- a/continuedev/src/continuedev/core/agent.py
+++ /dev/null
@@ -1,180 +0,0 @@
-import traceback
-import time
-from typing import Callable, Coroutine, List
-from ..models.filesystem_edit import FileEditWithFullContents
-from ..libs.llm import LLM
-from .observation import Observation
-from ..server.ide_protocol import AbstractIdeProtocolServer
-from ..libs.util.queue import AsyncSubscriptionQueue
-from ..models.main import ContinueBaseModel
-from .main import Policy, History, FullState, Step, HistoryNode
-from ..libs.steps.core.core import ReversibleStep, ManualEditStep, UserInputStep
-from .sdk import ContinueSDK
-
-
-class Autopilot(ContinueBaseModel):
-    llm: LLM
-    policy: Policy
-    ide: AbstractIdeProtocolServer
-    history: History = History.from_empty()
-    _on_update_callbacks: List[Callable[[FullState], None]] = []
-
-    _active: bool = False
-    _should_halt: bool = False
-    _main_user_input_queue: List[str] = []
-
-    _user_input_queue = AsyncSubscriptionQueue()
-
-    class Config:
-        arbitrary_types_allowed = True
-
-    def get_full_state(self) -> FullState:
-        return FullState(history=self.history, active=self._active, user_input_queue=self._main_user_input_queue)
-
-    def on_update(self, callback: Callable[["FullState"], None]):
-        """Subscribe to changes to state"""
-        self._on_update_callbacks.append(callback)
-
-    def update_subscribers(self):
-        full_state = self.get_full_state()
-        for callback in self._on_update_callbacks:
-            callback(full_state)
-
-    def __get_step_params(self, step: "Step"):
-        return ContinueSDK(autopilot=self, llm=self.llm.with_system_message(step.system_message))
-
-    def give_user_input(self, input: str, index: int):
-        self._user_input_queue.post(index, input)
-
-    async def wait_for_user_input(self) -> str:
-        self._active = False
-        self.update_subscribers()
-        user_input = await self._user_input_queue.get(self.history.current_index)
-        self._active = True
-        self.update_subscribers()
-        return user_input
-
-    _manual_edits_buffer: List[FileEditWithFullContents] = []
-
-    async def reverse_to_index(self, index: int):
-        try:
-            while self.history.get_current_index() >= index:
-                current_step = self.history.get_current().step
-                self.history.step_back()
-                if issubclass(current_step.__class__, ReversibleStep):
-                    await current_step.reverse(self.__get_step_params(current_step))
-
-                self.update_subscribers()
-        except Exception as e:
-            print(e)
-
-    def handle_manual_edits(self, edits: List[FileEditWithFullContents]):
-        for edit in edits:
-            self._manual_edits_buffer.append(edit)
-            # TODO: You're storing a lot of unecessary data here. Can compress into EditDiffs on the spot, and merge.
-            # self._manual_edits_buffer = merge_file_edit(self._manual_edits_buffer, edit)
-
-    def handle_traceback(self, traceback: str):
-        raise NotImplementedError
-
-    _step_depth: int = 0
-
-    async def _run_singular_step(self, step: "Step", is_future_step: bool = False) -> Coroutine[Observation, None, None]:
-        if not is_future_step:
-            # Check manual edits buffer, clear out if needed by creating a ManualEditStep
-            if len(self._manual_edits_buffer) > 0:
-                manualEditsStep = ManualEditStep.from_sequence(
-                    self._manual_edits_buffer)
-                self._manual_edits_buffer = []
-                await self._run_singular_step(manualEditsStep)
-
-        # Update history - do this first so we get top-first tree ordering
-        self.history.add_node(HistoryNode(
-            step=step, observation=None, depth=self._step_depth))
-
-        # Run step
-        self._step_depth += 1
-        observation = await step(self.__get_step_params(step))
-        self._step_depth -= 1
-
-        # Add observation to history
-        self.history.get_current().observation = observation
-
-        # Update its description
-        step._set_description(await step.describe(self.llm))
-
-        # Call all subscribed callbacks
-        self.update_subscribers()
-
-        return observation
-
-    async def run_from_step(self, step: "Step"):
-        # if self._active:
-        #     raise RuntimeError("Autopilot is already running")
-        self._active = True
-
-        next_step = step
-        is_future_step = False
-        while not (next_step is None or self._should_halt):
-            try:
-                if is_future_step:
-                    # If future step, then we are replaying and need to delete the step from history so it can be replaced
-                    self.history.remove_current_and_substeps()
-
-                observation = await self._run_singular_step(next_step, is_future_step)
-                if next_step := self.policy.next(self.history):
-                    is_future_step = False
-                elif next_step := self.history.take_next_step():
-                    is_future_step = True
-                else:
-                    next_step = None
-
-            except Exception as e:
-                print(
-                    f"Error while running step: \n{''.join(traceback.format_tb(e.__traceback__))}\n{e}")
-                next_step = None
-
-        self._active = False
-
-        # Doing this so active can make it to the frontend after steps are done. But want better state syncing tools
-        for callback in self._on_update_callbacks:
-            callback(None)
-
-    async def run_from_observation(self, observation: Observation):
-        next_step = self.policy.next(self.history)
-        await self.run_from_step(next_step)
-
-    async def run_policy(self):
-        first_step = self.policy.next(self.history)
-        await self.run_from_step(first_step)
-
-    async def _request_halt(self):
-        if self._active:
-            self._should_halt = True
-            while self._active:
-                time.sleep(0.1)
-        self._should_halt = False
-        return None
-
-    async def accept_user_input(self, user_input: str):
-        self._main_user_input_queue.append(user_input)
-        self.update_subscribers()
-
-        if len(self._main_user_input_queue) > 1:
-            return
-
-        # await self._request_halt()
-        # Just run the step that takes user input, and
-        # then up to the policy to decide how to deal with it.
-        self._main_user_input_queue.pop(0)
-        self.update_subscribers()
-        await self.run_from_step(UserInputStep(user_input=user_input))
-
-        while len(self._main_user_input_queue) > 0:
-            await self.run_from_step(UserInputStep(
-                user_input=self._main_user_input_queue.pop(0)))
-
-    async def accept_refinement_input(self, user_input: str, index: int):
-        await self._request_halt()
-        await self.reverse_to_index(index)
-        await self.run_from_step(UserInputStep(user_input=user_input))
diff --git a/continuedev/src/continuedev/core/autopilot.py b/continuedev/src/continuedev/core/autopilot.py
index 85f65dc3..db06c975 100644
--- a/continuedev/src/continuedev/core/autopilot.py
+++ b/continuedev/src/continuedev/core/autopilot.py
@@ -3,7 +3,7 @@ import time
 from typing import Callable, Coroutine, List
 from ..models.filesystem_edit import FileEditWithFullContents
 from ..libs.llm import LLM
-from .observation import Observation
+from .observation import Observation, InternalErrorObservation
 from ..server.ide_protocol import AbstractIdeProtocolServer
 from ..libs.util.queue import AsyncSubscriptionQueue
 from ..models.main import ContinueBaseModel
@@ -77,6 +77,11 @@ class Autopilot(ContinueBaseModel):
 
     _step_depth: int = 0
 
+    async def retry_at_index(self, index: int):
+        last_step = self.history.pop_last_step()
+        await self.update_subscribers()
+        await self._run_singular_step(last_step)
+
     async def _run_singular_step(self, step: "Step", is_future_step: bool = False) -> Coroutine[Observation, None, None]:
         capture_event(
             'step run', {'step_name': step.name, 'params': step.dict()})
@@ -96,14 +101,28 @@ class Autopilot(ContinueBaseModel):
         # Call all subscribed callbacks
         await self.update_subscribers()
 
-        # Run step
+        # Try to run step and handle errors
         self._step_depth += 1
-        observation = await step(ContinueSDK(self))
+
+        try:
+            observation = await step(ContinueSDK(self))
+        except Exception as e:
+            # Attach an InternalErrorObservation to the step and unhide it.
+            error_string = '\n\n'.join(
+                traceback.format_tb(e.__traceback__)) + f"\n\n{e.__repr__()}"
+            print(
+                f"Error while running step: \n{error_string}\n{e}")
+
+            observation = InternalErrorObservation(
+                error=error_string)
+            step.hide = False
+
         self._step_depth -= 1
 
         # Add observation to history
         self.history.get_last_at_depth(
             self._step_depth, include_current=True).observation = observation
+        await self.update_subscribers()
 
         # Update its description
         async def update_description():
@@ -122,22 +141,17 @@ class Autopilot(ContinueBaseModel):
         next_step = step
         is_future_step = False
         while not (next_step is None or self._should_halt):
-            try:
-                if is_future_step:
-                    # If future step, then we are replaying and need to delete the step from history so it can be replaced
-                    self.history.remove_current_and_substeps()
-
-                observation = await self._run_singular_step(next_step, is_future_step)
-                if next_step := self.policy.next(self.history):
-                    is_future_step = False
-                elif next_step := self.history.take_next_step():
-                    is_future_step = True
-                else:
-                    next_step = None
-
-            except Exception as e:
-                print(
-                    f"Error while running step: \n{''.join(traceback.format_tb(e.__traceback__))}\n{e}")
+            if is_future_step:
+                # If future step, then we are replaying and need to delete the step from history so it can be replaced
+                self.history.remove_current_and_substeps()
+
+            await self._run_singular_step(next_step, is_future_step)
+
+            if next_step := self.policy.next(self.history):
+                is_future_step = False
+            elif next_step := self.history.take_next_step():
+                is_future_step = True
+            else:
                 next_step = None
 
         self._active = False
diff --git a/continuedev/src/continuedev/core/main.py b/continuedev/src/continuedev/core/main.py
index a2336671..b2b97bae 100644
--- a/continuedev/src/continuedev/core/main.py
+++ b/continuedev/src/continuedev/core/main.py
@@ -67,6 +67,13 @@ class History(ContinueBaseModel):
             return None
         return state.observation
 
+    def pop_last_step(self) -> Union[HistoryNode, None]:
+        if self.current_index < 0:
+            return None
+        node = self.timeline.pop(self.current_index)
+        self.current_index -= 1
+        return node.step
+
     @classmethod
     def from_empty(cls):
         return cls(timeline=[], current_index=-1)
@@ -118,7 +125,7 @@ class Step(ContinueBaseModel):
         if self._description is not None:
             d["description"] = self._description
         else:
-            d["description"] = self.name
+            d["description"] = "`Description loading...`"
         return d
 
     @validator("name", pre=True, always=True)
diff --git a/continuedev/src/continuedev/core/observation.py b/continuedev/src/continuedev/core/observation.py
index fef04311..b6117236 100644
--- a/continuedev/src/continuedev/core/observation.py
+++ b/continuedev/src/continuedev/core/observation.py
@@ -33,3 +33,7 @@ class TextObservation(Observation):
         if v is None:
             return ""
         return v
+
+
+class InternalErrorObservation(Observation):
+    error: str
diff --git a/continuedev/src/continuedev/libs/llm/hf_inference_api.py b/continuedev/src/continuedev/libs/llm/hf_inference_api.py
index 83852d27..734da160 100644
--- a/continuedev/src/continuedev/libs/llm/hf_inference_api.py
+++ b/continuedev/src/continuedev/libs/llm/hf_inference_api.py
@@ -22,4 +22,11 @@ class HuggingFaceInferenceAPI(LLM):
                 "return_full_text": False,
             }
         })
-        return response.json()[0]["generated_text"]
+        data = response.json()
+
+        # Error if the response is not a list
+        if not isinstance(data, list):
+            raise Exception(
+                "Hugging Face returned an error response: \n\n", data)
+
+        return data[0]["generated_text"]
diff --git a/continuedev/src/continuedev/server/gui.py b/continuedev/src/continuedev/server/gui.py
index 3d1a5a82..b873a88f 100644
--- a/continuedev/src/continuedev/server/gui.py
+++ b/continuedev/src/continuedev/server/gui.py
@@ -75,6 +75,8 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
                 self.on_refinement_input(data["input"], data["index"])
             elif message_type == "reverse_to_index":
                 self.on_reverse_to_index(data["index"])
+            elif message_type == "retry_at_index":
+                self.on_retry_at_index(data["index"])
         except Exception as e:
             print(e)
 
@@ -100,6 +102,10 @@ class GUIProtocolServer(AbstractGUIProtocolServer):
         asyncio.create_task(
             self.session.autopilot.accept_refinement_input(input, index))
 
+    def on_retry_at_index(self, index: int):
+        asyncio.create_task(
+            self.session.autopilot.retry_at_index(index))
+
 
 @router.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket, session: Session = Depends(websocket_session)):
diff --git a/continuedev/src/continuedev/server/gui_protocol.py b/continuedev/src/continuedev/server/gui_protocol.py
index e32d80ef..287f9e3b 100644
--- a/continuedev/src/continuedev/server/gui_protocol.py
+++ b/continuedev/src/continuedev/server/gui_protocol.py
@@ -26,3 +26,7 @@ class AbstractGUIProtocolServer(ABC):
     @abstractmethod
     async def send_state_update(self, state: dict):
         """Send a state update to the client"""
+
+    @abstractmethod
+    def on_retry_at_index(self, index: int):
+        """Called when the user requests a retry at a previous index"""