Skip to content

Commit

Permalink
feat (translator): enhance chat response handling with retries and st…
Browse files Browse the repository at this point in the history
…reaming
  • Loading branch information
7shi committed Dec 20, 2024
1 parent b9ae66b commit 7377cef
Showing 1 changed file with 28 additions and 6 deletions.
34 changes: 28 additions & 6 deletions pdf2zh/translator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down

0 comments on commit 7377cef

Please sign in to comment.