Add the new o1 OpenAI model

mahaloz · Sep 13, 2024 · 1071c2a · 1071c2a
1 parent b14764c
commit 1071c2a
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 4 deletions.
diff --git a/dailalib/__init__.py b/dailalib/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "3.6.0"
+__version__ = "3.7.0"
 
 from .api import AIAPI, LiteLLMAIAPI
 from libbs.api import DecompilerInterface

diff --git a/dailalib/api/litellm/litellm_api.py b/dailalib/api/litellm/litellm_api.py
@@ -9,7 +9,11 @@
 class LiteLLMAIAPI(AIAPI):
     prompts_by_name = []
     DEFAULT_MODEL = "gpt-4o"
+    OPENAI_MODELS = {"gpt-4", "gpt-4o", "gpt-4-turbo", "gpt-3.5-turbo", "o1-mini", "o1-preview"}
     MODEL_TO_TOKENS = {
+        # TODO: update the token values for o1
+        "o1-mini": 8_000,
+        "o1-preview": 8_000,
         "gpt-4-turbo": 128_000,
         "gpt-4": 8_000,
         "gpt-4o": 8_000,
@@ -115,7 +119,7 @@ def fit_decompilation_to_token_max(decompilation: str, delta_step=10, model=DEFA
     def api_key(self):
         if not self._api_key:
             return None
-        elif "gpt" in self.model:
+        elif self.model in self.OPENAI_MODELS:
             return os.getenv("OPENAI_API_KEY", None)
         elif "claude" in self.model:
             return os.getenv("ANTHROPIC_API_KEY", None)
@@ -130,7 +134,7 @@ def api_key(self):
     def api_key(self, value):
         self._api_key = value
         if self._api_key:
-            if "gpt" in self.model:
+            if self.model in self.OPENAI_MODELS:
                 os.environ["OPENAI_API_KEY"] = self._api_key
             elif "claude" in self.model:
                 os.environ["ANTHROPIC_API_KEY"] = self._api_key

diff --git a/setup.cfg b/setup.cfg
@@ -14,7 +14,7 @@ long_description_content_type = text/markdown
 
 [options]
 install_requires =
-    litellm
+    litellm>=1.44.27
     tiktoken
     Jinja2
     libbs>=1.22.0