chore: refine litellm

iuiaoin · Aug 9, 2023 · aa1a35e · aa1a35e
1 parent 473e273
commit aa1a35e
Showing 1 changed file with 7 additions and 13 deletions.
diff --git a/bot/bot.py b/bot/bot.py
@@ -1,4 +1,4 @@
-# import litellm
+import litellm
 from common.context import Context
 from config import conf
 from common.singleton import singleton
@@ -9,28 +9,22 @@
 class Bot:
     def __init__(self):
         use_azure_chatgpt = conf().get("use_azure_chatgpt", False)
-        # model = conf().get("model", "gpt-3.5-turbo")
+        model = conf().get("model", "gpt-3.5-turbo")
         if use_azure_chatgpt:
             from bot.azure_chatgpt import AzureChatGPTBot
 
             self.bot = AzureChatGPTBot()
+
+        elif model in litellm.open_ai_chat_completion_models:
+            from bot.chatgpt import ChatGPTBot
+
+            self.bot = ChatGPTBot()
         else:
             # see litellm supported models here:
             # https://litellm.readthedocs.io/en/latest/supported/
             from bot.litellm import LiteLLMChatGPTBot
 
             self.bot = LiteLLMChatGPTBot()
 
-        # elif model in litellm.open_ai_chat_completion_models:
-        #     from bot.chatgpt import ChatGPTBot
-
-        #     self.bot = ChatGPTBot()
-        # else:
-        #     # see litellm supported models here:
-        #     # https://litellm.readthedocs.io/en/latest/supported/
-        #     from bot.litellm import LiteLLMChatGPTBot
-
-        #     self.bot = LiteLLMChatGPTBot()
-
     def reply(self, context: Context) -> Reply:
         return self.bot.reply(context)