AddorUpdate

microsoft · Nov 9, 2024 · 8bb3c09 · 8bb3c09
1 parent e19563b
commit 8bb3c09
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 20 deletions.
diff --git a/dotnet/samples/Hello/HelloAgentState/Program.cs b/dotnet/samples/Hello/HelloAgentState/Program.cs
@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Program.cs
 
+using System.Text.Json;
 using Microsoft.AutoGen.Abstractions;
 using Microsoft.AutoGen.Agents;
 using Microsoft.Extensions.DependencyInjection;
@@ -22,9 +23,9 @@ public class HelloAgent(
         [FromKeyedServices("EventTypes")] EventTypes typeRegistry) : ConsoleAgent(
             context,
             typeRegistry),
-            ISayHello,
             IHandle<NewMessageReceived>,
-            IHandle<ConversationClosed>
+            IHandle<ConversationClosed>,
+            IHandle<Shutdown>
     {
         private AgentState? State { get; set; }
         public async Task Handle(NewMessageReceived item)
@@ -34,11 +35,15 @@ public async Task Handle(NewMessageReceived item)
             {
                 Message = response
             }.ToCloudEvent(this.AgentId.Key);
-            var entry = "We said hello to " + item.Message;
+            Dictionary <string, string> state = new()
+            {
+                { "data", "We said hello to " + item.Message },
+                { "workflow", "Active" }
+            };
             await StoreAsync(new AgentState
             {
                 AgentId = this.AgentId,
-                TextData = entry
+                TextData = JsonSerializer.Serialize(state)
             }).ConfigureAwait(false);
             await PublishEventAsync(evt).ConfigureAwait(false);
             var goodbye = new ConversationClosed
@@ -47,30 +52,45 @@ await StoreAsync(new AgentState
                 UserMessage = "Goodbye"
             }.ToCloudEvent(this.AgentId.Key);
             await PublishEventAsync(goodbye).ConfigureAwait(false);
+            // send the shutdown message
+            await PublishEventAsync(new Shutdown { Message = this.AgentId.Key }.ToCloudEvent(this.AgentId.Key)).ConfigureAwait(false);
+
         }
         public async Task Handle(ConversationClosed item)
         {
             State = await ReadAsync<AgentState>(this.AgentId).ConfigureAwait(false);
-            var read = State?.TextData ?? "No state data found";
-            var goodbye = $"{read}\n*********************  {item.UserId} said {item.UserMessage}  ************************";
+            var state = JsonSerializer.Deserialize<Dictionary<string, string>>(State.TextData) ?? new Dictionary<string, string> { { "data", "No state data found" } };
+            var goodbye = $"\nState: {state}\n*********************  {item.UserId} said {item.UserMessage}  ************************";
             var evt = new Output
             {
                 Message = goodbye
             }.ToCloudEvent(this.AgentId.Key);
-            await PublishEventAsync(evt).ConfigureAwait(false);
-            //sleep
-            await Task.Delay(10000).ConfigureAwait(false);
-            await AgentsApp.ShutdownAsync().ConfigureAwait(false);
-
+            await PublishEventAsync(evt).ConfigureAwait(true);
+            state["workflow"] = "Complete";
+            await StoreAsync(new AgentState
+            {
+                AgentId = this.AgentId,
+                TextData = JsonSerializer.Serialize(state)
+            }).ConfigureAwait(false);
+        }
+        public async Task Handle(Shutdown item)
+        {
+            string? workflow = null;
+            // make sure the workflow is finished
+            while (workflow != "Complete")
+            {
+                State = await ReadAsync<AgentState>(this.AgentId).ConfigureAwait(true);
+                var state = JsonSerializer.Deserialize<Dictionary<string, string>>(State?.TextData ?? "{}") ?? new Dictionary<string, string>();
+                workflow = state["workflow"];
+                await Task.Delay(1000).ConfigureAwait(false);
+            }
+            // now we can shut down...
+            await AgentsApp.ShutdownAsync().ConfigureAwait(true);
         }
         public async Task<string> SayHello(string ask)
         {
             var response = $"\n\n\n\n***************Hello {ask}**********************\n\n\n\n";
             return response;
         }
     }
-    public interface ISayHello
-    {
-        public Task<string> SayHello(string ask);
-    }
 }
diff --git a/dotnet/src/Microsoft.AutoGen/Agents/Services/AgentWorker.cs b/dotnet/src/Microsoft.AutoGen/Agents/Services/AgentWorker.cs
@@ -66,11 +66,8 @@ public ValueTask SendResponseAsync(RpcResponse response, CancellationToken cance
     public ValueTask StoreAsync(AgentState value, CancellationToken cancellationToken = default)
     {
         var agentId = value.AgentId ?? throw new InvalidOperationException("AgentId is required when saving AgentState.");
-        var response = _agentStates.TryAdd(agentId.ToString(), value);
-        if (!response)
-        {
-            throw new InvalidOperationException($"Error saving AgentState for AgentId {agentId}.");
-        }
+        // add or update _agentStates with the new state
+        var response = _agentStates.AddOrUpdate(agentId.ToString(), value, (key, oldValue) => value);
         return ValueTask.CompletedTask;
     }
     public ValueTask<AgentState> ReadAsync(AgentId agentId, CancellationToken cancellationToken = default)

diff --git a/protos/agent_events.proto b/protos/agent_events.proto
@@ -35,3 +35,6 @@ message ConversationClosed {
   string user_id = 1;
   string user_message = 2;
 }
+message Shutdown {
+      string message = 1;
+}