From 5954675d284131054d9496cea74c702459bd0d18 Mon Sep 17 00:00:00 2001
From: Arsenii Shatokhin <arseny9795@gmail.com>
Date: Mon, 12 Feb 2024 08:21:02 +0400
Subject: [PATCH] Improved error handling in thread for long running agent
 conversations

---
 agency_swarm/threads/thread.py | 25 ++++++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/agency_swarm/threads/thread.py b/agency_swarm/threads/thread.py
index a698e8d2..380cb465 100644
--- a/agency_swarm/threads/thread.py
+++ b/agency_swarm/threads/thread.py
@@ -57,6 +57,7 @@ def get_completion(self, message: str, message_files=None, yield_messages=True,
             assistant_id=recipient_agent.id,
         )
 
+        run_failed = False
         while True:
             self.await_run_completion()
 
@@ -94,6 +95,12 @@ def get_completion(self, message: str, message_files=None, yield_messages=True,
                     )
                 except BadRequestError as e:
                     if 'Runs in status "expired"' in e.message:
+                        self.client.beta.threads.messages.create(
+                            thread_id=self.thread.id,
+                            role="user",
+                            content="Please repeat the exact same function calls again."
+                        )
+
                         self.run = self.client.beta.threads.runs.create(
                             thread_id=self.thread.id,
                             assistant_id=recipient_agent.id,
@@ -101,6 +108,14 @@ def get_completion(self, message: str, message_files=None, yield_messages=True,
 
                         self.await_run_completion()
 
+                        if self.run.status != "requires_action":
+                            raise Exception("Run Failed. Error: ", self.run.last_error)
+
+                        # change tool call ids
+                        tool_calls = self.run.required_action.submit_tool_outputs.tool_calls
+                        for i, tool_call in enumerate(tool_calls):
+                            tool_outputs[i]["tool_call_id"] = tool_call.id
+
                         self.run = self.client.beta.threads.runs.submit_tool_outputs(
                             thread_id=self.thread.id,
                             run_id=self.run.id,
@@ -110,7 +125,15 @@ def get_completion(self, message: str, message_files=None, yield_messages=True,
                         raise e
             # error
             elif self.run.status == "failed":
-                raise Exception("Run Failed. Error: ", self.run.last_error)
+                # retry run 1 time
+                if not run_failed and "something went wrong" in self.run.last_error:
+                    self.run = self.client.beta.threads.runs.create(
+                        thread_id=self.thread.id,
+                        assistant_id=recipient_agent.id,
+                    )
+                    run_failed = True
+                else:
+                    raise Exception("Run Failed. Error: ", self.run.last_error)
             # return assistant message
             else:
                 messages = self.client.beta.threads.messages.list(