.NET: Add additional error handling to existing providers. (#1837)

* Add additional error handling to existing providers. * Add additional information when logging mem0 messages. * Fix formatting.
2026-06-16 21:04:09 +08:00 · 2025-11-04 17:09:45 +00:00
parent ada5b83c80
commit 8b4aa1ebb5
5 changed files with 217 additions and 78 deletions
@@ -131,31 +131,62 @@ public sealed class Mem0Provider : AIContextProvider
            Environment.NewLine,
            context.RequestMessages.Where(m => !string.IsNullOrWhiteSpace(m.Text)).Select(m => m.Text));

-        var memories = (await this._client.SearchAsync(
-            this.ApplicationId,
-            this.AgentId,
-            this.ThreadId,
-            this.UserId,
-            queryText,
-            cancellationToken).ConfigureAwait(false)).ToList();
-
-        var contextInstructions = memories.Count == 0
-            ? null
-            : $"{this._contextPrompt}\n{string.Join(Environment.NewLine, memories)}";
-
-        if (this._logger is not null)
+        try
        {
-            this._logger.LogInformation("Mem0AIContextProvider retrieved {Count} memories.", memories.Count);
-            if (contextInstructions is not null)
+            var memories = (await this._client.SearchAsync(
+                this.ApplicationId,
+                this.AgentId,
+                this.ThreadId,
+                this.UserId,
+                queryText,
+                cancellationToken).ConfigureAwait(false)).ToList();
+
+            var outputMessageText = memories.Count == 0
+                ? null
+                : $"{this._contextPrompt}\n{string.Join(Environment.NewLine, memories)}";
+
+            if (this._logger is not null)
            {
-                this._logger.LogTrace("Mem0AIContextProvider instructions: {Instructions}", contextInstructions);
+                this._logger.LogInformation(
+                    "Mem0AIContextProvider: Retrieved {Count} memories. ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', ThreadId: '{ThreadId}', UserId: '{UserId}'",
+                    memories.Count,
+                    this.ApplicationId,
+                    this.AgentId,
+                    this.ThreadId,
+                    this.UserId);
+                if (outputMessageText is not null)
+                {
+                    this._logger.LogTrace(
+                        "Mem0AIContextProvider: Search Results\nInput:{Input}\nOutput:{MessageText}\nApplicationId: '{ApplicationId}', AgentId: '{AgentId}', ThreadId: '{ThreadId}', UserId: '{UserId}'",
+                        queryText,
+                        outputMessageText,
+                        this.ApplicationId,
+                        this.AgentId,
+                        this.ThreadId,
+                        this.UserId);
+                }
            }
-        }

-        return new AIContext
+            return new AIContext
+            {
+                Messages = [new ChatMessage(ChatRole.User, outputMessageText)]
+            };
+        }
+        catch (ArgumentException)
        {
-            Messages = [new ChatMessage(ChatRole.User, contextInstructions)]
-        };
+            throw;
+        }
+        catch (Exception ex)
+        {
+            this._logger?.LogError(
+                ex,
+                "Mem0AIContextProvider: Failed to search Mem0 for memories due to error. ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', ThreadId: '{ThreadId}', UserId: '{UserId}'",
+                this.ApplicationId,
+                this.AgentId,
+                this.ThreadId,
+                this.UserId);
+            return new AIContext();
+        }
    }

    /// <inheritdoc />
@@ -166,12 +197,20 @@ public sealed class Mem0Provider : AIContextProvider
            return; // Do not update memory on failed invocations.
        }

-        // Persist request and response messages after invocation.
-        await this.PersistMessagesAsync(context.RequestMessages, cancellationToken).ConfigureAwait(false);
-
-        if (context.ResponseMessages is not null)
+        try
        {
-            await this.PersistMessagesAsync(context.ResponseMessages, cancellationToken).ConfigureAwait(false);
+            // Persist request and response messages after invocation.
+            await this.PersistMessagesAsync(context.RequestMessages.Concat(context.ResponseMessages ?? []), cancellationToken).ConfigureAwait(false);
+        }
+        catch (Exception ex)
+        {
+            this._logger?.LogError(
+                ex,
+                "Mem0AIContextProvider: Failed to send messages to Mem0 due to error. ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', ThreadId: '{ThreadId}', UserId: '{UserId}'",
+                this.ApplicationId,
+                this.AgentId,
+                this.ThreadId,
+                this.UserId);
        }
    }

