From d6c3f0ea0f97d5f059427cd41e2efb764c54cb1a Mon Sep 17 00:00:00 2001
From: Vivian Fang <hi@vivi.sh>
Date: Tue, 31 Oct 2023 12:12:07 -0700
Subject: [PATCH] Gracefully catch errors when running agent.step() (#216)

* Gracefully catch errors when running agent.step()

* Modify retry message
---
 memgpt/main.py | 43 ++++++++++++++++++++++++++-----------------
 1 file changed, 26 insertions(+), 17 deletions(-)

diff --git a/memgpt/main.py b/memgpt/main.py
index ecbc0ed0..57ea9e8b 100644
--- a/memgpt/main.py
+++ b/memgpt/main.py
@@ -7,6 +7,7 @@ import glob
 import os
 import sys
 import pickle
+import traceback
 
 import questionary
 import typer
@@ -512,24 +513,32 @@ async def run_agent_loop(memgpt_agent, first, no_verify=False, cfg=None, legacy=
 
         skip_next_user_input = False
 
-        with console.status("[bold cyan]Thinking...") as status:
-            (
-                new_messages,
-                heartbeat_request,
-                function_failed,
-                token_warning,
-            ) = await memgpt_agent.step(user_message, first_message=False, skip_verify=no_verify)
+        while True:
+            try:
+                with console.status("[bold cyan]Thinking...") as status:
+                    (
+                        new_messages,
+                        heartbeat_request,
+                        function_failed,
+                        token_warning,
+                    ) = await memgpt_agent.step(user_message, first_message=False, skip_verify=no_verify)
 
-            # Skip user inputs if there's a memory warning, function execution failed, or the agent asked for control
-            if token_warning:
-                user_message = system.get_token_limit_warning()
-                skip_next_user_input = True
-            elif function_failed:
-                user_message = system.get_heartbeat(constants.FUNC_FAILED_HEARTBEAT_MESSAGE)
-                skip_next_user_input = True
-            elif heartbeat_request:
-                user_message = system.get_heartbeat(constants.REQ_HEARTBEAT_MESSAGE)
-                skip_next_user_input = True
+                    # Skip user inputs if there's a memory warning, function execution failed, or the agent asked for control
+                    if token_warning:
+                        user_message = system.get_token_limit_warning()
+                        skip_next_user_input = True
+                    elif function_failed:
+                        user_message = system.get_heartbeat(constants.FUNC_FAILED_HEARTBEAT_MESSAGE)
+                        skip_next_user_input = True
+                    elif heartbeat_request:
+                        user_message = system.get_heartbeat(constants.REQ_HEARTBEAT_MESSAGE)
+                        skip_next_user_input = True
+            except Exception as e:
+                print("An exception ocurred when running agent.step(): ")
+                traceback.print_exc()
+                retry = await questionary.confirm("Retry agent.step()?").ask_async()
+                if not retry:
+                    break
 
         counter += 1