From e2d069ad9cf893f479313b08ef49495ec2ca07ef Mon Sep 17 00:00:00 2001
From: Tuchuanhuhuhu <gzblog@hdu.edu.cn>
Date: Tue, 5 Mar 2024 11:58:06 +0800
Subject: [PATCH] bugfix: retry doesn't work when Claude API errors out (such
 as when overloaded)

---
 modules/models/Claude.py | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/modules/models/Claude.py b/modules/models/Claude.py
index 3e3a8817..5893a246 100644
--- a/modules/models/Claude.py
+++ b/modules/models/Claude.py
@@ -74,16 +74,19 @@ def get_answer_stream_iter(self):
         system_prompt = self.system_prompt
         history = self._get_claude_style_history()
 
-        with self.claude_client.messages.stream(
-            model=self.model_name,
-            max_tokens=self.max_generation_token,
-            messages=history,
-            system=system_prompt,
-        ) as stream:
-            partial_text = ""
-            for text in stream.text_stream:
-                partial_text += text
-                yield partial_text
+        try:
+            with self.claude_client.messages.stream(
+                model=self.model_name,
+                max_tokens=self.max_generation_token,
+                messages=history,
+                system=system_prompt,
+            ) as stream:
+                partial_text = ""
+                for text in stream.text_stream:
+                    partial_text += text
+                    yield partial_text
+        except Exception as e:
+            yield i18n(GENERAL_ERROR_MSG) + ": " + str(e)
 
     def get_answer_at_once(self):
         system_prompt = self.system_prompt