Skip to content

Commit 8616a63

Browse files
mattgotteinermattgotteiner
committed
Add embedding dimensions
1 parent 529b4a2 commit 8616a63

File tree

13 files changed

+111
-11
lines changed

13 files changed

+111
-11
lines changed

‎app/backend/Services/ReadRetrieveReadChatService.cs

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,19 +33,27 @@ public ReadRetrieveReadChatService(
3333
kernelBuilder = kernelBuilder.AddOpenAIChatCompletion(deployment, client);
3434

3535
var embeddingModelName = configuration["OpenAiEmbeddingDeployment"];
36+
int embeddingModelDimensions;
37+
if (!int.TryParse(configuration["AzureOpenAiEmbeddingModelDimensions"], out embeddingModelDimensions)) {
38+
embeddingModelDimensions = 1536;
39+
}
3640
ArgumentNullException.ThrowIfNullOrWhiteSpace(embeddingModelName);
37-
kernelBuilder = kernelBuilder.AddOpenAITextEmbeddingGeneration(embeddingModelName, client);
41+
kernelBuilder = kernelBuilder.AddOpenAITextEmbeddingGeneration(embeddingModelName, client, dimensions: embeddingModelDimensions);
3842
}
3943
else
4044
{
4145
var deployedModelName = configuration["AzureOpenAiChatGptDeployment"];
4246
ArgumentNullException.ThrowIfNullOrWhiteSpace(deployedModelName);
4347
var embeddingModelName = configuration["AzureOpenAiEmbeddingDeployment"];
48+
int embeddingModelDimensions;
49+
if (!int.TryParse(configuration["AzureOpenAiEmbeddingModelDimensions"], out embeddingModelDimensions)) {
50+
embeddingModelDimensions = 1536;
51+
}
4452
if (!string.IsNullOrEmpty(embeddingModelName))
4553
{
4654
var endpoint = configuration["AzureOpenAiServiceEndpoint"];
4755
ArgumentNullException.ThrowIfNullOrWhiteSpace(endpoint);
48-
kernelBuilder = kernelBuilder.AddAzureOpenAITextEmbeddingGeneration(embeddingModelName, endpoint, tokenCredential ?? new DefaultAzureCredential());
56+
kernelBuilder = kernelBuilder.AddAzureOpenAITextEmbeddingGeneration(embeddingModelName, endpoint, tokenCredential ?? new DefaultAzureCredential(), dimensions: embeddingModelDimensions);
4957
kernelBuilder = kernelBuilder.AddAzureOpenAIChatCompletion(deployedModelName, endpoint, tokenCredential ?? new DefaultAzureCredential());
5058
}
5159
}

‎app/functions/EmbedFunctions/Program.cs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
// Copyright (c) Microsoft. All rights reserved.
22

33
using Azure.AI.OpenAI;
4+
using Microsoft.Extensions.Logging.Abstractions;
45

