feat(api): add text embeddings dimensions param (#1103)

openai · Jan 25, 2024 · 20d97ef · 20d97ef
1 parent 62cd463
commit 20d97ef
Show file tree

Hide file tree

Showing 5 changed files with 62 additions and 18 deletions.
diff --git a/src/openai/resources/chat/completions.py b/src/openai/resources/chat/completions.py
@@ -46,6 +46,8 @@ def create(
         model: Union[
             str,
             Literal[
+                "gpt-4-0125-preview",
+                "gpt-4-turbo-preview",
                 "gpt-4-1106-preview",
                 "gpt-4-vision-preview",
                 "gpt-4",
@@ -152,7 +154,8 @@ def create(
               [See more information about frequency and presence penalties.](https://platform.openai.com/docs/guides/text-generation/parameter-details)
 
           response_format: An object specifying the format that the model must output. Compatible with
-              `gpt-4-1106-preview` and `gpt-3.5-turbo-1106`.
+              [GPT-4 Turbo](https://platform.openai.com/docs/models/gpt-4-and-gpt-4-turbo) and
+              `gpt-3.5-turbo-1106`.
 
               Setting to `{ "type": "json_object" }` enables JSON mode, which guarantees the
               message the model generates is valid JSON.
@@ -232,6 +235,8 @@ def create(
         model: Union[
             str,
             Literal[
+                "gpt-4-0125-preview",
+                "gpt-4-turbo-preview",
                 "gpt-4-1106-preview",
                 "gpt-4-vision-preview",
                 "gpt-4",
@@ -345,7 +350,8 @@ def create(
               [See more information about frequency and presence penalties.](https://platform.openai.com/docs/guides/text-generation/parameter-details)
 
           response_format: An object specifying the format that the model must output. Compatible with
-              `gpt-4-1106-preview` and `gpt-3.5-turbo-1106`.
+              [GPT-4 Turbo](https://platform.openai.com/docs/models/gpt-4-and-gpt-4-turbo) and
+              `gpt-3.5-turbo-1106`.
 
               Setting to `{ "type": "json_object" }` enables JSON mode, which guarantees the
               message the model generates is valid JSON.
@@ -418,6 +424,8 @@ def create(
         model: Union[
             str,
             Literal[
+                "gpt-4-0125-preview",
+                "gpt-4-turbo-preview",
                 "gpt-4-1106-preview",
                 "gpt-4-vision-preview",
                 "gpt-4",
@@ -531,7 +539,8 @@ def create(
               [See more information about frequency and presence penalties.](https://platform.openai.com/docs/guides/text-generation/parameter-details)
 
           response_format: An object specifying the format that the model must output. Compatible with
-              `gpt-4-1106-preview` and `gpt-3.5-turbo-1106`.
+              [GPT-4 Turbo](https://platform.openai.com/docs/models/gpt-4-and-gpt-4-turbo) and
+              `gpt-3.5-turbo-1106`.
 
               Setting to `{ "type": "json_object" }` enables JSON mode, which guarantees the
               message the model generates is valid JSON.
@@ -604,6 +613,8 @@ def create(
         model: Union[
             str,
             Literal[
+                "gpt-4-0125-preview",
+                "gpt-4-turbo-preview",
                 "gpt-4-1106-preview",
                 "gpt-4-vision-preview",
                 "gpt-4",
@@ -698,6 +709,8 @@ async def create(
         model: Union[
             str,
             Literal[
+                "gpt-4-0125-preview",
+                "gpt-4-turbo-preview",
                 "gpt-4-1106-preview",
                 "gpt-4-vision-preview",
                 "gpt-4",
@@ -804,7 +817,8 @@ async def create(
               [See more information about frequency and presence penalties.](https://platform.openai.com/docs/guides/text-generation/parameter-details)
 
           response_format: An object specifying the format that the model must output. Compatible with
-              `gpt-4-1106-preview` and `gpt-3.5-turbo-1106`.
+              [GPT-4 Turbo](https://platform.openai.com/docs/models/gpt-4-and-gpt-4-turbo) and
+              `gpt-3.5-turbo-1106`.
 
               Setting to `{ "type": "json_object" }` enables JSON mode, which guarantees the
               message the model generates is valid JSON.
@@ -884,6 +898,8 @@ async def create(
         model: Union[
             str,
             Literal[
+                "gpt-4-0125-preview",
+                "gpt-4-turbo-preview",
                 "gpt-4-1106-preview",
                 "gpt-4-vision-preview",
                 "gpt-4",
@@ -997,7 +1013,8 @@ async def create(
               [See more information about frequency and presence penalties.](https://platform.openai.com/docs/guides/text-generation/parameter-details)
 
           response_format: An object specifying the format that the model must output. Compatible with
-              `gpt-4-1106-preview` and `gpt-3.5-turbo-1106`.
+              [GPT-4 Turbo](https://platform.openai.com/docs/models/gpt-4-and-gpt-4-turbo) and
+              `gpt-3.5-turbo-1106`.
 
               Setting to `{ "type": "json_object" }` enables JSON mode, which guarantees the
               message the model generates is valid JSON.
@@ -1070,6 +1087,8 @@ async def create(
         model: Union[
             str,
             Literal[
+                "gpt-4-0125-preview",
+                "gpt-4-turbo-preview",
                 "gpt-4-1106-preview",
                 "gpt-4-vision-preview",
                 "gpt-4",
@@ -1183,7 +1202,8 @@ async def create(
               [See more information about frequency and presence penalties.](https://platform.openai.com/docs/guides/text-generation/parameter-details)
 
           response_format: An object specifying the format that the model must output. Compatible with
-              `gpt-4-1106-preview` and `gpt-3.5-turbo-1106`.
+              [GPT-4 Turbo](https://platform.openai.com/docs/models/gpt-4-and-gpt-4-turbo) and
+              `gpt-3.5-turbo-1106`.
 
               Setting to `{ "type": "json_object" }` enables JSON mode, which guarantees the
               message the model generates is valid JSON.
@@ -1256,6 +1276,8 @@ async def create(
         model: Union[
             str,
             Literal[
+                "gpt-4-0125-preview",
+                "gpt-4-turbo-preview",
                 "gpt-4-1106-preview",
                 "gpt-4-vision-preview",
                 "gpt-4",

diff --git a/src/openai/resources/embeddings.py b/src/openai/resources/embeddings.py
@@ -36,7 +36,8 @@ def create(
         self,
         *,
         input: Union[str, List[str], List[int], List[List[int]]],
-        model: Union[str, Literal["text-embedding-ada-002"]],
+        model: Union[str, Literal["text-embedding-ada-002", "text-embedding-3-small", "text-embedding-3-large"]],
+        dimensions: int | NotGiven = NOT_GIVEN,
         encoding_format: Literal["float", "base64"] | NotGiven = NOT_GIVEN,
         user: str | NotGiven = NOT_GIVEN,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
@@ -64,6 +65,9 @@ def create(
               [Model overview](https://platform.openai.com/docs/models/overview) for
               descriptions of them.
 
+          dimensions: The number of dimensions the resulting output embeddings should have. Only
+              supported in `text-embedding-3` and later models.
+
           encoding_format: The format to return the embeddings in. Can be either `float` or
               [`base64`](https://pypi.org/project/pybase64/).
 
@@ -83,6 +87,7 @@ def create(
             "input": input,
             "model": model,
             "user": user,
+            "dimensions": dimensions,
             "encoding_format": encoding_format,
         }
         if not is_given(encoding_format) and has_numpy():
@@ -132,7 +137,8 @@ async def create(
         self,
         *,
         input: Union[str, List[str], List[int], List[List[int]]],
-        model: Union[str, Literal["text-embedding-ada-002"]],
+        model: Union[str, Literal["text-embedding-ada-002", "text-embedding-3-small", "text-embedding-3-large"]],
+        dimensions: int | NotGiven = NOT_GIVEN,
         encoding_format: Literal["float", "base64"] | NotGiven = NOT_GIVEN,
         user: str | NotGiven = NOT_GIVEN,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
@@ -160,6 +166,9 @@ async def create(
               [Model overview](https://platform.openai.com/docs/models/overview) for
               descriptions of them.
 
+          dimensions: The number of dimensions the resulting output embeddings should have. Only
+              supported in `text-embedding-3` and later models.
+
           encoding_format: The format to return the embeddings in. Can be either `float` or
               [`base64`](https://pypi.org/project/pybase64/).
 
@@ -179,6 +188,7 @@ async def create(
             "input": input,
             "model": model,
             "user": user,
+            "dimensions": dimensions,
             "encoding_format": encoding_format,
         }
         if not is_given(encoding_format) and has_numpy():

diff --git a/src/openai/types/chat/completion_create_params.py b/src/openai/types/chat/completion_create_params.py
@@ -32,6 +32,8 @@ class CompletionCreateParamsBase(TypedDict, total=False):
         Union[
             str,
             Literal[
+                "gpt-4-0125-preview",
+                "gpt-4-turbo-preview",
                 "gpt-4-1106-preview",
                 "gpt-4-vision-preview",
                 "gpt-4",
@@ -133,7 +135,9 @@ class CompletionCreateParamsBase(TypedDict, total=False):
     response_format: ResponseFormat
     """An object specifying the format that the model must output.
 
-    Compatible with `gpt-4-1106-preview` and `gpt-3.5-turbo-1106`.
+    Compatible with
+    [GPT-4 Turbo](https://platform.openai.com/docs/models/gpt-4-and-gpt-4-turbo) and
+    `gpt-3.5-turbo-1106`.
 
     Setting to `{ "type": "json_object" }` enables JSON mode, which guarantees the
     message the model generates is valid JSON.

diff --git a/src/openai/types/embedding_create_params.py b/src/openai/types/embedding_create_params.py
@@ -20,7 +20,7 @@ class EmbeddingCreateParams(TypedDict, total=False):
     for counting tokens.
     """
 
-    model: Required[Union[str, Literal["text-embedding-ada-002"]]]
+    model: Required[Union[str, Literal["text-embedding-ada-002", "text-embedding-3-small", "text-embedding-3-large"]]]
     """ID of the model to use.
 
     You can use the
@@ -30,6 +30,12 @@ class EmbeddingCreateParams(TypedDict, total=False):
     descriptions of them.
     """
 
+    dimensions: int
+    """The number of dimensions the resulting output embeddings should have.
+
+    Only supported in `text-embedding-3` and later models.
+    """
+
     encoding_format: Literal["float", "base64"]
     """The format to return the embeddings in.
 

diff --git a/tests/api_resources/test_embeddings.py b/tests/api_resources/test_embeddings.py
@@ -21,15 +21,16 @@ class TestEmbeddings:
     def test_method_create(self, client: OpenAI) -> None:
         embedding = client.embeddings.create(
             input="The quick brown fox jumped over the lazy dog",
-            model="text-embedding-ada-002",
+            model="text-embedding-3-small",
         )
         assert_matches_type(CreateEmbeddingResponse, embedding, path=["response"])
 
     @parametrize
     def test_method_create_with_all_params(self, client: OpenAI) -> None:
         embedding = client.embeddings.create(
             input="The quick brown fox jumped over the lazy dog",
-            model="text-embedding-ada-002",
+            model="text-embedding-3-small",
+            dimensions=1,
             encoding_format="float",
             user="user-1234",
         )
@@ -39,7 +40,7 @@ def test_method_create_with_all_params(self, client: OpenAI) -> None:
     def test_raw_response_create(self, client: OpenAI) -> None:
         response = client.embeddings.with_raw_response.create(
             input="The quick brown fox jumped over the lazy dog",
-            model="text-embedding-ada-002",
+            model="text-embedding-3-small",
         )
 
         assert response.is_closed is True
@@ -51,7 +52,7 @@ def test_raw_response_create(self, client: OpenAI) -> None:
     def test_streaming_response_create(self, client: OpenAI) -> None:
         with client.embeddings.with_streaming_response.create(
             input="The quick brown fox jumped over the lazy dog",
-            model="text-embedding-ada-002",
+            model="text-embedding-3-small",
         ) as response:
             assert not response.is_closed
             assert response.http_request.headers.get("X-Stainless-Lang") == "python"
@@ -69,15 +70,16 @@ class TestAsyncEmbeddings:
     async def test_method_create(self, async_client: AsyncOpenAI) -> None:
         embedding = await async_client.embeddings.create(
             input="The quick brown fox jumped over the lazy dog",
-            model="text-embedding-ada-002",
+            model="text-embedding-3-small",
         )
         assert_matches_type(CreateEmbeddingResponse, embedding, path=["response"])
 
     @parametrize
     async def test_method_create_with_all_params(self, async_client: AsyncOpenAI) -> None:
         embedding = await async_client.embeddings.create(
             input="The quick brown fox jumped over the lazy dog",
-            model="text-embedding-ada-002",
+            model="text-embedding-3-small",
+            dimensions=1,
             encoding_format="float",
             user="user-1234",
         )
@@ -87,7 +89,7 @@ async def test_method_create_with_all_params(self, async_client: AsyncOpenAI) ->
     async def test_raw_response_create(self, async_client: AsyncOpenAI) -> None:
         response = await async_client.embeddings.with_raw_response.create(
             input="The quick brown fox jumped over the lazy dog",
-            model="text-embedding-ada-002",
+            model="text-embedding-3-small",
         )
 
         assert response.is_closed is True
@@ -99,7 +101,7 @@ async def test_raw_response_create(self, async_client: AsyncOpenAI) -> None:
     async def test_streaming_response_create(self, async_client: AsyncOpenAI) -> None:
         async with async_client.embeddings.with_streaming_response.create(
             input="The quick brown fox jumped over the lazy dog",
-            model="text-embedding-ada-002",
+            model="text-embedding-3-small",
         ) as response:
             assert not response.is_closed
             assert response.http_request.headers.get("X-Stainless-Lang") == "python"