yujunjun/autogen
1
0
Fork 0
mirror of https://github.com/microsoft/autogen.git synced 2025-10-14 01:18:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: save/load test for dotnet agents (#5284)

Browse Source
This commit is contained in:
Griffin Bassman 2025-02-06 17:09:26 -05:00 committed by GitHub
parent 25f26a338b
commit c8e4ad8242
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 230 additions and 140 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
dotnet/src/Microsoft.AutoGen/Contracts/AgentProxy.cs
View File

@ -1,6 +1,8 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// AgentProxy.cs
using System.Text.Json;
namespace Microsoft.AutoGen.Contracts;
/// <summary>
@ -55,7 +57,7 @@ public class AgentProxy(AgentId agentId, IAgentRuntime runtime)
/// </summary>
/// <param name="state">A dictionary representing the state of the agent. Must be JSON serializable.</param>
/// <returns>A task representing the asynchronous operation.</returns>
public ValueTask LoadStateAsync(IDictionary<string, object> state)
public ValueTask LoadStateAsync(IDictionary<string, JsonElement> state)
{
return this.runtime.LoadAgentStateAsync(this.Id, state);
}
@ -64,7 +66,7 @@ public class AgentProxy(AgentId agentId, IAgentRuntime runtime)
/// Saves the state of the agent. The result must be JSON serializable.
/// </summary>
/// <returns>A task representing the asynchronous operation, returning a dictionary containing the saved state.</returns>
public ValueTask<IDictionary<string, object>> SaveStateAsync()
public ValueTask<IDictionary<string, JsonElement>> SaveStateAsync()
{
return this.runtime.SaveAgentStateAsync(this.Id);
}

2
dotnet/src/Microsoft.AutoGen/Contracts/IAgentRuntime.cs
View File

@ -1,7 +1,7 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// IAgentRuntime.cs
using StateDict = System.Collections.Generic.IDictionary<string, object>;
using StateDict = System.Collections.Generic.IDictionary<string, System.Text.Json.JsonElement>;
namespace Microsoft.AutoGen.Contracts;

2
dotnet/src/Microsoft.AutoGen/Contracts/ISaveState.cs
View File

@ -1,7 +1,7 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// ISaveState.cs
using StateDict = System.Collections.Generic.IDictionary<string, object>;
using StateDict = System.Collections.Generic.IDictionary<string, System.Text.Json.JsonElement>;
namespace Microsoft.AutoGen.Contracts;

39
dotnet/src/Microsoft.AutoGen/Core.Grpc/GrpcAgentRuntime.cs
View File

