Rename UseSessionId to UseStickySessions + refactoring

microsoft · Nov 12, 2024 · 32d3793 · 32d3793
1 parent 475e584
commit 32d3793
Show file tree

Hide file tree

Showing 10 changed files with 36 additions and 27 deletions.
diff --git a/applications/tests/Evaluation.Tests/appsettings.json b/applications/tests/Evaluation.Tests/appsettings.json
@@ -81,7 +81,7 @@
         // 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
-        "UseSessionId": false
+        "UseStickySessions": false
       },
       "OpenAI": {
         // Name of the model used to generate text (text completion or chat completion)

diff --git a/examples/002-dotnet-Serverless/appsettings.json b/examples/002-dotnet-Serverless/appsettings.json
@@ -81,7 +81,7 @@
         // 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
-        "UseSessionId": false
+        "UseStickySessions": false
       },
       "OpenAI": {
         // Name of the model used to generate text (text completion or chat completion)

diff --git a/examples/210-KM-without-builder/appsettings.json b/examples/210-KM-without-builder/appsettings.json
@@ -259,7 +259,7 @@
         // 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
-        "UseSessionId": false
+        "UseStickySessions": false
       },
       "AzureAIDocIntel": {
         // "APIKey" or "AzureIdentity".

diff --git a/examples/401-evaluation/appsettings.json b/examples/401-evaluation/appsettings.json
@@ -81,7 +81,7 @@
         // 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
-        "UseSessionId": false
+        "UseStickySessions": false
       },
       "OpenAI": {
         // Name of the model used to generate text (text completion or chat completion)

diff --git a/extensions/AzureAISearch/AzureAISearch.FunctionalTests/appsettings.json b/extensions/AzureAISearch/AzureAISearch.FunctionalTests/appsettings.json
@@ -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)

diff --git a/extensions/AzureAISearch/AzureAISearch/AzureAISearchConfig.cs b/extensions/AzureAISearch/AzureAISearch/AzureAISearchConfig.cs
@@ -43,7 +43,7 @@ public enum AuthTypes
     ///
     /// See https://learn.microsoft.com/rest/api/searchservice/documents/search-post?view=rest-searchservice-2024-07-01&tabs=HTTP#request-body
     /// </summary>
-    public bool UseSessionId { get; set; } = false;
+    public bool UseStickySessions { get; set; } = false;
 
     public void SetCredential(TokenCredential credential)
     {

diff --git a/extensions/AzureAISearch/AzureAISearch/AzureAISearchMemory.cs b/extensions/AzureAISearch/AzureAISearch/AzureAISearchMemory.cs
@@ -34,7 +34,7 @@ public class AzureAISearchMemory : IMemoryDb, IMemoryDbUpsertBatch
     private readonly ITextEmbeddingGenerator _embeddingGenerator;
     private readonly ILogger<AzureAISearchMemory> _log;
     private readonly bool _useHybridSearch;
-    private readonly bool _useSessionId;
+    private readonly bool _useStickySessions;
 
     /// <summary>
     /// Create a new instance
@@ -50,7 +50,7 @@ public AzureAISearchMemory(
         this._embeddingGenerator = embeddingGenerator;
         this._log = (loggerFactory ?? DefaultLogger.Factory).CreateLogger<AzureAISearchMemory>();
         this._useHybridSearch = config.UseHybridSearch;
-        this._useSessionId = config.UseSessionId;
+        this._useStickySessions = config.UseStickySessions;
 
         if (string.IsNullOrEmpty(config.Endpoint))
         {
@@ -192,13 +192,12 @@ await client.IndexDocumentsAsync(
                 FilterMode = VectorFilterMode.PreFilter
             }
         };
-        this.ApplyCommonSearchOptions(options, withEmbeddings, filters);
+        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);
+            this._log.LogDebug("KNearestNeighborsCount: {0}", limit);
         }
 
         Response<SearchResults<AzureAISearchMemoryRecord>>? searchResult = null;
@@ -246,14 +245,7 @@ public async IAsyncEnumerable<MemoryRecord> GetListAsync(
     {
         var client = this.GetSearchClient(index);
 
-        SearchOptions options = new();
-        this.ApplyCommonSearchOptions(options, withEmbeddings, filters);
-
-        if (limit > 0)
-        {
-            options.Size = limit;
-            this._log.LogDebug("Max results: {0}", limit);
-        }
+        SearchOptions options = this.PrepareSearchOptions(null, withEmbeddings, filters, limit);
 
         Response<SearchResults<AzureAISearchMemoryRecord>>? searchResult = null;
         try
@@ -601,14 +593,20 @@ at Azure.Search.Documents.SearchClient.SearchInternal[T](SearchOptions options,
         return indexSchema;
     }
 
-    private void ApplyCommonSearchOptions(
-        SearchOptions options,
+    private SearchOptions PrepareSearchOptions(
+        SearchOptions? options,
         bool withEmbeddings,
-        ICollection<MemoryFilter>? filters = null)
+        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);
@@ -633,10 +631,19 @@ private void ApplyCommonSearchOptions(
         //     Size = limit
         // };
 
-        if (this._useSessionId)
+        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)

diff --git a/service/Service/appsettings.json b/service/Service/appsettings.json
@@ -291,7 +291,7 @@
         // 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
-        "UseSessionId": false
+        "UseStickySessions": false
       },
       "AzureAIDocIntel": {
         // "APIKey" or "AzureIdentity".

diff --git a/service/tests/Core.FunctionalTests/appsettings.json b/service/tests/Core.FunctionalTests/appsettings.json
@@ -12,7 +12,8 @@
         "Auth": "AzureIdentity",
         "Endpoint": "https://<...>",
         "APIKey": "",
-        "UseHybridSearch": false
+        "UseHybridSearch": false,
+        "UseStickySessions": true
       },
       "LlamaSharp": {
         "TextModel": {

diff --git a/tools/InteractiveSetup/Services/AzureAISearch.cs b/tools/InteractiveSetup/Services/AzureAISearch.cs
@@ -22,7 +22,7 @@ public static void Setup(Context ctx, bool force = false)
                 { "Auth", "ApiKey" },
                 { "APIKey", "" },
                 { "UseHybridSearch", false },
-                { "UseSessionId", false }
+                { "UseStickySessions", false }
             };
         }
 
@@ -46,6 +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]["UseSessionId"] = SetupUI.AskBoolean("Use session ID (yes/no)?", (bool)config["UseSessionId"]));
+        AppSettings.Change(x => x.Services[ServiceName]["UseStickySessions"] = SetupUI.AskBoolean("Use sticky sessions (yes/no)?", (bool)config["UseStickySessions"]));
     }
 }