Skip to content

Commit

Permalink
.Net: Use AsynEnumerable search interface for azure ai search. (#9090)
Browse files Browse the repository at this point in the history
### Motivation and Context

AzureAISearch has two ways of accessing search results.

### Description

- Switching to IAsyncEnumerable interface to improve performance.
- Switch to AzureOpenAI for integration test embeddings to remove need
for key management.
- Fix small bug in powershell module to create azure ai search test
infra

#9077

### Contribution Checklist

<!-- Before submitting this PR, please make sure: -->

- [x] The code builds clean without any errors or warnings
- [x] The PR follows the [SK Contribution
Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md)
and the [pre-submission formatting
script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts)
raises no violations
- [x] All unit tests pass, and I have added new tests where possible
- [x] I didn't break anyone 😄
  • Loading branch information
westey-m authored Oct 4, 2024
1 parent cc8fe45 commit fba0aba
Show file tree
Hide file tree
Showing 3 changed files with 49 additions and 21 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -446,7 +446,7 @@ public Task<VectorSearchResults<TRecord>> VectorizableTextSearchAsync(string sea
/// <param name="searchOptions">The options controlling the behavior of the search operation.</param>
/// <param name="includeVectors">A value indicating whether to include vectors in the result or not.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
/// <returns></returns>
/// <returns>The mapped search results.</returns>
private async Task<VectorSearchResults<TRecord>> SearchAndMapToDataModelAsync(
string? searchText,
SearchOptions searchOptions,
Expand All @@ -462,23 +462,14 @@ private async Task<VectorSearchResults<TRecord>> SearchAndMapToDataModelAsync(
OperationName,
() => this._searchClient.SearchAsync<JsonObject>(searchText, searchOptions, cancellationToken)).ConfigureAwait(false);

var mappedJsonObjectResults = jsonObjectResults.Value.GetResults().Select(result =>
{
var document = VectorStoreErrorHandler.RunModelConversion(
DatabaseName,
this._collectionName,
OperationName,
() => this._options.JsonObjectCustomMapper!.MapFromStorageToDataModel(result.Document, new() { IncludeVectors = includeVectors }));
return new VectorSearchResult<TRecord>(document, result.Score);
});

return new VectorSearchResults<TRecord>(mappedJsonObjectResults.ToAsyncEnumerable()) { TotalCount = jsonObjectResults.Value.TotalCount };
var mappedJsonObjectResults = this.MapSearchResultsAsync(jsonObjectResults.Value.GetResultsAsync(), OperationName, includeVectors);
return new VectorSearchResults<TRecord>(mappedJsonObjectResults) { TotalCount = jsonObjectResults.Value.TotalCount };
}

// Execute search and map using the built in Azure AI Search mapper.
Response<SearchResults<TRecord>> results = await this.RunOperationAsync(OperationName, () => this._searchClient.SearchAsync<TRecord>(searchText, searchOptions, cancellationToken)).ConfigureAwait(false);
var mappedResults = results.Value.GetResults().Select(result => new VectorSearchResult<TRecord>(result.Document, result.Score));
return new VectorSearchResults<TRecord>(mappedResults.ToAsyncEnumerable()) { TotalCount = results.Value.TotalCount };
var mappedResults = this.MapSearchResultsAsync(results.Value.GetResultsAsync());
return new VectorSearchResults<TRecord>(mappedResults) { TotalCount = results.Value.TotalCount };
}

/// <summary>
Expand Down Expand Up @@ -515,6 +506,39 @@ private Task<Response<IndexDocumentsResult>> MapToStorageModelAndUploadDocumentA
() => this._searchClient.UploadDocumentsAsync<TRecord>(records, innerOptions, cancellationToken));
}

