feat: Add text-embedding-3-small and text-embedding-3-large in model …

…catalog (#310)
davidmigloz · Jan 26, 2024 · fda1602 · fda1602
1 parent f5a7886
commit fda1602
Show file tree

Hide file tree

Showing 5 changed files with 34 additions and 20 deletions.
diff --git a/packages/openai_dart/lib/src/generated/schema/create_embedding_request.dart b/packages/openai_dart/lib/src/generated/schema/create_embedding_request.dart
@@ -66,6 +66,10 @@ class CreateEmbeddingRequest with _$CreateEmbeddingRequest {
 enum EmbeddingModels {
   @JsonValue('text-embedding-ada-002')
   textEmbeddingAda002,
+  @JsonValue('text-embedding-3-small')
+  textEmbedding3Small,
+  @JsonValue('text-embedding-3-large')
+  textEmbedding3Large,
 }
 
 // ==========================================

diff --git a/packages/openai_dart/lib/src/generated/schema/schema.g.dart b/packages/openai_dart/lib/src/generated/schema/schema.g.dart
diff --git a/packages/openai_dart/oas/openapi_curated.yaml b/packages/openai_dart/oas/openapi_curated.yaml
@@ -2154,15 +2154,15 @@ components:
         model:
           title: EmbeddingModel
           description: *model_description
-          example: "text-embedding-ada-002"
+          example: "text-embedding-3-small"
           anyOf:
             - type: string
               description: The ID of the model to use for this request.
             - type: string
               title: EmbeddingModels
               description: |
                 Available completion models. Mind that the list may not be exhaustive nor up-to-date.
-              enum: [ "text-embedding-ada-002" ]
+              enum: [ "text-embedding-ada-002", "text-embedding-3-small", "text-embedding-3-large" ]
         input:
           title: EmbeddingInput
           description: |

diff --git a/packages/openai_dart/oas/openapi_official.yaml b/packages/openai_dart/oas/openapi_official.yaml
@@ -6740,11 +6740,11 @@ components:
           x-oaiExpandable: true
         model:
           description: *model_description
-          example: "text-embedding-ada-002"
+          example: "text-embedding-3-small"
           anyOf:
             - type: string
             - type: string
-              enum: ["text-embedding-ada-002"]
+              enum: ["text-embedding-ada-002", "text-embedding-3-small", "text-embedding-3-large"]
           x-oaiTypeLabel: string
         encoding_format:
           description: "The format to return the embeddings in. Can be either `float` or [`base64`](https://pypi.org/project/pybase64/)."

diff --git a/packages/openai_dart/test/openai_client_embeddings_test.dart b/packages/openai_dart/test/openai_client_embeddings_test.dart
@@ -21,26 +21,34 @@ void main() {
     });
 
     test('Test call embeddings API', () async {
-      const request = CreateEmbeddingRequest(
-        model: EmbeddingModel.model(EmbeddingModels.textEmbeddingAda002),
-        input: EmbeddingInput.string(
-          'The food was delicious and the waiter...',
-        ),
-      );
-      final res = await client.createEmbedding(request: request);
-      expect(res.data, hasLength(1));
-      expect(res.data.first.index, 0);
-      expect(res.data.first.embeddingVector, hasLength(1536));
-      expect(res.data.first.object, EmbeddingObject.embedding);
-      expect(res.model, startsWith('text-embedding-ada-002'));
-      expect(res.object, CreateEmbeddingResponseObject.list);
-      expect(res.usage?.promptTokens, greaterThan(0));
-      expect(res.usage?.totalTokens, greaterThan(0));
+      final models = [
+        (EmbeddingModels.textEmbeddingAda002, 1536),
+        (EmbeddingModels.textEmbedding3Small, 1536),
+        (EmbeddingModels.textEmbedding3Large, 3072),
+      ];
+
+      for (final (model, modelDim) in models) {
+        final request = CreateEmbeddingRequest(
+          model: EmbeddingModel.model(model),
+          input: const EmbeddingInput.string(
+            'The food was delicious and the waiter...',
+          ),
+        );
+        final res = await client.createEmbedding(request: request);
+        expect(res.data, hasLength(1));
+        expect(res.data.first.index, 0);
+        expect(res.data.first.embeddingVector, hasLength(modelDim));
+        expect(res.data.first.object, EmbeddingObject.embedding);
+        expect(res.model.replaceAll('-', ''), model.name.toLowerCase());
+        expect(res.object, CreateEmbeddingResponseObject.list);
+        expect(res.usage?.promptTokens, greaterThan(0));
+        expect(res.usage?.totalTokens, greaterThan(0));
+      }
     });
 
     test('Test call embeddings API with encoding base64', () async {
       const request = CreateEmbeddingRequest(
-        model: EmbeddingModel.model(EmbeddingModels.textEmbeddingAda002),
+        model: EmbeddingModel.model(EmbeddingModels.textEmbedding3Small),
         input: EmbeddingInput.string(
           'The food was delicious and the waiter...',
         ),