Add GPT-4 support (#11)

Sinaptik-AI · Apr 30, 2023 · beac78d · beac78d
1 parent 12bbc2c
commit beac78d
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/pandasai/llm/openai.py b/pandasai/llm/openai.py
@@ -12,6 +12,16 @@
 class OpenAI(LLM):
     """OpenAI LLM"""
 
+    _supported_chat_models = [
+        "gpt-4",
+        "gpt-4-0314",
+        "gpt-4-32k",
+        "gpt-4-32k-0314",
+        "gpt-3.5-turbo",
+        "gpt-3.5-turbo-0301",
+    ]
+    _supported_completion_models = ["text-davinci-003"]
+
     api_token: str
     model: str = "gpt-3.5-turbo"
     temperature: float = 0
@@ -93,9 +103,9 @@ def chat_completion(self, value: str) -> str:
         return response["choices"][0]["message"]["content"]
 
     def call(self, instruction: str, value: str) -> str:
-        if self.model == "text-davinci-003":
+        if self.model in self._supported_completion_models:
             response = self.completion(str(instruction) + str(value))
-        elif self.model == "gpt-3.5-turbo":
+        elif self.model in self._supported_chat_models:
             response = self.chat_completion(str(instruction) + str(value))
         else:
             raise UnsupportedOpenAIModelError("Unsupported model")