/// <summary>
/// Map the search results from <see cref="SearchResult{JsonObject}"/> to <see cref="VectorSearchResults{TRecord}"/> objects using the configured mapper type.
/// </summary>
/// <param name="results">The search results to map.</param>
/// <param name="operationName">The name of the current operation for telemetry purposes.</param>
/// <param name="includeVectors">A value indicating whether to include vectors in the resultset or not.</param>
/// <returns>The mapped results.</returns>
private async IAsyncEnumerable<VectorSearchResult<TRecord>> MapSearchResultsAsync(IAsyncEnumerable<SearchResult<JsonObject>> results, string operationName, bool includeVectors)
{
await foreach (var result in results.ConfigureAwait(false))
{
var document = VectorStoreErrorHandler.RunModelConversion(
DatabaseName,
this._collectionName,
operationName,
() => this._options.JsonObjectCustomMapper!.MapFromStorageToDataModel(result.Document, new() { IncludeVectors = includeVectors }));
yield return new VectorSearchResult<TRecord>(document, result.Score);
}
}

/// <summary>
/// Map the search results from <see cref="SearchResult{TRecord}"/> to <see cref="VectorSearchResults{TRecord}"/> objects.
/// </summary>
/// <param name="results">The search results to map.</param>
/// <returns>The mapped results.</returns>
private async IAsyncEnumerable<VectorSearchResult<TRecord>> MapSearchResultsAsync(IAsyncEnumerable<SearchResult<TRecord>> results)
{
await foreach (var result in results.ConfigureAwait(false))
{
yield return new VectorSearchResult<TRecord>(result.Document, result.Score);
}
}

/// <summary>
/// Convert the public <see cref="GetRecordOptions"/> options model to the Azure AI Search <see cref="GetDocumentOptions"/> options model.
/// </summary>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,10 @@

using System;
using System.Threading.Tasks;
using Azure.Identity;
using Microsoft.Extensions.Configuration;
using Microsoft.SemanticKernel.Connectors.AzureAISearch;
using Microsoft.SemanticKernel.Connectors.OpenAI;
using Microsoft.SemanticKernel.Connectors.AzureOpenAI;
using Microsoft.SemanticKernel.Data;
using SemanticKernel.IntegrationTests.Data;
using SemanticKernel.IntegrationTests.TestSettings;
Expand Down Expand Up @@ -69,11 +70,14 @@ public override Task<ITextSearch> CreateTextSearchAsync()
{
if (this.VectorStore is null)
{
OpenAIConfiguration? openAIConfiguration = this.Configuration.GetSection("OpenAIEmbeddings").Get<OpenAIConfiguration>();
Assert.NotNull(openAIConfiguration);
Assert.NotEmpty(openAIConfiguration.ModelId);
Assert.NotEmpty(openAIConfiguration.ApiKey);
this.EmbeddingGenerator = new OpenAITextEmbeddingGenerationService(openAIConfiguration.ModelId, openAIConfiguration.ApiKey);
AzureOpenAIConfiguration? azureOpenAIConfiguration = this.Configuration.GetSection("AzureOpenAIEmbeddings").Get<AzureOpenAIConfiguration>();
Assert.NotNull(azureOpenAIConfiguration);
Assert.NotEmpty(azureOpenAIConfiguration.DeploymentName);
Assert.NotEmpty(azureOpenAIConfiguration.Endpoint);
this.EmbeddingGenerator = new AzureOpenAITextEmbeddingGenerationService(
azureOpenAIConfiguration.DeploymentName,
azureOpenAIConfiguration.Endpoint,
new AzureCliCredential());

this.VectorStore = new AzureAISearchVectorStore(fixture.SearchIndexClient);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ function New-AzureAISearchIntegrationInfra($overrideResourceGroup = $resourceGro
}

# Create the ai search service if it doesn't exist.
$service = Get-AzSearchService -ResourceGroupName $resourceGroup -Name $aiSearchResourceName
$service = Get-AzSearchService -ResourceGroupName $overrideResourceGroup -Name $overrideAISearchResourceName
if (-not $service) {
Write-Host "Service does not exist, creating '$overrideAISearchResourceName' ..."
New-AzSearchService -ResourceGroupName $overrideResourceGroup -Name $overrideAISearchResourceName -Sku "Basic" -Location "North Europe" -PartitionCount 1 -ReplicaCount 1 -HostingMode Default
Expand Down

0 comments on commit fba0aba

Please sign in to comment.