Skip to content

Commit

Permalink
Support Azure AI Search Sticky Sessions (#877)
Browse files Browse the repository at this point in the history 
## Motivation and Context (Why the change? What's the scenario?)

The `UseSessionId` property on the `SearchOptions` can help improve the
relevance of search results when multiple replicas are configured. This
should be the case in any production deployment as two replicas are
required for high availability of read-only workloads (queries), and
three or more for high availability of read-write workloads (queries and
indexing).

> A value to be used to create a sticky session, which can help getting
more consistent results. As long as the same sessionId is used, a
best-effort attempt will be made to target the same replica set. Be wary
that reusing the same sessionID values repeatedly can interfere with the
load balancing of the requests across replicas and adversely affect the
performance of the search service. The value used as sessionId cannot
start with a '_' character.


https://learn.microsoft.com/en-us/rest/api/searchservice/documents/search-post?view=rest-searchservice-2024-07-01&tabs=HTTP#request-body

Using this setting should also provide more consistent results from the
`AzureAISearchMemory.GetListAsync` method. Without an `orderby`
parameter results are sorted by score, and when `*` is used as the
search query, all results have the same score of `1`. This increases the
likelihood of getting inconsistent results when the search query hits
different replicas.
---------

Co-authored-by: Devis Lucato <dluc@users.noreply.github.com>
  • Loading branch information
@alexmg @dluc
alexmg and dluc authored Nov 12, 2024
1 parent fb52c86 commit 490c013
Show file tree
Hide file tree
Showing 12 changed files with 131 additions and 49 deletions.
13 changes: 12 additions & 1 deletion applications/tests/Evaluation.Tests/appsettings.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,18 @@
"APIKey": "",
// Hybrid search is not enabled by default. Note that when using hybrid search
// relevance scores are different, usually lower, than when using just vector search
"UseHybridSearch": false
"UseHybridSearch": false,
// Helps improve relevance score consistency for search services with multiple replicas by
// attempting to route a given request to the same replica for that session. Use this when
// favoring consistent scoring over lower latency. Can adversely affect performance.
//
// Whether to use sticky sessions, which can help getting more consistent results.
// When using sticky sessions, a best-effort attempt will be made to target the same replica set.
// Be wary that reusing the same replica repeatedly can interfere with the load balancing of
// the requests across replicas and adversely affect the performance of the search service.
//
// See https://learn.microsoft.com/rest/api/searchservice/documents/search-post?view=rest-searchservice-2024-07-01&tabs=HTTP#request-body
"UseStickySessions": false
},
"OpenAI": {
// Name of the model used to generate text (text completion or chat completion)
Expand Down
13 changes: 12 additions & 1 deletion examples/002-dotnet-Serverless/appsettings.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,18 @@
"APIKey": "",
// Hybrid search is not enabled by default. Note that when using hybrid search
// relevance scores are different, usually lower, than when using just vector search
"UseHybridSearch": false
"UseHybridSearch": false,
// Helps improve relevance score consistency for search services with multiple replicas by
// attempting to route a given request to the same replica for that session. Use this when
// favoring consistent scoring over lower latency. Can adversely affect performance.
//
// Whether to use sticky sessions, which can help getting more consistent results.
// When using sticky sessions, a best-effort attempt will be made to target the same replica set.
// Be wary that reusing the same replica repeatedly can interfere with the load balancing of
// the requests across replicas and adversely affect the performance of the search service.
//
// See https://learn.microsoft.com/rest/api/searchservice/documents/search-post?view=rest-searchservice-2024-07-01&tabs=HTTP#request-body
"UseStickySessions": false
},
"OpenAI": {
// Name of the model used to generate text (text completion or chat completion)
Expand Down
13 changes: 12 additions & 1 deletion examples/210-KM-without-builder/appsettings.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -248,7 +248,18 @@
"APIKey": "",
// Hybrid search is not enabled by default. Note that when using hybrid search
// relevance scores are different, usually lower, than when using just vector search
"UseHybridSearch": false
"UseHybridSearch": false,
// Helps improve relevance score consistency for search services with multiple replicas by
// attempting to route a given request to the same replica for that session. Use this when
// favoring consistent scoring over lower latency. Can adversely affect performance.
//
// Whether to use sticky sessions, which can help getting more consistent results.
// When using sticky sessions, a best-effort attempt will be made to target the same replica set.
// Be wary that reusing the same replica repeatedly can interfere with the load balancing of
// the requests across replicas and adversely affect the performance of the search service.
//
// See https://learn.microsoft.com/rest/api/searchservice/documents/search-post?view=rest-searchservice-2024-07-01&tabs=HTTP#request-body
"UseStickySessions": false
},
"AzureAIDocIntel": {
// "APIKey" or "AzureIdentity".
Expand Down
13 changes: 12 additions & 1 deletion examples/401-evaluation/appsettings.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,18 @@
"APIKey": "",
// Hybrid search is not enabled by default. Note that when using hybrid search
// relevance scores are different, usually lower, than when using just vector search
"UseHybridSearch": false
"UseHybridSearch": false,
// Helps improve relevance score consistency for search services with multiple replicas by
// attempting to route a given request to the same replica for that session. Use this when
// favoring consistent scoring over lower latency. Can adversely affect performance.
//
// Whether to use sticky sessions, which can help getting more consistent results.
// When using sticky sessions, a best-effort attempt will be made to target the same replica set.
// Be wary that reusing the same replica repeatedly can interfere with the load balancing of
// the requests across replicas and adversely affect the performance of the search service.
//
// See https://learn.microsoft.com/rest/api/searchservice/documents/search-post?view=rest-searchservice-2024-07-01&tabs=HTTP#request-body
"UseStickySessions": false
},
"OpenAI": {
// Name of the model used to generate text (text completion or chat completion)
Expand Down
4 changes: 2 additions & 2 deletions extensions/AzureAISearch/AzureAISearch.FunctionalTests/DefaultTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,14 +14,14 @@ public class DefaultTests : BaseFunctionalTestCase
public DefaultTests(IConfiguration cfg, ITestOutputHelper output) : base(cfg, output)
{
Assert.False(string.IsNullOrEmpty(this.AzureAiSearchConfig.Endpoint));
Assert.False(string.IsNullOrEmpty(this.AzureAiSearchConfig.APIKey));
Assert.False(this.AzureAiSearchConfig.Auth == AzureAISearchConfig.AuthTypes.APIKey && string.IsNullOrEmpty(this.AzureAiSearchConfig.APIKey));
Assert.False(string.IsNullOrEmpty(this.OpenAiConfig.APIKey));

this._memory = new KernelMemoryBuilder()
.With(new KernelMemoryConfig { DefaultIndexName = "default4tests" })
.WithSearchClientConfig(new SearchClientConfig { EmptyAnswer = NotFound })
.WithOpenAI(this.OpenAiConfig)
.WithAzureAISearchMemoryDb(this.AzureAiSearchConfig.Endpoint, this.AzureAiSearchConfig.APIKey)
.WithAzureAISearchMemoryDb(this.AzureAiSearchConfig)
.Build<MemoryServerless>();
}

Expand Down
3 changes: 2 additions & 1 deletion extensions/AzureAISearch/AzureAISearch.FunctionalTests/appsettings.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,8 @@
// using the env vars AZURE_TENANT_ID, AZURE_CLIENT_ID, AZURE_CLIENT_SECRET.
"Auth": "AzureIdentity",
"Endpoint": "https://<...>",
"APIKey": ""
"APIKey": "",
"UseStickySessions": true
},
"OpenAI": {
// Name of the model used to generate text (text completion or chat completion)
Expand Down
14 changes: 14 additions & 0 deletions extensions/AzureAISearch/AzureAISearch/AzureAISearchConfig.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,20 @@ public enum AuthTypes
/// </summary>
public bool UseHybridSearch { get; set; } = false;

/// <summary>
/// Helps improve relevance score consistency for search services with multiple replicas by
/// attempting to route a given request to the same replica for that session. Use this when
/// favoring consistent scoring over lower latency. Can adversely affect performance.
///
/// Whether to use sticky sessions, which can help getting more consistent results.
/// When using sticky sessions, a best-effort attempt will be made to target the same replica set.
/// Be wary that reusing the same replica repeatedly can interfere with the load balancing of
/// the requests across replicas and adversely affect the performance of the search service.
///
/// See https://learn.microsoft.com/rest/api/searchservice/documents/search-post?view=rest-searchservice-2024-07-01&amp;tabs=HTTP#request-body
/// </summary>
public bool UseStickySessions { get; set; } = false;

public void SetCredential(TokenCredential credential)
{
this.Auth = AuthTypes.ManualTokenCredential;
Expand Down
88 changes: 48 additions & 40 deletions extensions/AzureAISearch/AzureAISearch/AzureAISearchMemory.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ public class AzureAISearchMemory : IMemoryDb, IMemoryDbUpsertBatch
private readonly ITextEmbeddingGenerator _embeddingGenerator;
private readonly ILogger<AzureAISearchMemory> _log;
private readonly bool _useHybridSearch;
private readonly bool _useStickySessions;

/// <summary>
/// Create a new instance
Expand All @@ -49,6 +50,7 @@ public AzureAISearchMemory(
this._embeddingGenerator = embeddingGenerator;
this._log = (loggerFactory ?? DefaultLogger.Factory).CreateLogger<AzureAISearchMemory>();
this._useHybridSearch = config.UseHybridSearch;
this._useStickySessions = config.UseStickySessions;

if (string.IsNullOrEmpty(config.Endpoint))
{
Expand Down Expand Up @@ -190,22 +192,12 @@ await client.IndexDocumentsAsync(
FilterMode = VectorFilterMode.PreFilter
}
};
DefineFieldsToSelect(options, withEmbeddings);
options = this.PrepareSearchOptions(options, withEmbeddings, filters, limit);

if (limit > 0)
{
vectorQuery.KNearestNeighborsCount = limit;
options.Size = limit;
this._log.LogDebug("KNearestNeighborsCount and max results: {0}", limit);
}

// Remove empty filters
filters = filters?.Where(f => !f.IsEmpty()).ToList();

if (filters is { Count: > 0 })
{
options.Filter = AzureAISearchFiltering.BuildSearchFilter(filters);
this._log.LogDebug("Filtering vectors, condition: {0}", options.Filter);
this._log.LogDebug("KNearestNeighborsCount: {0}", limit);
}

Response<SearchResults<AzureAISearchMemoryRecord>>? searchResult = null;
Expand Down Expand Up @@ -253,33 +245,7 @@ public async IAsyncEnumerable<MemoryRecord> GetListAsync(
{
var client = this.GetSearchClient(index);

SearchOptions options = new();
DefineFieldsToSelect(options, withEmbeddings);

if (limit > 0)
{
options.Size = limit;
this._log.LogDebug("Max results: {0}", limit);
}

// Remove empty filters
filters = filters?.Where(f => !f.IsEmpty()).ToList();

if (filters is { Count: > 0 })
{
options.Filter = AzureAISearchFiltering.BuildSearchFilter(filters);
this._log.LogDebug("Filtering vectors, condition: {0}", options.Filter);
}

// See: https://learn.microsoft.com/azure/search/search-query-understand-collection-filters
// fieldValue = fieldValue.Replace("'", "''", StringComparison.Ordinal);
// var options = new SearchOptions
// {
// Filter = fieldIsCollection
// ? $"{fieldName}/any(s: s eq '{fieldValue}')"
// : $"{fieldName} eq '{fieldValue}')",
// Size = limit
// };
SearchOptions options = this.PrepareSearchOptions(null, withEmbeddings, filters, limit);

Response<SearchResults<AzureAISearchMemoryRecord>>? searchResult = null;
try
Expand Down Expand Up @@ -627,15 +593,57 @@ at Azure.Search.Documents.SearchClient.SearchInternal[T](SearchOptions options,
return indexSchema;
}

private static void DefineFieldsToSelect(SearchOptions options, bool withEmbeddings)
private SearchOptions PrepareSearchOptions(
SearchOptions? options,
bool withEmbeddings,
ICollection<MemoryFilter>? filters = null,
int limit = 1)
{
options ??= new SearchOptions();

// Define which fields to fetch
options.Select.Add(AzureAISearchMemoryRecord.IdField);
options.Select.Add(AzureAISearchMemoryRecord.TagsField);
options.Select.Add(AzureAISearchMemoryRecord.PayloadField);

// Embeddings are fetched only when needed, to reduce latency and cost
if (withEmbeddings)
{
options.Select.Add(AzureAISearchMemoryRecord.VectorField);
}

// Remove empty filters
filters = filters?.Where(f => !f.IsEmpty()).ToList();

if (filters is { Count: > 0 })
{
options.Filter = AzureAISearchFiltering.BuildSearchFilter(filters);
this._log.LogDebug("Filtering vectors, condition: {0}", options.Filter);
}

// See: https://learn.microsoft.com/azure/search/search-query-understand-collection-filters
// fieldValue = fieldValue.Replace("'", "''", StringComparison.Ordinal);
// var options = new SearchOptions
// {
// Filter = fieldIsCollection
// ? $"{fieldName}/any(s: s eq '{fieldValue}')"
// : $"{fieldName} eq '{fieldValue}')",
// Size = limit
// };

if (limit > 0)
{
options.Size = limit;
this._log.LogDebug("Max results: {0}", limit);
}

// Decide whether to use a sticky session for the current request
if (this._useStickySessions)
{
options.SessionId = Guid.NewGuid().ToString("N");
}

return options;
}

private static double ScoreToCosineSimilarity(double score)
Expand Down
13 changes: 12 additions & 1 deletion service/Service/appsettings.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -280,7 +280,18 @@
"APIKey": "",
// Hybrid search is not enabled by default. Note that when using hybrid search
// relevance scores are different, usually lower, than when using just vector search
"UseHybridSearch": false
"UseHybridSearch": false,
// Helps improve relevance score consistency for search services with multiple replicas by
// attempting to route a given request to the same replica for that session. Use this when
// favoring consistent scoring over lower latency. Can adversely affect performance.
//
// Whether to use sticky sessions, which can help getting more consistent results.
// When using sticky sessions, a best-effort attempt will be made to target the same replica set.
// Be wary that reusing the same replica repeatedly can interfere with the load balancing of
// the requests across replicas and adversely affect the performance of the search service.
//
// See https://learn.microsoft.com/rest/api/searchservice/documents/search-post?view=rest-searchservice-2024-07-01&tabs=HTTP#request-body
"UseStickySessions": false
},
"AzureAIDocIntel": {
// "APIKey" or "AzureIdentity".
Expand Down
1 change: 1 addition & 0 deletions service/tests/Abstractions.UnitTests/Diagnostics/SensitiveDataLoggerTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ public sealed class SensitiveDataLoggerTests : IDisposable
private const string DotNetEnvVar = "DOTNET_ENVIRONMENT";

[Fact]
[Trait("Category", "UnitTest")]
public void ItIsDisabledByDefault()
{
// Assert
Expand Down
3 changes: 2 additions & 1 deletion service/tests/Core.FunctionalTests/appsettings.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,8 @@
"Auth": "AzureIdentity",
"Endpoint": "https://<...>",
"APIKey": "",
"UseHybridSearch": false
"UseHybridSearch": false,
"UseStickySessions": true
},
"LlamaSharp": {
"TextModel": {
Expand Down
2 changes: 2 additions & 0 deletions tools/InteractiveSetup/Services/AzureAISearch.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ public static void Setup(Context ctx, bool force = false)
{ "Auth", "ApiKey" },
{ "APIKey", "" },
{ "UseHybridSearch", false },
{ "UseStickySessions", false }
};
}

Expand All @@ -45,5 +46,6 @@ public static void Setup(Context ctx, bool force = false)

AppSettings.Change(x => x.Services[ServiceName]["Endpoint"] = SetupUI.AskOpenQuestion("Azure AI Search <endpoint>", config["Endpoint"].ToString()));
AppSettings.Change(x => x.Services[ServiceName]["UseHybridSearch"] = SetupUI.AskBoolean("Use hybrid search (yes/no)?", (bool)config["UseHybridSearch"]));
AppSettings.Change(x => x.Services[ServiceName]["UseStickySessions"] = SetupUI.AskBoolean("Use sticky sessions (yes/no)?", (bool)config["UseStickySessions"]));
}
}

0 comments on commit 490c013

Please sign in to comment.