From 1d99ebf3750556af08597e6a4381c834d4d2af4a Mon Sep 17 00:00:00 2001
From: mkXultra <electrincgoat@gmail.com>
Date: Thu, 22 Aug 2024 20:22:46 +0900
Subject: [PATCH] support: gpt4o

---
 book_maker/cli.py                             |  5 ++++-
 book_maker/translator/__init__.py             |  1 +
 .../translator/chatgptapi_translator.py       | 19 +++++++++++++++++++
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/book_maker/cli.py b/book_maker/cli.py
index 3e87360..f42db4f 100644
--- a/book_maker/cli.py
+++ b/book_maker/cli.py
@@ -331,7 +331,7 @@ So you are close to reaching the limit. You have to choose your own value, there
     translate_model = MODEL_DICT.get(options.model)
     assert translate_model is not None, "unsupported model"
     API_KEY = ""
-    if options.model in ["openai", "chatgptapi", "gpt4", "gpt4omini"]:
+    if options.model in ["openai", "chatgptapi", "gpt4", "gpt4omini", "gpt4o"]:
         if OPENAI_API_KEY := (
             options.openai_key
             or env.get(
@@ -449,6 +449,7 @@ So you are close to reaching the limit. You have to choose your own value, there
             "chatgptapi",
             "gpt4",
             "gpt4omini",
+            "gpt4o",
         ], "only support chatgptapi for deployment_id"
         if not options.api_base:
             raise ValueError("`api_base` must be provided when using `deployment_id`")
@@ -471,6 +472,8 @@ So you are close to reaching the limit. You have to choose your own value, there
         e.translate_model.set_gpt4_models()
     if options.model == "gpt4omini":
         e.translate_model.set_gpt4omini_models()
+    if options.model == "gpt4o":
+        e.translate_model.set_gpt4o_models()
     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 d7e84cb..8f848d5 100644
--- a/book_maker/translator/__init__.py
+++ b/book_maker/translator/__init__.py
@@ -14,6 +14,7 @@ MODEL_DICT = {
     "chatgptapi": ChatGPTAPI,
     "gpt4": ChatGPTAPI,
     "gpt4omini": ChatGPTAPI,
+    "gpt4o": ChatGPTAPI,
     "google": Google,
     "caiyun": Caiyun,
     "deepl": DeepL,
diff --git a/book_maker/translator/chatgptapi_translator.py b/book_maker/translator/chatgptapi_translator.py
index c97121d..cad3086 100644
--- a/book_maker/translator/chatgptapi_translator.py
+++ b/book_maker/translator/chatgptapi_translator.py
@@ -42,6 +42,12 @@ GPT4oMINI_MODEL_LIST = [
     "gpt-4o-mini",
     "gpt-4o-mini-2024-07-18",
 ]
+GPT4o_MODEL_LIST = [
+    "gpt-4o",
+    "gpt-4o-2024-05-13",
+    "gpt-4o-2024-08-06",
+    "chatgpt-4o-latest"
+]
 
 
 class ChatGPTAPI(Base):
@@ -403,6 +409,19 @@ class ChatGPTAPI(Base):
             model_list = list(set(my_model_list) & set(GPT4oMINI_MODEL_LIST))
             print(f"Using model list {model_list}")
             self.model_list = cycle(model_list)
+            self.model_list = cycle(["plamo-beta"])
+
+    def set_gpt4o_models(self):
+        # for issue #375 azure can not use model list
+        if self.deployment_id:
+            self.model_list = cycle(["gpt-4o"])
+        else:
+            my_model_list = [
+                i["id"] for i in self.openai_client.models.list().model_dump()["data"]
+            ]
+            model_list = list(set(my_model_list) & set(GPT4o_MODEL_LIST))
+            print(f"Using model list {model_list}")
+            self.model_list = cycle(model_list)
 
     def set_model_list(self, model_list):
         model_list = list(set(model_list))