feat: add gpt-4o&gpt-4o-mini

context-labs · Jul 29, 2024 · 2138bbc · 2138bbc
1 parent bddf7cc
commit 2138bbc
Show file tree

Hide file tree

Showing 5 changed files with 66 additions and 0 deletions.
diff --git a/src/cli/commands/index/selectModel.ts b/src/cli/commands/index/selectModel.ts
@@ -26,6 +26,18 @@ export const selectModel = (
         getMaxPromptLength(prompts, LLMModels.GPT432k)
     ) {
       return models[LLMModels.GPT432k];
+    } else if (
+      llms.includes(LLMModels.GPT4o) &&
+      models[LLMModels.GPT4o].maxLength >
+        getMaxPromptLength(prompts, LLMModels.GPT4o)
+    ) {
+      return models[LLMModels.GPT4o];
+    } else if (
+      llms.includes(LLMModels.GPT4omini) &&
+      models[LLMModels.GPT4omini].maxLength >
+        getMaxPromptLength(prompts, LLMModels.GPT4omini)
+    ) {
+      return models[LLMModels.GPT4omini];
     } else {
       return null;
     }
@@ -42,6 +54,18 @@ export const selectModel = (
           getMaxPromptLength(prompts, LLMModels.GPT432k)
       ) {
         return models[LLMModels.GPT432k];
+      } else if (
+        llms.includes(LLMModels.GPT4o) &&
+        models[LLMModels.GPT4o].maxLength >
+          getMaxPromptLength(prompts, LLMModels.GPT4o)
+      ) {
+        return models[LLMModels.GPT4o];
+      } else if (
+        llms.includes(LLMModels.GPT4omini) &&
+        models[LLMModels.GPT4omini].maxLength >
+          getMaxPromptLength(prompts, LLMModels.GPT4omini)
+      ) {
+        return models[LLMModels.GPT4omini];
       } else {
         return null;
       }

diff --git a/src/cli/commands/init/index.ts b/src/cli/commands/init/index.ts
@@ -111,6 +111,10 @@ export const init = async (
           name: 'GPT-3.5 Turbo, GPT-4 8K (Early Access), GPT-4 32K (Early Access)',
           value: [LLMModels.GPT3, LLMModels.GPT4, LLMModels.GPT432k],
         },
+        {
+          name: 'GPT-4o, GPT-4o-mini',
+          value: [LLMModels.GPT4o, LLMModels.GPT4omini],
+        },
       ],
     },
     {

diff --git a/src/cli/commands/user/index.ts b/src/cli/commands/user/index.ts
@@ -63,6 +63,10 @@ export const user = async (
           name: 'GPT-3.5 Turbo, GPT-4 8K (Early Access), GPT-4 32K (Early Access)',
           value: [LLMModels.GPT3, LLMModels.GPT4, LLMModels.GPT432k],
         },
+        {
+          name: 'GPT-4o, GPT-4o-mini',
+          value: [LLMModels.GPT4o, LLMModels.GPT4omini],
+        },
       ],
     },
   ];

diff --git a/src/cli/utils/LLMUtil.ts b/src/cli/utils/LLMUtil.ts
@@ -50,6 +50,38 @@ export const models: Record<LLMModels, LLMModelDetails> = {
     failed: 0,
     total: 0,
   },
+  [LLMModels.GPT4o]: {
+    name: LLMModels.GPT4o,
+    inputCostPer1KTokens: 0.005,
+    outputCostPer1KTokens: 0.015,
+    maxLength: 4096,
+    llm: new OpenAIChat({
+      temperature: 0.1,
+      openAIApiKey: process.env.OPENAI_API_KEY,
+      modelName: LLMModels.GPT4o,
+    }),
+    inputTokens: 0,
+    outputTokens: 0,
+    succeeded: 0,
+    failed: 0,
+    total: 0,
+  },
+  [LLMModels.GPT4omini]: {
+    name: LLMModels.GPT4omini,
+    inputCostPer1KTokens: 0.00015,
+    outputCostPer1KTokens: 0.0006,
+    maxLength: 16384,
+    llm: new OpenAIChat({
+      temperature: 0.1,
+      openAIApiKey: process.env.OPENAI_API_KEY,
+      modelName: LLMModels.GPT4omini,
+    }),
+    inputTokens: 0,
+    outputTokens: 0,
+    succeeded: 0,
+    failed: 0,
+    total: 0,
+  },
 };
 
 export const printModelDetails = (models: LLMModelDetails[]): void => {

diff --git a/src/types.ts b/src/types.ts
@@ -85,6 +85,8 @@ export enum LLMModels {
   GPT3 = 'gpt-3.5-turbo',
   GPT4 = 'gpt-4',
   GPT432k = 'gpt-4-32k',
+  GPT4o = 'gpt-4o',
+  GPT4omini = 'gpt-4o-mini',
 }
 
 export type LLMModelDetails = {