diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/CheckpointManager.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/CheckpointManager.cs index 9334158b5b..c50283e728 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/CheckpointManager.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/CheckpointManager.cs @@ -20,7 +20,7 @@ public sealed class CheckpointManager : ICheckpointManager return new CheckpointManagerImpl(marshaller, store); } - private CheckpointManager(ICheckpointManager impl) + internal CheckpointManager(ICheckpointManager impl) { this._impl = impl; } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs index 965a23a22f..8fa289d7ea 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs @@ -3,8 +3,6 @@ using System; using System.Collections.Concurrent; using System.Collections.Generic; -using System.Diagnostics; -using System.Linq; using System.Runtime.CompilerServices; using System.Text.Json; using System.Threading; @@ -16,18 +14,19 @@ namespace Microsoft.Agents.AI.Workflows; internal sealed class WorkflowHostAgent : AIAgent { - private readonly Workflow> _workflow; + private readonly Workflow _workflow; private readonly string? _id; + private readonly CheckpointManager? _checkpointManager; private readonly ConcurrentDictionary _assignedRunIds = []; - private readonly Dictionary _runningWorkflows = []; - public WorkflowHostAgent(Workflow> workflow, string? id = null, string? name = null) + public WorkflowHostAgent(Workflow> workflow, string? id = null, string? name = null, CheckpointManager? checkpointManager = null) { this._workflow = Throw.IfNull(workflow); this._id = id; this.Name = name; + this._checkpointManager = checkpointManager; } public override string? Name { get; } @@ -45,59 +44,10 @@ internal sealed class WorkflowHostAgent : AIAgent return result; } - public override AgentThread GetNewThread() => new WorkflowThread(this.Id, this.Name, this.GenerateNewId()); + public override AgentThread GetNewThread() => new WorkflowThread(this._workflow, this.GenerateNewId(), this._checkpointManager); public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null) - => new WorkflowThread(serializedThread, jsonSerializerOptions); - - private async - IAsyncEnumerable InvokeStageAsync( - WorkflowThread conversation, - [EnumeratorCancellation] CancellationToken cancellationToken = default) - { - string runId = conversation.RunId; - List messages = conversation.MessageStore.GetFromBookmark().ToList(); - - try - { - // technically there is a race condition here between assigning the ID, and checking if it exists - // in the case of new threads. - if (!this._runningWorkflows.TryGetValue(runId, out StreamingRun? run)) - { - run = await InProcessExecution.StreamAsync(this._workflow, messages, cancellationToken: cancellationToken) - .ConfigureAwait(false); - this._runningWorkflows[runId] = run; - } - else - { - bool sentMessages = await run.TrySendMessageAsync(messages).ConfigureAwait(false); - Debug.Assert(sentMessages, "Hosted workflow is required to take List as input."); - } - - await run.TrySendMessageAsync(new TurnToken(emitEvents: true)).ConfigureAwait(false); - await foreach (WorkflowEvent evt in run.WatchStreamAsync(blockOnPendingRequest: false, cancellationToken) - .ConfigureAwait(false) - .WithCancellation(cancellationToken)) - { - switch (evt) - { - case AgentRunUpdateEvent agentUpdate: - yield return agentUpdate.Update; - break; - case RequestInfoEvent requestInfo: - FunctionCallContent fcContent = requestInfo.Request.ToFunctionCall(); - AgentRunResponseUpdate update = conversation.CreateUpdate(fcContent); - yield return update; - break; - } - } - } - finally - { - // Do we want to try to undo the step, and not update the bookmark? - conversation.MessageStore.UpdateBookmark(); - } - } + => new WorkflowThread(this._workflow, serializedThread, this._checkpointManager, jsonSerializerOptions); private async ValueTask UpdateThreadAsync(IEnumerable messages, AgentThread? thread = null, CancellationToken cancellationToken = default) { @@ -122,14 +72,14 @@ internal sealed class WorkflowHostAgent : AIAgent WorkflowThread workflowThread = await this.UpdateThreadAsync(messages, thread, cancellationToken).ConfigureAwait(false); MessageMerger merger = new(); - await foreach (AgentRunResponseUpdate update in this.InvokeStageAsync(workflowThread, cancellationToken) - .ConfigureAwait(false) - .WithCancellation(cancellationToken)) + await foreach (AgentRunResponseUpdate update in workflowThread.InvokeStageAsync(cancellationToken) + .ConfigureAwait(false) + .WithCancellation(cancellationToken)) { merger.AddUpdate(update); } - return merger.ComputeMerged(workflowThread.ResponseId, this.Id, this.Name); + return merger.ComputeMerged(workflowThread.LastResponseId!, this.Id, this.Name); } public override async @@ -140,9 +90,9 @@ internal sealed class WorkflowHostAgent : AIAgent [EnumeratorCancellation] CancellationToken cancellationToken = default) { WorkflowThread workflowThread = await this.UpdateThreadAsync(messages, thread, cancellationToken).ConfigureAwait(false); - await foreach (AgentRunResponseUpdate update in this.InvokeStageAsync(workflowThread, cancellationToken) - .ConfigureAwait(false) - .WithCancellation(cancellationToken)) + await foreach (AgentRunResponseUpdate update in workflowThread.InvokeStageAsync(cancellationToken) + .ConfigureAwait(false) + .WithCancellation(cancellationToken)) { yield return update; } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowMessageStore.cs index 1fffb502fc..39c83bcadf 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowMessageStore.cs @@ -1,11 +1,11 @@ // Copyright (c) Microsoft. All rights reserved. -using System; using System.Collections.Generic; using System.Text.Json; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; +using Microsoft.Shared.Diagnostics; namespace Microsoft.Agents.AI.Workflows; @@ -18,22 +18,22 @@ internal sealed class WorkflowMessageStore : ChatMessageStore { } - public WorkflowMessageStore(JsonElement serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null) + public WorkflowMessageStore(StoreState state) { - if (serializedStoreState.ValueKind is not JsonValueKind.Object) - { - throw new ArgumentException("The provided JsonElement must be a json object", nameof(serializedStoreState)); - } + this.ImportStoreState(Throw.IfNull(state)); + } - StoreState? state = - serializedStoreState.Deserialize( - AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(StoreState))) as StoreState; + private void ImportStoreState(StoreState state, bool clearMessages = false) + { + if (clearMessages) + { + this._chatMessages.Clear(); + } if (state?.Messages is not null) { this._chatMessages.AddRange(state.Messages); } - this._bookmark = state?.Bookmark ?? 0; } @@ -66,13 +66,11 @@ internal sealed class WorkflowMessageStore : ChatMessageStore public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) { - StoreState state = new() - { - Bookmark = this._bookmark, - Messages = this._chatMessages, - }; + StoreState state = this.ExportStoreState(); return JsonSerializer.SerializeToElement(state, WorkflowsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(StoreState))); } + + internal StoreState ExportStoreState() => new() { Bookmark = this._bookmark, Messages = this._chatMessages }; } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowThread.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowThread.cs index 8db2eabe93..61dc0ab337 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowThread.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowThread.cs @@ -1,7 +1,13 @@ // Copyright (c) Microsoft. All rights reserved. using System; +using System.Collections.Generic; +using System.Linq; +using System.Runtime.CompilerServices; using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Agents.AI.Workflows.Checkpointing; using Microsoft.Extensions.AI; using Microsoft.Shared.Diagnostics; @@ -9,26 +15,71 @@ namespace Microsoft.Agents.AI.Workflows; internal sealed class WorkflowThread : AgentThread { - public WorkflowThread(string workflowId, string? workflowName, string runId) + private readonly CheckpointManager _checkpointManager; + private readonly InMemoryCheckpointManager? _inMemoryCheckpointManager; + private readonly Workflow _workflow; + + public WorkflowThread(Workflow workflow, string runId, CheckpointManager? checkpointManager = null) { - this.MessageStore = new(); + this._workflow = Throw.IfNull(workflow); + + // If the user provided an external checkpoint manager, use that, otherwise rely on an in-memory one. + // TODO: Implement persist-only-last functionality for in-memory checkpoint manager, to avoid unbounded + // memory growth. + this._checkpointManager = checkpointManager ?? new(this._inMemoryCheckpointManager = new()); + this.RunId = Throw.IfNullOrEmpty(runId); + this.MessageStore = new WorkflowMessageStore(); } - public WorkflowThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null) + public WorkflowThread(Workflow workflow, JsonElement serializedThread, CheckpointManager? checkpointManager = null, JsonSerializerOptions? jsonSerializerOptions = null) { - throw new NotImplementedException("Pending Checkpointing work."); + this._workflow = Throw.IfNull(workflow); + + JsonMarshaller marshaller = new(jsonSerializerOptions); + ThreadState threadState = marshaller.Marshal(serializedThread); + + this._inMemoryCheckpointManager = threadState.CheckpointManager; + if (this._inMemoryCheckpointManager is not null && checkpointManager is not null) + { + // The thread was externalized with an in-memory checkpoint manager, but the caller is providing an external one. + throw new ArgumentException("Cannot provide an external checkpoint manager when deserializing a thread that " + + "was serialized with an in-memory checkpoint manager.", nameof(checkpointManager)); + } + else if (this._inMemoryCheckpointManager is null && checkpointManager is null) + { + // The thread was externalized without an in-memory checkpoint manager, and the caller is not providing an external one. + throw new ArgumentException("An external checkpoint manager must be provided when deserializing a thread that " + + "was serialized without an in-memory checkpoint manager.", nameof(checkpointManager)); + } + else + { + this._checkpointManager = checkpointManager ?? new(this._inMemoryCheckpointManager!); + } + + this.RunId = threadState.RunId; + this.LastCheckpoint = threadState.LastCheckpoint; + this.MessageStore = new WorkflowMessageStore(threadState.MessageStoreState); } - public string RunId { get; } - public int Halts { get; } + public CheckpointInfo? LastCheckpoint { get; set; } - public string ResponseId => $"{this.RunId}@{this.Halts}"; + protected override Task MessagesReceivedAsync(IEnumerable newMessages, CancellationToken cancellationToken = default) + => this.MessageStore.AddMessagesAsync(newMessages, cancellationToken); public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - => throw new NotImplementedException("Pending Checkpointing work."); + { + JsonMarshaller marshaller = new(jsonSerializerOptions); + ThreadState info = new( + this.RunId, + this.LastCheckpoint, + this.MessageStore.ExportStoreState(), + this._inMemoryCheckpointManager); - public AgentRunResponseUpdate CreateUpdate(params AIContent[] parts) + return marshaller.Marshal(info); + } + + public AgentRunResponseUpdate CreateUpdate(string responseId, params AIContent[] parts) { Throw.IfNullOrEmpty(parts); @@ -36,6 +87,8 @@ internal sealed class WorkflowThread : AgentThread { CreatedAt = DateTimeOffset.UtcNow, MessageId = Guid.NewGuid().ToString("N"), + Role = ChatRole.Assistant, + ResponseId = responseId }; this.MessageStore.AddMessages(update.ToChatMessage()); @@ -43,6 +96,89 @@ internal sealed class WorkflowThread : AgentThread return update; } + private async ValueTask> CreateOrResumeRunAsync(List messages, CancellationToken cancellationToken = default) + { + if (this.LastCheckpoint is not null) + { + Checkpointed checkpointed = + await InProcessExecution.ResumeStreamAsync(this._workflow, + this.LastCheckpoint, + this._checkpointManager, + this.RunId, + cancellationToken) + .ConfigureAwait(false); + + await checkpointed.Run.TrySendMessageAsync(messages).ConfigureAwait(false); + return checkpointed; + } + + return await InProcessExecution.StreamAsync(this._workflow, + messages, + this._checkpointManager, + this.RunId, + cancellationToken) + .ConfigureAwait(false); + } + + internal async + IAsyncEnumerable InvokeStageAsync( + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + try + { + this.LastResponseId = Guid.NewGuid().ToString("N"); + List messages = this.MessageStore.GetFromBookmark().ToList(); + +#pragma warning disable CA2007 // Analyzer misfiring and not seeing .ConfigureAwait(false) below. + await using Checkpointed checkpointed = + await this.CreateOrResumeRunAsync(messages, cancellationToken).ConfigureAwait(false); +#pragma warning restore CA2007 + + StreamingRun run = checkpointed.Run; + await run.TrySendMessageAsync(new TurnToken(emitEvents: true)).ConfigureAwait(false); + await foreach (WorkflowEvent evt in run.WatchStreamAsync(blockOnPendingRequest: false, cancellationToken) + .ConfigureAwait(false) + .WithCancellation(cancellationToken)) + { + switch (evt) + { + case AgentRunUpdateEvent agentUpdate: + yield return agentUpdate.Update; + break; + case RequestInfoEvent requestInfo: + FunctionCallContent fcContent = requestInfo.Request.ToFunctionCall(); + AgentRunResponseUpdate update = this.CreateUpdate(this.LastResponseId, fcContent); + yield return update; + break; + case SuperStepCompletedEvent stepCompleted: + this.LastCheckpoint = stepCompleted.CompletionInfo?.Checkpoint; + break; + } + } + } + finally + { + // Do we want to try to undo the step, and not update the bookmark? + this.MessageStore.UpdateBookmark(); + } + } + + public string? LastResponseId { get; set; } + + public string RunId { get; } + /// public WorkflowMessageStore MessageStore { get; } + + internal sealed class ThreadState( + string runId, + CheckpointInfo? lastCheckpoint, + WorkflowMessageStore.StoreState messageStoreState, + InMemoryCheckpointManager? checkpointManager = null) + { + public string RunId { get; } = runId; + public CheckpointInfo? LastCheckpoint { get; } = lastCheckpoint; + public WorkflowMessageStore.StoreState MessageStoreState { get; } = messageStoreState; + public InMemoryCheckpointManager? CheckpointManager { get; } = checkpointManager; + } } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowsJsonUtilities.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowsJsonUtilities.cs index 33cf002517..f64c1a6b80 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowsJsonUtilities.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowsJsonUtilities.cs @@ -7,7 +7,6 @@ using System.Text.Json.Serialization; using Microsoft.Agents.AI.Workflows.Checkpointing; using Microsoft.Agents.AI.Workflows.Execution; using Microsoft.Extensions.AI; -using static Microsoft.Agents.AI.Workflows.WorkflowMessageStore; namespace Microsoft.Agents.AI.Workflows; @@ -80,7 +79,8 @@ internal static partial class WorkflowsJsonUtilities [JsonSerializable(typeof(EdgeConnection))] // Workflow-as-Agent - [JsonSerializable(typeof(StoreState))] + [JsonSerializable(typeof(WorkflowMessageStore.StoreState))] + [JsonSerializable(typeof(WorkflowThread.ThreadState))] // Message Types [JsonSerializable(typeof(ChatMessage))] @@ -91,7 +91,7 @@ internal static partial class WorkflowsJsonUtilities // Event Types //[JsonSerializable(typeof(WorkflowEvent))] // Currently cannot be serialized because it includes Exceptions. - // We'll need a way to marshal this correct in the AgentRuntime case. + // We'll need a way to marshal this correctly in the AgentRuntime case. // For now this is okay, because we never serialize WorkflowEvents into // checkpoints. [JsonSerializable(typeof(JsonElement))]