.NET: Add AzureStorageBlob integration as AgentThreadStore support #1893
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,157 @@ | ||
| // Copyright (c) Microsoft. All rights reserved. | ||
|
|
||
| using System; | ||
| using System.IO; | ||
| using System.Text; | ||
| using System.Text.Json; | ||
| using System.Threading; | ||
| using System.Threading.Tasks; | ||
| using Azure; | ||
| using Azure.Storage.Blobs; | ||
| using Azure.Storage.Blobs.Models; | ||
| using Microsoft.Shared.Diagnostics; | ||
|
|
||
| namespace Microsoft.Agents.AI.Hosting.AzureStorage.Blob; | ||
|
|
||
| internal sealed class AzureBlobAgentThreadStore : AgentThreadStore | ||
| { | ||
| private static readonly BlobOpenWriteOptions s_uploadJsonOptions = new() | ||
| { | ||
| HttpHeaders = new() | ||
| { | ||
| ContentType = "application/json" | ||
| } | ||
| }; | ||
|
|
||
| private readonly BlobContainerClient _containerClient; | ||
| private readonly AzureBlobAgentThreadStoreOptions _options; | ||
| private bool _containerInitialized; | ||
|
|
||
| /// <summary> | ||
| /// Initializes a new instance of the <see cref="AzureBlobAgentThreadStore"/> class using a <see cref="BlobContainerClient"/>. | ||
| /// </summary> | ||
| /// <param name="containerClient">The blob container client to use for storage operations.</param> | ||
| /// <param name="options">Optional configuration options. If <see langword="null"/>, default options will be used.</param> | ||
| /// <exception cref="ArgumentNullException"><paramref name="containerClient"/> is <see langword="null"/>.</exception> | ||
| public AzureBlobAgentThreadStore(BlobContainerClient containerClient, AzureBlobAgentThreadStoreOptions? options = null) | ||
| { | ||
| this._containerClient = containerClient ?? throw new ArgumentNullException(nameof(containerClient)); | ||
| this._options = options ?? new AzureBlobAgentThreadStoreOptions(); | ||
| } | ||
|
|
||
| /// <inheritdoc/> | ||
| public override async ValueTask SaveThreadAsync( | ||
| AIAgent agent, | ||
| string conversationId, | ||
| AgentThread thread, | ||
| CancellationToken cancellationToken = default) | ||
| { | ||
| Throw.IfNull(agent); | ||
| Throw.IfNull(conversationId); | ||
| Throw.IfNull(thread); | ||
|
|
||
| await this.EnsureContainerExistsAsync(cancellationToken).ConfigureAwait(false); | ||
|
|
||
| var blobName = this.GetBlobName(agent.Id, conversationId); | ||
| var blobClient = this._containerClient.GetBlobClient(blobName); | ||
|
|
||
| JsonElement serializedThread = thread.Serialize(); | ||
| using Stream stream = await blobClient.OpenWriteAsync(overwrite: true, s_uploadJsonOptions, cancellationToken).ConfigureAwait(false); | ||
| using Utf8JsonWriter writer = new(stream); | ||
|
|
||
| serializedThread.WriteTo(writer); | ||
| await writer.FlushAsync(cancellationToken).ConfigureAwait(false); | ||
| } | ||
|
|
||
| /// <inheritdoc/> | ||
| public override async ValueTask<AgentThread> GetThreadAsync( | ||
| AIAgent agent, | ||
| string conversationId, | ||
| CancellationToken cancellationToken = default) | ||
| { | ||
| Throw.IfNull(agent); | ||
| Throw.IfNull(conversationId); | ||
|
|
||
| await this.EnsureContainerExistsAsync(cancellationToken).ConfigureAwait(false); | ||
|
|
||
| var blobName = this.GetBlobName(agent.Id, conversationId); | ||
| var blobClient = this._containerClient.GetBlobClient(blobName); | ||
|
|
||
| try | ||
| { | ||
| var stream = await blobClient.OpenReadAsync(cancellationToken: cancellationToken).ConfigureAwait(false); | ||
| var jsonDoc = await JsonDocument.ParseAsync(stream, cancellationToken: cancellationToken).ConfigureAwait(false); | ||
| var serializedThread = jsonDoc.RootElement; | ||
|
|
||
| return agent.DeserializeThread(serializedThread); | ||
| } | ||
| catch (RequestFailedException ex) when (ex.Status == 404) | ||
| { | ||
| // Blob doesn't exist, return a new thread | ||
| return agent.GetNewThread(); | ||
| } | ||
| } | ||
|
|
||
| /// <summary> | ||
| /// Ensures that the blob container exists, creating it if necessary. | ||
| /// </summary> | ||
| private async Task EnsureContainerExistsAsync(CancellationToken cancellationToken) | ||
| { | ||
| if (this._containerInitialized) | ||
| { | ||
| return; | ||
| } | ||
|
|
||
| if (!this._containerInitialized) | ||
| { | ||
| if (this._options.CreateContainerIfNotExists) | ||
| { | ||
| await this._containerClient.CreateIfNotExistsAsync(cancellationToken: cancellationToken).ConfigureAwait(false); | ||
| } | ||
|
|
||
| this._containerInitialized = true; | ||
| } | ||
DeagleGross marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| } | ||
|
|
||
| /// <summary> | ||
| /// Generates the blob name for a given agent and conversation. | ||
| /// </summary> | ||
| // internal for testing | ||
| internal string GetBlobName(string agentId, string conversationId) | ||
| { | ||
| string sanitizedAgentId = this.SanitizeBlobNameSegment(agentId); | ||
| string sanitizedConversationId = this.SanitizeBlobNameSegment(conversationId); | ||
| string baseName = $"{sanitizedAgentId}/{sanitizedConversationId}.json"; | ||
|
|
||
| return string.IsNullOrEmpty(this._options.BlobNamePrefix) | ||
| ? baseName | ||
| : $"{this._options.BlobNamePrefix.TrimEnd('/')}/{baseName}"; | ||
| } | ||
|
|
||
| /// <summary> | ||
| /// Sanitizes a string to be safe for use in blob names. | ||
| /// </summary> | ||
| private string SanitizeBlobNameSegment(string input) | ||
| { | ||
| if (string.IsNullOrWhiteSpace(input)) | ||
| { | ||
| return "default"; | ||
| } | ||
|
|
||
| // Replace invalid characters with underscore | ||
| StringBuilder builder = new(input.Length); | ||
| foreach (char c in input) | ||
| { | ||
| if (char.IsLetterOrDigit(c) || c == '-' || c == '_' || c == '.') | ||
| { | ||
| builder.Append(c); | ||
| } | ||
| else | ||
| { | ||
| builder.Append('_'); | ||
| } | ||
| } | ||
|
|
||
| return builder.ToString(); | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,26 @@ | ||
| // Copyright (c) Microsoft. All rights reserved. | ||
|
|
||
| namespace Microsoft.Agents.AI.Hosting.AzureStorage.Blob; | ||
|
|
||
| /// <summary> | ||
| /// Configuration options for <see cref="AzureBlobAgentThreadStore"/>. | ||
| /// </summary> | ||
| public sealed class AzureBlobAgentThreadStoreOptions | ||
| { | ||
| /// <summary> | ||
| /// Gets or sets a value indicating whether to automatically create the container if it doesn't exist. | ||
| /// </summary> | ||
| /// <remarks> | ||
| /// Defaults to <see langword="true"/>. | ||
| /// </remarks> | ||
| public bool CreateContainerIfNotExists { get; set; } = true; | ||
|
|
||
| /// <summary> | ||
| /// Gets or sets the blob name prefix to use for organizing threads. | ||
| /// </summary> | ||
| /// <remarks> | ||
| /// This can be used to namespace threads within a container. | ||
| /// For example, setting this to "prod/" will store all blobs under a "prod/" prefix. | ||
| /// </remarks> | ||
| public string? BlobNamePrefix { get; set; } | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,41 @@ | ||
| // Copyright (c) Microsoft. All rights reserved. | ||
|
|
||
| using System; | ||
| using Azure.Storage.Blobs; | ||
| using Microsoft.Agents.AI.Hosting.AzureStorage.Blob; | ||
|
|
||
| namespace Microsoft.Agents.AI.Hosting; | ||
|
|
||
| /// <summary> | ||
| /// Provides extension methods for configuring <see cref="AIAgent"/>. | ||
| /// </summary> | ||
| public static partial class HostedAgentBuilderExtensions | ||
| { | ||
|
Contributor
Author
There was a problem hiding this comment. not sure if we need other API here, but felt like |
||
| /// <summary> | ||
| /// Configures the host agent builder to use an Azure Blob thread store for agent thread management. | ||
| /// </summary> | ||
| /// <param name="builder">The host agent builder to configure with the Azure blob thread store.</param> | ||
| /// <param name="containerClient">The blob container client to use for storage operations.</param> | ||
| /// <param name="options">Optional configuration options for the blob thread store.</param> | ||
| /// <returns>The same <paramref name="builder"/> instance, configured to use Azure blob thread store.</returns> | ||
| public static IHostedAgentBuilder WithAzureBlobThreadStore(this IHostedAgentBuilder builder, BlobContainerClient containerClient, AzureBlobAgentThreadStoreOptions? options = null) | ||
| => WithAzureBlobThreadStore(builder, sp => containerClient, options); | ||
|
|
||
| /// <summary> | ||
| /// Configures the agent builder to use Azure Blob Storage as the thread store for agent state persistence. | ||
| /// </summary> | ||
| /// <param name="builder">The agent builder to configure with Azure Blob thread store support.</param> | ||
| /// <param name="createBlobContainer">A factory function that provides a configured BlobContainerClient instance for accessing the Azure Blob | ||
| /// container used to store thread data.</param> | ||
| /// <param name="options">Optional settings for customizing the Azure Blob thread store behavior. If null, default options are used.</param> | ||
| /// <returns>The same agent builder instance, configured to use Azure Blob Storage for thread persistence.</returns> | ||
| public static IHostedAgentBuilder WithAzureBlobThreadStore( | ||
| this IHostedAgentBuilder builder, | ||
| Func<IServiceProvider, BlobContainerClient> createBlobContainer, | ||
| AzureBlobAgentThreadStoreOptions? options = null) | ||
| => builder.WithThreadStore((sp, key) => | ||
| { | ||
| var blobContainer = createBlobContainer(sp); | ||
| return new AzureBlobAgentThreadStore(blobContainer, options); | ||
| }); | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,27 @@ | ||
| <Project Sdk="Microsoft.NET.Sdk"> | ||
|
|
||
| <PropertyGroup> | ||
| <TargetFrameworks>$(ProjectsCoreTargetFrameworks)</TargetFrameworks> | ||
| <TargetFrameworks Condition="'$(Configuration)' == 'Debug'">$(ProjectsDebugCoreTargetFrameworks)</TargetFrameworks> | ||
| <VersionSuffix>preview</VersionSuffix> | ||
| <InjectSharedThrow>true</InjectSharedThrow> | ||
|
|
||
| <!-- NuGet Package Settings --> | ||
| <Title>Microsoft Agent Framework AzureStorage</Title> | ||
| <Description>Provides AzureStorage integration with Microsoft Agent Framework.</Description> | ||
| </PropertyGroup> | ||
|
|
||
| <Import Project="$(RepoRoot)/dotnet/nuget/nuget-package.props" /> | ||
|
|
||
| <ItemGroup> | ||
| <PackageReference Include="Azure.Storage.Blobs" /> | ||
| </ItemGroup> | ||
|
|
||
| <ItemGroup> | ||
| <ProjectReference Include="..\Microsoft.Agents.AI.Hosting\Microsoft.Agents.AI.Hosting.csproj" /> | ||
| </ItemGroup> | ||
|
|
||
| <ItemGroup> | ||
| <InternalsVisibleTo Include="Microsoft.Agents.AI.Hosting.AzureStorage.UnitTests" /> | ||
| </ItemGroup> | ||
| </Project> |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.