mirror of
https://github.com/microsoft/autogen.git
synced 2025-10-14 09:29:40 +00:00
feat: save/load test for dotnet agents (#5284)
This commit is contained in:
parent
25f26a338b
commit
c8e4ad8242
@ -1,6 +1,8 @@
|
|||||||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||||
// AgentProxy.cs
|
// AgentProxy.cs
|
||||||
|
|
||||||
|
using System.Text.Json;
|
||||||
|
|
||||||
namespace Microsoft.AutoGen.Contracts;
|
namespace Microsoft.AutoGen.Contracts;
|
||||||
|
|
||||||
/// <summary>
|
/// <summary>
|
||||||
@ -55,7 +57,7 @@ public class AgentProxy(AgentId agentId, IAgentRuntime runtime)
|
|||||||
/// </summary>
|
/// </summary>
|
||||||
/// <param name="state">A dictionary representing the state of the agent. Must be JSON serializable.</param>
|
/// <param name="state">A dictionary representing the state of the agent. Must be JSON serializable.</param>
|
||||||
/// <returns>A task representing the asynchronous operation.</returns>
|
/// <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);
|
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.
|
/// Saves the state of the agent. The result must be JSON serializable.
|
||||||
/// </summary>
|
/// </summary>
|
||||||
/// <returns>A task representing the asynchronous operation, returning a dictionary containing the saved state.</returns>
|
/// <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);
|
return this.runtime.SaveAgentStateAsync(this.Id);
|
||||||
}
|
}
|
||||||
|
@ -1,7 +1,7 @@
|
|||||||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||||
// IAgentRuntime.cs
|
// 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;
|
namespace Microsoft.AutoGen.Contracts;
|
||||||
|
|
||||||
|
@ -1,7 +1,7 @@
|
|||||||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||||
// ISaveState.cs
|
// 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;
|
namespace Microsoft.AutoGen.Contracts;
|
||||||
|
|
||||||
|
@ -2,6 +2,7 @@
|
|||||||
// GrpcAgentRuntime.cs
|
// GrpcAgentRuntime.cs
|
||||||
|
|
||||||
using System.Collections.Concurrent;
|
using System.Collections.Concurrent;
|
||||||
|
using System.Text.Json;
|
||||||
using Grpc.Core;
|
using Grpc.Core;
|
||||||
using Microsoft.AutoGen.Contracts;
|
using Microsoft.AutoGen.Contracts;
|
||||||
using Microsoft.AutoGen.Protobuf;
|
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)
|
public ValueTask<Contracts.AgentId> GetAgentAsync(string agent, string key = "default", bool lazy = true)
|
||||||
=> this.GetAgentAsync(new Contracts.AgentId(agent, key), lazy);
|
=> 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);
|
IHostableAgent agent = await this._agentsContainer.EnsureAgentAsync(agentId);
|
||||||
return await agent.SaveStateAsync();
|
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);
|
IHostableAgent agent = await this._agentsContainer.EnsureAgentAsync(agentId);
|
||||||
await agent.LoadStateAsync(state);
|
await agent.LoadStateAsync(state);
|
||||||
@ -375,37 +376,41 @@ public sealed class GrpcAgentRuntime : IHostedService, IAgentRuntime, IMessageSi
|
|||||||
return ValueTask.FromResult(new AgentProxy(agentId, this));
|
return ValueTask.FromResult(new AgentProxy(agentId, this));
|
||||||
}
|
}
|
||||||
|
|
||||||
public async ValueTask<IDictionary<string, object>> SaveStateAsync()
|
public async ValueTask LoadStateAsync(IDictionary<string, JsonElement> state)
|
||||||
{
|
|
||||||
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)
|
|
||||||
{
|
{
|
||||||
HashSet<AgentType> registeredTypes = this._agentsContainer.RegisteredAgentTypes;
|
HashSet<AgentType> registeredTypes = this._agentsContainer.RegisteredAgentTypes;
|
||||||
|
|
||||||
foreach (var agentIdStr in state.Keys)
|
foreach (var agentIdStr in state.Keys)
|
||||||
{
|
{
|
||||||
Contracts.AgentId agentId = Contracts.AgentId.FromStr(agentIdStr);
|
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))
|
if (registeredTypes.Contains(agentId.Type))
|
||||||
{
|
{
|
||||||
IHostableAgent agent = await this._agentsContainer.EnsureAgentAsync(agentId);
|
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)
|
public async ValueTask OnMessageAsync(Message message, CancellationToken cancellation = default)
|
||||||
{
|
{
|
||||||
switch (message.MessageCase)
|
switch (message.MessageCase)
|
||||||
|
@ -4,6 +4,7 @@
|
|||||||
|
|
||||||
using System.Diagnostics;
|
using System.Diagnostics;
|
||||||
using System.Reflection;
|
using System.Reflection;
|
||||||
|
using System.Text.Json;
|
||||||
using Microsoft.AutoGen.Contracts;
|
using Microsoft.AutoGen.Contracts;
|
||||||
using Microsoft.Extensions.Logging;
|
using Microsoft.Extensions.Logging;
|
||||||
|
|
||||||
@ -92,11 +93,11 @@ public abstract class BaseAgent : IAgent, IHostableAgent
|
|||||||
return null;
|
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;
|
return ValueTask.CompletedTask;
|
||||||
}
|
}
|
||||||
|
@ -3,6 +3,7 @@
|
|||||||
|
|
||||||
using System.Collections.Concurrent;
|
using System.Collections.Concurrent;
|
||||||
using System.Diagnostics;
|
using System.Diagnostics;
|
||||||
|
using System.Text.Json;
|
||||||
using Microsoft.AutoGen.Contracts;
|
using Microsoft.AutoGen.Contracts;
|
||||||
using Microsoft.Extensions.Hosting;
|
using Microsoft.Extensions.Hosting;
|
||||||
|
|
||||||
@ -12,7 +13,7 @@ public sealed class InProcessRuntime : IAgentRuntime, IHostedService
|
|||||||
{
|
{
|
||||||
public bool DeliverToSelf { get; set; } //= false;
|
public bool DeliverToSelf { get; set; } //= false;
|
||||||
|
|
||||||
Dictionary<AgentId, IHostableAgent> agentInstances = new();
|
internal Dictionary<AgentId, IHostableAgent> agentInstances = new();
|
||||||
Dictionary<string, ISubscriptionDefinition> subscriptions = new();
|
Dictionary<string, ISubscriptionDefinition> subscriptions = new();
|
||||||
Dictionary<AgentType, Func<AgentId, IAgentRuntime, ValueTask<IHostableAgent>>> agentFactories = new();
|
Dictionary<AgentType, Func<AgentId, IAgentRuntime, ValueTask<IHostableAgent>>> agentFactories = new();
|
||||||
|
|
||||||
@ -152,13 +153,13 @@ public sealed class InProcessRuntime : IAgentRuntime, IHostedService
|
|||||||
return agent.Metadata;
|
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);
|
IHostableAgent agent = await this.EnsureAgentAsync(agentId);
|
||||||
await agent.LoadStateAsync(state);
|
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);
|
IHostableAgent agent = await this.EnsureAgentAsync(agentId);
|
||||||
return await agent.SaveStateAsync();
|
return await agent.SaveStateAsync();
|
||||||
@ -187,16 +188,21 @@ public sealed class InProcessRuntime : IAgentRuntime, IHostedService
|
|||||||
return ValueTask.CompletedTask;
|
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)
|
foreach (var agentIdStr in state.Keys)
|
||||||
{
|
{
|
||||||
AgentId agentId = AgentId.FromStr(agentIdStr);
|
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))
|
if (this.agentFactories.ContainsKey(agentId.Type))
|
||||||
{
|
{
|
||||||
IHostableAgent agent = await this.EnsureAgentAsync(agentId);
|
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)
|
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;
|
return state;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -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")]
|
@ -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");
|
|
||||||
}
|
|
||||||
}
|
|
@ -54,7 +54,7 @@ public class AgentTests()
|
|||||||
return ValueTask.FromResult(agent);
|
return ValueTask.FromResult(agent);
|
||||||
});
|
});
|
||||||
|
|
||||||
// Ensure the agent is actually created
|
// Ensure the agent id is registered
|
||||||
AgentId agentId = await runtime.GetAgentAsync("MyAgent", lazy: false);
|
AgentId agentId = await runtime.GetAgentAsync("MyAgent", lazy: false);
|
||||||
|
|
||||||
// Validate agent ID
|
// Validate agent ID
|
||||||
@ -146,25 +146,4 @@ public class AgentTests()
|
|||||||
|
|
||||||
Assert.True(agent.ReceivedItems.Count == 1);
|
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");
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
@ -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");
|
||||||
|
}
|
||||||
|
}
|
@ -1,6 +1,7 @@
|
|||||||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||||
// TestAgent.cs
|
// TestAgent.cs
|
||||||
|
|
||||||
|
using System.Text.Json;
|
||||||
using Microsoft.AutoGen.Contracts;
|
using Microsoft.AutoGen.Contracts;
|
||||||
using Microsoft.Extensions.Logging;
|
using Microsoft.Extensions.Logging;
|
||||||
|
|
||||||
@ -59,7 +60,7 @@ public class TestAgent(AgentId id,
|
|||||||
/// Key: source
|
/// Key: source
|
||||||
/// Value: message
|
/// Value: message
|
||||||
/// </summary>
|
/// </summary>
|
||||||
private readonly Dictionary<string, object> _receivedMessages = new();
|
protected Dictionary<string, object> _receivedMessages = new();
|
||||||
public Dictionary<string, object> ReceivedMessages => _receivedMessages;
|
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>
|
/// <summary>
|
||||||
/// The test agent showing an agent that subscribes to itself.
|
/// The test agent showing an agent that subscribes to itself.
|
||||||
/// </summary>
|
/// </summary>
|
||||||
|
Loading…
x
Reference in New Issue
Block a user