Optimize pipeline for fatal/non-retriable errors (#898)

* In case of fatal non retriable errors, stop retrying and immediately
move queue msgs to poison queue.
* New `KernelMemoryException.IsTransient` optional property to explicit
if an error is transient or not.
* Use `KernelMemoryException.IsTransient` in Azure AI Doc Intelligence
when an image cannot be processed, for instance when the image is too
big (note: max size depends on SKU, see service public docs).
* Use `KernelMemoryException.IsTransient` in AI clients, don't retry if
the HTTP error cannot be recovered without code/config changes
* Use `KernelMemoryException.IsTransient` in AI Mime type detection,
don't retry processing unsupported files.
* Change `IQueue` and `IPipelineStepHandler` signatures to support
failure types, as opposed to just a boolean `false`.

Author: dluc

examples/201-dotnet-serverless-custom-handler/Program.cs

@@ -47,7 +47,7 @@ public MyHandler(
     public string StepName { get; }
 
     /// <inheritdoc />
-    public async Task<(bool success, DataPipeline updatedPipeline)> InvokeAsync(
+    public async Task<(ReturnType returnType, DataPipeline updatedPipeline)> InvokeAsync(
         DataPipeline pipeline, CancellationToken cancellationToken = default)
     {
         /* ... your custom ...
@@ -64,6 +64,6 @@ public MyHandler(
         // Remove this - here only to avoid build errors
         await Task.Delay(0, cancellationToken).ConfigureAwait(false);
 
-        return (true, pipeline);
+        return (ReturnType.Success, pipeline);
     }
 }

examples/202-dotnet-custom-handler-as-a-service/MyHandler.cs

@@ -38,7 +38,7 @@ public Task StopAsync(CancellationToken cancellationToken = default)
     }
 
     /// <inheritdoc />
-    public async Task<(bool success, DataPipeline updatedPipeline)> InvokeAsync(DataPipeline pipeline, CancellationToken cancellationToken = default)
+    public async Task<(ReturnType returnType, DataPipeline updatedPipeline)> InvokeAsync(DataPipeline pipeline, CancellationToken cancellationToken = default)
     {
         /* ... your custom ...
          * ... handler ...
@@ -49,6 +49,6 @@ public Task StopAsync(CancellationToken cancellationToken = default)
         // Remove this - here only to avoid build errors
         await Task.Delay(0, cancellationToken).ConfigureAwait(false);
 
-        return (true, pipeline);
+        return (ReturnType.Success, pipeline);
     }
 }

extensions/Anthropic/Client/RawAnthropicClient.cs

@@ -64,7 +64,8 @@ internal async IAsyncEnumerable<StreamingResponseMessage> CallClaudeStreamingAsy
         if (!response.IsSuccessStatusCode)
         {
             var responseError = await response.Content.ReadAsStringAsync(cancellationToken).ConfigureAwait(false);
-            throw new KernelMemoryException($"Failed to send request: {response.StatusCode} - {responseError}");
+            throw new KernelMemoryException($"Failed to send request: {response.StatusCode} - {responseError}",
+                isTransient: response.StatusCode.IsTransientError());
         }
 
         var responseStream = await response.Content.ReadAsStreamAsync(cancellationToken).ConfigureAwait(false);

extensions/AzureAIDocIntel/AzureAIDocIntelEngine.cs

@@ -58,12 +58,19 @@ public AzureAIDocIntelEngine(
     ///<inheritdoc/>
     public async Task<string> ExtractTextFromImageAsync(Stream imageContent, CancellationToken cancellationToken = default)
     {
-        // Start the OCR operation
-        var operation = await this._recognizerClient.AnalyzeDocumentAsync(WaitUntil.Completed, "prebuilt-read", imageContent, cancellationToken: cancellationToken).ConfigureAwait(false);
+        try
+        {
+            // Start the OCR operation
+            var operation = await this._recognizerClient.AnalyzeDocumentAsync(WaitUntil.Completed, "prebuilt-read", imageContent, cancellationToken: cancellationToken).ConfigureAwait(false);
 
-        // Wait for the result
-        Response<AnalyzeResult> operationResponse = await operation.WaitForCompletionAsync(cancellationToken).ConfigureAwait(false);
+            // Wait for the result
+            Response<AnalyzeResult> operationResponse = await operation.WaitForCompletionAsync(cancellationToken).ConfigureAwait(false);
 
-        return operationResponse.Value.Content;
+            return operationResponse.Value.Content;
+        }
+        catch (RequestFailedException e)
+        {
+            throw new AzureAIDocIntelException(e.Message, e, isTransient: HttpErrors.IsTransientError(e.Status));
+        }
     }
 }

extensions/AzureAIDocIntel/AzureAIDocIntelException.cs

@@ -0,0 +1,26 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+
+namespace Microsoft.KernelMemory.DataFormats.AzureAIDocIntel;
+
+public class AzureAIDocIntelException : KernelMemoryException
+{
+    /// <inheritdoc />
+    public AzureAIDocIntelException(bool? isTransient = null)
+    {
+        this.IsTransient = isTransient;
+    }
+
+    /// <inheritdoc />
+    public AzureAIDocIntelException(string message, bool? isTransient = null) : base(message)
+    {
+        this.IsTransient = isTransient;
+    }
+
+    /// <inheritdoc />
+    public AzureAIDocIntelException(string message, Exception? innerException, bool? isTransient = null) : base(message, innerException)
+    {
+        this.IsTransient = isTransient;
+    }
+}

extensions/AzureAISearch/AzureAISearch/AzureAISearchMemoryException.cs

@@ -7,17 +7,21 @@ namespace Microsoft.KernelMemory.MemoryDb.AzureAISearch;
 public class AzureAISearchMemoryException : KernelMemoryException
 {
     /// <inheritdoc />
-    public AzureAISearchMemoryException()
+    public AzureAISearchMemoryException(bool? isTransient = null)
     {
+        this.IsTransient = isTransient;
     }
 
     /// <inheritdoc />
-    public AzureAISearchMemoryException(string? message) : base(message)
+    public AzureAISearchMemoryException(string message, bool? isTransient = null) : base(message)
     {
+        this.IsTransient = isTransient;
     }
 
     /// <inheritdoc />
-    public AzureAISearchMemoryException(string? message, Exception? innerException) : base(message, innerException)
+    public AzureAISearchMemoryException(string message, Exception? innerException, bool? isTransient = null) : base(message, innerException)
     {
+        this.IsTransient = isTransient;
     }
 }
+

extensions/AzureAISearch/AzureAISearch/Internals/MemoryDbSchema.cs

@@ -13,41 +13,41 @@ public void Validate(bool vectorSizeRequired = false)
     {
         if (this.Fields.Count == 0)
         {
-            throw new KernelMemoryException("The schema is empty");
+            throw new AzureAISearchMemoryException("The schema is empty", isTransient: false);
         }
 
         if (this.Fields.All(x => x.Type != MemoryDbField.FieldType.Vector))
         {
-            throw new KernelMemoryException("The schema doesn't contain a vector field");
+            throw new AzureAISearchMemoryException("The schema doesn't contain a vector field", isTransient: false);
         }
 
         int keys = this.Fields.Count(x => x.IsKey);
         switch (keys)
         {
             case 0:
-                throw new KernelMemoryException("The schema doesn't contain a key field");
+                throw new AzureAISearchMemoryException("The schema doesn't contain a key field", isTransient: false);
             case > 1:
-                throw new KernelMemoryException("The schema cannot contain more than one key");
+                throw new AzureAISearchMemoryException("The schema cannot contain more than one key", isTransient: false);
         }
 
         if (vectorSizeRequired && this.Fields.Any(x => x is { Type: MemoryDbField.FieldType.Vector, VectorSize: 0 }))
         {
-            throw new KernelMemoryException("Vector fields must have a size greater than zero defined");
+            throw new AzureAISearchMemoryException("Vector fields must have a size greater than zero defined", isTransient: false);
         }
 
         if (this.Fields.Any(x => x is { Type: MemoryDbField.FieldType.Bool, IsKey: true }))
         {
-            throw new KernelMemoryException("Boolean fields cannot be used as unique keys");
+            throw new AzureAISearchMemoryException("Boolean fields cannot be used as unique keys", isTransient: false);
         }
 
         if (this.Fields.Any(x => x is { Type: MemoryDbField.FieldType.ListOfStrings, IsKey: true }))
         {
-            throw new KernelMemoryException("Collection fields cannot be used as unique keys");
+            throw new AzureAISearchMemoryException("Collection fields cannot be used as unique keys", isTransient: false);
         }
 
         if (this.Fields.Any(x => x is { Type: MemoryDbField.FieldType.Vector, IsKey: true }))
         {
-            throw new KernelMemoryException("Vector fields cannot be used as unique keys");
+            throw new AzureAISearchMemoryException("Vector fields cannot be used as unique keys", isTransient: false);
         }
     }
 }

extensions/AzureOpenAI/AzureOpenAI/AzureOpenAIException.cs

@@ -0,0 +1,26 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+
+namespace Microsoft.KernelMemory.AI.AzureOpenAI;
+
+public class AzureOpenAIException : KernelMemoryException
+{
+    /// <inheritdoc />
+    public AzureOpenAIException(bool? isTransient = null)
+    {
+        this.IsTransient = isTransient;
+    }
+
+    /// <inheritdoc />
+    public AzureOpenAIException(string message, bool? isTransient = null) : base(message)
+    {
+        this.IsTransient = isTransient;
+    }
+
+    /// <inheritdoc />
+    public AzureOpenAIException(string message, Exception? innerException, bool? isTransient = null) : base(message, innerException)
+    {
+        this.IsTransient = isTransient;
+    }
+}

extensions/AzureOpenAI/AzureOpenAI/AzureOpenAITextEmbeddingGenerator.cs

@@ -12,6 +12,7 @@
 using Microsoft.KernelMemory.AI.AzureOpenAI.Internals;
 using Microsoft.KernelMemory.AI.OpenAI;
 using Microsoft.KernelMemory.Diagnostics;
+using Microsoft.SemanticKernel;
 using Microsoft.SemanticKernel.AI.Embeddings;
 using Microsoft.SemanticKernel.Connectors.AzureOpenAI;
 
@@ -121,15 +122,29 @@ public IReadOnlyList<string> GetTokens(string text)
     public Task<Embedding> GenerateEmbeddingAsync(string text, CancellationToken cancellationToken = default)
     {
         this._log.LogTrace("Generating embedding");
-        return this._client.GenerateEmbeddingAsync(text, cancellationToken);
+        try
+        {
+            return this._client.GenerateEmbeddingAsync(text, cancellationToken);
+        }
+        catch (HttpOperationException e)
+        {
+            throw new AzureOpenAIException(e.Message, e, isTransient: e.StatusCode.IsTransientError());
+        }
     }
 
     /// <inheritdoc/>
     public async Task<Embedding[]> GenerateEmbeddingBatchAsync(IEnumerable<string> textList, CancellationToken cancellationToken = default)
     {
         var list = textList.ToList();
         this._log.LogTrace("Generating embeddings, batch size '{0}'", list.Count);
-        IList<ReadOnlyMemory<float>> embeddings = await this._client.GenerateEmbeddingsAsync(list, cancellationToken: cancellationToken).ConfigureAwait(false);
-        return embeddings.Select(e => new Embedding(e)).ToArray();
+        try
+        {
+            IList<ReadOnlyMemory<float>> embeddings = await this._client.GenerateEmbeddingsAsync(list, cancellationToken: cancellationToken).ConfigureAwait(false);
+            return embeddings.Select(e => new Embedding(e)).ToArray();
+        }
+        catch (HttpOperationException e)
+        {
+            throw new AzureOpenAIException(e.Message, e, isTransient: e.StatusCode.IsTransientError());
+        }
     }
 }

extensions/AzureOpenAI/AzureOpenAI/AzureOpenAITextGenerator.cs

@@ -5,6 +5,7 @@
 using System.Net.Http;
 using System.Runtime.CompilerServices;
 using System.Threading;
+using System.Threading.Tasks;
 using Azure.AI.OpenAI;
 using Microsoft.Extensions.Logging;
 using Microsoft.KernelMemory.AI.AzureOpenAI.Internals;
@@ -140,8 +141,17 @@ public async IAsyncEnumerable<string> GenerateTextAsync(
         }
 
         this._log.LogTrace("Sending chat message generation request");
-        IAsyncEnumerable<StreamingTextContent> result = this._client.GetStreamingTextContentsAsync(prompt, skOptions, cancellationToken: cancellationToken);
-        await foreach (StreamingTextContent x in result)
+        IAsyncEnumerable<StreamingTextContent> result;
+        try
+        {
+            result = this._client.GetStreamingTextContentsAsync(prompt, skOptions, cancellationToken: cancellationToken);
+        }
+        catch (HttpOperationException e)
+        {
+            throw new AzureOpenAIException(e.Message, e, isTransient: e.StatusCode.IsTransientError());
+        }
+
+        await foreach (StreamingTextContent x in result.WithCancellation(cancellationToken))
         {
             if (x.Text == null) { continue; }