@@ -22,7 +22,7 @@ namespace Microsoft.Agents.AI.Data;
 /// <para>
 /// The provider supports two behaviors controlled via <see cref="TextSearchProviderOptions.SearchTime"/>:
 /// <list type="bullet">
-/// <item><description><see cref="TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke"/> – Automatically performs a search prior to every AI invocation and injects results as additional instructions.</description></item>
+/// <item><description><see cref="TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke"/> – Automatically performs a search prior to every AI invocation and injects results as additional messages.</description></item>
 /// <item><description><see cref="TextSearchProviderOptions.TextSearchBehavior.OnDemandFunctionCalling"/> – Exposes a function tool that the model may invoke to retrieve contextual information when needed.</description></item>
 /// </list>
 /// </para>
@@ -122,12 +122,13 @@ public sealed class TextSearchProvider : AIContextProvider
        if (this._options.SearchTime != TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke)
        {
            // Expose the search tool for on-demand invocation.
-            return new AIContext { Tools = this._tools }; // No automatic instructions injection.
+            return new AIContext { Tools = this._tools }; // No automatic message injection.
        }

        // Aggregate text from memory + current request messages.
        var sbInput = new StringBuilder();
-        foreach (var messageText in this._recentMessagesText)
+        var requestMessagesText = context.RequestMessages.Where(x => !string.IsNullOrWhiteSpace(x?.Text)).Select(x => x.Text);
+        foreach (var messageText in this._recentMessagesText.Concat(requestMessagesText))
        {
            if (sbInput.Length > 0)
            {
@@ -136,38 +137,35 @@ public sealed class TextSearchProvider : AIContextProvider
            sbInput.Append(messageText);
        }

-        foreach (var message in context.RequestMessages)
-        {
-            if (!string.IsNullOrWhiteSpace(message?.Text))
-            {
-                if (sbInput.Length > 0)
-                {
-                    sbInput.Append('\n');
-                }
-                sbInput.Append(message.Text);
-            }
-        }
        string input = sbInput.ToString();

-        // Search
-        var results = await this._searchAsync(input, cancellationToken).ConfigureAwait(false);
-        IList<TextSearchResult> materialized = results as IList<TextSearchResult> ?? results.ToList();
-        if (materialized.Count == 0)
+        try
        {
-            this._logger?.LogWarning("TextSearchProvider: No search results found.");
+            // Search
+            var results = await this._searchAsync(input, cancellationToken).ConfigureAwait(false);
+            IList<TextSearchResult> materialized = results as IList<TextSearchResult> ?? results.ToList();
+            this._logger?.LogInformation("TextSearchProvider: Retrieved {Count} search results.", materialized.Count);
+
+            if (materialized.Count == 0)
+            {
+                return new AIContext();
+            }
+
+            // Format search results
+            string formatted = this.FormatResults(materialized);
+
+            this._logger?.LogTrace("TextSearchProvider: Search Results\nInput:{Input}\nOutput:{MessageText}", input, formatted);
+
+            return new AIContext
+            {
+                Messages = [new ChatMessage(ChatRole.User, formatted) { AdditionalProperties = new AdditionalPropertiesDictionary() { ["IsTextSearchProviderOutput"] = true } }]
+            };
+        }
+        catch (Exception ex)
+        {
+            this._logger?.LogError(ex, "TextSearchProvider: Failed to search for data due to error");
            return new AIContext();
        }
-
-        // Format search results
-        string formatted = this.FormatResults(materialized);
-
-        this._logger?.LogInformation("TextSearchProvider: Retrieved {Count} search results.", materialized.Count);
-        this._logger?.LogTrace("TextSearchProvider Input:{Input}\nContext Instructions:{Instructions}", input, formatted);
-
-        return new AIContext
-        {
-            Messages = [new ChatMessage(ChatRole.User, formatted) { AdditionalProperties = new AdditionalPropertiesDictionary() { ["IsTextSearchProviderOutput"] = true } }]
-        };
    }

    /// <inheritdoc />
