diff --git a/book_maker/translator/claude_translator.py b/book_maker/translator/claude_translator.py index 5ca99e1..94b153d 100644 --- a/book_maker/translator/claude_translator.py +++ b/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 @@ class Claude(Base): **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 @@ class Claude(Base): 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 diff --git a/setup.py b/setup.py index 25a8484..d5ab3e3 100644 --- a/setup.py +++ b/setup.py @@ -14,6 +14,7 @@ packages = [ "google-generativeai", "langdetect", "backoff", + "anthropic", ]