@ -2,6 +2,7 @@
// GrpcAgentRuntime.cs
using System.Collections.Concurrent;
using System.Text.Json;
using Grpc.Core;
using Microsoft.AutoGen.Contracts;
using Microsoft.AutoGen.Protobuf;
@ -319,13 +320,13 @@ public sealed class GrpcAgentRuntime : IHostedService, IAgentRuntime, IMessageSi
public ValueTask<Contracts.AgentId> GetAgentAsync(string agent, string key = "default", bool lazy = true)
=> this.GetAgentAsync(new Contracts.AgentId(agent, key), lazy);
public async ValueTask<IDictionary<string, object>> SaveAgentStateAsync(Contracts.AgentId agentId)
public async ValueTask<IDictionary<string, JsonElement>> SaveAgentStateAsync(Contracts.AgentId agentId)
{
IHostableAgent agent = await this._agentsContainer.EnsureAgentAsync(agentId);
return await agent.SaveStateAsync();
}
public async ValueTask LoadAgentStateAsync(Contracts.AgentId agentId, IDictionary<string, object> state)
public async ValueTask LoadAgentStateAsync(Contracts.AgentId agentId, IDictionary<string, JsonElement> state)
{
IHostableAgent agent = await this._agentsContainer.EnsureAgentAsync(agentId);
await agent.LoadStateAsync(state);
@ -375,37 +376,41 @@ public sealed class GrpcAgentRuntime : IHostedService, IAgentRuntime, IMessageSi
return ValueTask.FromResult(new AgentProxy(agentId, this));
}
public async ValueTask<IDictionary<string, object>> SaveStateAsync()
{
Dictionary<string, object> state = new();
foreach (var agent in this._agentsContainer.LiveAgents)
{
state[agent.Id.ToString()] = await agent.SaveStateAsync();
}
return state;
}
public async ValueTask LoadStateAsync(IDictionary<string, object> state)
public async ValueTask LoadStateAsync(IDictionary<string, JsonElement> state)
{
HashSet<AgentType> registeredTypes = this._agentsContainer.RegisteredAgentTypes;
foreach (var agentIdStr in state.Keys)
{
Contracts.AgentId agentId = Contracts.AgentId.FromStr(agentIdStr);
if (state[agentIdStr] is not IDictionary<string, object> agentStateDict)
if (state[agentIdStr].ValueKind != JsonValueKind.Object)
{
throw new Exception($"Agent state for {agentId} is not a {typeof(IDictionary<string, object>)}: {state[agentIdStr].GetType()}");
throw new Exception($"Agent state for {agentId} is not a valid JSON object.");
}
var agentState = JsonSerializer.Deserialize<IDictionary<string, JsonElement>>(state[agentIdStr].GetRawText())
?? throw new Exception($"Failed to deserialize state for {agentId}.");
if (registeredTypes.Contains(agentId.Type))
{
IHostableAgent agent = await this._agentsContainer.EnsureAgentAsync(agentId);
await agent.LoadStateAsync(agentStateDict);
await agent.LoadStateAsync(agentState);
}
}
}
public async ValueTask<IDictionary<string, JsonElement>> SaveStateAsync()
{
Dictionary<string, JsonElement> state = new();
foreach (var agent in this._agentsContainer.LiveAgents)
{
var agentState = await agent.SaveStateAsync();
state[agent.Id.ToString()] = JsonSerializer.SerializeToElement(agentState);
}
return state;
}
public async ValueTask OnMessageAsync(Message message, CancellationToken cancellation = default)
{
switch (message.MessageCase)

7
dotnet/src/Microsoft.AutoGen/Core/BaseAgent.cs
View File

@ -4,6 +4,7 @@
using System.Diagnostics;
using System.Reflection;
using System.Text.Json;
using Microsoft.AutoGen.Contracts;
using Microsoft.Extensions.Logging;
@ -92,11 +93,11 @@ public abstract class BaseAgent : IAgent, IHostableAgent
return null;
}
public virtual ValueTask<IDictionary<string, object>> SaveStateAsync()
public virtual ValueTask<IDictionary<string, JsonElement>> SaveStateAsync()
{
return ValueTask.FromResult<IDictionary<string, object>>(new Dictionary<string, object>());
return ValueTask.FromResult<IDictionary<string, JsonElement>>(new Dictionary<string, JsonElement>());
}
public virtual ValueTask LoadStateAsync(IDictionary<string, object> state)
public virtual ValueTask LoadStateAsync(IDictionary<string, JsonElement> state)
{
return ValueTask.CompletedTask;
}

26
dotnet/src/Microsoft.AutoGen/Core/InProcessRuntime.cs
View File

