From 5ad87bca4fe27e26b449f3f710b3d9a3a8a17cee Mon Sep 17 00:00:00 2001
From: Risin <5331402+risin42@users.noreply.github.com>
Date: Tue, 15 Oct 2024 18:40:44 +0900
Subject: [PATCH] feat: add support for --interval option to gemini

---
 book_maker/cli.py                          | 7 +++++++
 book_maker/loader/epub_loader.py           | 2 ++
 book_maker/translator/gemini_translator.py | 6 ++++--
 3 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/book_maker/cli.py b/book_maker/cli.py
index d926e0e..4f1437f 100644
--- a/book_maker/cli.py
+++ b/book_maker/cli.py
@@ -316,6 +316,12 @@ So you are close to reaching the limit. You have to choose your own value, there
         action="store_true",
         help="Use pre-generated batch translations to create files. Run with --batch first before using this option",
     )
+    parser.add_argument(
+        "--interval",
+        type=float,
+        default=0.01,
+        help="Request interval in seconds (e.g., 0.1 for 100ms). Currently only supported for Gemini models. Deafult: 0.01",
+    )
 
     options = parser.parse_args()
 
@@ -422,6 +428,7 @@ So you are close to reaching the limit. You have to choose your own value, there
         single_translate=options.single_translate,
         context_flag=options.context_flag,
         temperature=options.temperature,
+        interval=options.interval,
     )
     # other options
     if options.allow_navigable_strings:
diff --git a/book_maker/loader/epub_loader.py b/book_maker/loader/epub_loader.py
index c2b1437..5f87c66 100644
--- a/book_maker/loader/epub_loader.py
+++ b/book_maker/loader/epub_loader.py
@@ -35,6 +35,7 @@ class EPUBBookLoader(BaseBookLoader):
         context_flag=False,
         temperature=1.0,
         context_paragraph_limit=0,
+        interval=0.01,
     ):
         self.epub_name = epub_name
         self.new_epub = epub.EpubBook()
@@ -45,6 +46,7 @@ class EPUBBookLoader(BaseBookLoader):
             context_flag=context_flag,
             context_paragraph_limit=context_paragraph_limit,
             temperature=temperature,
+            interval=interval,
             **prompt_config_to_kwargs(prompt_config),
         )
         self.is_test = is_test
diff --git a/book_maker/translator/gemini_translator.py b/book_maker/translator/gemini_translator.py
index da8e7f8..fad32ae 100644
--- a/book_maker/translator/gemini_translator.py
+++ b/book_maker/translator/gemini_translator.py
@@ -43,9 +43,11 @@ class Gemini(Base):
         key,
         language,
         temperature=1.0,
+        interval=0.01,
         **kwargs,
     ) -> None:
         super().__init__(key, language)
+        self.interval = interval
         generation_config["temperature"] = temperature
         model = genai.GenerativeModel(
             model_name="gemini-pro",
@@ -90,8 +92,8 @@ class Gemini(Base):
             self.convo.history = self.convo.history[2:]
 
         print("[bold green]" + re.sub("\n{3,}", "\n\n", t_text) + "[/bold green]")
-        # for limit
-        time.sleep(0.5)
+        # for rate limit(RPM)
+        time.sleep(self.interval)
         if num:
             t_text = str(num) + "\n" + t_text
         return t_text