Merge branch 'release/4.4.1'

Author: nicholasbarlow

azure-pipelines.yml

@@ -62,4 +62,5 @@ steps:
       deleteIntegrationEnvironmentScriptFilePath: "./build/integration-test.ps1"
       deleteIntegrationEnvironmentScriptArguments: "-Action TearDown"
       publishCodeCoverage: true
-      useGitVersionDotNetTool: true
+      useGitVersionDotNetTool: true
+      publishToDevFeed: true

docs/4.4.1-release-notes.md

@@ -0,0 +1,3 @@
+# Fixes
+- Reduce TooManyRequests exceptions thrown by Azure OpenAI 
+- Use distributed locks to control rate limits across pods

src/ExternalSearch.Providers.AzureOpenAI/AzureOpenAIExternalSearchProvider.cs

@@ -78,7 +78,12 @@ public IEnumerable<IExternalSearchQueryResult> ExecuteSearch(ExecutionContext co
         IDictionary<string, object> config, IProvider provider)
     {
         var jobData = new AzureOpenAIExternalSearchJobData(config);
-        return InternalExecuteSearch(context, query, jobData);
+
+        return ActionExtensions.ExecuteWithRetry(
+            () => InternalExecuteSearch(context, query, jobData).ToArray(), // important we need to materialize the enumerable for ExecuteWithRetry to work
+            retryCount: 1000,
+            isTransient: ex => ex.ToString().Contains("TooManyRequests")    // the core will retry but only 3 times
+        );
     }
 
     public IEnumerable<Clue> BuildClues(ExecutionContext context, IExternalSearchQuery query,
@@ -228,7 +233,7 @@ public ConnectionVerificationResult VerifyConnection(ExecutionContext context,
             return new ConnectionVerificationResult(false, "Prompt must contain at least one output. eg, {output:vocabulary:product.description}");
         }
 
-        var deploymentSupportsCompletion = DeploymentSupportsCompletion(context, deploymentName);
+        var deploymentSupportsCompletion = DeploymentSupportsCompletion(context, deploymentName, baseUrl);
 
         prompt = "Hello";
 
@@ -253,10 +258,8 @@ public ConnectionVerificationResult VerifyConnection(ExecutionContext context,
         }
     }
 
