diff --git a/book_maker/cli.py b/book_maker/cli.py index ee91108..c5b823d 100644 --- a/book_maker/cli.py +++ b/book_maker/cli.py @@ -375,7 +375,7 @@ So you are close to reaching the limit. You have to choose your own value, there API_KEY = options.deepl_key or env.get("BBM_DEEPL_API_KEY") if not API_KEY: raise Exception("Please provide deepl key") - elif options.model == "claude": + elif options.model.startswith("claude"): API_KEY = options.claude_key or env.get("BBM_CLAUDE_API_KEY") if not API_KEY: raise Exception("Please provide claude key") @@ -494,6 +494,8 @@ So you are close to reaching the limit. You have to choose your own value, there e.translate_model.set_gpt4omini_models() if options.model == "gpt4o": e.translate_model.set_gpt4o_models() + if options.model.startswith("claude-"): + e.translate_model.set_claude_model(options.model) if options.block_size > 0: e.block_size = options.block_size if options.batch_flag: diff --git a/book_maker/translator/__init__.py b/book_maker/translator/__init__.py index 003dbb8..2e55bcc 100644 --- a/book_maker/translator/__init__.py +++ b/book_maker/translator/__init__.py @@ -21,6 +21,11 @@ MODEL_DICT = { "deepl": DeepL, "deeplfree": DeepLFree, "claude": Claude, + "claude-3-5-sonnet-latest": Claude, + "claude-3-5-sonnet-20241022": Claude, + "claude-3-5-sonnet-20240620": Claude, + "claude-3-5-haiku-latest": Claude, + "claude-3-5-haiku-20241022": Claude, "gemini": Gemini, "geminipro": Gemini, "groq": GroqClient, diff --git a/book_maker/translator/claude_translator.py b/book_maker/translator/claude_translator.py index f21da6f..3b6417d 100644 --- a/book_maker/translator/claude_translator.py +++ b/book_maker/translator/claude_translator.py @@ -30,6 +30,9 @@ class Claude(Base): def rotate_key(self): pass + def set_claude_model(self, model_name): + self.model = model_name + def translate(self, text): print(text) self.rotate_key()