feat: support for Custom API (#354)

feat: support for Custom API

README-CN.md

@@ -69,6 +69,9 @@ python3 make_book.py --book_name test_books/animal_farm.epub --model deepl --dee
 # Use the Claude model with Japanese
 python3 make_book.py --book_name test_books/animal_farm.epub --model claude --claude_key ${claude_key} --language ja
 
+# Use the CustomAPI model with Japanese
+python3 make_book.py --book_name test_books/animal_farm.epub --model customapi --custom_api ${custom_api} --language ja
+
 # Translate contents in <div> and <p>
 python3 make_book.py --book_name test_books/animal_farm.epub --translate-tags div,p
 

README.md

@@ -87,6 +87,9 @@ python3 make_book.py --book_name test_books/animal_farm.epub --model deepl --dee
 # Use the Claude model with Japanese
 python3 make_book.py --book_name test_books/animal_farm.epub --model claude --claude_key ${claude_key} --language ja
 
+# Use the CustomAPI model with Japanese
+python3 make_book.py --book_name test_books/animal_farm.epub --model customapi --custom_api ${custom_api} --language ja
+
 # Translate contents in <div> and <p>
 python3 make_book.py --book_name test_books/animal_farm.epub --translate-tags div,p
 

book_maker/cli.py

@@ -99,6 +99,13 @@ def main():
         help="you can find claude key from here (https://console.anthropic.com/account/keys)",
     )
 
+    parser.add_argument(
+        "--custom_api",
+        dest="custom_api",
+        type=str,
+        help="you should build your own translation api",
+    )
+
     parser.add_argument(
         "--test",
         dest="test",
@@ -296,6 +303,10 @@ So you are close to reaching the limit. You have to choose your own value, there
         API_KEY = options.claude_key or env.get("BBM_CLAUDE_API_KEY")
         if not API_KEY:
             raise Exception("Please provide claude key")
+    elif options.model == "customapi":
+        API_KEY = options.custom_api or env.get("BBM_CUSTOM_API")
+        if not API_KEY:
+            raise Exception("Please provide custom translate api")
     else:
         API_KEY = ""
 

book_maker/translator/__init__.py

@@ -6,6 +6,7 @@ from book_maker.translator.google_translator import Google
 from book_maker.translator.gpt3_translator import GPT3
 from book_maker.translator.gpt4_translator import GPT4
 from book_maker.translator.claude_translator import Claude
+from book_maker.translator.custom_api_translator import CustomAPI
 
 MODEL_DICT = {
     "chatgptapi": ChatGPTAPI,
@@ -16,5 +17,6 @@ MODEL_DICT = {
     "deeplfree": DeepLFree,
     "gpt4": GPT4,
     "claude": Claude,
+    "customapi": CustomAPI
     # add more here
 }

book_maker/translator/custom_api_translator.py

@@ -0,0 +1,31 @@
+from .base_translator import Base
+import re
+import json
+import requests
+import time
+from rich import print
+
+
+class CustomAPI(Base):
+    """
+    Custom API translator
+    """
+
+    def __init__(self, custom_api, language, **kwargs) -> None:
+        super().__init__(custom_api, language)
+        self.language = language
+        self.custom_api = custom_api
+
+    def rotate_key(self):
+        pass
+
+    def translate(self, text):
+        print(text)
+        custom_api = self.custom_api
+        data = {"text": text, "source_lang": self.language, "target_lang": "auto"}
+        post_data = json.dumps(data)
+        r = requests.post(url=custom_api, data=post_data, timeout=10).text
+        t_text = json.loads(r)["data"]
+        print("[bold green]" + re.sub("\n{3,}", "\n\n", t_text) + "[/bold green]")
+        time.sleep(5)
+        return t_text

docs/model_lang.md

@@ -2,7 +2,7 @@
 ## Models
 `-m, --model <Model>` <br>
 
-Currently `bbook_maker` supports these models: `chatgptapi` , `gpt3` , `google` , `caiyun` , `deepl` , `deeplfree` , `gpt4` , `claude` .
+Currently `bbook_maker` supports these models: `chatgptapi` , `gpt3` , `google` , `caiyun` , `deepl` , `deeplfree` , `gpt4` , `claude` , `customapi`.
 Default model is `chatgptapi` . 
 
 ### OPENAI models
@@ -96,6 +96,10 @@ Support [Claude](https://console.anthropic.com/docs) model. Use `--model claude
     bbook_maker --book_name test_books/animal_farm.epub --model claude --claude_key ${claude_key}
             
 
+### Custom API
+Support CustomAPI model. Use `--model customapi --custom_api ${custom_api}` .
+
+    bbook_maker --book_name test_books/animal_farm.epub --model customapi --custom_api ${custom_api}  
 
 ### Google