@@ -240,16 +238,16 @@ public sealed class TextSearchProvider : AIContextProvider
    {
        var results = await this._searchAsync(userQuestion, cancellationToken).ConfigureAwait(false);
        IList<TextSearchResult> materialized = results as IList<TextSearchResult> ?? results.ToList();
-        string formatted = this.FormatResults(materialized);
+        string outputText = this.FormatResults(materialized);

        this._logger?.LogInformation("TextSearchProvider: Retrieved {Count} search results.", materialized.Count);
-        this._logger?.LogTrace("TextSearchProvider Input:{UserQuestion}\nContext Instructions:{Instructions}", userQuestion, formatted);
+        this._logger?.LogTrace("TextSearchProvider Input:{UserQuestion}\nOutput:{MessageText}", userQuestion, outputText);

-        return formatted;
+        return outputText;
    }

    /// <summary>
-    /// Formats search results into an instructions string for model consumption.
+    /// Formats search results into an output string for model consumption.
    /// </summary>
    /// <param name="results">The results.</param>
    /// <returns>Formatted string (may be empty).</returns>
@@ -30,12 +30,12 @@ public sealed class TextSearchProviderOptions
    public string? FunctionToolDescription { get; set; }

    /// <summary>
-    /// Gets or sets the context prompt prefixed to automatically injected results.
+    /// Gets or sets the context prompt prefixed to results.
    /// </summary>
    public string? ContextPrompt { get; set; }

    /// <summary>
-    /// Gets or sets the instruction appended after automatically injected results to request citations.
+    /// Gets or sets the instruction appended after results to request citations.
    /// </summary>
    public string? CitationsPrompt { get; set; }