@ -3,6 +3,7 @@
using System.Collections.Concurrent;
using System.Diagnostics;
using System.Text.Json;
using Microsoft.AutoGen.Contracts;
using Microsoft.Extensions.Hosting;
@ -12,7 +13,7 @@ public sealed class InProcessRuntime : IAgentRuntime, IHostedService
{
public bool DeliverToSelf { get; set; } //= false;
Dictionary<AgentId, IHostableAgent> agentInstances = new();
internal Dictionary<AgentId, IHostableAgent> agentInstances = new();
Dictionary<string, ISubscriptionDefinition> subscriptions = new();
Dictionary<AgentType, Func<AgentId, IAgentRuntime, ValueTask<IHostableAgent>>> agentFactories = new();
@ -152,13 +153,13 @@ public sealed class InProcessRuntime : IAgentRuntime, IHostedService
return agent.Metadata;
}
public async ValueTask LoadAgentStateAsync(AgentId agentId, IDictionary<string, object> state)
public async ValueTask LoadAgentStateAsync(AgentId agentId, IDictionary<string, JsonElement> state)
{
IHostableAgent agent = await this.EnsureAgentAsync(agentId);
await agent.LoadStateAsync(state);
}
public async ValueTask<IDictionary<string, object>> SaveAgentStateAsync(AgentId agentId)
public async ValueTask<IDictionary<string, JsonElement>> SaveAgentStateAsync(AgentId agentId)
{
IHostableAgent agent = await this.EnsureAgentAsync(agentId);
return await agent.SaveStateAsync();
@ -187,16 +188,21 @@ public sealed class InProcessRuntime : IAgentRuntime, IHostedService
return ValueTask.CompletedTask;
}
public async ValueTask LoadStateAsync(IDictionary<string, object> state)
public async ValueTask LoadStateAsync(IDictionary<string, JsonElement> state)
{
foreach (var agentIdStr in state.Keys)
{
AgentId agentId = AgentId.FromStr(agentIdStr);
if (state[agentIdStr] is not IDictionary<string, object> agentState)
if (state[agentIdStr].ValueKind != JsonValueKind.Object)
{
throw new Exception($"Agent state for {agentId} is not a {typeof(IDictionary<string, object>)}: {state[agentIdStr].GetType()}");
throw new Exception($"Agent state for {agentId} is not a valid JSON object.");
}
// Deserialize before using
var agentState = JsonSerializer.Deserialize<IDictionary<string, JsonElement>>(state[agentIdStr].GetRawText())
?? throw new Exception($"Failed to deserialize state for {agentId}.");
if (this.agentFactories.ContainsKey(agentId.Type))
{
IHostableAgent agent = await this.EnsureAgentAsync(agentId);
@ -205,14 +211,14 @@ public sealed class InProcessRuntime : IAgentRuntime, IHostedService
}
}
public async ValueTask<IDictionary<string, object>> SaveStateAsync()
public async ValueTask<IDictionary<string, JsonElement>> SaveStateAsync()
{
Dictionary<string, object> state = new();
Dictionary<string, JsonElement> state = new();
foreach (var agentId in this.agentInstances.Keys)
{
state[agentId.ToString()] = await this.agentInstances[agentId].SaveStateAsync();
var agentState = await this.agentInstances[agentId].SaveStateAsync();
state[agentId.ToString()] = JsonSerializer.SerializeToElement(agentState);
}
return state;
}

6
dotnet/src/Microsoft.AutoGen/Core/Properties/AssemblyInfo.cs Normal file
View File

@ -0,0 +1,6 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// AssemblyInfo.cs
using System.Runtime.CompilerServices;
[assembly: InternalsVisibleTo("Microsoft.AutoGen.Core.Tests, PublicKey=0024000004800000940000000602000000240000525341310004000001000100f1d038d0b85ae392ad72011df91e9343b0b5df1bb8080aa21b9424362d696919e0e9ac3a8bca24e283e10f7a569c6f443e1d4e3ebc84377c87ca5caa562e80f9932bf5ea91b7862b538e13b8ba91c7565cf0e8dfeccfea9c805ae3bda044170ecc7fc6f147aeeac422dd96aeb9eb1f5a5882aa650efe2958f2f8107d2038f2ab")]

83
dotnet/test/Microsoft.AutoGen.Core.Tests/AgentRuntimeTests.cs
View File

@ -1,83 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// AgentRuntimeTests.cs
using FluentAssertions;
using Microsoft.AutoGen.Contracts;
using Microsoft.Extensions.Logging;
using Xunit;
namespace Microsoft.AutoGen.Core.Tests;
[Trait("Category", "UnitV2")]
public class AgentRuntimeTests()
{
// Agent will not deliver to self will success when runtime.DeliverToSelf is false (default)
[Fact]
public async Task RuntimeAgentPublishToSelfDefaultNoSendTest()
{
var runtime = new InProcessRuntime();
await runtime.StartAsync();
Logger<BaseAgent> logger = new(new LoggerFactory());
SubscribedSelfPublishAgent agent = null!;
await runtime.RegisterAgentFactoryAsync("MyAgent", (id, runtime) =>
{
agent = new SubscribedSelfPublishAgent(id, runtime, logger);
return ValueTask.FromResult(agent);
});
// Ensure the agent is actually created
AgentId agentId = await runtime.GetAgentAsync("MyAgent", lazy: false);
// Validate agent ID
agentId.Should().Be(agent.Id, "Agent ID should match the registered agent");
await runtime.RegisterImplicitAgentSubscriptionsAsync<SubscribedSelfPublishAgent>("MyAgent");
var topicType = "TestTopic";
await runtime.PublishMessageAsync("SelfMessage", new TopicId(topicType)).ConfigureAwait(true);
await runtime.RunUntilIdleAsync();
// Agent has default messages and could not publish to self
agent.Text.Source.Should().Be("DefaultTopic");
agent.Text.Content.Should().Be("DefaultContent");
}
// Agent delivery to self will success when runtime.DeliverToSelf is true
[Fact]
public async Task RuntimeAgentPublishToSelfDeliverToSelfTrueTest()
{
var runtime = new InProcessRuntime();
runtime.DeliverToSelf = true;
await runtime.StartAsync();
Logger<BaseAgent> logger = new(new LoggerFactory());
SubscribedSelfPublishAgent agent = null!;
await runtime.RegisterAgentFactoryAsync("MyAgent", (id, runtime) =>
{
agent = new SubscribedSelfPublishAgent(id, runtime, logger);
return ValueTask.FromResult(agent);
});
// Ensure the agent is actually created
AgentId agentId = await runtime.GetAgentAsync("MyAgent", lazy: false);
// Validate agent ID
agentId.Should().Be(agent.Id, "Agent ID should match the registered agent");
await runtime.RegisterImplicitAgentSubscriptionsAsync<SubscribedSelfPublishAgent>("MyAgent");
var topicType = "TestTopic";
await runtime.PublishMessageAsync("SelfMessage", new TopicId(topicType)).ConfigureAwait(true);
await runtime.RunUntilIdleAsync();
// Agent sucessfully published to self
agent.Text.Source.Should().Be("TestTopic");
agent.Text.Content.Should().Be("SelfMessage");
}
}

23
dotnet/test/Microsoft.AutoGen.Core.Tests/AgentTests.cs
View File

@ -54,7 +54,7 @@ public class AgentTests()
return ValueTask.FromResult(agent);
});
// Ensure the agent is actually created
// Ensure the agent id is registered
AgentId agentId = await runtime.GetAgentAsync("MyAgent", lazy: false);
// Validate agent ID
@ -146,25 +146,4 @@ public class AgentTests()
Assert.True(agent.ReceivedItems.Count == 1);
}
[Fact]
public async Task AgentShouldSaveStateCorrectlyTest()
{
var runtime = new InProcessRuntime();
await runtime.StartAsync();
Logger<BaseAgent> logger = new(new LoggerFactory());
TestAgent agent = new TestAgent(new AgentId("TestType", "TestKey"), runtime, logger);
var state = await agent.SaveStateAsync();
// Ensure state is a dictionary
state.Should().NotBeNull();
state.Should().BeOfType<Dictionary<string, object>>();
state.Should().BeEmpty("Default SaveStateAsync should return an empty dictionary.");
// Add a sample value and verify it updates correctly
state["testKey"] = "testValue";
state.Should().ContainKey("testKey").WhoseValue.Should().Be("testValue");
}
}

