Merge pull request #305 from 7shi/fix-ollama

Enhanced Ollama Response Handling with Retries and Streaming

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):