fix: #391

Signed-off-by: yihong0618 <zouzou0208@gmail.com>

book_maker/translator/claude_translator.py

@@ -1,6 +1,7 @@
 import re
-import requests
+import time
 from rich import print
+from anthropic import Anthropic
 
 from .base_translator import Base
 
@@ -16,23 +17,9 @@ def __init__(
         **kwargs,
     ) -> None:
         super().__init__(key, language)
-        self.api_url = (
-            f"{api_base}v1/complete"
-            if api_base
-            else "https://api.anthropic.com/v1/complete"
-        )
-        self.headers = {
-            "Content-Type": "application/json",
-            "x-api-key": key,
-        }
-        self.data = {
-            "prompt": "",
-            "model": "claude-v1.3",
-            "max_tokens_to_sample": 1024,
-            "temperature": temperature,
-            "stop_sequences": ["\n\nHuman:"],
-        }
-        self.session = requests.session()
+        self.api_url = f"{api_base}" if api_base else "https://api.anthropic.com"
+        self.client = Anthropic(base_url=api_base, api_key=key, timeout=20)
+
         self.language = language
         self.prompt_template = (
             prompt_template
@@ -45,14 +32,19 @@ def rotate_key(self):
     def translate(self, text):
         print(text)
         self.rotate_key()
-        self.data["prompt"] = self.prompt_template.format(
+        prompt = self.prompt_template.format(
             text=text,
             language=self.language,
         )
-        r = self.session.post(self.api_url, headers=self.headers, json=self.data)
-        if not r.ok:
-            return text
-        t_text = r.json().get("completion").strip()
+        message = [{"role": "user", "content": prompt}]
+        r = self.client.messages.create(
+            max_tokens=4096,
+            messages=message,
+            model="claude-3-haiku-20240307",  # default it for now
+        )
+        t_text = r.content[0].text
+        # api limit rate and spider rule
+        time.sleep(1)
 
         print("[bold green]" + re.sub("\n{3,}", "\n\n", t_text) + "[/bold green]")
         return t_text

setup.py

@@ -14,6 +14,7 @@
     "google-generativeai",
     "langdetect",
     "backoff",
+    "anthropic",
 ]