56
var host = new HostBuilder()
67
.ConfigureServices(services =>
@@ -65,17 +66,24 @@ uri is not null
6566

6667
OpenAIClient? openAIClient = null;
6768
string? embeddingModelName = null;
69+
int embeddingModelDimensions = -1;
6870

6971
if (useAOAI)
7072
{
7173
var openaiEndPoint = Environment.GetEnvironmentVariable("AZURE_OPENAI_ENDPOINT") ?? throw new ArgumentNullException("AZURE_OPENAI_ENDPOINT is null");
7274
embeddingModelName = Environment.GetEnvironmentVariable("AZURE_OPENAI_EMBEDDING_DEPLOYMENT") ?? throw new ArgumentNullException("AZURE_OPENAI_EMBEDDING_DEPLOYMENT is null");
75+
if (!int.TryParse(Environment.GetEnvironmentVariable("AZURE_OPENAI_EMBEDDING_MODEL_DIMENSIONS"), out embeddingModelDimensions)) {
76+
embeddingModelDimensions = 1536;
77+
}
7378
openAIClient = new OpenAIClient(new Uri(openaiEndPoint), new DefaultAzureCredential());
7479
}
7580
else
7681
{
7782
embeddingModelName = Environment.GetEnvironmentVariable("OPENAI_EMBEDDING_DEPLOYMENT") ?? throw new ArgumentNullException("OPENAI_EMBEDDING_DEPLOYMENT is null");
7883
var openaiKey = Environment.GetEnvironmentVariable("OPENAI_API_KEY") ?? throw new ArgumentNullException("OPENAI_API_KEY is null");
84+
if (!int.TryParse(Environment.GetEnvironmentVariable("AZURE_OPENAI_EMBEDDING_MODEL_DIMENSIONS"), out embeddingModelDimensions)) {
85+
embeddingModelDimensions = 1536;
86+
}
7987
openAIClient = new OpenAIClient(openaiKey);
8088
}
8189

@@ -94,6 +102,7 @@ uri is not null
94102
return new AzureSearchEmbedService(
95103
openAIClient: openAIClient,
96104
embeddingModelName: embeddingModelName,
105+
embeddingModelDimensions: embeddingModelDimensions,
97106
searchClient: searchClient,
98107
searchIndexName: searchIndexName,
99108
searchIndexClient: searchIndexClient,
@@ -108,6 +117,7 @@ uri is not null
108117
return new AzureSearchEmbedService(
109118
openAIClient: openAIClient,
110119
embeddingModelName: embeddingModelName,
120+
embeddingModelDimensions: embeddingModelDimensions,
111121
searchClient: searchClient,
112122
searchIndexName: searchIndexName,
113123
searchIndexClient: searchIndexClient,

‎app/map-env.ps1

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,7 @@ foreach ($line in (& azd env get-values)) {
66
if ($line -match "([^=]+)=(.*)") {
77
$key = $matches[1]
88
$value = $matches[2] -replace '^"|"$'
9-
[Environment]::SetEnvironmentVariable(
10-
$key, $value, [System.EnvironmentVariableTarget]::User)
9+
Set-Item "env:$key" $value
1110
}
1211
}
1312

‎app/prepdocs/PrepareDocs/AppOptions.cs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ internal record class AppOptions(
1313
string? AzureOpenAIServiceEndpoint,
1414
string? SearchIndexName,
1515
string? EmbeddingModelName,
16+
int EmbeddingModelDimensions,
1617
bool Remove,
1718
bool RemoveAll,
1819
string? FormRecognizerServiceEndpoint,

‎app/prepdocs/PrepareDocs/Program.Clients.cs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,12 +26,14 @@ private static Task<AzureSearchEmbedService> GetAzureSearchEmbedService(AppOptio
2626
var blobContainerClient = await GetCorpusBlobContainerClientAsync(o);
2727
var openAIClient = await GetOpenAIClientAsync(o);
2828
var embeddingModelName = o.EmbeddingModelName ?? throw new ArgumentNullException(nameof(o.EmbeddingModelName));
29+
var embeddingModelDimensions = o.EmbeddingModelDimensions;
2930
var searchIndexName = o.SearchIndexName ?? throw new ArgumentNullException(nameof(o.SearchIndexName));
3031
var computerVisionService = await GetComputerVisionServiceAsync(o);
3132

3233
return new AzureSearchEmbedService(
3334
openAIClient: openAIClient,
3435
embeddingModelName: embeddingModelName,
36+
embeddingModelDimensions: embeddingModelDimensions,
3537
searchClient: searchClient,
3638
searchIndexName: searchIndexName,
3739
searchIndexClient: searchIndexClient,

‎app/prepdocs/PrepareDocs/Program.Options.cs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,9 @@ internal static partial class Program
3232
private static readonly Option<string> s_embeddingModelName =
3333
new(name: "--embeddingmodel", description: "Optional. Name of the Azure AI Search embedding model to use for embedding content in the search index (will be created if it doesn't exist)");
3434

35+
private static readonly Option<int> s_embeddingModelDimensions =
36+
new(name: "--embeddingmodeldimensions", description: "Optional. Name of the Azure AI Search embedding model dimensions to use for embedding content in the search index");
37+
3538
private static readonly Option<bool> s_remove =
3639
new(name: "--remove", description: "Remove references to this document from blob storage and the search index");
3740

@@ -63,6 +66,7 @@ internal static partial class Program
6366
s_searchIndexName,
6467
s_azureOpenAIService,
6568
s_embeddingModelName,
69+
s_embeddingModelDimensions,
6670
s_remove,
6771
s_removeAll,
6872
s_formRecognizerServiceEndpoint,
@@ -81,6 +85,7 @@ internal static partial class Program
8185
SearchIndexName: context.ParseResult.GetValueForOption(s_searchIndexName),
8286
AzureOpenAIServiceEndpoint: context.ParseResult.GetValueForOption(s_azureOpenAIService),
8387
EmbeddingModelName: context.ParseResult.GetValueForOption(s_embeddingModelName),
88+
EmbeddingModelDimensions: context.ParseResult.GetValueForOption(s_embeddingModelDimensions),
8489
Remove: context.ParseResult.GetValueForOption(s_remove),
8590
RemoveAll: context.ParseResult.GetValueForOption(s_removeAll),
8691
FormRecognizerServiceEndpoint: context.ParseResult.GetValueForOption(s_formRecognizerServiceEndpoint),

‎app/shared/Shared/Services/AzureSearchEmbedService.cs

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
public sealed partial class AzureSearchEmbedService(
1919
OpenAIClient openAIClient,
2020
string embeddingModelName,
21+
int embeddingModelDimensions,
2122
SearchClient searchClient,
2223
string searchIndexName,
2324
SearchIndexClient searchIndexClient,
@@ -135,7 +136,7 @@ public async Task CreateSearchIndexAsync(string searchIndexName, CancellationTok
135136
new SimpleField("sourcefile", SearchFieldDataType.String) { IsFacetable = true },
136137
new SearchField("embedding", SearchFieldDataType.Collection(SearchFieldDataType.Single))
137138
{
138-
VectorSearchDimensions = 1536,
139+
VectorSearchDimensions = embeddingModelDimensions,
139140
IsSearchable = true,
140141
VectorSearchProfileName = vectorSearchProfile,
141142
}
@@ -449,7 +450,7 @@ private async Task IndexSectionsAsync(IEnumerable<Section> sections)
449450
var batch = new IndexDocumentsBatch<SearchDocument>();
450451
foreach (var section in sections)
451452
{
452-
var embeddings = await openAIClient.GetEmbeddingsAsync(new Azure.AI.OpenAI.EmbeddingsOptions(embeddingModelName, [section.Content.Replace('\r', ' ')]));
453+
var embeddings = await openAIClient.GetEmbeddingsAsync(new Azure.AI.OpenAI.EmbeddingsOptions(embeddingModelName, [section.Content.Replace('\r', ' ')]) { Dimensions = embeddingModelDimensions });
453454
var embedding = embeddings.Value.Data.FirstOrDefault()?.Embedding.ToArray() ?? [];
454455
batch.Actions.Add(new IndexDocumentsAction<SearchDocument>(
455456
IndexActionType.MergeOrUpload,

‎app/tests/MinimalApi.Tests/AzureDocumentSearchServiceTest.cs

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,34 @@ public async Task QueryDocumentsTestTextOnlyAsync()
3737
records.Count().Should().Be(3);
3838
}
3939

40+
[EnvironmentVariablesFact("AZURE_SEARCH_INDEX", "AZURE_SEARCH_SERVICE_ENDPOINT", "AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_EMBEDDING_DEPLOYMENT", "AZURE_OPENAI_EMBEDDING_MODEL_DIMENSIONS")]
41+
public async Task QueryDocumentsTestEmbeddingOnlyDimensionsAsync()
42+
{
43+
var index = Environment.GetEnvironmentVariable("AZURE_SEARCH_INDEX") ?? throw new InvalidOperationException();
44+
var searchServceEndpoint = Environment.GetEnvironmentVariable("AZURE_SEARCH_SERVICE_ENDPOINT") ?? throw new InvalidOperationException();
45+
var openAiEndpoint = Environment.GetEnvironmentVariable("AZURE_OPENAI_ENDPOINT") ?? throw new InvalidOperationException();
46+
var openAiEmbeddingDeployment = Environment.GetEnvironmentVariable("AZURE_OPENAI_EMBEDDING_DEPLOYMENT") ?? throw new InvalidOperationException();
47+
var openAiEmbeddingModelDimensions = int.Parse(Environment.GetEnvironmentVariable("AZURE_OPENAI_EMBEDDING_MODEL_DIMENSIONS") ?? "");
48+
var openAIClient = new OpenAIClient(new Uri(openAiEndpoint), new DefaultAzureCredential());
49+
var query = "What is included in my Northwind Health Plus plan that is not in standard?";
50+
var embeddingResponse = await openAIClient.GetEmbeddingsAsync(new EmbeddingsOptions(openAiEmbeddingDeployment, [query]) { Dimensions = openAiEmbeddingModelDimensions });
51+
var embedding = embeddingResponse.Value.Data.First().Embedding;
52+
var searchClient = new SearchClient(new Uri(searchServceEndpoint), index, new DefaultAzureCredential());
53+
var service = new AzureSearchService(searchClient);
54+
55+
// query only
56+
var option = new RequestOverrides
57+
{
58+
RetrievalMode = RetrievalMode.Vector,
59+
Top = 3,
60+
SemanticCaptions = true,
61+
SemanticRanker = true,
62+
};
63+
64+
var records = await service.QueryDocumentsAsync(query: query, embedding: embedding.ToArray(), overrides: option);
65+
records.Count().Should().Be(3);
66+
}
67+
4068
[EnvironmentVariablesFact("AZURE_SEARCH_INDEX", "AZURE_SEARCH_SERVICE_ENDPOINT", "AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_EMBEDDING_DEPLOYMENT")]
4169
public async Task QueryDocumentsTestEmbeddingOnlyAsync()
4270
{

0 commit comments

Comments
 (0)