o1 openai models

kharvd · Sep 14, 2024 · 1c9291c · 1c9291c
1 parent d90b1b9
commit 1c9291c
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/gptcli/assistant.py b/gptcli/assistant.py
@@ -72,6 +72,7 @@ def get_completion_provider(model: str) -> CompletionProvider:
         or model.startswith("ft:gpt")
         or model.startswith("oai-compat:")
         or model.startswith("chatgpt")
+        or model.startswith("o1")
     ):
         return OpenAICompletionProvider()
     elif model.startswith("claude"):

diff --git a/gptcli/providers/openai.py b/gptcli/providers/openai.py
@@ -122,6 +122,15 @@ def complete(
     "response": 0.600 / 1_000_000,
 }
 
+O_1_PREVIEW_PRICE_PER_TOKEN: Pricing = {
+    "prompt": 15.0 / 1_000_000,
+    "response": 60.0 / 1_000_000,
+}
+
+O_1_MINI_PRICE_PER_TOKEN: Pricing = {
+    "prompt": 3.0 / 1_000_000,
+    "response": 12.0 / 1_000_000,
+}
 
 def gpt_pricing(model: str) -> Optional[Pricing]:
     if model.startswith("gpt-3.5-turbo-16k"):
@@ -140,5 +149,9 @@ def gpt_pricing(model: str) -> Optional[Pricing]:
         return GPT_4_TURBO_PRICE_PER_TOKEN
     elif model.startswith("gpt-4"):
         return GPT_4_PRICE_PER_TOKEN
+    elif model.startswith("o1-preview"):
+        return O_1_PREVIEW_PRICE_PER_TOKEN
+    elif model.startswith("o1-mini"):
+        return O_1_MINI_PRICE_PER_TOKEN
     else:
         return None