Skip to content

Commit 490c013

Browse files
alexmgdluc
alexmg
and
dluc
authored
Support Azure AI Search Sticky Sessions (#877)
## 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 <[email protected]>
1 parent fb52c86 commit 490c013

File tree

12 files changed

+131
-49
lines changed

12 files changed

+131
-49
lines changed

‎applications/tests/Evaluation.Tests/appsettings.json‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,18 @@
7070
"APIKey": "",
7171
// Hybrid search is not enabled by default. Note that when using hybrid search
7272
// relevance scores are different, usually lower, than when using just vector search
73-
"UseHybridSearch": false
73+
"UseHybridSearch": false,
74+
// Helps improve relevance score consistency for search services with multiple replicas by
75+
// attempting to route a given request to the same replica for that session. Use this when
76+
// favoring consistent scoring over lower latency. Can adversely affect performance.
77+
//
78+
// Whether to use sticky sessions, which can help getting more consistent results.
79+
// When using sticky sessions, a best-effort attempt will be made to target the same replica set.
80+
// Be wary that reusing the same replica repeatedly can interfere with the load balancing of
81+
// the requests across replicas and adversely affect the performance of the search service.
82+
//
83+
// See https://learn.microsoft.com/rest/api/searchservice/documents/search-post?view=rest-searchservice-2024-07-01&tabs=HTTP#request-body
84+
"UseStickySessions": false
7485
},
7586
"OpenAI": {
7687
// Name of the model used to generate text (text completion or chat completion)

‎examples/002-dotnet-Serverless/appsettings.json‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,18 @@
7070
"APIKey": "",
7171
// Hybrid search is not enabled by default. Note that when using hybrid search
7272
// relevance scores are different, usually lower, than when using just vector search
73-
"UseHybridSearch": false
73+
"UseHybridSearch": false,
74+
// Helps improve relevance score consistency for search services with multiple replicas by
75+
// attempting to route a given request to the same replica for that session. Use this when
76+
// favoring consistent scoring over lower latency. Can adversely affect performance.
77+
//
78+
// Whether to use sticky sessions, which can help getting more consistent results.
79+
// When using sticky sessions, a best-effort attempt will be made to target the same replica set.
80+
// Be wary that reusing the same replica repeatedly can interfere with the load balancing of
81+
// the requests across replicas and adversely affect the performance of the search service.
82+
//
83+
// See https://learn.microsoft.com/rest/api/searchservice/documents/search-post?view=rest-searchservice-2024-07-01&tabs=HTTP#request-body
84+
"UseStickySessions": false
7485
},
7586
"OpenAI": {
7687
// Name of the model used to generate text (text completion or chat completion)

‎examples/210-KM-without-builder/appsettings.json‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -248,7 +248,18 @@
248248
"APIKey": "",
249249
// Hybrid search is not enabled by default. Note that when using hybrid search
250250
// relevance scores are different, usually lower, than when using just vector search
251-
"UseHybridSearch": false
251+
"UseHybridSearch": false,
252+
// Helps improve relevance score consistency for search services with multiple replicas by
253+
// attempting to route a given request to the same replica for that session. Use this when
254+
// favoring consistent scoring over lower latency. Can adversely affect performance.
255+
//
256+
// Whether to use sticky sessions, which can help getting more consistent results.
257+
// When using sticky sessions, a best-effort attempt will be made to target the same replica set.
258+
// Be wary that reusing the same replica repeatedly can interfere with the load balancing of
259+
// the requests across replicas and adversely affect the performance of the search service.
260+
//
261+
// See https://learn.microsoft.com/rest/api/searchservice/documents/search-post?view=rest-searchservice-2024-07-01&tabs=HTTP#request-body
262+
"UseStickySessions": false
252263
},
253264
"AzureAIDocIntel": {
254265
// "APIKey" or "AzureIdentity".

‎examples/401-evaluation/appsettings.json‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,18 @@
7070
"APIKey": "",
7171
// Hybrid search is not enabled by default. Note that when using hybrid search
7272
// relevance scores are different, usually lower, than when using just vector search
73-
"UseHybridSearch": false
73+
"UseHybridSearch": false,
74+
// Helps improve relevance score consistency for search services with multiple replicas by
75+
// attempting to route a given request to the same replica for that session. Use this when
76+
// favoring consistent scoring over lower latency. Can adversely affect performance.
77+
//
78+
// Whether to use sticky sessions, which can help getting more consistent results.
79+
// When using sticky sessions, a best-effort attempt will be made to target the same replica set.
80+
// Be wary that reusing the same replica repeatedly can interfere with the load balancing of
81+
// the requests across replicas and adversely affect the performance of the search service.
82+
//
83+
// See https://learn.microsoft.com/rest/api/searchservice/documents/search-post?view=rest-searchservice-2024-07-01&tabs=HTTP#request-body
84+
"UseStickySessions": false
7485
},
7586
"OpenAI": {
7687
// Name of the model used to generate text (text completion or chat completion)

‎extensions/AzureAISearch/AzureAISearch.FunctionalTests/DefaultTests.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,14 +14,14 @@ public class DefaultTests : BaseFunctionalTestCase
1414
public DefaultTests(IConfiguration cfg, ITestOutputHelper output) : base(cfg, output)
1515
{
1616
Assert.False(string.IsNullOrEmpty(this.AzureAiSearchConfig.Endpoint));
17-
Assert.False(string.IsNullOrEmpty(this.AzureAiSearchConfig.APIKey));
17+
Assert.False(this.AzureAiSearchConfig.Auth == AzureAISearchConfig.AuthTypes.APIKey && string.IsNullOrEmpty(this.AzureAiSearchConfig.APIKey));
1818
Assert.False(string.IsNullOrEmpty(this.OpenAiConfig.APIKey));
1919

2020
this._memory = new KernelMemoryBuilder()
2121
.With(new KernelMemoryConfig { DefaultIndexName = "default4tests" })
2222
.WithSearchClientConfig(new SearchClientConfig { EmptyAnswer = NotFound })
2323
.WithOpenAI(this.OpenAiConfig)
24-
.WithAzureAISearchMemoryDb(this.AzureAiSearchConfig.Endpoint, this.AzureAiSearchConfig.APIKey)
24+
.WithAzureAISearchMemoryDb(this.AzureAiSearchConfig)
2525
.Build<MemoryServerless>();
2626
}
2727

‎extensions/AzureAISearch/AzureAISearch.FunctionalTests/appsettings.json‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,8 @@
1111
// using the env vars AZURE_TENANT_ID, AZURE_CLIENT_ID, AZURE_CLIENT_SECRET.
1212
"Auth": "AzureIdentity",
1313
"Endpoint": "https://<...>",
14-
"APIKey": ""
14+
"APIKey": "",
15+
"UseStickySessions": true
1516
},
1617
"OpenAI": {
1718
// Name of the model used to generate text (text completion or chat completion)

‎extensions/AzureAISearch/AzureAISearch/AzureAISearchConfig.cs‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,20 @@ public enum AuthTypes
3131
/// </summary>
3232
public bool UseHybridSearch { get; set; } = false;
3333

34+
/// <summary>
35+
/// Helps improve relevance score consistency for search services with multiple replicas by
36+
/// attempting to route a given request to the same replica for that session. Use this when
37+
/// favoring consistent scoring over lower latency. Can adversely affect performance.
38+
///
39+
/// Whether to use sticky sessions, which can help getting more consistent results.
40+
/// When using sticky sessions, a best-effort attempt will be made to target the same replica set.
41+
/// Be wary that reusing the same replica repeatedly can interfere with the load balancing of
42+
/// the requests across replicas and adversely affect the performance of the search service.
43+
///
44+
/// See https://learn.microsoft.com/rest/api/searchservice/documents/search-post?view=rest-searchservice-2024-07-01&amp;tabs=HTTP#request-body
45+
/// </summary>
46+
public bool UseStickySessions { get; set; } = false;
47+
3448
public void SetCredential(TokenCredential credential)
3549
{
3650
this.Auth = AuthTypes.ManualTokenCredential;

‎extensions/AzureAISearch/AzureAISearch/AzureAISearchMemory.cs‎

Lines changed: 48 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ public class AzureAISearchMemory : IMemoryDb, IMemoryDbUpsertBatch
3434
private readonly ITextEmbeddingGenerator _embeddingGenerator;
3535
private readonly ILogger<AzureAISearchMemory> _log;
3636
private readonly bool _useHybridSearch;
37+
private readonly bool _useStickySessions;
3738

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

5355
if (string.IsNullOrEmpty(config.Endpoint))
5456
{
@@ -190,22 +192,12 @@ await client.IndexDocumentsAsync(
190192
FilterMode = VectorFilterMode.PreFilter
191193
}
192194
};
193-
DefineFieldsToSelect(options, withEmbeddings);
195+
options = this.PrepareSearchOptions(options, withEmbeddings, filters, limit);
194196

195197
if (limit > 0)
196198
{
197199
vectorQuery.KNearestNeighborsCount = limit;
198-
options.Size = limit;
199-
this._log.LogDebug("KNearestNeighborsCount and max results: {0}", limit);
200-
}
201-
202-
// Remove empty filters
203-
filters = filters?.Where(f => !f.IsEmpty()).ToList();
204-
205-
if (filters is { Count: > 0 })
206-
{
207-
options.Filter = AzureAISearchFiltering.BuildSearchFilter(filters);
208-
this._log.LogDebug("Filtering vectors, condition: {0}", options.Filter);
200+
this._log.LogDebug("KNearestNeighborsCount: {0}", limit);
209201
}
210202

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

256-
SearchOptions options = new();
257-
DefineFieldsToSelect(options, withEmbeddings);
258-
259-
if (limit > 0)
260-
{
261-
options.Size = limit;
262-
this._log.LogDebug("Max results: {0}", limit);
263-
}
264-
265-
// Remove empty filters
266-
filters = filters?.Where(f => !f.IsEmpty()).ToList();
267-
268-
if (filters is { Count: > 0 })
269-
{
270-
options.Filter = AzureAISearchFiltering.BuildSearchFilter(filters);
271-
this._log.LogDebug("Filtering vectors, condition: {0}", options.Filter);
272-
}
273-
274-
// See: https://learn.microsoft.com/azure/search/search-query-understand-collection-filters
275-
// fieldValue = fieldValue.Replace("'", "''", StringComparison.Ordinal);
276-
// var options = new SearchOptions
277-
// {
278-
// Filter = fieldIsCollection
279-
// ? $"{fieldName}/any(s: s eq '{fieldValue}')"
280-
// : $"{fieldName} eq '{fieldValue}')",
281-
// Size = limit
282-
// };
248+
SearchOptions options = this.PrepareSearchOptions(null, withEmbeddings, filters, limit);
283249

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

630-
private static void DefineFieldsToSelect(SearchOptions options, bool withEmbeddings)
596+
private SearchOptions PrepareSearchOptions(
597+
SearchOptions? options,
598+
bool withEmbeddings,
599+
ICollection<MemoryFilter>? filters = null,
600+
int limit = 1)
631601
{
602+
options ??= new SearchOptions();
603+
604+
// Define which fields to fetch
632605
options.Select.Add(AzureAISearchMemoryRecord.IdField);
633606
options.Select.Add(AzureAISearchMemoryRecord.TagsField);
634607
options.Select.Add(AzureAISearchMemoryRecord.PayloadField);
608+
609+
// Embeddings are fetched only when needed, to reduce latency and cost
635610
if (withEmbeddings)
636611
{
637612
options.Select.Add(AzureAISearchMemoryRecord.VectorField);
638613
}
614+
615+
// Remove empty filters
616+
filters = filters?.Where(f => !f.IsEmpty()).ToList();
617+
618+
if (filters is { Count: > 0 })
619+
{
620+
options.Filter = AzureAISearchFiltering.BuildSearchFilter(filters);
621+
this._log.LogDebug("Filtering vectors, condition: {0}", options.Filter);
622+
}
623+
624+
// See: https://learn.microsoft.com/azure/search/search-query-understand-collection-filters
625+
// fieldValue = fieldValue.Replace("'", "''", StringComparison.Ordinal);
626+
// var options = new SearchOptions
627+
// {
628+
// Filter = fieldIsCollection
629+
// ? $"{fieldName}/any(s: s eq '{fieldValue}')"
630+
// : $"{fieldName} eq '{fieldValue}')",
631+
// Size = limit
632+
// };
633+
634+
if (limit > 0)
635+
{
636+
options.Size = limit;
637+
this._log.LogDebug("Max results: {0}", limit);
638+
}
639+
640+
// Decide whether to use a sticky session for the current request
641+
if (this._useStickySessions)
642+
{
643+
options.SessionId = Guid.NewGuid().ToString("N");
644+
}
645+
646+
return options;
639647
}
640648

641649
private static double ScoreToCosineSimilarity(double score)

‎service/Service/appsettings.json‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -280,7 +280,18 @@
280280
"APIKey": "",
281281
// Hybrid search is not enabled by default. Note that when using hybrid search
282282
// relevance scores are different, usually lower, than when using just vector search
283-
"UseHybridSearch": false
283+
"UseHybridSearch": false,
284+
// Helps improve relevance score consistency for search services with multiple replicas by
285+
// attempting to route a given request to the same replica for that session. Use this when
286+
// favoring consistent scoring over lower latency. Can adversely affect performance.
287+
//
288+
// Whether to use sticky sessions, which can help getting more consistent results.
289+
// When using sticky sessions, a best-effort attempt will be made to target the same replica set.
290+
// Be wary that reusing the same replica repeatedly can interfere with the load balancing of
291+
// the requests across replicas and adversely affect the performance of the search service.
292+
//
293+
// See https://learn.microsoft.com/rest/api/searchservice/documents/search-post?view=rest-searchservice-2024-07-01&tabs=HTTP#request-body
294+
"UseStickySessions": false
284295
},
285296
"AzureAIDocIntel": {
286297
// "APIKey" or "AzureIdentity".

‎service/tests/Abstractions.UnitTests/Diagnostics/SensitiveDataLoggerTests.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ public sealed class SensitiveDataLoggerTests : IDisposable
1010
private const string DotNetEnvVar = "DOTNET_ENVIRONMENT";
1111

1212
[Fact]
13+
[Trait("Category", "UnitTest")]
1314
public void ItIsDisabledByDefault()
1415
{
1516
// Assert

0 commit comments

Comments
 (0)