Add support for using new ada models with different dimensions (#1378)

* update reqs

* Add parameters for ada 3

* Update readme

* Update TOC

* Fix tests and mocks

* Mypy fixes

* Addressing feedback (more testing still needed)

* More readme notes

* Cast openaidimensions to int

* Support batch for ada3

* right model name

* Note about regions

* Add constants for tests model name and dimensions

* Default to int

* Typing error

* Fix env var name

* Undo unneeded parameter move

* Dont specify dimensons for old models

* typing

.azdo/pipelines/azure-dev.yml

@@ -70,6 +70,7 @@ steps:
       AZURE_OPENAI_EMB_DEPLOYMENT: $(AZURE_OPENAI_EMB_DEPLOYMENT)
       AZURE_OPENAI_EMB_DEPLOYMENT_CAPACITY: $(AZURE_OPENAI_EMB_DEPLOYMENT_CAPACITY)
       AZURE_OPENAI_EMB_DEPLOYMENT_VERSION: $(AZURE_OPENAI_EMB_DEPLOYMENT_VERSION)
+      AZURE_OPENAI_EMB_DIMENSIONS: $(AZURE_OPENAI_EMB_DIMENSIONS)
       OPENAI_HOST: $(OPENAI_HOST)
       OPENAI_API_KEY: $(OPENAI_API_KEY)
       OPENAI_ORGANIZATION: $(OPENAI_ORGANIZATION)

.github/workflows/azure-dev.yml

@@ -61,6 +61,7 @@ jobs:
       AZURE_OPENAI_EMB_DEPLOYMENT: ${{ vars.AZURE_OPENAI_EMB_DEPLOYMENT }}
       AZURE_OPENAI_EMB_DEPLOYMENT_CAPACITY: ${{ vars.AZURE_OPENAI_EMB_DEPLOYMENT_CAPACITY }}
       AZURE_OPENAI_EMB_DEPLOYMENT_VERSION: ${{ vars.AZURE_OPENAI_EMB_DEPLOYMENT_VERSION }}
+      AZURE_OPENAI_EMB_DIMENSIONS: ${{ vars.AZURE_OPENAI_EMB_DIMENSIONS }}
       OPENAI_HOST: ${{ vars.OPENAI_HOST }}
       OPENAI_API_KEY: ${{ vars.OPENAI_API_KEY }}
       OPENAI_ORGANIZATION: ${{ vars.OPENAI_ORGANIZATION }}

app/backend/app.py

@@ -222,6 +222,7 @@ async def setup_clients():
     OPENAI_HOST = os.getenv("OPENAI_HOST", "azure")
     OPENAI_CHATGPT_MODEL = os.environ["AZURE_OPENAI_CHATGPT_MODEL"]
     OPENAI_EMB_MODEL = os.getenv("AZURE_OPENAI_EMB_MODEL_NAME", "text-embedding-ada-002")
+    OPENAI_EMB_DIMENSIONS = int(os.getenv("AZURE_OPENAI_EMB_DIMENSIONS", 1536))
     # Used with Azure OpenAI deployments
     AZURE_OPENAI_SERVICE = os.getenv("AZURE_OPENAI_SERVICE")
     AZURE_OPENAI_GPT4V_DEPLOYMENT = os.environ.get("AZURE_OPENAI_GPT4V_DEPLOYMENT")
@@ -345,6 +346,7 @@ async def setup_clients():
         chatgpt_deployment=AZURE_OPENAI_CHATGPT_DEPLOYMENT,
         embedding_model=OPENAI_EMB_MODEL,
         embedding_deployment=AZURE_OPENAI_EMB_DEPLOYMENT,
+        embedding_dimensions=OPENAI_EMB_DIMENSIONS,
         sourcepage_field=KB_FIELDS_SOURCEPAGE,
         content_field=KB_FIELDS_CONTENT,
         query_language=AZURE_SEARCH_QUERY_LANGUAGE,
@@ -365,6 +367,7 @@ async def setup_clients():
             gpt4v_model=AZURE_OPENAI_GPT4V_MODEL,
             embedding_model=OPENAI_EMB_MODEL,
             embedding_deployment=AZURE_OPENAI_EMB_DEPLOYMENT,
+            embedding_dimensions=OPENAI_EMB_DIMENSIONS,
             sourcepage_field=KB_FIELDS_SOURCEPAGE,
             content_field=KB_FIELDS_CONTENT,
             query_language=AZURE_SEARCH_QUERY_LANGUAGE,
@@ -382,6 +385,7 @@ async def setup_clients():
             gpt4v_model=AZURE_OPENAI_GPT4V_MODEL,
             embedding_model=OPENAI_EMB_MODEL,
             embedding_deployment=AZURE_OPENAI_EMB_DEPLOYMENT,
+            embedding_dimensions=OPENAI_EMB_DIMENSIONS,
             sourcepage_field=KB_FIELDS_SOURCEPAGE,
             content_field=KB_FIELDS_CONTENT,
             query_language=AZURE_SEARCH_QUERY_LANGUAGE,
@@ -396,6 +400,7 @@ async def setup_clients():
         chatgpt_deployment=AZURE_OPENAI_CHATGPT_DEPLOYMENT,
         embedding_model=OPENAI_EMB_MODEL,
         embedding_deployment=AZURE_OPENAI_EMB_DEPLOYMENT,
+        embedding_dimensions=OPENAI_EMB_DIMENSIONS,
         sourcepage_field=KB_FIELDS_SOURCEPAGE,
         content_field=KB_FIELDS_CONTENT,
         query_language=AZURE_SEARCH_QUERY_LANGUAGE,

app/backend/approaches/approach.py

@@ -1,7 +1,17 @@
 import os
 from abc import ABC
 from dataclasses import dataclass
-from typing import Any, AsyncGenerator, Awaitable, Callable, List, Optional, Union, cast
+from typing import (
+    Any,
+    AsyncGenerator,
+    Awaitable,
+    Callable,
+    List,
+    Optional,
+    TypedDict,
+    Union,
+    cast,
+)
 from urllib.parse import urljoin
 
 import aiohttp
@@ -90,6 +100,7 @@ def __init__(
         query_speller: Optional[str],
         embedding_deployment: Optional[str],  # Not needed for non-Azure OpenAI or for retrieval_mode="text"
         embedding_model: str,
+        embedding_dimensions: int,
         openai_host: str,
         vision_endpoint: str,
         vision_token_provider: Callable[[], Awaitable[str]],
@@ -101,6 +112,7 @@ def __init__(
         self.query_speller = query_speller
         self.embedding_deployment = embedding_deployment
         self.embedding_model = embedding_model
+        self.embedding_dimensions = embedding_dimensions
         self.openai_host = openai_host
         self.vision_endpoint = vision_endpoint
         self.vision_token_provider = vision_token_provider
@@ -204,10 +216,23 @@ def get_citation(self, sourcepage: str, use_image_citation: bool) -> str:
             return sourcepage
 
     async def compute_text_embedding(self, q: str):
+        SUPPORTED_DIMENSIONS_MODEL = {
+            "text-embedding-ada-002": False,
+            "text-embedding-3-small": True,
+            "text-embedding-3-large": True,
+        }
+
+        class ExtraArgs(TypedDict, total=False):
+            dimensions: int
+
+        dimensions_args: ExtraArgs = (
+            {"dimensions": self.embedding_dimensions} if SUPPORTED_DIMENSIONS_MODEL[self.embedding_model] else {}
+        )
         embedding = await self.openai_client.embeddings.create(
             # Azure OpenAI takes the deployment name as the model name
             model=self.embedding_deployment if self.embedding_deployment else self.embedding_model,
             input=q,
+            **dimensions_args,
         )
         query_vector = embedding.data[0].embedding
         return VectorizedQuery(vector=query_vector, k_nearest_neighbors=50, fields="embedding")

app/backend/approaches/chatreadretrieveread.py

@@ -32,6 +32,7 @@ def __init__(
         chatgpt_deployment: Optional[str],  # Not needed for non-Azure OpenAI
         embedding_deployment: Optional[str],  # Not needed for non-Azure OpenAI or for retrieval_mode="text"
         embedding_model: str,
+        embedding_dimensions: int,
         sourcepage_field: str,
         content_field: str,
         query_language: str,
@@ -44,6 +45,7 @@ def __init__(
         self.chatgpt_deployment = chatgpt_deployment
         self.embedding_deployment = embedding_deployment
         self.embedding_model = embedding_model
+        self.embedding_dimensions = embedding_dimensions
         self.sourcepage_field = sourcepage_field
         self.content_field = content_field
         self.query_language = query_language

app/backend/approaches/chatreadretrievereadvision.py

@@ -35,6 +35,7 @@ def __init__(
         gpt4v_model: str,
         embedding_deployment: Optional[str],  # Not needed for non-Azure OpenAI or for retrieval_mode="text"
         embedding_model: str,
+        embedding_dimensions: int,
         sourcepage_field: str,
         content_field: str,
         query_language: str,
@@ -50,6 +51,7 @@ def __init__(
         self.gpt4v_model = gpt4v_model
         self.embedding_deployment = embedding_deployment
         self.embedding_model = embedding_model
+        self.embedding_dimensions = embedding_dimensions
         self.sourcepage_field = sourcepage_field
         self.content_field = content_field
         self.query_language = query_language

app/backend/approaches/retrievethenread.py

@@ -52,6 +52,7 @@ def __init__(
         chatgpt_deployment: Optional[str],  # Not needed for non-Azure OpenAI
         embedding_model: str,
         embedding_deployment: Optional[str],  # Not needed for non-Azure OpenAI or for retrieval_mode="text"
+        embedding_dimensions: int,
         sourcepage_field: str,
         content_field: str,
         query_language: str,
@@ -63,6 +64,7 @@ def __init__(
         self.auth_helper = auth_helper
         self.chatgpt_model = chatgpt_model
         self.embedding_model = embedding_model
+        self.embedding_dimensions = embedding_dimensions
         self.chatgpt_deployment = chatgpt_deployment
         self.embedding_deployment = embedding_deployment
         self.sourcepage_field = sourcepage_field

app/backend/approaches/retrievethenreadvision.py

@@ -48,6 +48,7 @@ def __init__(
         gpt4v_model: str,
         embedding_deployment: Optional[str],  # Not needed for non-Azure OpenAI or for retrieval_mode="text"
         embedding_model: str,
+        embedding_dimensions: int,
         sourcepage_field: str,
         content_field: str,
         query_language: str,
@@ -61,6 +62,7 @@ def __init__(
         self.auth_helper = auth_helper
         self.embedding_model = embedding_model
         self.embedding_deployment = embedding_deployment
+        self.embedding_dimensions = embedding_dimensions
         self.sourcepage_field = sourcepage_field
         self.content_field = content_field
         self.gpt4v_deployment = gpt4v_deployment

azure.yaml

@@ -51,6 +51,7 @@ pipeline:
       - AZURE_OPENAI_EMB_DEPLOYMENT
       - AZURE_OPENAI_EMB_DEPLOYMENT_CAPACITY
       - AZURE_OPENAI_EMB_DEPLOYMENT_VERSION
+      - AZURE_OPENAI_EMB_DIMENSIONS
       - OPENAI_HOST
       - OPENAI_API_KEY
       - OPENAI_ORGANIZATION

docs/deploy_features.md

@@ -5,6 +5,7 @@ This document covers optional features that can be enabled in the deployed Azure
 You should typically enable these features before running `azd up`. Once you've set them, return to the [deployment steps](../README.md#deploying).
 
 * [Using GPT-4](#using-gpt-4)
+* [Using text-embedding-3 models](#using-text-embedding-3-models)
 * [Enabling GPT-4 Turbo with Vision](#enabling-gpt-4-turbo-with-vision)
 * [Enabling Integrated Vectorization](#enabling-integrated-vectorization)
 * [Enabling authentication](#enabling-authentication)
@@ -58,6 +59,43 @@ Execute the following commands inside your terminal:
 >
 > Note that this does not delete your GPT-4 deployment; it just makes your application create a new or reuse an old GPT 3.5 deployment. If you want to delete it, you can go to your Azure OpenAI studio and do so.
 
+## Using text-embedding-3 models
+
+By default, the deployed Azure web app uses the `text-embedding-ada-002` embedding model. If you want to use one of the text-embedding-3 models, you can do so by following these steps:
+
+1. Run one of the following commands to set the desired model:
+
+    ```shell
+    azd env set AZURE_OPENAI_EMB_MODEL_NAME text-embedding-3-small
+    ```
+
+    ```shell
+    azd env set AZURE_OPENAI_EMB_MODEL_NAME text-embedding-3-large
+    ```
+
+2. Specify the desired dimensions of the model: (from 256-3072, model dependent)
+
+    ```shell
+    azd env set AZURE_OPENAI_EMB_DIMENSIONS 256
+    ```
+
+3. Set the model version to "1" (the only version as of March 2024):
+
+    ```shell
+    azd env set AZURE_OPENAI_EMB_DEPLOYMENT_VERSION 1
+    ```
+
+3. When prompted during `azd up`, make sure to select a region for the OpenAI resource group location that supports the text-embedding-3 models. There are [limited regions available](https://learn.microsoft.com/azure/ai-services/openai/concepts/models#embeddings-models).
+
+If you have already deployed:
+
+* You'll need to change the deployment name by running `azd env set AZURE_OPENAI_EMB_DEPLOYMENT <new-deployment-name>`
+* You'll need to create a new index, and re-index all of the data using the new model. You can either delete the current index in the Azure Portal, or create an index with a different name by running `azd env set AZURE_SEARCH_INDEX new-index-name`. When you next run `azd up`, the new index will be created and the data will be re-indexed.
+* If your OpenAI resource is not in one of the supported regions, you should delete `openAiResourceGroupLocation` from `.azure/YOUR-ENV-NAME/config.json`. When running `azd up`, you will be prompted to select a new region.
+
+> ![NOTE]
+> The text-embedding-3 models are not currently supported by the integrated vectorization feature.
+
 ## Enabling GPT-4 Turbo with Vision
 
 This section covers the integration of GPT-4 Vision with Azure AI Search. Learn how to enhance your search capabilities with the power of image and text indexing, enabling advanced search functionalities over diverse document types. For a detailed guide on setup and usage, visit our [Enabling GPT-4 Turbo with Vision](docs/gpt4v.md) page.
@@ -73,6 +111,8 @@ To enable integrated vectorization with this sample:
 3. Run `azd up` to update system and user roles
 4. You can view the resources such as the indexer and skillset in Azure Portal and monitor the status of the vectorization process.
 
+This feature is not currently compatible with GPT4-vision or the newer text-embedding-3 models.
+
 ## Enabling authentication
 
 By default, the deployed Azure web app will have no authentication or access restrictions enabled, meaning anyone with routable network access to the web app can chat with your indexed data. If you'd like to automatically setup authentication and user login as part of the `azd up` process, see [this guide](./login_and_acl.md).

infra/main.bicep

@@ -102,11 +102,13 @@ param embeddingModelName string = ''
 param embeddingDeploymentName string = ''
 param embeddingDeploymentVersion string = ''
 param embeddingDeploymentCapacity int = 0
+param embeddingDimensions int = 0
 var embedding = {
   modelName: !empty(embeddingModelName) ? embeddingModelName : 'text-embedding-ada-002'
   deploymentName: !empty(embeddingDeploymentName) ? embeddingDeploymentName : 'embedding'
   deploymentVersion: !empty(embeddingDeploymentVersion) ? embeddingDeploymentVersion : '2'
   deploymentCapacity: embeddingDeploymentCapacity != 0 ? embeddingDeploymentCapacity : 30
+  dimensions: embeddingDimensions != 0 ? embeddingDimensions : 1536
 }
 
 param gpt4vModelName string = 'gpt-4'
@@ -263,6 +265,7 @@ module backend 'core/host/appservice.bicep' = {
       AZURE_OPENAI_CUSTOM_URL: azureOpenAiCustomUrl
       AZURE_OPENAI_API_VERSION: azureOpenAiApiVersion
       AZURE_OPENAI_EMB_MODEL_NAME: embedding.modelName
+      AZURE_OPENAI_EMB_DIMENSIONS: embedding.dimensions
       AZURE_OPENAI_CHATGPT_MODEL: chatGpt.modelName
       AZURE_OPENAI_GPT4V_MODEL: gpt4vModelName
       // Specific to Azure OpenAI

infra/main.parameters.json

@@ -113,6 +113,9 @@
     "embeddingDeploymentCapacity":{
       "value": "${AZURE_OPENAI_EMB_DEPLOYMENT_CAPACITY}"
     },
+    "embeddingDimensions": {
+      "value": "${AZURE_OPENAI_EMB_DIMENSIONS}"
+    },
     "openAiHost": {
       "value": "${OPENAI_HOST=azure}"
     },

scripts/prepdocs.ps1

@@ -47,6 +47,10 @@ if ($env:USE_VECTORS -eq $false) {
   $disableVectorsArg="--novectors"
 }
 
+if ($env:AZURE_OPENAI_EMB_DIMENSIONS) {
+  $openaiDimensionsArg = "--openaidimensions $env:AZURE_OPENAI_EMB_DIMENSIONS"
+}
+
 if ($env:USE_LOCAL_PDF_PARSER -eq $true) {
   $localPdfParserArg = "--localpdfparser"
 }
@@ -71,7 +75,7 @@ $argumentList = "./scripts/prepdocs.py $dataArg --verbose " + `
 "--storageaccount $env:AZURE_STORAGE_ACCOUNT --container $env:AZURE_STORAGE_CONTAINER --storageresourcegroup $env:AZURE_STORAGE_RESOURCE_GROUP " + `
 "--searchservice $env:AZURE_SEARCH_SERVICE --index $env:AZURE_SEARCH_INDEX " + `
 "$searchAnalyzerNameArg $searchSecretNameArg " + `
-"--openaihost `"$env:OPENAI_HOST`" --openaimodelname `"$env:AZURE_OPENAI_EMB_MODEL_NAME`" " + `
+"--openaihost `"$env:OPENAI_HOST`" --openaimodelname `"$env:AZURE_OPENAI_EMB_MODEL_NAME`" $openaiDimensionsArg " + `
 "--openaiservice `"$env:AZURE_OPENAI_SERVICE`" --openaideployment `"$env:AZURE_OPENAI_EMB_DEPLOYMENT`" " + `
 "--openaikey `"$env:OPENAI_API_KEY`" --openaiorg `"$env:OPENAI_ORGANIZATION`" " + `
 "--documentintelligenceservice $env:AZURE_DOCUMENTINTELLIGENCE_SERVICE " + `

scripts/prepdocs.py

@@ -122,6 +122,7 @@ def setup_embeddings_service(
     openai_model_name: str,
     openai_service: str,
     openai_deployment: str,
+    openai_dimensions: int,
     openai_key: Union[str, None],
     openai_org: Union[str, None],
     disable_vectors: bool = False,
@@ -139,6 +140,7 @@ def setup_embeddings_service(
             open_ai_service=openai_service,
             open_ai_deployment=openai_deployment,
             open_ai_model_name=openai_model_name,
+            open_ai_dimensions=openai_dimensions,
             credential=azure_open_ai_credential,
             disable_batch=disable_batch_vectors,
         )
@@ -147,6 +149,7 @@ def setup_embeddings_service(
             raise ValueError("OpenAI key is required when using the non-Azure OpenAI API")
         return OpenAIEmbeddingService(
             open_ai_model_name=openai_model_name,
+            open_ai_dimensions=openai_dimensions,
             credential=openai_key,
             organization=openai_org,
             disable_batch=disable_batch_vectors,
@@ -308,6 +311,13 @@ async def main(strategy: Strategy, setup_index: bool = True):
     parser.add_argument(
         "--openaimodelname", help="Name of the Azure OpenAI embedding model ('text-embedding-ada-002' recommended)"
     )
+    parser.add_argument(
+        "--openaidimensions",
+        required=False,
+        default=1536,
+        type=int,
+        help="Dimensions for the embedding model (defaults to 1536 for 'text-embedding-ada-002')",
+    )
     parser.add_argument(
         "--novectors",
         action="store_true",
@@ -434,6 +444,7 @@ async def main(strategy: Strategy, setup_index: bool = True):
         openai_model_name=args.openaimodelname,
         openai_service=args.openaiservice,
         openai_deployment=args.openaideployment,
+        openai_dimensions=args.openaidimensions,
         openai_key=clean_key_if_exists(args.openaikey),
         openai_org=args.openaiorg,
         disable_vectors=args.novectors,

scripts/prepdocs.sh

@@ -48,6 +48,10 @@ if [ "$USE_VECTORS" = false ]; then
   disableVectorsArg="--novectors"
 fi
 
+if [ -n "$AZURE_OPENAI_EMB_DIMENSIONS" ]; then
+  openAiDimensionsArg="--openaidimensions $AZURE_OPENAI_EMB_DIMENSIONS"
+fi
+
 if [ "$USE_LOCAL_PDF_PARSER" = true ]; then
   localPdfParserArg="--localpdfparser"
 fi
@@ -69,7 +73,7 @@ fi
 --storageaccount "$AZURE_STORAGE_ACCOUNT" --container "$AZURE_STORAGE_CONTAINER" --storageresourcegroup $AZURE_STORAGE_RESOURCE_GROUP \
 --searchservice "$AZURE_SEARCH_SERVICE" --index "$AZURE_SEARCH_INDEX" \
 $searchAnalyzerNameArg $searchSecretNameArg \
---openaihost "$OPENAI_HOST" --openaimodelname "$AZURE_OPENAI_EMB_MODEL_NAME" \
+--openaihost "$OPENAI_HOST" --openaimodelname "$AZURE_OPENAI_EMB_MODEL_NAME" $openAiDimensionsArg \
 --openaiservice "$AZURE_OPENAI_SERVICE" --openaideployment "$AZURE_OPENAI_EMB_DEPLOYMENT"  \
 --openaikey "$OPENAI_API_KEY" --openaiorg "$OPENAI_ORGANIZATION" \
 --documentintelligenceservice "$AZURE_DOCUMENTINTELLIGENCE_SERVICE" \