141
dotnet/test/Microsoft.AutoGen.Core.Tests/InProcessRuntimeTests.cs Normal file
View File

@ -0,0 +1,141 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// InProcessRuntimeTests.cs
using System.Text.Json;
using FluentAssertions;
using Microsoft.AutoGen.Contracts;
using Microsoft.Extensions.Logging;
using Xunit;
namespace Microsoft.AutoGen.Core.Tests;
[Trait("Category", "UnitV2")]
public class InProcessRuntimeTests()
{
// Agent will not deliver to self will success when runtime.DeliverToSelf is false (default)
[Fact]
public async Task RuntimeAgentPublishToSelfDefaultNoSendTest()
{
var runtime = new InProcessRuntime();
await runtime.StartAsync();
Logger<BaseAgent> logger = new(new LoggerFactory());
SubscribedSelfPublishAgent agent = null!;
await runtime.RegisterAgentFactoryAsync("MyAgent", (id, runtime) =>
{
agent = new SubscribedSelfPublishAgent(id, runtime, logger);
return ValueTask.FromResult(agent);
});
// Ensure the agent is actually created
AgentId agentId = await runtime.GetAgentAsync("MyAgent", lazy: false);
// Validate agent ID
agentId.Should().Be(agent.Id, "Agent ID should match the registered agent");
await runtime.RegisterImplicitAgentSubscriptionsAsync<SubscribedSelfPublishAgent>("MyAgent");
var topicType = "TestTopic";
await runtime.PublishMessageAsync("SelfMessage", new TopicId(topicType)).ConfigureAwait(true);
await runtime.RunUntilIdleAsync();
// Agent has default messages and could not publish to self
agent.Text.Source.Should().Be("DefaultTopic");
agent.Text.Content.Should().Be("DefaultContent");
}
// Agent delivery to self will success when runtime.DeliverToSelf is true
[Fact]
public async Task RuntimeAgentPublishToSelfDeliverToSelfTrueTest()
{
var runtime = new InProcessRuntime();
runtime.DeliverToSelf = true;
await runtime.StartAsync();
Logger<BaseAgent> logger = new(new LoggerFactory());
SubscribedSelfPublishAgent agent = null!;
await runtime.RegisterAgentFactoryAsync("MyAgent", (id, runtime) =>
{
agent = new SubscribedSelfPublishAgent(id, runtime, logger);
return ValueTask.FromResult(agent);
});
// Ensure the agent is actually created
AgentId agentId = await runtime.GetAgentAsync("MyAgent", lazy: false);
// Validate agent ID
agentId.Should().Be(agent.Id, "Agent ID should match the registered agent");
await runtime.RegisterImplicitAgentSubscriptionsAsync<SubscribedSelfPublishAgent>("MyAgent");
var topicType = "TestTopic";
await runtime.PublishMessageAsync("SelfMessage", new TopicId(topicType)).ConfigureAwait(true);
await runtime.RunUntilIdleAsync();
// Agent sucessfully published to self
agent.Text.Source.Should().Be("TestTopic");
agent.Text.Content.Should().Be("SelfMessage");
}
[Fact]
public async Task RuntimeShouldSaveLoadStateCorrectlyTest()
{
// Create a runtime and register an agent
var runtime = new InProcessRuntime();
await runtime.StartAsync();
Logger<BaseAgent> logger = new(new LoggerFactory());
SubscribedSaveLoadAgent agent = null!;
await runtime.RegisterAgentFactoryAsync("MyAgent", (id, runtime) =>
{
agent = new SubscribedSaveLoadAgent(id, runtime, logger);
return ValueTask.FromResult(agent);
});
// Get agent ID and instantiate agent by publishing
AgentId agentId = await runtime.GetAgentAsync("MyAgent", lazy: true);
await runtime.RegisterImplicitAgentSubscriptionsAsync<SubscribedSaveLoadAgent>("MyAgent");
var topicType = "TestTopic";
await runtime.PublishMessageAsync(new TextMessage { Source = topicType, Content = "test" }, new TopicId(topicType)).ConfigureAwait(true);
await runtime.RunUntilIdleAsync();
agent.ReceivedMessages.Any().Should().BeTrue("Agent should receive messages when subscribed.");
// Save the state
var savedState = await runtime.SaveStateAsync();
// Ensure saved state contains the agent's state
savedState.Should().ContainKey(agentId.ToString());
// Ensure the agent's state is stored as a valid JSON object
savedState[agentId.ToString()].ValueKind.Should().Be(JsonValueKind.Object, "Agent state should be stored as a JSON object");
// Serialize and Deserialize the state to simulate persistence
string json = JsonSerializer.Serialize(savedState);
json.Should().NotBeNullOrEmpty("Serialized state should not be empty");
var deserializedState = JsonSerializer.Deserialize<IDictionary<string, JsonElement>>(json)
?? throw new Exception("Deserialized state is unexpectedly null");
deserializedState.Should().ContainKey(agentId.ToString());
// Start new runtime and restore the state
var newRuntime = new InProcessRuntime();
await newRuntime.StartAsync();
await newRuntime.RegisterAgentFactoryAsync("MyAgent", (id, runtime) =>
{
agent = new SubscribedSaveLoadAgent(id, runtime, logger);
return ValueTask.FromResult(agent);
});
await newRuntime.RegisterImplicitAgentSubscriptionsAsync<SubscribedSaveLoadAgent>("MyAgent");
// Show that no agent instances exist in the new runtime
newRuntime.agentInstances.Count.Should().Be(0, "Agent should be registered in the new runtime");
// Load the state into the new runtime and show that agent is now instantiated
await newRuntime.LoadStateAsync(deserializedState);
newRuntime.agentInstances.Count.Should().Be(1, "Agent should be registered in the new runtime");
newRuntime.agentInstances.Should().ContainKey(agentId, "Agent should be loaded into the new runtime");
}
}