@@ -85,7 +85,7 @@ public sealed class TextSearchProviderOptions
    public enum TextSearchBehavior
    {
        /// <summary>
-        /// Execute search prior to each invocation and inject results as instructions.
+        /// Execute search prior to each invocation and inject results as a message.
        /// </summary>
        BeforeAIInvoke,

@@ -9,6 +9,7 @@ using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
 using Microsoft.Extensions.Logging;
+using Moq;

 namespace Microsoft.Agents.AI.Mem0.UnitTests;

@@ -17,13 +18,23 @@ namespace Microsoft.Agents.AI.Mem0.UnitTests;
 /// </summary>
 public sealed class Mem0ProviderTests : IDisposable
 {
+    private readonly Mock<ILogger<Mem0Provider>> _loggerMock;
+    private readonly Mock<ILoggerFactory> _loggerFactoryMock;
    private readonly RecordingHandler _handler = new();
    private readonly HttpClient _httpClient;
-    private readonly ILoggerFactory _loggerFactory = LoggerFactory.Create(b => b.AddProvider(new NullLoggerProvider()));
    private bool _disposed;

    public Mem0ProviderTests()
    {
+        this._loggerMock = new();
+        this._loggerFactoryMock = new();
+        this._loggerFactoryMock
+            .Setup(f => f.CreateLogger(It.IsAny<string>()))
+            .Returns(this._loggerMock.Object);
+        this._loggerFactoryMock
+            .Setup(f => f.CreateLogger(typeof(Mem0Provider).FullName!))
+            .Returns(this._loggerMock.Object);
+
        this._httpClient = new HttpClient(this._handler)
        {
            BaseAddress = new Uri("https://localhost/")
@@ -80,7 +91,7 @@ public sealed class Mem0ProviderTests : IDisposable
            ThreadId = "thread",
            UserId = "user"
        };
-        var sut = new Mem0Provider(this._httpClient, options, this._loggerFactory);
+        var sut = new Mem0Provider(this._httpClient, options, this._loggerFactoryMock.Object);
        var invokingContext = new AIContextProvider.InvokingContext(new[] { new ChatMessage(ChatRole.User, "What is my name?") });

        // Act
@@ -99,6 +110,24 @@ public sealed class Mem0ProviderTests : IDisposable
        var contextMessage = Assert.Single(aiContext.Messages);
        Assert.Equal(ChatRole.User, contextMessage.Role);
        Assert.Contains("Name is Caoimhe", contextMessage.Text);
+
+        this._loggerMock.Verify(
+            l => l.Log(
+                LogLevel.Information,
+                It.IsAny<EventId>(),
+                It.Is<It.IsAnyType>((v, t) => v.ToString()!.Contains("Mem0AIContextProvider: Retrieved 1 memories.")),
+                It.IsAny<Exception>(),
+                It.IsAny<Func<It.IsAnyType, Exception?, string>>()),
+            Times.Once);
+
+        this._loggerMock.Verify(
+            l => l.Log(
+                LogLevel.Trace,
+                It.IsAny<EventId>(),
+                It.Is<It.IsAnyType>((v, t) => v.ToString()!.Contains("Mem0AIContextProvider: Search Results\nInput:What is my name?\nOutput")),
+                It.IsAny<Exception>(),
+                It.IsAny<Func<It.IsAnyType, Exception?, string>>()),
+            Times.Once);
    }

    [Fact]
@@ -156,6 +185,39 @@ public sealed class Mem0ProviderTests : IDisposable
        Assert.Empty(this._handler.Requests);
    }

+    [Fact]
+    public async Task InvokedAsync_ShouldNotThrow_WhenStorageFailsAsync()
+    {
+        // Arrange
+        var options = new Mem0ProviderOptions { ApplicationId = "a", AgentId = "b", ThreadId = "c", UserId = "d" };
+        var sut = new Mem0Provider(this._httpClient, options, this._loggerFactoryMock.Object);
+        this._handler.EnqueueEmptyInternalServerError();
+
+        var requestMessages = new List<ChatMessage>
+        {
+            new(ChatRole.User, "User text"),
+            new(ChatRole.System, "System text"),
+            new(ChatRole.Tool, "Tool text should be ignored")
+        };
+        var responseMessages = new List<ChatMessage>
+        {
+            new(ChatRole.Assistant, "Assistant text")
+        };
+
+        // Act
+        await sut.InvokedAsync(new AIContextProvider.InvokedContext(requestMessages, aiContextProviderMessages: null) { ResponseMessages = responseMessages });
+
+        // Assert
+        this._loggerMock.Verify(
+            l => l.Log(
+                LogLevel.Error,
+                It.IsAny<EventId>(),
+                It.Is<It.IsAnyType>((v, t) => v.ToString()!.Contains("Mem0AIContextProvider: Failed to send messages to Mem0 due to error")),
+                It.IsAny<Exception>(),
+                It.IsAny<Func<It.IsAnyType, Exception?, string>>()),
+            Times.Once);
+    }
+
    [Fact]
    public async Task ClearStoredMemoriesAsync_SendsDeleteWithQueryAsync()
    {
@@ -267,6 +329,30 @@ public sealed class Mem0ProviderTests : IDisposable
        await Assert.ThrowsAsync<ArgumentException>(() => sut.InvokingAsync(ctx).AsTask());
    }

+    [Fact]
+    public async Task InvokingAsync_ShouldNotThrow_WhenSearchFailsAsync()
+    {
+        // Arrange
+        var options = new Mem0ProviderOptions { ApplicationId = "app" };
+        var provider = new Mem0Provider(this._httpClient, options, loggerFactory: this._loggerFactoryMock.Object);
+        var invokingContext = new AIContextProvider.InvokingContext(new[] { new ChatMessage(ChatRole.User, "Q?") });
+
+        // Act
+        var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None);
+
+        // Assert
+        Assert.Null(aiContext.Messages);
+        Assert.Null(aiContext.Tools);
+        this._loggerMock.Verify(
+            l => l.Log(
+                LogLevel.Error,
+                It.IsAny<EventId>(),
+                It.Is<It.IsAnyType>((v, t) => v.ToString()!.Contains("Mem0AIContextProvider: Failed to search Mem0 for memories due to error")),
+                It.IsAny<Exception>(),
+                It.IsAny<Func<It.IsAnyType, Exception?, string>>()),
+            Times.Once);
+    }
+
    private static bool ContainsOrdinal(string source, string value) => source.IndexOf(value, StringComparison.Ordinal) >= 0;

    public void Dispose()
