From 7377cef859b999f696a824872a190e2e4d2c4051 Mon Sep 17 00:00:00 2001
From: 7shi <7shi@live.jp>
Date: Fri, 20 Dec 2024 17:59:35 +0900
Subject: [PATCH] feat (translator): enhance chat response handling with
 retries and streaming

---
 pdf2zh/translator.py | 34 ++++++++++++++++++++++++++++------
 1 file changed, 28 insertions(+), 6 deletions(-)

diff --git a/pdf2zh/translator.py b/pdf2zh/translator.py
index 5046f6e3..c514bec6 100644
--- a/pdf2zh/translator.py
+++ b/pdf2zh/translator.py
@@ -216,12 +216,34 @@ def translate(self, text):
         print(len(self.prompt(text, self.prompttext)))
         print(self.prompt(text, self.prompttext)[0])
         print(self.prompt(text, self.prompttext)[1])
-        response = self.client.chat(
-            model=self.model,
-            options=self.options,
-            messages=self.prompt(text, self.prompttext),
-        )
-        return response["message"]["content"].strip()
+        maxlen = max(2000, len(text) * 3)
+        for model in self.model.split(";"):
+            for i in range(2):
+                if i:
+                    print("[Retry]")
+                response = ""
+                try:
+                    stream = self.client.chat(
+                        model=model,
+                        options=self.options,
+                        messages=self.prompt(text, self.prompttext),
+                        stream=True,
+                    )
+                    for chunk in stream:
+                        chunk = chunk["message"]["content"]
+                        print(chunk, end="", flush=True)
+                        response += chunk
+                        if len(response) > maxlen:
+                            raise Exception("Response too long")
+                    if not response.endswith("\n"):
+                        print()
+                    return response.strip()
+                except Exception as e:
+                    print()
+                    print(e)
+        print("[Aborted.]")
+        # if translation fails after multiple retries, return the original text to prevent hang-up
+        return text
 
 
 class OpenAITranslator(BaseTranslator):