Add the full list of supported models (#11)

platformatic · Apr 10, 2024 · b9cddee · b9cddee
1 parent cafc995
commit b9cddee
Show file tree

Hide file tree

Showing 6 changed files with 49 additions and 13 deletions.
diff --git a/config.d.ts b/config.d.ts
@@ -226,13 +226,34 @@ export interface AiWarpConfig {
   aiProvider:
     | {
         openai: {
-          model: "gpt-3.5-turbo" | "gpt-4";
+          model:
+            | "gpt-4-0125-preview"
+            | "gpt-4-turbo-preview"
+            | "gpt-4-1106-preview"
+            | "gpt-4-vision-preview"
+            | "gpt-4-1106-vision-preview"
+            | "gpt-4"
+            | "gpt-4-0613"
+            | "gpt-4-32k"
+            | "gpt-4-32k-0613"
+            | "gpt-3.5-turbo-0125"
+            | "gpt-3.5-turbo"
+            | "gpt-3.5-turbo-1106"
+            | "gpt-3.5-turbo-instruct"
+            | "gpt-3.5-turbo-16k"
+            | "gpt-3.5-turbo-0613"
+            | "gpt-3.5-turbo-16k-0613";
           apiKey: string;
         };
       }
     | {
         mistral: {
-          model: "mistral-tiny";
+          model:
+            | "open-mistral-7b"
+            | "open-mixtral-8x7b"
+            | "mistral-small-latest"
+            | "mistral-medium-latest"
+            | "mistral-large-latest";
           apiKey: string;
         };
       };

diff --git a/lib/schema.ts b/lib/schema.ts
@@ -18,9 +18,22 @@ const aiWarpSchema = {
                 model: {
                   type: 'string',
                   enum: [
+                    'gpt-4-0125-preview',
+                    'gpt-4-turbo-preview',
+                    'gpt-4-1106-preview',
+                    'gpt-4-vision-preview',
+                    'gpt-4-1106-vision-preview',
+                    'gpt-4',
+                    'gpt-4-0613',
+                    'gpt-4-32k',
+                    'gpt-4-32k-0613',
+                    'gpt-3.5-turbo-0125',
                     'gpt-3.5-turbo',
-                    'gpt-4'
-                    // TODO: fill
+                    'gpt-3.5-turbo-1106',
+                    'gpt-3.5-turbo-instruct',
+                    'gpt-3.5-turbo-16k',
+                    'gpt-3.5-turbo-0613',
+                    'gpt-3.5-turbo-16k-0613'
                   ]
                 },
                 apiKey: { type: 'string' }
@@ -40,8 +53,11 @@ const aiWarpSchema = {
                 model: {
                   type: 'string',
                   enum: [
-                    'mistral-tiny'
-                    // TODO: fill
+                    'open-mistral-7b',
+                    'open-mixtral-8x7b',
+                    'mistral-small-latest',
+                    'mistral-medium-latest',
+                    'mistral-large-latest'
                   ]
                 },
                 apiKey: { type: 'string' }

diff --git a/tests/e2e/api.test.ts b/tests/e2e/api.test.ts
@@ -28,7 +28,7 @@ const providers: Provider[] = [
     name: 'Mistral',
     config: {
       mistral: {
-        model: 'mistral-tiny',
+        model: 'open-mistral-7b',
         apiKey: ''
       }
     }

diff --git a/tests/types/schema.test-d.ts b/tests/types/schema.test-d.ts
@@ -1,7 +1,6 @@
 import { expectAssignable } from 'tsd'
 import { AiWarpConfig } from '../../config'
 
-// TODO: add more when more models are added
 expectAssignable<AiWarpConfig['aiProvider']>({
   openai: {
     model: 'gpt-3.5-turbo',
@@ -18,7 +17,7 @@ expectAssignable<AiWarpConfig['aiProvider']>({
 
 expectAssignable<AiWarpConfig['aiProvider']>({
   mistral: {
-    model: 'mistral-tiny',
+    model: 'open-mistral-7b',
     apiKey: ''
   }
 })
@@ -35,7 +34,7 @@ expectAssignable<AiWarpConfig>({
 expectAssignable<AiWarpConfig>({
   aiProvider: {
     mistral: {
-      model: 'mistral-tiny',
+      model: 'open-mistral-7b',
       apiKey: ''
     }
   }

diff --git a/tests/unit/ai-providers.test.ts b/tests/unit/ai-providers.test.ts
@@ -10,7 +10,7 @@ const expectedStreamBody = buildExpectedStreamBodyString()
 
 const providers: AiProvider[] = [
   new OpenAiProvider('gpt-3.5-turbo', ''),
-  new MistralProvider('mistral-tiny', '')
+  new MistralProvider('open-mistral-7b', '')
 ]
 
 for (const provider of providers) {

diff --git a/tests/unit/generator.test.ts b/tests/unit/generator.test.ts
@@ -93,14 +93,14 @@ describe('AiWarpGenerator', () => {
 
     generator.setConfig({
       aiProvider: 'mistral',
-      aiModel: 'mistral-tiny'
+      aiModel: 'open-mistral-7b'
     })
     await generator.run()
 
     configFile = JSON.parse(await readFile(join(dir, 'platformatic.json'), 'utf8'))
     assert.deepStrictEqual(configFile.aiProvider, {
       mistral: {
-        model: 'mistral-tiny',
+        model: 'open-mistral-7b',
         apiKey: '{PLT_MISTRAL_API_KEY}'
       }
     })