@@ -275,7 +361,6 @@ public sealed class Mem0ProviderTests : IDisposable
        {
            this._httpClient.Dispose();
            this._handler.Dispose();
-            this._loggerFactory.Dispose();
            this._disposed = true;
        }
    }
@@ -309,18 +394,7 @@ public sealed class Mem0ProviderTests : IDisposable
        }

        public void EnqueueEmptyOk() => this._responses.Enqueue(new HttpResponseMessage(System.Net.HttpStatusCode.OK));
-    }

-    private sealed class NullLoggerProvider : ILoggerProvider
-    {
-        public ILogger CreateLogger(string categoryName) => new NullLogger();
-        public void Dispose() { }
-
-        private sealed class NullLogger : ILogger
-        {
-            public IDisposable? BeginScope<TState>(TState state) where TState : notnull => null;
-            public bool IsEnabled(LogLevel logLevel) => false;
-            public void Log<TState>(LogLevel logLevel, EventId eventId, TState state, Exception? exception, Func<TState, Exception?, string> formatter) { }
-        }
+        public void EnqueueEmptyInternalServerError() => this._responses.Enqueue(new HttpResponseMessage(System.Net.HttpStatusCode.InternalServerError));
    }
 }
@@ -116,7 +116,7 @@ public sealed class TextSearchProviderTests
                l => l.Log(
                    LogLevel.Trace,
                    It.IsAny<EventId>(),
-                    It.Is<It.IsAnyType>((v, t) => v.ToString()!.Contains("TextSearchProvider Input:Sample user question?\nAdditional part\nContext Instructions")),
+                    It.Is<It.IsAnyType>((v, t) => v.ToString()!.Contains("TextSearchProvider: Search Results\nInput:Sample user question?\nAdditional part\nOutput")),
                    It.IsAny<Exception?>(),
                    It.IsAny<Func<It.IsAnyType, Exception?, string>>()),
                Times.AtLeastOnce);
@@ -150,6 +150,29 @@ public sealed class TextSearchProviderTests
        Assert.Equal(expectedDescription, tool.Description);
    }

+    [Fact]
+    public async Task InvokingAsync_ShouldNotThrow_WhenSearchFailsAsync()
+    {
+        // Arrange
+        var provider = new TextSearchProvider(this.FailingSearchAsync, loggerFactory: this._loggerFactoryMock.Object);
+        var invokingContext = new AIContextProvider.InvokingContext(new[] { new ChatMessage(ChatRole.User, "Q?") });
+
+        // Act
+        var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None);
+
+        // Assert
+        Assert.Null(aiContext.Messages);
+        Assert.Null(aiContext.Tools);
+        this._loggerMock.Verify(
+            l => l.Log(
+                LogLevel.Error,
+                It.IsAny<EventId>(),
+                It.Is<It.IsAnyType>((v, t) => v.ToString()!.Contains("TextSearchProvider: Failed to search for data due to error")),
+                It.IsAny<Exception>(),
+                It.IsAny<Func<It.IsAnyType, Exception?, string>>()),
+            Times.AtLeastOnce);
+    }
+
    [Theory]
    [InlineData(null, null)]
    [InlineData("Custom context prompt", "Custom citations prompt")]
@@ -619,6 +642,11 @@ public sealed class TextSearchProviderTests
        return Task.FromResult<IEnumerable<TextSearchProvider.TextSearchResult>>([]);
    }

+    private Task<IEnumerable<TextSearchProvider.TextSearchResult>> FailingSearchAsync(string input, CancellationToken ct)
+    {
+        throw new InvalidOperationException("Search Failed");
+    }
+
    private sealed class RawPayload
    {
        public string Id { get; set; } = string.Empty;