-    private bool DeploymentSupportsCompletion(ExecutionContext executionContext, string deploymentName)
+    private bool DeploymentSupportsCompletion(ExecutionContext executionContext, string deploymentName, string baseUrl)
     {
-        var baseUrl = executionContext.Organization.Settings.GetValue("OpenAiBaseUrl", "OpenAiBaseUrl", "");
-
         var deploymentSupportsCompletionCacheKey = $"{nameof(DeploymentSupportsCompletion)}_{executionContext.Organization.Id}_{baseUrl}_{deploymentName}";
         if (_cache.TryGetValue(deploymentSupportsCompletionCacheKey, out var cached) && cached != null)
         {
@@ -412,37 +415,61 @@ private IEnumerable<IExternalSearchQueryResult> InternalExecuteSearch(ExecutionC
 
         using (context.Log.BeginScope("{0} {1}: query {2}", GetType().Name, "ExecuteSearch", query))
         {
-            context.Log.LogTrace("Starting external search for Id: '{Id}' QueryKey: '{QueryKey}'", query.Id, query.QueryKey);
-
             var prompt = query.QueryParameters["prompt"].Single();
             var deploymentName = query.QueryParameters["deploymentName"].Single();
+            var baseUrl = context.Organization.Settings.GetValue("OpenAiBaseUrl", "OpenAiBaseUrl", "");
 
-            var deploymentSupportsCompletion = DeploymentSupportsCompletion(context, deploymentName);
-
-            var response = deploymentSupportsCompletion
-                ? QueryInternalUsingCompletionApi(context, deploymentName, prompt)
-                : QueryInternalUsingChatApi(context, deploymentName, prompt);
-
-            JObject jsonResponse;
-
-            try
+            using (new DistributedLock(context, $"{nameof(AzureOpenAIExternalSearchProvider)}_{baseUrl}_{deploymentName}", exclusive: false))
             {
-                jsonResponse = JObject.Parse(response);
-            }
-            catch
-            {
-                prompt +=
-                    "\n\nImportant: A prior attempt to answer this question resulted in malformed JSON. Please retry and verify that the output adheres strictly to the specified JSON format. The response must consist solely of valid JSON, as it will be programmatically processed.";
+                context.Log.LogTrace("Starting external search for Id: '{Id}' QueryKey: '{QueryKey}'", query.Id, query.QueryKey);
+
+                var deploymentSupportsCompletion = DeploymentSupportsCompletion(context, deploymentName, baseUrl);
 
-                response = deploymentSupportsCompletion
+                var response = deploymentSupportsCompletion
                     ? QueryInternalUsingCompletionApi(context, deploymentName, prompt)
                     : QueryInternalUsingChatApi(context, deploymentName, prompt);
 
-                jsonResponse = JObject.Parse(response);
+                JObject jsonResponse;
+
+                try
+                {
+                    jsonResponse = ParseResponse(response);
+                }
+                catch(Exception ex)
+                {
+                    context.Log.LogDebug(ex, $"Failed to parse json response from AI. Will try again. response was:\n{response}");
+
+                    prompt +=
+                        "\n\nImportant: A prior attempt to answer this question resulted in malformed JSON. Please retry and verify that the output adheres strictly to the specified JSON format. The response must consist solely of valid JSON, as it will be programmatically processed.";
+
+                    response = deploymentSupportsCompletion
+                        ? QueryInternalUsingCompletionApi(context, deploymentName, prompt)
+                        : QueryInternalUsingChatApi(context, deploymentName, prompt);
+
+                    try
+                    {
+                        jsonResponse = ParseResponse(response);
+                    }
+                    catch (Exception ex2)
+                    {
+                        throw new Exception($"Unable to deserialize the response generated by AI. The generated response was:\n{response}", ex2);
+                    }
+                }
+
+                yield return new ExternalSearchQueryResult<JObject>(query, jsonResponse);
             }
+        }
+    }
 
-            yield return new ExternalSearchQueryResult<JObject>(query, jsonResponse);
+    private static JObject ParseResponse(string response)
+    {
+        var m = Regex.Match(response, "(?<json>{.+})", RegexOptions.Singleline);
+        if (m.Success)
+        {
+            response = m.Groups["json"].Value;
         }
+
+        return JObject.Parse(response);
     }
 
     private string QueryInternalUsingCompletionApi(ExecutionContext executionContext, string deploymentName, string prompt, bool logError = true)
@@ -470,7 +497,8 @@ private string QueryInternalUsingCompletionApi(ExecutionContext executionContext
 
         if (response.StatusCode == HttpStatusCode.TooManyRequests)
         {
-            Thread.Sleep(2000); // while developing we observed that the error message says to retry in 2s
+            WaitDueToTooManyRequests(executionContext, deploymentName, response, baseUrl);
+
             throw new Exception($"Too many requests - Call to openai returned HTTP {response.StatusCode}"); // hack the message must start with 'Too many requests' for the core to retry
         }
 
@@ -523,7 +551,8 @@ private string QueryInternalUsingChatApi(ExecutionContext executionContext, stri
 
         if (response.StatusCode == HttpStatusCode.TooManyRequests)
         {
-            Thread.Sleep(2000); // while developing we observed that the error message says to retry in 2s
+            WaitDueToTooManyRequests(executionContext, deploymentName, response, baseUrl);
+
             throw new Exception($"Too many requests - Call to openai returned HTTP {response.StatusCode}"); // hack the message must start with 'Too many requests' for the core to retry
         }
 
@@ -555,6 +584,36 @@ private string QueryInternalUsingChatApi(ExecutionContext executionContext, stri
         return content.TrimEnd();
     }
 
+    private static void WaitDueToTooManyRequests(ExecutionContext executionContext, string deploymentName, IRestResponse response, string baseUrl)
+    {
+        var responseAt = DateTime.Now;
+
+        // Try to get Retry-After from headers (Different model is having different request per minute)
+        var retryAfterHeader = response.Headers
+            .FirstOrDefault(h => h.Name.Equals("Retry-After", StringComparison.OrdinalIgnoreCase));
+
+        if (!int.TryParse(retryAfterHeader?.Value?.ToString(), out var waitSeconds))
+        {
+            waitSeconds = 2;
+        }
+
+        using (new DistributedLock(executionContext, $"{nameof(AzureOpenAIExternalSearchProvider)}_{baseUrl}_{deploymentName}", exclusive: true))
+        {
+            var timeAlreadyWaited = (int)DateTime.Now.Subtract(responseAt).TotalMilliseconds;
+            var waitTime = waitSeconds * 1000 - timeAlreadyWaited;
+
+            if (waitTime > 0)
+            {
+                executionContext.Log.Log(LogLevel.Debug, $"Sleeping thread for {waitTime}ms due to TooManyRequest response received");
+                Thread.Sleep(waitTime);
+            }
+            else
+            {
+                executionContext.Log.Log(LogLevel.Debug, $"TooManyRequest response received however no need to sleep as another thread already blocked");
+            }
+        }
+    }
+
     public override IPreviewImage GetPrimaryEntityPreviewImage(ExecutionContext context,
         IExternalSearchQueryResult result, IExternalSearchRequest request)
     {

src/ExternalSearch.Providers.AzureOpenAI/DistributedLock.cs

@@ -0,0 +1,52 @@
+using System;
+using System.Collections.Generic;
+using System.Data;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+using CluedIn.Core;
+using CluedIn.Core.DataStore;
+using Microsoft.Data.SqlClient;
+
+namespace CluedIn.ExternalSearch.Providers.AzureOpenAI
+{
+    internal class DistributedLock : IDisposable
+    {
+        private readonly SqlConnection _connection;
+        private readonly SqlTransaction _transaction;
+
+        public DistributedLock(ExecutionContext context, string lockName, bool exclusive = false)
+        {
+            var connectionStringKey = context.ApplicationContext.System.DataShards.GetDataShard(DataShardType.Locking).ReadConnectionString;
+            var connectionString = context.ApplicationContext.System.ConnectionStrings.GetConnectionString(connectionStringKey);
+
+            _connection = new SqlConnection(connectionString);
+            _connection.Open();
+            _transaction = _connection.BeginTransaction();
+            using var cmd = new SqlCommand("sp_getapplock", _connection);
+            cmd.CommandType = CommandType.StoredProcedure;
+            cmd.Transaction = _transaction;
+            cmd.CommandTimeout = 0;
+
+            cmd.Parameters.Add(new SqlParameter("@Resource", SqlDbType.NVarChar, 255) { Value = lockName });
+            cmd.Parameters.Add(new SqlParameter("@LockMode", SqlDbType.NVarChar, 32) { Value = exclusive ? "Exclusive" : "Shared" });
+            cmd.Parameters.Add(new SqlParameter("@LockTimeout", SqlDbType.Int) { Value = -1 });
+            cmd.Parameters.Add(new SqlParameter("@Result", SqlDbType.Int) { Direction = ParameterDirection.ReturnValue });
+
+            cmd.ExecuteNonQuery();
+
+            var queryResult = (int)cmd.Parameters["@Result"].Value;
+
+            if (queryResult < 0)
+            {
+                throw new ApplicationException($"sp_getapplock returned {queryResult}");    // should never happen as @LockTimeout is -1
+            }
+        }
+
+        public void Dispose()
+        {
+            _transaction.Rollback();
+            _connection.Close();
+        }
+    }
+}