35
dotnet/test/Microsoft.AutoGen.Core.Tests/TestAgent.cs
View File

@ -1,6 +1,7 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// TestAgent.cs
using System.Text.Json;
using Microsoft.AutoGen.Contracts;
using Microsoft.Extensions.Logging;
@ -59,7 +60,7 @@ public class TestAgent(AgentId id,
/// Key: source
/// Value: message
/// </summary>
private readonly Dictionary<string, object> _receivedMessages = new();
protected Dictionary<string, object> _receivedMessages = new();
public Dictionary<string, object> ReceivedMessages => _receivedMessages;
}
@ -73,6 +74,38 @@ public class SubscribedAgent : TestAgent
}
}
[TypeSubscription("TestTopic")]
public class SubscribedSaveLoadAgent : TestAgent
{
public SubscribedSaveLoadAgent(AgentId id,
IAgentRuntime runtime,
Logger<BaseAgent>? logger = null) : base(id, runtime, logger)
{
}
public override ValueTask<IDictionary<string, JsonElement>> SaveStateAsync()
{
var jsonSafeDictionary = _receivedMessages.ToDictionary(
kvp => kvp.Key,
kvp => JsonSerializer.SerializeToElement(kvp.Value) // Convert each object to JsonElement
);
return ValueTask.FromResult<IDictionary<string, JsonElement>>(jsonSafeDictionary);
}
public override ValueTask LoadStateAsync(IDictionary<string, JsonElement> state)
{
_receivedMessages.Clear();
foreach (var kvp in state)
{
_receivedMessages[kvp.Key] = kvp.Value.Deserialize<object>() ?? throw new Exception($"Failed to deserialize key: {kvp.Key}");
}
return ValueTask.CompletedTask;
}
}
/// <summary>
/// The test agent showing an agent that subscribes to itself.
/// </summary>