diff --git a/.github/workflows/dotnet-build-and-test.yml b/.github/workflows/dotnet-build-and-test.yml index 97d1d60b3a..09a6caaf61 100644 --- a/.github/workflows/dotnet-build-and-test.yml +++ b/.github/workflows/dotnet-build-and-test.yml @@ -124,6 +124,16 @@ jobs: popd rm -rf "$TEMP_DIR" + # Start Cosmos DB Emulator for Cosmos-based unit tests (only on Windows) + - name: Start Azure Cosmos DB Emulator + if: runner.os == 'Windows' + shell: pwsh + run: | + Write-Host "Launching Azure Cosmos DB Emulator" + Import-Module "$env:ProgramFiles\Azure Cosmos DB Emulator\PSModules\Microsoft.Azure.CosmosDB.Emulator" + Start-CosmosDbEmulator -NoUI -Key "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==" + echo "COSMOS_EMULATOR_AVAILABLE=true" >> $env:GITHUB_ENV + - name: Run Unit Tests shell: bash run: | @@ -143,6 +153,10 @@ jobs: echo "Skipping $project - does not support target framework ${{ matrix.targetFramework }} (supports: $target_frameworks)" fi done + env: + # Cosmos DB Emulator connection settings + COSMOSDB_ENDPOINT: https://localhost:8081 + COSMOSDB_KEY: C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw== - name: Log event name and matrix integration-tests shell: bash @@ -181,6 +195,9 @@ jobs: fi done env: + # Cosmos DB Emulator connection settings + COSMOSDB_ENDPOINT: https://localhost:8081 + COSMOSDB_KEY: C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw== # OpenAI Models OpenAI__ApiKey: ${{ secrets.OPENAI__APIKEY }} OpenAI__ChatModelId: ${{ vars.OPENAI__CHATMODELID }} @@ -200,7 +217,7 @@ jobs: # Generate test reports and check coverage - name: Generate test reports if: matrix.targetFramework == env.COVERAGE_FRAMEWORK - uses: danielpalme/ReportGenerator-GitHub-Action@5.4.18 + uses: danielpalme/ReportGenerator-GitHub-Action@5.5.0 with: reports: "./TestResults/Coverage/**/coverage.cobertura.xml" targetdir: "./TestResults/Reports" diff --git a/dotnet/Directory.Build.props b/dotnet/Directory.Build.props index b862c85072..54a125a13b 100644 --- a/dotnet/Directory.Build.props +++ b/dotnet/Directory.Build.props @@ -8,7 +8,7 @@ true latest enable - $(NoWarn);NU5128 + $(NoWarn);NU5128;CS8002 true net10.0;net9.0;net8.0 $(TargetFrameworksCore);netstandard2.0;net472 diff --git a/dotnet/Directory.Packages.props b/dotnet/Directory.Packages.props index 62793f39fd..68ea321902 100644 --- a/dotnet/Directory.Packages.props +++ b/dotnet/Directory.Packages.props @@ -25,6 +25,10 @@ + + + + @@ -129,6 +133,7 @@ + diff --git a/dotnet/agent-framework-dotnet.slnx b/dotnet/agent-framework-dotnet.slnx index f1efae7ba8..00763b09c8 100644 --- a/dotnet/agent-framework-dotnet.slnx +++ b/dotnet/agent-framework-dotnet.slnx @@ -343,6 +343,7 @@ + @@ -383,6 +384,7 @@ + diff --git a/dotnet/nuget/nuget-package.props b/dotnet/nuget/nuget-package.props index 2282c9ce13..ca719be243 100644 --- a/dotnet/nuget/nuget-package.props +++ b/dotnet/nuget/nuget-package.props @@ -2,9 +2,9 @@ 1.0.0 - $(VersionPrefix)-$(VersionSuffix).251114.1 - $(VersionPrefix)-preview.251114.1 - 1.0.0-preview.251114.1 + $(VersionPrefix)-$(VersionSuffix).251125.1 + $(VersionPrefix)-preview.251125.1 + 1.0.0-preview.251125.1 Debug;Release;Publish true diff --git a/dotnet/samples/GettingStarted/AgentProviders/Agent_With_AzureAIAgentsPersistent/README.md b/dotnet/samples/GettingStarted/AgentProviders/Agent_With_AzureAIAgentsPersistent/README.md index 9e981de64b..d6b5497601 100644 --- a/dotnet/samples/GettingStarted/AgentProviders/Agent_With_AzureAIAgentsPersistent/README.md +++ b/dotnet/samples/GettingStarted/AgentProviders/Agent_With_AzureAIAgentsPersistent/README.md @@ -1,3 +1,13 @@ +# Classic Foundry Agents + +This sample demonstrates how to create an agent using the classic Foundry Agents experience. + +# Classic vs New Foundry Agents + +Below is a comparison between the classic and new Foundry Agents approaches: + +[Migration Guide](https://learn.microsoft.com/en-us/azure/ai-foundry/agents/how-to/migrate?view=foundry) + # Prerequisites Before you begin, ensure you have the following prerequisites: diff --git a/dotnet/samples/GettingStarted/AgentProviders/Agent_With_AzureAIProject/Program.cs b/dotnet/samples/GettingStarted/AgentProviders/Agent_With_AzureAIProject/Program.cs index dd4a011e4d..2c2b9d1969 100644 --- a/dotnet/samples/GettingStarted/AgentProviders/Agent_With_AzureAIProject/Program.cs +++ b/dotnet/samples/GettingStarted/AgentProviders/Agent_With_AzureAIProject/Program.cs @@ -10,35 +10,34 @@ using Microsoft.Agents.AI; var endpoint = Environment.GetEnvironmentVariable("AZURE_FOUNDRY_PROJECT_ENDPOINT") ?? throw new InvalidOperationException("AZURE_FOUNDRY_PROJECT_ENDPOINT is not set."); var deploymentName = Environment.GetEnvironmentVariable("AZURE_FOUNDRY_PROJECT_DEPLOYMENT_NAME") ?? "gpt-4o-mini"; -const string JokerInstructions = "You are good at telling jokes."; const string JokerName = "JokerAgent"; // Get a client to create/retrieve/delete server side agents with Azure Foundry Agents. var aiProjectClient = new AIProjectClient(new Uri(endpoint), new AzureCliCredential()); // Define the agent you want to create. (Prompt Agent in this case) -var agentVersionCreationOptions = new AgentVersionCreationOptions(new PromptAgentDefinition(model: deploymentName) { Instructions = JokerInstructions }); +var agentVersionCreationOptions = new AgentVersionCreationOptions(new PromptAgentDefinition(model: deploymentName) { Instructions = "You are good at telling jokes." }); // Azure.AI.Agents SDK creates and manages agent by name and versions. // You can create a server side agent version with the Azure.AI.Agents SDK client below. -var agentVersion = aiProjectClient.Agents.CreateAgentVersion(agentName: JokerName, options: agentVersionCreationOptions); +var createdAgentVersion = aiProjectClient.Agents.CreateAgentVersion(agentName: JokerName, options: agentVersionCreationOptions); // Note: // agentVersion.Id = ":", // agentVersion.Version = , // agentVersion.Name = -// You can retrieve an AIAgent for a already created server side agent version. -AIAgent jokerAgentV1 = aiProjectClient.GetAIAgent(agentVersion); +// You can retrieve an AIAgent for an already created server side agent version. +AIAgent existingJokerAgent = aiProjectClient.GetAIAgent(createdAgentVersion); -// You can also create another AIAgent version (V2) by providing the same name with a different definition. -AIAgent jokerAgentV2 = aiProjectClient.CreateAIAgent(name: JokerName, model: deploymentName, instructions: JokerInstructions + "V2"); +// You can also create another AIAgent version by providing the same name with a different definition. +AIAgent newJokerAgent = aiProjectClient.CreateAIAgent(name: JokerName, model: deploymentName, instructions: "You are extremely hilarious at telling jokes."); // You can also get the AIAgent latest version just providing its name. AIAgent jokerAgentLatest = aiProjectClient.GetAIAgent(name: JokerName); -var latestVersion = jokerAgentLatest.GetService()!; +var latestAgentVersion = jokerAgentLatest.GetService()!; // The AIAgent version can be accessed via the GetService method. -Console.WriteLine($"Latest agent version id: {latestVersion.Id}"); +Console.WriteLine($"Latest agent version id: {latestAgentVersion.Id}"); // Once you have the AIAgent, you can invoke it like any other AIAgent. AgentThread thread = jokerAgentLatest.GetNewThread(); @@ -47,5 +46,5 @@ Console.WriteLine(await jokerAgentLatest.RunAsync("Tell me a joke about a pirate // This will use the same thread to continue the conversation. Console.WriteLine(await jokerAgentLatest.RunAsync("Now tell me a joke about a cat and a dog using last joke as the anchor.", thread)); -// Cleanup by agent name removes both agent versions created (jokerAgentV1 + jokerAgentV2). -aiProjectClient.Agents.DeleteAgent(jokerAgentV1.Name); +// Cleanup by agent name removes both agent versions created. +aiProjectClient.Agents.DeleteAgent(existingJokerAgent.Name); diff --git a/dotnet/samples/GettingStarted/AgentProviders/Agent_With_AzureAIProject/README.md b/dotnet/samples/GettingStarted/AgentProviders/Agent_With_AzureAIProject/README.md index 9e981de64b..7e4a28f6a1 100644 --- a/dotnet/samples/GettingStarted/AgentProviders/Agent_With_AzureAIProject/README.md +++ b/dotnet/samples/GettingStarted/AgentProviders/Agent_With_AzureAIProject/README.md @@ -1,3 +1,13 @@ +# New Foundry Agents + +This sample demonstrates how to create an agent using the new Foundry Agents experience. + +# Classic vs New Foundry Agents + +Below is a comparison between the classic and new Foundry Agents approaches: + +[Migration Guide](https://learn.microsoft.com/en-us/azure/ai-foundry/agents/how-to/migrate?view=foundry) + # Prerequisites Before you begin, ensure you have the following prerequisites: diff --git a/dotnet/samples/GettingStarted/FoundryAgents/FoundryAgents_Step01.1_Basics/Program.cs b/dotnet/samples/GettingStarted/FoundryAgents/FoundryAgents_Step01.1_Basics/Program.cs index 3c374d799f..9a7ee0736a 100644 --- a/dotnet/samples/GettingStarted/FoundryAgents/FoundryAgents_Step01.1_Basics/Program.cs +++ b/dotnet/samples/GettingStarted/FoundryAgents/FoundryAgents_Step01.1_Basics/Program.cs @@ -11,19 +11,17 @@ using Microsoft.Extensions.AI; string endpoint = Environment.GetEnvironmentVariable("AZURE_FOUNDRY_PROJECT_ENDPOINT") ?? throw new InvalidOperationException("AZURE_FOUNDRY_PROJECT_ENDPOINT is not set."); string deploymentName = Environment.GetEnvironmentVariable("AZURE_FOUNDRY_PROJECT_DEPLOYMENT_NAME") ?? "gpt-4o-mini"; -const string JokerInstructionsV1 = "You are good at telling jokes."; -const string JokerInstructionsV2 = "You are extremely hilarious at telling jokes."; const string JokerName = "JokerAgent"; // Get a client to create/retrieve/delete server side agents with Azure Foundry Agents. AIProjectClient aiProjectClient = new(new Uri(endpoint), new AzureCliCredential()); // Define the agent you want to create. (Prompt Agent in this case) -AgentVersionCreationOptions options = new(new PromptAgentDefinition(model: deploymentName) { Instructions = JokerInstructionsV1 }); +AgentVersionCreationOptions options = new(new PromptAgentDefinition(model: deploymentName) { Instructions = "You are good at telling jokes." }); // Azure.AI.Agents SDK creates and manages agent by name and versions. // You can create a server side agent version with the Azure.AI.Agents SDK client below. -AgentVersion agentVersion = aiProjectClient.Agents.CreateAgentVersion(agentName: JokerName, options); +AgentVersion createdAgentVersion = aiProjectClient.Agents.CreateAgentVersion(agentName: JokerName, options); // Note: // agentVersion.Id = ":", @@ -31,20 +29,20 @@ AgentVersion agentVersion = aiProjectClient.Agents.CreateAgentVersion(agentName: // agentVersion.Name = // You can retrieve an AIAgent for an already created server side agent version. -AIAgent jokerAgentV1 = aiProjectClient.GetAIAgent(agentVersion); +AIAgent existingJokerAgent = aiProjectClient.GetAIAgent(createdAgentVersion); -// You can also create another AIAgent version (V2) by providing the same name with a different definition/instruction. -AIAgent jokerAgentV2 = aiProjectClient.CreateAIAgent(name: JokerName, model: deploymentName, instructions: JokerInstructionsV2); +// You can also create another AIAgent version by providing the same name with a different definition/instruction. +AIAgent newJokerAgent = aiProjectClient.CreateAIAgent(name: JokerName, model: deploymentName, instructions: "You are extremely hilarious at telling jokes."); // You can also get the AIAgent latest version by just providing its name. AIAgent jokerAgentLatest = aiProjectClient.GetAIAgent(name: JokerName); -AgentVersion latestVersion = jokerAgentLatest.GetService()!; +AgentVersion latestAgentVersion = jokerAgentLatest.GetService()!; // The AIAgent version can be accessed via the GetService method. -Console.WriteLine($"Latest agent version id: {latestVersion.Id}"); +Console.WriteLine($"Latest agent version id: {latestAgentVersion.Id}"); // Once you have the AIAgent, you can invoke it like any other AIAgent. Console.WriteLine(await jokerAgentLatest.RunAsync("Tell me a joke about a pirate.")); -// Cleanup by agent name removes both agent versions created (jokerAgentV1 + jokerAgentV2). -await aiProjectClient.Agents.DeleteAgentAsync(jokerAgentV1.Name); +// Cleanup by agent name removes both agent versions created. +await aiProjectClient.Agents.DeleteAgentAsync(existingJokerAgent.Name); diff --git a/dotnet/samples/GettingStarted/FoundryAgents/README.md b/dotnet/samples/GettingStarted/FoundryAgents/README.md index 9369f5b34e..daeb2db8df 100644 --- a/dotnet/samples/GettingStarted/FoundryAgents/README.md +++ b/dotnet/samples/GettingStarted/FoundryAgents/README.md @@ -6,6 +6,15 @@ of Azure Foundry Agents and can be used with Azure Foundry as the AI provider. These samples showcase how to work with agents managed through Azure Foundry, including agent creation, versioning, multi-turn conversations, and advanced features like code interpretation and computer use. +## Classic vs New Foundry Agents + +> [!NOTE] +> Recently, Azure Foundry introduced a new and improved experience for creating and managing AI agents, which is the target of these samples. + +For more information about the previous classic agents and for what's new in Foundry Agents, see the [Foundry Agents migration documentation](https://learn.microsoft.com/en-us/azure/ai-foundry/agents/how-to/migrate?view=foundry). + +For a sample demonstrating how to use classic Foundry Agents, see the following: [Agent with Azure AI Persistent](../AgentProviders/Agent_With_AzureAIAgentsPersistent/README.md). + ## Getting started with Foundry Agents prerequisites Before you begin, ensure you have the following prerequisites: diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs new file mode 100644 index 0000000000..fff7f56fa5 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs @@ -0,0 +1,688 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; +using System.Linq; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Azure.Core; +using Microsoft.Azure.Cosmos; +using Microsoft.Extensions.AI; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Agents.AI; + +/// +/// Provides a Cosmos DB implementation of the abstract class. +/// +[RequiresUnreferencedCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with trimming.")] +[RequiresDynamicCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with NativeAOT.")] +public sealed class CosmosChatMessageStore : ChatMessageStore, IDisposable +{ + private readonly CosmosClient _cosmosClient; + private readonly Container _container; + private readonly bool _ownsClient; + private bool _disposed; + + // Hierarchical partition key support + private readonly string? _tenantId; + private readonly string? _userId; + private readonly PartitionKey _partitionKey; + private readonly bool _useHierarchicalPartitioning; + + /// + /// Cached JSON serializer options for .NET 9.0 compatibility. + /// + private static readonly JsonSerializerOptions s_defaultJsonOptions = CreateDefaultJsonOptions(); + + private static JsonSerializerOptions CreateDefaultJsonOptions() + { + var options = new JsonSerializerOptions(); +#if NET9_0_OR_GREATER + // Configure TypeInfoResolver for .NET 9.0 to enable JSON serialization + options.TypeInfoResolver = new System.Text.Json.Serialization.Metadata.DefaultJsonTypeInfoResolver(); +#endif + return options; + } + + /// + /// Gets or sets the maximum number of messages to return in a single query batch. + /// Default is 100 for optimal performance. + /// + public int MaxItemCount { get; set; } = 100; + + /// + /// Gets or sets the maximum number of items per transactional batch operation. + /// Default is 100, maximum allowed by Cosmos DB is 100. + /// + public int MaxBatchSize { get; set; } = 100; + + /// + /// Gets or sets the maximum number of messages to retrieve from the store. + /// This helps prevent exceeding LLM context windows in long conversations. + /// Default is null (no limit). When set, only the most recent messages are returned. + /// + public int? MaxMessagesToRetrieve { get; set; } + + /// + /// Gets or sets the Time-To-Live (TTL) in seconds for messages. + /// Default is 86400 seconds (24 hours). Set to null to disable TTL. + /// + public int? MessageTtlSeconds { get; set; } = 86400; + + /// + /// Gets the conversation ID associated with this message store. + /// + public string ConversationId { get; init; } + + /// + /// Gets the database ID associated with this message store. + /// + public string DatabaseId { get; init; } + + /// + /// Gets the container ID associated with this message store. + /// + public string ContainerId { get; init; } + + /// + /// Internal primary constructor used by all public constructors. + /// + /// The instance to use for Cosmos DB operations. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The unique identifier for this conversation thread. + /// Whether this instance owns the CosmosClient and should dispose it. + /// Optional tenant identifier for hierarchical partitioning. + /// Optional user identifier for hierarchical partitioning. + internal CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId, string conversationId, bool ownsClient, string? tenantId = null, string? userId = null) + { + this._cosmosClient = Throw.IfNull(cosmosClient); + this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); + this.ConversationId = Throw.IfNullOrWhitespace(conversationId); + this.DatabaseId = databaseId; + this.ContainerId = containerId; + this._ownsClient = ownsClient; + + // Initialize partitioning mode + this._tenantId = tenantId; + this._userId = userId; + this._useHierarchicalPartitioning = tenantId != null && userId != null; + + this._partitionKey = this._useHierarchicalPartitioning + ? new PartitionKeyBuilder() + .Add(tenantId!) + .Add(userId!) + .Add(conversationId) + .Build() + : new PartitionKey(conversationId); + } + + /// + /// Initializes a new instance of the class using a connection string. + /// + /// The Cosmos DB connection string. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(string connectionString, string databaseId, string containerId) + : this(connectionString, databaseId, containerId, Guid.NewGuid().ToString("N")) + { + } + + /// + /// Initializes a new instance of the class using a connection string. + /// + /// The Cosmos DB connection string. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The unique identifier for this conversation thread. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(string connectionString, string databaseId, string containerId, string conversationId) + : this(new CosmosClient(Throw.IfNullOrWhitespace(connectionString)), databaseId, containerId, conversationId, ownsClient: true) + { + } + + /// + /// Initializes a new instance of the class using TokenCredential for authentication. + /// + /// The Cosmos DB account endpoint URI. + /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId) + : this(accountEndpoint, tokenCredential, databaseId, containerId, Guid.NewGuid().ToString("N")) + { + } + + /// + /// Initializes a new instance of the class using a TokenCredential for authentication. + /// + /// The Cosmos DB account endpoint URI. + /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The unique identifier for this conversation thread. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId, string conversationId) + : this(new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)), databaseId, containerId, conversationId, ownsClient: true) + { + } + + /// + /// Initializes a new instance of the class using an existing . + /// + /// The instance to use for Cosmos DB operations. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// Thrown when is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId) + : this(cosmosClient, databaseId, containerId, Guid.NewGuid().ToString("N")) + { + } + + /// + /// Initializes a new instance of the class using an existing . + /// + /// The instance to use for Cosmos DB operations. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The unique identifier for this conversation thread. + /// Thrown when is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId, string conversationId) + : this(cosmosClient, databaseId, containerId, conversationId, ownsClient: false) + { + } + + /// + /// Initializes a new instance of the class using a connection string with hierarchical partition keys. + /// + /// The Cosmos DB connection string. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The tenant identifier for hierarchical partitioning. + /// The user identifier for hierarchical partitioning. + /// The session identifier for hierarchical partitioning. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(string connectionString, string databaseId, string containerId, string tenantId, string userId, string sessionId) + : this(new CosmosClient(Throw.IfNullOrWhitespace(connectionString)), databaseId, containerId, Throw.IfNullOrWhitespace(sessionId), ownsClient: true, Throw.IfNullOrWhitespace(tenantId), Throw.IfNullOrWhitespace(userId)) + { + } + + /// + /// Initializes a new instance of the class using a TokenCredential for authentication with hierarchical partition keys. + /// + /// The Cosmos DB account endpoint URI. + /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The tenant identifier for hierarchical partitioning. + /// The user identifier for hierarchical partitioning. + /// The session identifier for hierarchical partitioning. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId, string tenantId, string userId, string sessionId) + : this(new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)), databaseId, containerId, Throw.IfNullOrWhitespace(sessionId), ownsClient: true, Throw.IfNullOrWhitespace(tenantId), Throw.IfNullOrWhitespace(userId)) + { + } + + /// + /// Initializes a new instance of the class using an existing with hierarchical partition keys. + /// + /// The instance to use for Cosmos DB operations. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The tenant identifier for hierarchical partitioning. + /// The user identifier for hierarchical partitioning. + /// The session identifier for hierarchical partitioning. + /// Thrown when is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId, string tenantId, string userId, string sessionId) + : this(cosmosClient, databaseId, containerId, Throw.IfNullOrWhitespace(sessionId), ownsClient: false, Throw.IfNullOrWhitespace(tenantId), Throw.IfNullOrWhitespace(userId)) + { + } + + /// + /// Creates a new instance of the class from previously serialized state. + /// + /// The instance to use for Cosmos DB operations. + /// A representing the serialized state of the message store. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// Optional settings for customizing the JSON deserialization process. + /// A new instance of initialized from the serialized state. + /// Thrown when is null. + /// Thrown when the serialized state cannot be deserialized. + public static CosmosChatMessageStore CreateFromSerializedState(CosmosClient cosmosClient, JsonElement serializedStoreState, string databaseId, string containerId, JsonSerializerOptions? jsonSerializerOptions = null) + { + Throw.IfNull(cosmosClient); + Throw.IfNullOrWhitespace(databaseId); + Throw.IfNullOrWhitespace(containerId); + + if (serializedStoreState.ValueKind is not JsonValueKind.Object) + { + throw new ArgumentException("Invalid serialized state", nameof(serializedStoreState)); + } + + var state = JsonSerializer.Deserialize(serializedStoreState, jsonSerializerOptions); + if (state?.ConversationIdentifier is not { } conversationId) + { + throw new ArgumentException("Invalid serialized state", nameof(serializedStoreState)); + } + + // Use the internal constructor with all parameters to ensure partition key logic is centralized + return state.UseHierarchicalPartitioning && state.TenantId != null && state.UserId != null + ? new CosmosChatMessageStore(cosmosClient, databaseId, containerId, conversationId, ownsClient: false, state.TenantId, state.UserId) + : new CosmosChatMessageStore(cosmosClient, databaseId, containerId, conversationId, ownsClient: false); + } + + /// + public override async Task> GetMessagesAsync(CancellationToken cancellationToken = default) + { +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (this._disposed) + { + throw new ObjectDisposedException(this.GetType().FullName); + } +#pragma warning restore CA1513 + + // Fetch most recent messages in descending order when limit is set, then reverse to ascending + var orderDirection = this.MaxMessagesToRetrieve.HasValue ? "DESC" : "ASC"; + var query = new QueryDefinition($"SELECT * FROM c WHERE c.conversationId = @conversationId AND c.type = @type ORDER BY c.timestamp {orderDirection}") + .WithParameter("@conversationId", this.ConversationId) + .WithParameter("@type", "ChatMessage"); + + var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions + { + PartitionKey = this._partitionKey, + MaxItemCount = this.MaxItemCount // Configurable query performance + }); + + var messages = new List(); + + while (iterator.HasMoreResults) + { + var response = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false); + + foreach (var document in response) + { + if (this.MaxMessagesToRetrieve.HasValue && messages.Count >= this.MaxMessagesToRetrieve.Value) + { + break; + } + + if (!string.IsNullOrEmpty(document.Message)) + { + var message = JsonSerializer.Deserialize(document.Message, s_defaultJsonOptions); + if (message != null) + { + messages.Add(message); + } + } + } + + if (this.MaxMessagesToRetrieve.HasValue && messages.Count >= this.MaxMessagesToRetrieve.Value) + { + break; + } + } + + // If we fetched in descending order (most recent first), reverse to ascending order + if (this.MaxMessagesToRetrieve.HasValue) + { + messages.Reverse(); + } + + return messages; + } + + /// + public override async Task AddMessagesAsync(IEnumerable messages, CancellationToken cancellationToken = default) + { + if (messages is null) + { + throw new ArgumentNullException(nameof(messages)); + } + +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (this._disposed) + { + throw new ObjectDisposedException(this.GetType().FullName); + } +#pragma warning restore CA1513 + + var messageList = messages as IReadOnlyCollection ?? messages.ToList(); + if (messageList.Count == 0) + { + return; + } + + // Use transactional batch for atomic operations + if (messageList.Count > 1) + { + await this.AddMessagesInBatchAsync(messageList, cancellationToken).ConfigureAwait(false); + } + else + { + await this.AddSingleMessageAsync(messageList.First(), cancellationToken).ConfigureAwait(false); + } + } + + /// + /// Adds multiple messages using transactional batch operations for atomicity. + /// + private async Task AddMessagesInBatchAsync(IReadOnlyCollection messages, CancellationToken cancellationToken) + { + var currentTimestamp = DateTimeOffset.UtcNow.ToUnixTimeSeconds(); + + // Process messages in optimal batch sizes + for (int i = 0; i < messages.Count; i += this.MaxBatchSize) + { + var batchMessages = messages.Skip(i).Take(this.MaxBatchSize).ToList(); + await this.ExecuteBatchOperationAsync(batchMessages, currentTimestamp, cancellationToken).ConfigureAwait(false); + } + } + + /// + /// Executes a single batch operation with enhanced error handling. + /// Cosmos SDK handles throttling (429) retries automatically. + /// + private async Task ExecuteBatchOperationAsync(List messages, long timestamp, CancellationToken cancellationToken) + { + // Create all documents upfront for validation and batch operation + var documents = new List(messages.Count); + foreach (var message in messages) + { + documents.Add(this.CreateMessageDocument(message, timestamp)); + } + + // Defensive check: Verify all messages share the same partition key values + // In hierarchical partitioning, this means same tenantId, userId, and sessionId + // In simple partitioning, this means same conversationId + if (documents.Count > 0) + { + if (this._useHierarchicalPartitioning) + { + // Verify all documents have matching hierarchical partition key components + var firstDoc = documents[0]; + if (!documents.All(d => d.TenantId == firstDoc.TenantId && d.UserId == firstDoc.UserId && d.SessionId == firstDoc.SessionId)) + { + throw new InvalidOperationException("All messages in a batch must share the same partition key values (tenantId, userId, sessionId)."); + } + } + else + { + // Verify all documents have matching conversationId + var firstConversationId = documents[0].ConversationId; + if (!documents.All(d => d.ConversationId == firstConversationId)) + { + throw new InvalidOperationException("All messages in a batch must share the same partition key value (conversationId)."); + } + } + } + + // All messages in this store share the same partition key by design + // Transactional batches require all items to share the same partition key + var batch = this._container.CreateTransactionalBatch(this._partitionKey); + + foreach (var document in documents) + { + batch.CreateItem(document); + } + + try + { + var response = await batch.ExecuteAsync(cancellationToken).ConfigureAwait(false); + if (!response.IsSuccessStatusCode) + { + throw new InvalidOperationException($"Batch operation failed with status: {response.StatusCode}. Details: {response.ErrorMessage}"); + } + } + catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.RequestEntityTooLarge) + { + // If batch is too large, split into smaller batches + if (messages.Count == 1) + { + // Can't split further, use single operation + await this.AddSingleMessageAsync(messages[0], cancellationToken).ConfigureAwait(false); + return; + } + + // Split the batch in half and retry + var midpoint = messages.Count / 2; + var firstHalf = messages.Take(midpoint).ToList(); + var secondHalf = messages.Skip(midpoint).ToList(); + + await this.ExecuteBatchOperationAsync(firstHalf, timestamp, cancellationToken).ConfigureAwait(false); + await this.ExecuteBatchOperationAsync(secondHalf, timestamp, cancellationToken).ConfigureAwait(false); + } + } + + /// + /// Adds a single message to the store. + /// + private async Task AddSingleMessageAsync(ChatMessage message, CancellationToken cancellationToken) + { + var document = this.CreateMessageDocument(message, DateTimeOffset.UtcNow.ToUnixTimeSeconds()); + + try + { + await this._container.CreateItemAsync(document, this._partitionKey, cancellationToken: cancellationToken).ConfigureAwait(false); + } + catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.RequestEntityTooLarge) + { + throw new InvalidOperationException( + "Message exceeds Cosmos DB's maximum item size limit of 2MB. " + + "Message ID: " + message.MessageId + ", Serialized size is too large. " + + "Consider reducing message content or splitting into smaller messages.", + ex); + } + } + + /// + /// Creates a message document with enhanced metadata. + /// + private CosmosMessageDocument CreateMessageDocument(ChatMessage message, long timestamp) + { + return new CosmosMessageDocument + { + Id = Guid.NewGuid().ToString(), + ConversationId = this.ConversationId, + Timestamp = timestamp, + MessageId = message.MessageId, + Role = message.Role.Value, + Message = JsonSerializer.Serialize(message, s_defaultJsonOptions), + Type = "ChatMessage", // Type discriminator + Ttl = this.MessageTtlSeconds, // Configurable TTL + // Include hierarchical metadata when using hierarchical partitioning + TenantId = this._useHierarchicalPartitioning ? this._tenantId : null, + UserId = this._useHierarchicalPartitioning ? this._userId : null, + SessionId = this._useHierarchicalPartitioning ? this.ConversationId : null + }; + } + + /// + public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) + { +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (this._disposed) + { + throw new ObjectDisposedException(this.GetType().FullName); + } +#pragma warning restore CA1513 + + var state = new StoreState + { + ConversationIdentifier = this.ConversationId, + TenantId = this._tenantId, + UserId = this._userId, + UseHierarchicalPartitioning = this._useHierarchicalPartitioning + }; + + var options = jsonSerializerOptions ?? s_defaultJsonOptions; + return JsonSerializer.SerializeToElement(state, options); + } + + /// + /// Gets the count of messages in this conversation. + /// This is an additional utility method beyond the base contract. + /// + /// The cancellation token. + /// The number of messages in the conversation. + public async Task GetMessageCountAsync(CancellationToken cancellationToken = default) + { +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (this._disposed) + { + throw new ObjectDisposedException(this.GetType().FullName); + } +#pragma warning restore CA1513 + + // Efficient count query + var query = new QueryDefinition("SELECT VALUE COUNT(1) FROM c WHERE c.conversationId = @conversationId AND c.Type = @type") + .WithParameter("@conversationId", this.ConversationId) + .WithParameter("@type", "ChatMessage"); + + var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions + { + PartitionKey = this._partitionKey + }); + + // COUNT queries always return a result + var response = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false); + return response.FirstOrDefault(); + } + + /// + /// Deletes all messages in this conversation. + /// This is an additional utility method beyond the base contract. + /// + /// The cancellation token. + /// The number of messages deleted. + public async Task ClearMessagesAsync(CancellationToken cancellationToken = default) + { +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (this._disposed) + { + throw new ObjectDisposedException(this.GetType().FullName); + } +#pragma warning restore CA1513 + + // Batch delete for efficiency + var query = new QueryDefinition("SELECT VALUE c.id FROM c WHERE c.conversationId = @conversationId AND c.Type = @type") + .WithParameter("@conversationId", this.ConversationId) + .WithParameter("@type", "ChatMessage"); + + var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions + { + PartitionKey = this._partitionKey, + MaxItemCount = this.MaxItemCount + }); + + var deletedCount = 0; + + while (iterator.HasMoreResults) + { + var response = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false); + var batch = this._container.CreateTransactionalBatch(this._partitionKey); + var batchItemCount = 0; + + foreach (var itemId in response) + { + if (!string.IsNullOrEmpty(itemId)) + { + batch.DeleteItem(itemId); + batchItemCount++; + deletedCount++; + } + } + + if (batchItemCount > 0) + { + await batch.ExecuteAsync(cancellationToken).ConfigureAwait(false); + } + } + + return deletedCount; + } + + /// + public void Dispose() + { + if (!this._disposed) + { + if (this._ownsClient) + { + this._cosmosClient?.Dispose(); + } + this._disposed = true; + } + } + + private sealed class StoreState + { + public string ConversationIdentifier { get; set; } = string.Empty; + public string? TenantId { get; set; } + public string? UserId { get; set; } + public bool UseHierarchicalPartitioning { get; set; } + } + + /// + /// Represents a document stored in Cosmos DB for chat messages. + /// + [SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated by Cosmos DB operations")] + private sealed class CosmosMessageDocument + { + [Newtonsoft.Json.JsonProperty("id")] + public string Id { get; set; } = string.Empty; + + [Newtonsoft.Json.JsonProperty("conversationId")] + public string ConversationId { get; set; } = string.Empty; + + [Newtonsoft.Json.JsonProperty("timestamp")] + public long Timestamp { get; set; } + + [Newtonsoft.Json.JsonProperty("messageId")] + public string? MessageId { get; set; } + + [Newtonsoft.Json.JsonProperty("role")] + public string? Role { get; set; } + + [Newtonsoft.Json.JsonProperty("message")] + public string Message { get; set; } = string.Empty; + + [Newtonsoft.Json.JsonProperty("type")] + public string Type { get; set; } = string.Empty; + + [Newtonsoft.Json.JsonProperty("ttl")] + public int? Ttl { get; set; } + + /// + /// Tenant ID for hierarchical partitioning scenarios (optional). + /// + [Newtonsoft.Json.JsonProperty("tenantId")] + public string? TenantId { get; set; } + + /// + /// User ID for hierarchical partitioning scenarios (optional). + /// + [Newtonsoft.Json.JsonProperty("userId")] + public string? UserId { get; set; } + + /// + /// Session ID for hierarchical partitioning scenarios (same as ConversationId for compatibility). + /// + [Newtonsoft.Json.JsonProperty("sessionId")] + public string? SessionId { get; set; } + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosCheckpointStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosCheckpointStore.cs new file mode 100644 index 0000000000..62987b1dfc --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosCheckpointStore.cs @@ -0,0 +1,279 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; +using System.Linq; +using System.Text.Json; +using System.Threading.Tasks; +using Azure.Core; +using Microsoft.Azure.Cosmos; +using Microsoft.Shared.Diagnostics; +using Newtonsoft.Json; +using Newtonsoft.Json.Linq; + +namespace Microsoft.Agents.AI.Workflows.Checkpointing; + +/// +/// Provides a Cosmos DB implementation of the abstract class. +/// +/// The type of objects to store as checkpoint values. +[RequiresUnreferencedCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with trimming.")] +[RequiresDynamicCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with NativeAOT.")] +public class CosmosCheckpointStore : JsonCheckpointStore, IDisposable +{ + private readonly CosmosClient _cosmosClient; + private readonly Container _container; + private readonly bool _ownsClient; + private bool _disposed; + + /// + /// Initializes a new instance of the class using a connection string. + /// + /// The Cosmos DB connection string. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosCheckpointStore(string connectionString, string databaseId, string containerId) + { + var cosmosClientOptions = new CosmosClientOptions(); + + this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(connectionString), cosmosClientOptions); + this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); + this._ownsClient = true; + } + + /// + /// Initializes a new instance of the class using a TokenCredential for authentication. + /// + /// The Cosmos DB account endpoint URI. + /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosCheckpointStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId) + { + var cosmosClientOptions = new CosmosClientOptions + { + SerializerOptions = new CosmosSerializationOptions + { + PropertyNamingPolicy = CosmosPropertyNamingPolicy.CamelCase + } + }; + + this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential), cosmosClientOptions); + this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); + this._ownsClient = true; + } + + /// + /// Initializes a new instance of the class using an existing . + /// + /// The instance to use for Cosmos DB operations. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// Thrown when is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosCheckpointStore(CosmosClient cosmosClient, string databaseId, string containerId) + { + this._cosmosClient = Throw.IfNull(cosmosClient); + + this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); + this._ownsClient = false; + } + + /// + /// Gets the identifier of the Cosmos DB database. + /// + public string DatabaseId => this._container.Database.Id; + + /// + /// Gets the identifier of the Cosmos DB container. + /// + public string ContainerId => this._container.Id; + + /// + public override async ValueTask CreateCheckpointAsync(string runId, JsonElement value, CheckpointInfo? parent = null) + { + if (string.IsNullOrWhiteSpace(runId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(runId)); + } + +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (this._disposed) + { + throw new ObjectDisposedException(this.GetType().FullName); + } +#pragma warning restore CA1513 + + var checkpointId = Guid.NewGuid().ToString("N"); + var checkpointInfo = new CheckpointInfo(runId, checkpointId); + + var document = new CosmosCheckpointDocument + { + Id = $"{runId}_{checkpointId}", + RunId = runId, + CheckpointId = checkpointId, + Value = JToken.Parse(value.GetRawText()), + ParentCheckpointId = parent?.CheckpointId, + Timestamp = DateTimeOffset.UtcNow.ToUnixTimeSeconds() + }; + + await this._container.CreateItemAsync(document, new PartitionKey(runId)).ConfigureAwait(false); + return checkpointInfo; + } + + /// + public override async ValueTask RetrieveCheckpointAsync(string runId, CheckpointInfo key) + { + if (string.IsNullOrWhiteSpace(runId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(runId)); + } + + if (key is null) + { + throw new ArgumentNullException(nameof(key)); + } + +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (this._disposed) + { + throw new ObjectDisposedException(this.GetType().FullName); + } +#pragma warning restore CA1513 + + var id = $"{runId}_{key.CheckpointId}"; + + try + { + var response = await this._container.ReadItemAsync(id, new PartitionKey(runId)).ConfigureAwait(false); + using var document = JsonDocument.Parse(response.Resource.Value.ToString()); + return document.RootElement.Clone(); + } + catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.NotFound) + { + throw new InvalidOperationException($"Checkpoint with ID '{key.CheckpointId}' for run '{runId}' not found."); + } + } + + /// + public override async ValueTask> RetrieveIndexAsync(string runId, CheckpointInfo? withParent = null) + { + if (string.IsNullOrWhiteSpace(runId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(runId)); + } + +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (this._disposed) + { + throw new ObjectDisposedException(this.GetType().FullName); + } +#pragma warning restore CA1513 + + QueryDefinition query = withParent == null + ? new QueryDefinition("SELECT c.runId, c.checkpointId FROM c WHERE c.runId = @runId ORDER BY c.timestamp ASC") + .WithParameter("@runId", runId) + : new QueryDefinition("SELECT c.runId, c.checkpointId FROM c WHERE c.runId = @runId AND c.parentCheckpointId = @parentCheckpointId ORDER BY c.timestamp ASC") + .WithParameter("@runId", runId) + .WithParameter("@parentCheckpointId", withParent.CheckpointId); + + var iterator = this._container.GetItemQueryIterator(query); + var checkpoints = new List(); + + while (iterator.HasMoreResults) + { + var response = await iterator.ReadNextAsync().ConfigureAwait(false); + checkpoints.AddRange(response.Select(r => new CheckpointInfo(r.RunId, r.CheckpointId))); + } + + return checkpoints; + } + + /// + public void Dispose() + { + this.Dispose(true); + GC.SuppressFinalize(this); + } + + /// + /// Releases the unmanaged resources used by the and optionally releases the managed resources. + /// + /// true to release both managed and unmanaged resources; false to release only unmanaged resources. + protected virtual void Dispose(bool disposing) + { + if (!this._disposed) + { + if (disposing && this._ownsClient) + { + this._cosmosClient?.Dispose(); + } + this._disposed = true; + } + } + + /// + /// Represents a checkpoint document stored in Cosmos DB. + /// + internal sealed class CosmosCheckpointDocument + { + [JsonProperty("id")] + public string Id { get; set; } = string.Empty; + + [JsonProperty("runId")] + public string RunId { get; set; } = string.Empty; + + [JsonProperty("checkpointId")] + public string CheckpointId { get; set; } = string.Empty; + + [JsonProperty("value")] + public JToken Value { get; set; } = JValue.CreateNull(); + + [JsonProperty("parentCheckpointId")] + public string? ParentCheckpointId { get; set; } + + [JsonProperty("timestamp")] + public long Timestamp { get; set; } + } + + /// + /// Represents the result of a checkpoint query. + /// + [SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated by Cosmos DB query deserialization")] + private sealed class CheckpointQueryResult + { + public string RunId { get; set; } = string.Empty; + public string CheckpointId { get; set; } = string.Empty; + } +} + +/// +/// Provides a non-generic Cosmos DB implementation of the abstract class. +/// +[RequiresUnreferencedCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with trimming.")] +[RequiresDynamicCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with NativeAOT.")] +public sealed class CosmosCheckpointStore : CosmosCheckpointStore +{ + /// + public CosmosCheckpointStore(string connectionString, string databaseId, string containerId) + : base(connectionString, databaseId, containerId) + { + } + + /// + public CosmosCheckpointStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId) + : base(accountEndpoint, tokenCredential, databaseId, containerId) + { + } + + /// + public CosmosCheckpointStore(CosmosClient cosmosClient, string databaseId, string containerId) + : base(cosmosClient, databaseId, containerId) + { + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs new file mode 100644 index 0000000000..4e3b66fd54 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs @@ -0,0 +1,95 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Diagnostics.CodeAnalysis; +using Azure.Identity; +using Microsoft.Azure.Cosmos; + +namespace Microsoft.Agents.AI; + +/// +/// Provides extension methods for integrating Cosmos DB chat message storage with the Agent Framework. +/// +public static class CosmosDBChatExtensions +{ + /// + /// Configures the agent to use Cosmos DB for message storage with connection string authentication. + /// + /// The chat client agent options to configure. + /// The Cosmos DB connection string. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The configured . + /// Thrown when is null. + /// Thrown when any string parameter is null or whitespace. + [RequiresUnreferencedCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with trimming.")] + [RequiresDynamicCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with NativeAOT.")] + public static ChatClientAgentOptions WithCosmosDBMessageStore( + this ChatClientAgentOptions options, + string connectionString, + string databaseId, + string containerId) + { + if (options is null) + { + throw new ArgumentNullException(nameof(options)); + } + + options.ChatMessageStoreFactory = context => new CosmosChatMessageStore(connectionString, databaseId, containerId); + return options; + } + + /// + /// Configures the agent to use Cosmos DB for message storage with managed identity authentication. + /// + /// The chat client agent options to configure. + /// The Cosmos DB account endpoint URI. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The configured . + /// Thrown when is null. + /// Thrown when any string parameter is null or whitespace. + [RequiresUnreferencedCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with trimming.")] + [RequiresDynamicCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with NativeAOT.")] + public static ChatClientAgentOptions WithCosmosDBMessageStoreUsingManagedIdentity( + this ChatClientAgentOptions options, + string accountEndpoint, + string databaseId, + string containerId) + { + if (options is null) + { + throw new ArgumentNullException(nameof(options)); + } + + options.ChatMessageStoreFactory = context => new CosmosChatMessageStore(accountEndpoint, new DefaultAzureCredential(), databaseId, containerId); + return options; + } + + /// + /// Configures the agent to use Cosmos DB for message storage with an existing . + /// + /// The chat client agent options to configure. + /// The instance to use for Cosmos DB operations. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The configured . + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + [RequiresUnreferencedCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with trimming.")] + [RequiresDynamicCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with NativeAOT.")] + public static ChatClientAgentOptions WithCosmosDBMessageStore( + this ChatClientAgentOptions options, + CosmosClient cosmosClient, + string databaseId, + string containerId) + { + if (options is null) + { + throw new ArgumentNullException(nameof(options)); + } + + options.ChatMessageStoreFactory = context => new CosmosChatMessageStore(cosmosClient, databaseId, containerId); + return options; + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs new file mode 100644 index 0000000000..9d8bc52e68 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs @@ -0,0 +1,218 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Diagnostics.CodeAnalysis; +using Azure.Identity; +using Microsoft.Agents.AI.Workflows.Checkpointing; +using Microsoft.Azure.Cosmos; + +namespace Microsoft.Agents.AI.Workflows; + +/// +/// Provides extension methods for integrating Cosmos DB checkpoint storage with the Agent Framework. +/// +public static class CosmosDBWorkflowExtensions +{ + /// + /// Creates a Cosmos DB checkpoint store using connection string authentication. + /// + /// The Cosmos DB connection string. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// A new instance of . + /// Thrown when any string parameter is null or whitespace. + [RequiresUnreferencedCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with trimming.")] + [RequiresDynamicCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with NativeAOT.")] + public static CosmosCheckpointStore CreateCheckpointStore( + string connectionString, + string databaseId, + string containerId) + { + if (string.IsNullOrWhiteSpace(connectionString)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(connectionString)); + } + + if (string.IsNullOrWhiteSpace(databaseId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + } + + if (string.IsNullOrWhiteSpace(containerId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + } + + return new CosmosCheckpointStore(connectionString, databaseId, containerId); + } + + /// + /// Creates a Cosmos DB checkpoint store using managed identity authentication. + /// + /// The Cosmos DB account endpoint URI. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// A new instance of . + /// Thrown when any string parameter is null or whitespace. + [RequiresUnreferencedCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with trimming.")] + [RequiresDynamicCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with NativeAOT.")] + public static CosmosCheckpointStore CreateCheckpointStoreUsingManagedIdentity( + string accountEndpoint, + string databaseId, + string containerId) + { + if (string.IsNullOrWhiteSpace(accountEndpoint)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(accountEndpoint)); + } + + if (string.IsNullOrWhiteSpace(databaseId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + } + + if (string.IsNullOrWhiteSpace(containerId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + } + + return new CosmosCheckpointStore(accountEndpoint, new DefaultAzureCredential(), databaseId, containerId); + } + + /// + /// Creates a Cosmos DB checkpoint store using an existing . + /// + /// The instance to use for Cosmos DB operations. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// A new instance of . + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + [RequiresUnreferencedCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with trimming.")] + [RequiresDynamicCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with NativeAOT.")] + public static CosmosCheckpointStore CreateCheckpointStore( + CosmosClient cosmosClient, + string databaseId, + string containerId) + { + if (cosmosClient is null) + { + throw new ArgumentNullException(nameof(cosmosClient)); + } + + if (string.IsNullOrWhiteSpace(databaseId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + } + + if (string.IsNullOrWhiteSpace(containerId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + } + + return new CosmosCheckpointStore(cosmosClient, databaseId, containerId); + } + + /// + /// Creates a generic Cosmos DB checkpoint store using connection string authentication. + /// + /// The type of objects to store as checkpoint values. + /// The Cosmos DB connection string. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// A new instance of . + /// Thrown when any string parameter is null or whitespace. + [RequiresUnreferencedCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with trimming.")] + [RequiresDynamicCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with NativeAOT.")] + public static CosmosCheckpointStore CreateCheckpointStore( + string connectionString, + string databaseId, + string containerId) + { + if (string.IsNullOrWhiteSpace(connectionString)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(connectionString)); + } + + if (string.IsNullOrWhiteSpace(databaseId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + } + + if (string.IsNullOrWhiteSpace(containerId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + } + + return new CosmosCheckpointStore(connectionString, databaseId, containerId); + } + + /// + /// Creates a generic Cosmos DB checkpoint store using managed identity authentication. + /// + /// The type of objects to store as checkpoint values. + /// The Cosmos DB account endpoint URI. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// A new instance of . + /// Thrown when any string parameter is null or whitespace. + [RequiresUnreferencedCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with trimming.")] + [RequiresDynamicCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with NativeAOT.")] + public static CosmosCheckpointStore CreateCheckpointStoreUsingManagedIdentity( + string accountEndpoint, + string databaseId, + string containerId) + { + if (string.IsNullOrWhiteSpace(accountEndpoint)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(accountEndpoint)); + } + + if (string.IsNullOrWhiteSpace(databaseId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + } + + if (string.IsNullOrWhiteSpace(containerId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + } + + return new CosmosCheckpointStore(accountEndpoint, new DefaultAzureCredential(), databaseId, containerId); + } + + /// + /// Creates a generic Cosmos DB checkpoint store using an existing . + /// + /// The type of objects to store as checkpoint values. + /// The instance to use for Cosmos DB operations. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// A new instance of . + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + [RequiresUnreferencedCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with trimming.")] + [RequiresDynamicCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with NativeAOT.")] + public static CosmosCheckpointStore CreateCheckpointStore( + CosmosClient cosmosClient, + string databaseId, + string containerId) + { + if (cosmosClient is null) + { + throw new ArgumentNullException(nameof(cosmosClient)); + } + + if (string.IsNullOrWhiteSpace(databaseId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + } + + if (string.IsNullOrWhiteSpace(containerId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + } + + return new CosmosCheckpointStore(cosmosClient, databaseId, containerId); + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/Microsoft.Agents.AI.CosmosNoSql.csproj b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/Microsoft.Agents.AI.CosmosNoSql.csproj new file mode 100644 index 0000000000..7e13ec5998 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/Microsoft.Agents.AI.CosmosNoSql.csproj @@ -0,0 +1,41 @@ + + + + $(TargetFrameworksCore) + Microsoft.Agents.AI + $(NoWarn);MEAI001 + preview + + + + true + true + true + true + true + true + + + + + + + Microsoft Agent Framework Cosmos DB NoSQL Integration + Provides Cosmos DB NoSQL implementations for Microsoft Agent Framework storage abstractions including ChatMessageStore and CheckpointStore. + + + + + + + + + + + + + + + + + diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/.editorconfig b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/.editorconfig new file mode 100644 index 0000000000..83e05f582a --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/.editorconfig @@ -0,0 +1,9 @@ +# EditorConfig overrides for Cosmos DB Unit Tests +# Multi-targeting (net472 + net9.0) causes false positives for IDE0005 (unnecessary using directives) + +root = false + +[*.cs] +# Suppress IDE0005 for this project - multi-targeting causes false positives +# These using directives ARE necessary but appear unnecessary in one target framework +dotnet_diagnostic.IDE0005.severity = none diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs new file mode 100644 index 0000000000..6f2a256206 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -0,0 +1,760 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Text.Json; +using System.Text.Json.Serialization.Metadata; +using System.Threading.Tasks; +using Azure.Core; +using Azure.Identity; +using Microsoft.Agents.AI; +using Microsoft.Azure.Cosmos; +using Microsoft.Extensions.AI; +using Xunit; + +namespace Microsoft.Agents.AI.CosmosNoSql.UnitTests; + +/// +/// Contains tests for . +/// +/// Test Modes: +/// - Default Mode: Cleans up all test data after each test run (deletes database) +/// - Preserve Mode: Keeps containers and data for inspection in Cosmos DB Emulator Data Explorer +/// +/// To enable Preserve Mode, set environment variable: COSMOS_PRESERVE_CONTAINERS=true +/// Example: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test +/// +/// In Preserve Mode, you can view the data in Cosmos DB Emulator Data Explorer at: +/// https://localhost:8081/_explorer/index.html +/// Database: AgentFrameworkTests +/// Container: ChatMessages +/// +/// Environment Variable Reference: +/// | Variable | Values | Description | +/// |----------|--------|-------------| +/// | COSMOS_PRESERVE_CONTAINERS | true / false | Controls whether to preserve test data after completion | +/// +/// Usage Examples: +/// - Run all tests in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/ +/// - Run specific test category in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/ --filter "Category=CosmosDB" +/// - Reset to cleanup mode: $env:COSMOS_PRESERVE_CONTAINERS=""; dotnet test tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/ +/// +[Collection("CosmosDB")] +public sealed class CosmosChatMessageStoreTests : IAsyncLifetime, IDisposable +{ + // Cosmos DB Emulator connection settings + private const string EmulatorEndpoint = "https://localhost:8081"; + private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; + private const string TestContainerId = "ChatMessages"; + private const string HierarchicalTestContainerId = "HierarchicalChatMessages"; + // Use unique database ID per test class instance to avoid conflicts +#pragma warning disable CA1802 // Use literals where appropriate + private static readonly string s_testDatabaseId = $"AgentFrameworkTests-ChatStore-{Guid.NewGuid():N}"; +#pragma warning restore CA1802 + + private string _connectionString = string.Empty; + private bool _emulatorAvailable; + private bool _preserveContainer; + private CosmosClient? _setupClient; // Only used for test setup/cleanup + + public async Task InitializeAsync() + { + // Check environment variable to determine if we should preserve containers + // Set COSMOS_PRESERVE_CONTAINERS=true to keep containers and data for inspection + this._preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase); + + this._connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey}"; + + try + { + // Only create CosmosClient for test setup - the actual tests will use connection string constructors + this._setupClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); + + // Test connection by attempting to create database + var databaseResponse = await this._setupClient.CreateDatabaseIfNotExistsAsync(s_testDatabaseId); + + // Create container for simple partitioning tests + await databaseResponse.Database.CreateContainerIfNotExistsAsync( + TestContainerId, + "/conversationId", + throughput: 400); + + // Create container for hierarchical partitioning tests with hierarchical partition key + var hierarchicalContainerProperties = new ContainerProperties(HierarchicalTestContainerId, new List { "/tenantId", "/userId", "/sessionId" }); + await databaseResponse.Database.CreateContainerIfNotExistsAsync( + hierarchicalContainerProperties, + throughput: 400); + + this._emulatorAvailable = true; + } + catch (Exception) + { + // Emulator not available, tests will be skipped + this._emulatorAvailable = false; + this._setupClient?.Dispose(); + this._setupClient = null; + } + } + + public async Task DisposeAsync() + { + if (this._setupClient != null && this._emulatorAvailable) + { + try + { + if (this._preserveContainer) + { + // Preserve mode: Don't delete the database/container, keep data for inspection + // This allows viewing data in the Cosmos DB Emulator Data Explorer + // No cleanup needed - data persists for debugging + } + else + { + // Clean mode: Delete the test database and all data + var database = this._setupClient.GetDatabase(s_testDatabaseId); + await database.DeleteAsync(); + } + } + catch (Exception ex) + { + // Ignore cleanup errors during test teardown + Console.WriteLine($"Warning: Cleanup failed: {ex.Message}"); + } + finally + { + this._setupClient.Dispose(); + } + } + } + + public void Dispose() + { + this._setupClient?.Dispose(); + GC.SuppressFinalize(this); + } + + private void SkipIfEmulatorNotAvailable() + { + // In CI: Skip if COSMOS_EMULATOR_AVAILABLE is not set to "true" + // Locally: Skip if emulator connection check failed + var ciEmulatorAvailable = string.Equals(Environment.GetEnvironmentVariable("COSMOS_EMULATOR_AVAILABLE"), "true", StringComparison.OrdinalIgnoreCase); + + Xunit.Skip.If(!ciEmulatorAvailable && !this._emulatorAvailable, "Cosmos DB Emulator is not available"); + } + + #region Constructor Tests + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithConnectionString_ShouldCreateInstance() + { + // Arrange & Act + this.SkipIfEmulatorNotAvailable(); + + // Act + using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, "test-conversation"); + + // Assert + Assert.NotNull(store); + Assert.Equal("test-conversation", store.ConversationId); + Assert.Equal(s_testDatabaseId, store.DatabaseId); + Assert.Equal(TestContainerId, store.ContainerId); + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithConnectionStringNoConversationId_ShouldCreateInstance() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + + // Act + using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId); + + // Assert + Assert.NotNull(store); + Assert.NotNull(store.ConversationId); + Assert.Equal(s_testDatabaseId, store.DatabaseId); + Assert.Equal(TestContainerId, store.ContainerId); + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithNullConnectionString_ShouldThrowArgumentException() + { + // Arrange & Act & Assert + Assert.Throws(() => + new CosmosChatMessageStore((string)null!, s_testDatabaseId, TestContainerId, "test-conversation")); + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithEmptyConversationId_ShouldThrowArgumentException() + { + // Arrange & Act & Assert + this.SkipIfEmulatorNotAvailable(); + + Assert.Throws(() => + new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, "")); + } + + #endregion + + #region AddMessagesAsync Tests + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + var conversationId = Guid.NewGuid().ToString(); + using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversationId); + var message = new ChatMessage(ChatRole.User, "Hello, world!"); + + // Act + await store.AddMessagesAsync([message]); + + // Wait a moment for eventual consistency + await Task.Delay(100); + + // Assert + var messages = await store.GetMessagesAsync(); + var messageList = messages.ToList(); + + // Simple assertion - if this fails, we know the deserialization is the issue + if (messageList.Count == 0) + { + // Let's check if we can find ANY items in the container for this conversation + var directQuery = new QueryDefinition("SELECT VALUE COUNT(1) FROM c WHERE c.conversationId = @conversationId") + .WithParameter("@conversationId", conversationId); + var countIterator = this._setupClient!.GetDatabase(s_testDatabaseId).GetContainer(TestContainerId) + .GetItemQueryIterator(directQuery, requestOptions: new QueryRequestOptions + { + PartitionKey = new PartitionKey(conversationId) + }); + + var countResponse = await countIterator.ReadNextAsync(); + var count = countResponse.FirstOrDefault(); + + // Debug: Let's see what the raw query returns + var rawQuery = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId") + .WithParameter("@conversationId", conversationId); + var rawIterator = this._setupClient!.GetDatabase(s_testDatabaseId).GetContainer(TestContainerId) + .GetItemQueryIterator(rawQuery, requestOptions: new QueryRequestOptions + { + PartitionKey = new PartitionKey(conversationId) + }); + + List rawResults = new(); + while (rawIterator.HasMoreResults) + { + var rawResponse = await rawIterator.ReadNextAsync(); + rawResults.AddRange(rawResponse); + } + + string rawJson = rawResults.Count > 0 ? Newtonsoft.Json.JsonConvert.SerializeObject(rawResults[0], Newtonsoft.Json.Formatting.Indented) : "null"; + Assert.Fail($"GetMessagesAsync returned 0 messages, but direct count query found {count} items for conversation {conversationId}. Raw document: {rawJson}"); + } + + Assert.Single(messageList); + Assert.Equal("Hello, world!", messageList[0].Text); + Assert.Equal(ChatRole.User, messageList[0].Role); + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public async Task AddMessagesAsync_WithMultipleMessages_ShouldAddAllMessagesAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + var conversationId = Guid.NewGuid().ToString(); + using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversationId); + var messages = new[] + { + new ChatMessage(ChatRole.User, "First message"), + new ChatMessage(ChatRole.Assistant, "Second message"), + new ChatMessage(ChatRole.User, "Third message") + }; + + // Act + await store.AddMessagesAsync(messages); + + // Assert + var retrievedMessages = await store.GetMessagesAsync(); + var messageList = retrievedMessages.ToList(); + Assert.Equal(3, messageList.Count); + Assert.Equal("First message", messageList[0].Text); + Assert.Equal("Second message", messageList[1].Text); + Assert.Equal("Third message", messageList[2].Text); + } + + #endregion + + #region GetMessagesAsync Tests + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public async Task GetMessagesAsync_WithNoMessages_ShouldReturnEmptyAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + + // Act + var messages = await store.GetMessagesAsync(); + + // Assert + Assert.Empty(messages); + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public async Task GetMessagesAsync_WithConversationIsolation_ShouldOnlyReturnMessagesForConversationAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + var conversation1 = Guid.NewGuid().ToString(); + var conversation2 = Guid.NewGuid().ToString(); + + using var store1 = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversation1); + using var store2 = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversation2); + + await store1.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message for conversation 1")]); + await store2.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message for conversation 2")]); + + // Act + var messages1 = await store1.GetMessagesAsync(); + var messages2 = await store2.GetMessagesAsync(); + + // Assert + var messageList1 = messages1.ToList(); + var messageList2 = messages2.ToList(); + Assert.Single(messageList1); + Assert.Single(messageList2); + Assert.Equal("Message for conversation 1", messageList1[0].Text); + Assert.Equal("Message for conversation 2", messageList2[0].Text); + } + + #endregion + + #region Integration Tests + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public async Task FullWorkflow_AddAndGet_ShouldWorkCorrectlyAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + var conversationId = $"test-conversation-{Guid.NewGuid():N}"; // Use unique conversation ID + using var originalStore = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversationId); + + var messages = new[] + { + new ChatMessage(ChatRole.System, "You are a helpful assistant."), + new ChatMessage(ChatRole.User, "Hello!"), + new ChatMessage(ChatRole.Assistant, "Hi there! How can I help you today?"), + new ChatMessage(ChatRole.User, "What's the weather like?"), + new ChatMessage(ChatRole.Assistant, "I'm sorry, I don't have access to current weather data.") + }; + + // Act 1: Add messages + await originalStore.AddMessagesAsync(messages); + + // Act 2: Verify messages were added + var retrievedMessages = await originalStore.GetMessagesAsync(); + var retrievedList = retrievedMessages.ToList(); + Assert.Equal(5, retrievedList.Count); + + // Act 3: Create new store instance for same conversation (test persistence) + using var newStore = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversationId); + var persistedMessages = await newStore.GetMessagesAsync(); + var persistedList = persistedMessages.ToList(); + + // Assert final state + Assert.Equal(5, persistedList.Count); + Assert.Equal("You are a helpful assistant.", persistedList[0].Text); + Assert.Equal("Hello!", persistedList[1].Text); + Assert.Equal("Hi there! How can I help you today?", persistedList[2].Text); + Assert.Equal("What's the weather like?", persistedList[3].Text); + Assert.Equal("I'm sorry, I don't have access to current weather data.", persistedList[4].Text); + } + + #endregion + + #region Disposal Tests + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public void Dispose_AfterUse_ShouldNotThrow() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + + // Act & Assert + store.Dispose(); // Should not throw + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public void Dispose_MultipleCalls_ShouldNotThrow() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + + // Act & Assert + store.Dispose(); // First call + store.Dispose(); // Second call - should not throw + } + + #endregion + + #region Hierarchical Partitioning Tests + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithHierarchicalConnectionString_ShouldCreateInstance() + { + // Arrange & Act + this.SkipIfEmulatorNotAvailable(); + + // Act + using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); + + // Assert + Assert.NotNull(store); + Assert.Equal("session-789", store.ConversationId); + Assert.Equal(s_testDatabaseId, store.DatabaseId); + Assert.Equal(HierarchicalTestContainerId, store.ContainerId); + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithHierarchicalEndpoint_ShouldCreateInstance() + { + // Arrange & Act + this.SkipIfEmulatorNotAvailable(); + + // Act + TokenCredential credential = new DefaultAzureCredential(); + using var store = new CosmosChatMessageStore(EmulatorEndpoint, credential, s_testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); + + // Assert + Assert.NotNull(store); + Assert.Equal("session-789", store.ConversationId); + Assert.Equal(s_testDatabaseId, store.DatabaseId); + Assert.Equal(HierarchicalTestContainerId, store.ContainerId); + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithHierarchicalCosmosClient_ShouldCreateInstance() + { + // Arrange & Act + this.SkipIfEmulatorNotAvailable(); + + using var cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); + using var store = new CosmosChatMessageStore(cosmosClient, s_testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); + + // Assert + Assert.NotNull(store); + Assert.Equal("session-789", store.ConversationId); + Assert.Equal(s_testDatabaseId, store.DatabaseId); + Assert.Equal(HierarchicalTestContainerId, store.ContainerId); + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithHierarchicalNullTenantId_ShouldThrowArgumentException() + { + // Arrange & Act & Assert + this.SkipIfEmulatorNotAvailable(); + + Assert.Throws(() => + new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, null!, "user-456", "session-789")); + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithHierarchicalEmptyUserId_ShouldThrowArgumentException() + { + // Arrange & Act & Assert + this.SkipIfEmulatorNotAvailable(); + + Assert.Throws(() => + new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, "tenant-123", "", "session-789")); + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithHierarchicalWhitespaceSessionId_ShouldThrowArgumentException() + { + // Arrange & Act & Assert + this.SkipIfEmulatorNotAvailable(); + + Assert.Throws(() => + new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", " ")); + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public async Task AddMessagesAsync_WithHierarchicalPartitioning_ShouldAddMessageWithMetadataAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + const string TenantId = "tenant-123"; + const string UserId = "user-456"; + const string SessionId = "session-789"; + // Test hierarchical partitioning constructor with connection string + using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); + var message = new ChatMessage(ChatRole.User, "Hello from hierarchical partitioning!"); + + // Act + await store.AddMessagesAsync([message]); + + // Wait a moment for eventual consistency + await Task.Delay(100); + + // Assert + var messages = await store.GetMessagesAsync(); + var messageList = messages.ToList(); + + Assert.Single(messageList); + Assert.Equal("Hello from hierarchical partitioning!", messageList[0].Text); + Assert.Equal(ChatRole.User, messageList[0].Role); + + // Verify that the document is stored with hierarchical partitioning metadata + var directQuery = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId AND c.type = @type") + .WithParameter("@conversationId", SessionId) + .WithParameter("@type", "ChatMessage"); + + var iterator = this._setupClient!.GetDatabase(s_testDatabaseId).GetContainer(HierarchicalTestContainerId) + .GetItemQueryIterator(directQuery, requestOptions: new QueryRequestOptions + { + PartitionKey = new PartitionKeyBuilder().Add(TenantId).Add(UserId).Add(SessionId).Build() + }); + + var response = await iterator.ReadNextAsync(); + var document = response.FirstOrDefault(); + + Assert.NotNull(document); + // The document should have hierarchical metadata + Assert.Equal(SessionId, (string)document!.conversationId); + Assert.Equal(TenantId, (string)document!.tenantId); + Assert.Equal(UserId, (string)document!.userId); + Assert.Equal(SessionId, (string)document!.sessionId); + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public async Task AddMessagesAsync_WithHierarchicalMultipleMessages_ShouldAddAllMessagesAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + const string TenantId = "tenant-batch"; + const string UserId = "user-batch"; + const string SessionId = "session-batch"; + // Test hierarchical partitioning constructor with connection string + using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); + var messages = new[] + { + new ChatMessage(ChatRole.User, "First hierarchical message"), + new ChatMessage(ChatRole.Assistant, "Second hierarchical message"), + new ChatMessage(ChatRole.User, "Third hierarchical message") + }; + + // Act + await store.AddMessagesAsync(messages); + + // Wait a moment for eventual consistency + await Task.Delay(100); + + // Assert + var retrievedMessages = await store.GetMessagesAsync(); + var messageList = retrievedMessages.ToList(); + + Assert.Equal(3, messageList.Count); + Assert.Equal("First hierarchical message", messageList[0].Text); + Assert.Equal("Second hierarchical message", messageList[1].Text); + Assert.Equal("Third hierarchical message", messageList[2].Text); + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public async Task GetMessagesAsync_WithHierarchicalPartitionIsolation_ShouldIsolateMessagesByUserIdAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + const string TenantId = "tenant-isolation"; + const string UserId1 = "user-1"; + const string UserId2 = "user-2"; + const string SessionId = "session-isolation"; + + // Different userIds create different hierarchical partitions, providing proper isolation + using var store1 = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId1, SessionId); + using var store2 = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId2, SessionId); + + // Add messages to both stores + await store1.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message from user 1")]); + await store2.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message from user 2")]); + + // Wait a moment for eventual consistency + await Task.Delay(100); + + // Act & Assert + var messages1 = await store1.GetMessagesAsync(); + var messageList1 = messages1.ToList(); + + var messages2 = await store2.GetMessagesAsync(); + var messageList2 = messages2.ToList(); + + // With true hierarchical partitioning, each user sees only their own messages + Assert.Single(messageList1); + Assert.Single(messageList2); + Assert.Equal("Message from user 1", messageList1[0].Text); + Assert.Equal("Message from user 2", messageList2[0].Text); + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public async Task SerializeDeserialize_WithHierarchicalPartitioning_ShouldPreserveStateAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + const string TenantId = "tenant-serialize"; + const string UserId = "user-serialize"; + const string SessionId = "session-serialize"; + + using var originalStore = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); + await originalStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Test serialization message")]); + + // Act - Serialize the store state + var serializedState = originalStore.Serialize(); + + // Create a new store from the serialized state + using var cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); + var serializerOptions = new JsonSerializerOptions + { + TypeInfoResolver = new DefaultJsonTypeInfoResolver() + }; + using var deserializedStore = CosmosChatMessageStore.CreateFromSerializedState(cosmosClient, serializedState, s_testDatabaseId, HierarchicalTestContainerId, serializerOptions); + + // Wait a moment for eventual consistency + await Task.Delay(100); + + // Assert - The deserialized store should have the same functionality + var messages = await deserializedStore.GetMessagesAsync(); + var messageList = messages.ToList(); + + Assert.Single(messageList); + Assert.Equal("Test serialization message", messageList[0].Text); + Assert.Equal(SessionId, deserializedStore.ConversationId); + Assert.Equal(s_testDatabaseId, deserializedStore.DatabaseId); + Assert.Equal(HierarchicalTestContainerId, deserializedStore.ContainerId); + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public async Task HierarchicalAndSimplePartitioning_ShouldCoexistAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + const string SessionId = "coexist-session"; + + // Create simple store using simple partitioning container and hierarchical store using hierarchical container + using var simpleStore = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, SessionId); + using var hierarchicalStore = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, "tenant-coexist", "user-coexist", SessionId); + + // Add messages to both + await simpleStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Simple partitioning message")]); + await hierarchicalStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Hierarchical partitioning message")]); + + // Wait a moment for eventual consistency + await Task.Delay(100); + + // Act & Assert + var simpleMessages = await simpleStore.GetMessagesAsync(); + var simpleMessageList = simpleMessages.ToList(); + + var hierarchicalMessages = await hierarchicalStore.GetMessagesAsync(); + var hierarchicalMessageList = hierarchicalMessages.ToList(); + + // Each should only see its own messages since they use different containers + Assert.Single(simpleMessageList); + Assert.Single(hierarchicalMessageList); + Assert.Equal("Simple partitioning message", simpleMessageList[0].Text); + Assert.Equal("Hierarchical partitioning message", hierarchicalMessageList[0].Text); + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public async Task MaxMessagesToRetrieve_ShouldLimitAndReturnMostRecentAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + const string ConversationId = "max-messages-test"; + + using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, ConversationId); + + // Add 10 messages + var messages = new List(); + for (int i = 1; i <= 10; i++) + { + messages.Add(new ChatMessage(ChatRole.User, $"Message {i}")); + await Task.Delay(10); // Small delay to ensure different timestamps + } + await store.AddMessagesAsync(messages); + + // Wait for eventual consistency + await Task.Delay(100); + + // Act - Set max to 5 and retrieve + store.MaxMessagesToRetrieve = 5; + var retrievedMessages = await store.GetMessagesAsync(); + var messageList = retrievedMessages.ToList(); + + // Assert - Should get the 5 most recent messages (6-10) in ascending order + Assert.Equal(5, messageList.Count); + Assert.Equal("Message 6", messageList[0].Text); + Assert.Equal("Message 7", messageList[1].Text); + Assert.Equal("Message 8", messageList[2].Text); + Assert.Equal("Message 9", messageList[3].Text); + Assert.Equal("Message 10", messageList[4].Text); + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public async Task MaxMessagesToRetrieve_Null_ShouldReturnAllMessagesAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + const string ConversationId = "max-messages-null-test"; + + using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, ConversationId); + + // Add 10 messages + var messages = new List(); + for (int i = 1; i <= 10; i++) + { + messages.Add(new ChatMessage(ChatRole.User, $"Message {i}")); + } + await store.AddMessagesAsync(messages); + + // Wait for eventual consistency + await Task.Delay(100); + + // Act - No limit set (default null) + var retrievedMessages = await store.GetMessagesAsync(); + var messageList = retrievedMessages.ToList(); + + // Assert - Should get all 10 messages + Assert.Equal(10, messageList.Count); + Assert.Equal("Message 1", messageList[0].Text); + Assert.Equal("Message 10", messageList[9].Text); + } + + #endregion +} diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs new file mode 100644 index 0000000000..dfa1f14221 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs @@ -0,0 +1,454 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Linq; +using System.Text.Json; +using System.Threading.Tasks; +using Microsoft.Agents.AI.Workflows; +using Microsoft.Agents.AI.Workflows.Checkpointing; +using Microsoft.Azure.Cosmos; +using Xunit; + +namespace Microsoft.Agents.AI.CosmosNoSql.UnitTests; + +/// +/// Contains tests for . +/// +/// Test Modes: +/// - Default Mode: Cleans up all test data after each test run (deletes database) +/// - Preserve Mode: Keeps containers and data for inspection in Cosmos DB Emulator Data Explorer +/// +/// To enable Preserve Mode, set environment variable: COSMOS_PRESERVE_CONTAINERS=true +/// Example: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test +/// +/// In Preserve Mode, you can view the data in Cosmos DB Emulator Data Explorer at: +/// https://localhost:8081/_explorer/index.html +/// Database: AgentFrameworkTests +/// Container: Checkpoints +/// +[Collection("CosmosDB")] +public class CosmosCheckpointStoreTests : IAsyncLifetime, IDisposable +{ + // Cosmos DB Emulator connection settings + private const string EmulatorEndpoint = "https://localhost:8081"; + private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; + private const string TestContainerId = "Checkpoints"; + // Use unique database ID per test class instance to avoid conflicts +#pragma warning disable CA1802 // Use literals where appropriate + private static readonly string s_testDatabaseId = $"AgentFrameworkTests-CheckpointStore-{Guid.NewGuid():N}"; +#pragma warning restore CA1802 + + private string _connectionString = string.Empty; + private CosmosClient? _cosmosClient; + private Database? _database; + private bool _emulatorAvailable; + private bool _preserveContainer; + + // JsonSerializerOptions configured for .NET 9+ compatibility + private static readonly JsonSerializerOptions s_jsonOptions = CreateJsonOptions(); + + private static JsonSerializerOptions CreateJsonOptions() + { + var options = new JsonSerializerOptions(); +#if NET9_0_OR_GREATER + options.TypeInfoResolver = new System.Text.Json.Serialization.Metadata.DefaultJsonTypeInfoResolver(); +#endif + return options; + } + + public async Task InitializeAsync() + { + // Check environment variable to determine if we should preserve containers + // Set COSMOS_PRESERVE_CONTAINERS=true to keep containers and data for inspection + this._preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase); + + this._connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey}"; + + try + { + this._cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); + + // Test connection by attempting to create database + this._database = await this._cosmosClient.CreateDatabaseIfNotExistsAsync(s_testDatabaseId); + await this._database.CreateContainerIfNotExistsAsync( + TestContainerId, + "/runId", + throughput: 400); + + this._emulatorAvailable = true; + } + catch (Exception ex) when (!(ex is OutOfMemoryException || ex is StackOverflowException || ex is AccessViolationException)) + { + // Emulator not available, tests will be skipped + this._emulatorAvailable = false; + this._cosmosClient?.Dispose(); + this._cosmosClient = null; + } + } + + public async Task DisposeAsync() + { + if (this._cosmosClient != null && this._emulatorAvailable) + { + try + { + if (this._preserveContainer) + { + // Preserve mode: Don't delete the database/container, keep data for inspection + // This allows viewing data in the Cosmos DB Emulator Data Explorer + // No cleanup needed - data persists for debugging + } + else + { + // Clean mode: Delete the test database and all data + await this._database!.DeleteAsync(); + } + } + catch (Exception ex) + { + // Ignore cleanup errors, but log for diagnostics + Console.WriteLine($"[DisposeAsync] Cleanup error: {ex.Message}\n{ex.StackTrace}"); + } + finally + { + this._cosmosClient.Dispose(); + } + } + } + + private void SkipIfEmulatorNotAvailable() + { + // In CI: Skip if COSMOS_EMULATOR_AVAILABLE is not set to "true" + // Locally: Skip if emulator connection check failed + var ciEmulatorAvailable = string.Equals(Environment.GetEnvironmentVariable("COSMOS_EMULATOR_AVAILABLE"), "true", StringComparison.OrdinalIgnoreCase); + + Xunit.Skip.If(!ciEmulatorAvailable && !this._emulatorAvailable, "Cosmos DB Emulator is not available"); + } + + #region Constructor Tests + + [SkippableFact] + public void Constructor_WithCosmosClient_SetsProperties() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + + // Act + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); + + // Assert + Assert.Equal(s_testDatabaseId, store.DatabaseId); + Assert.Equal(TestContainerId, store.ContainerId); + } + + [SkippableFact] + public void Constructor_WithConnectionString_SetsProperties() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + + // Act + using var store = new CosmosCheckpointStore(this._connectionString, s_testDatabaseId, TestContainerId); + + // Assert + Assert.Equal(s_testDatabaseId, store.DatabaseId); + Assert.Equal(TestContainerId, store.ContainerId); + } + + [SkippableFact] + public void Constructor_WithNullCosmosClient_ThrowsArgumentNullException() + { + // Act & Assert + Assert.Throws(() => + new CosmosCheckpointStore((CosmosClient)null!, s_testDatabaseId, TestContainerId)); + } + + [SkippableFact] + public void Constructor_WithNullConnectionString_ThrowsArgumentException() + { + // Act & Assert + Assert.Throws(() => + new CosmosCheckpointStore((string)null!, s_testDatabaseId, TestContainerId)); + } + + #endregion + + #region Checkpoint Operations Tests + + [SkippableFact] + public async Task CreateCheckpointAsync_NewCheckpoint_CreatesSuccessfullyAsync() + { + this.SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test checkpoint" }, s_jsonOptions); + + // Act + var checkpointInfo = await store.CreateCheckpointAsync(runId, checkpointValue); + + // Assert + Assert.NotNull(checkpointInfo); + Assert.Equal(runId, checkpointInfo.RunId); + Assert.NotNull(checkpointInfo.CheckpointId); + Assert.NotEmpty(checkpointInfo.CheckpointId); + } + + [SkippableFact] + public async Task RetrieveCheckpointAsync_ExistingCheckpoint_ReturnsCorrectValueAsync() + { + this.SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + var originalData = new { message = "Hello, World!", timestamp = DateTimeOffset.UtcNow }; + var checkpointValue = JsonSerializer.SerializeToElement(originalData, s_jsonOptions); + + // Act + var checkpointInfo = await store.CreateCheckpointAsync(runId, checkpointValue); + var retrievedValue = await store.RetrieveCheckpointAsync(runId, checkpointInfo); + + // Assert + Assert.Equal(JsonValueKind.Object, retrievedValue.ValueKind); + Assert.True(retrievedValue.TryGetProperty("message", out var messageProp)); + Assert.Equal("Hello, World!", messageProp.GetString()); + } + + [SkippableFact] + public async Task RetrieveCheckpointAsync_NonExistentCheckpoint_ThrowsInvalidOperationExceptionAsync() + { + this.SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + var fakeCheckpointInfo = new CheckpointInfo(runId, "nonexistent-checkpoint"); + + // Act & Assert + await Assert.ThrowsAsync(() => + store.RetrieveCheckpointAsync(runId, fakeCheckpointInfo).AsTask()); + } + + [SkippableFact] + public async Task RetrieveIndexAsync_EmptyStore_ReturnsEmptyCollectionAsync() + { + this.SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + + // Act + var index = await store.RetrieveIndexAsync(runId); + + // Assert + Assert.NotNull(index); + Assert.Empty(index); + } + + [SkippableFact] + public async Task RetrieveIndexAsync_WithCheckpoints_ReturnsAllCheckpointsAsync() + { + this.SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); + + // Create multiple checkpoints + var checkpoint1 = await store.CreateCheckpointAsync(runId, checkpointValue); + var checkpoint2 = await store.CreateCheckpointAsync(runId, checkpointValue); + var checkpoint3 = await store.CreateCheckpointAsync(runId, checkpointValue); + + // Act + var index = (await store.RetrieveIndexAsync(runId)).ToList(); + + // Assert + Assert.Equal(3, index.Count); + Assert.Contains(index, c => c.CheckpointId == checkpoint1.CheckpointId); + Assert.Contains(index, c => c.CheckpointId == checkpoint2.CheckpointId); + Assert.Contains(index, c => c.CheckpointId == checkpoint3.CheckpointId); + } + + [SkippableFact] + public async Task CreateCheckpointAsync_WithParent_CreatesHierarchyAsync() + { + this.SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); + + // Act + var parentCheckpoint = await store.CreateCheckpointAsync(runId, checkpointValue); + var childCheckpoint = await store.CreateCheckpointAsync(runId, checkpointValue, parentCheckpoint); + + // Assert + Assert.NotEqual(parentCheckpoint.CheckpointId, childCheckpoint.CheckpointId); + Assert.Equal(runId, parentCheckpoint.RunId); + Assert.Equal(runId, childCheckpoint.RunId); + } + + [SkippableFact] + public async Task RetrieveIndexAsync_WithParentFilter_ReturnsFilteredResultsAsync() + { + this.SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); + + // Create parent and child checkpoints + var parent = await store.CreateCheckpointAsync(runId, checkpointValue); + var child1 = await store.CreateCheckpointAsync(runId, checkpointValue, parent); + var child2 = await store.CreateCheckpointAsync(runId, checkpointValue, parent); + + // Create an orphan checkpoint + var orphan = await store.CreateCheckpointAsync(runId, checkpointValue); + + // Act + var allCheckpoints = (await store.RetrieveIndexAsync(runId)).ToList(); + var childrenOfParent = (await store.RetrieveIndexAsync(runId, parent)).ToList(); + + // Assert + Assert.Equal(4, allCheckpoints.Count); // parent + 2 children + orphan + Assert.Equal(2, childrenOfParent.Count); // only children + + Assert.Contains(childrenOfParent, c => c.CheckpointId == child1.CheckpointId); + Assert.Contains(childrenOfParent, c => c.CheckpointId == child2.CheckpointId); + Assert.DoesNotContain(childrenOfParent, c => c.CheckpointId == parent.CheckpointId); + Assert.DoesNotContain(childrenOfParent, c => c.CheckpointId == orphan.CheckpointId); + } + + #endregion + + #region Run Isolation Tests + + [SkippableFact] + public async Task CheckpointOperations_DifferentRuns_IsolatesDataAsync() + { + this.SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); + var runId1 = Guid.NewGuid().ToString(); + var runId2 = Guid.NewGuid().ToString(); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); + + // Act + var checkpoint1 = await store.CreateCheckpointAsync(runId1, checkpointValue); + var checkpoint2 = await store.CreateCheckpointAsync(runId2, checkpointValue); + + var index1 = (await store.RetrieveIndexAsync(runId1)).ToList(); + var index2 = (await store.RetrieveIndexAsync(runId2)).ToList(); + + // Assert + Assert.Single(index1); + Assert.Single(index2); + Assert.Equal(checkpoint1.CheckpointId, index1[0].CheckpointId); + Assert.Equal(checkpoint2.CheckpointId, index2[0].CheckpointId); + Assert.NotEqual(checkpoint1.CheckpointId, checkpoint2.CheckpointId); + } + + #endregion + + #region Error Handling Tests + + [SkippableFact] + public async Task CreateCheckpointAsync_WithNullRunId_ThrowsArgumentExceptionAsync() + { + this.SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); + + // Act & Assert + await Assert.ThrowsAsync(() => + store.CreateCheckpointAsync(null!, checkpointValue).AsTask()); + } + + [SkippableFact] + public async Task CreateCheckpointAsync_WithEmptyRunId_ThrowsArgumentExceptionAsync() + { + this.SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); + + // Act & Assert + await Assert.ThrowsAsync(() => + store.CreateCheckpointAsync("", checkpointValue).AsTask()); + } + + [SkippableFact] + public async Task RetrieveCheckpointAsync_WithNullCheckpointInfo_ThrowsArgumentNullExceptionAsync() + { + this.SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + + // Act & Assert + await Assert.ThrowsAsync(() => + store.RetrieveCheckpointAsync(runId, null!).AsTask()); + } + + #endregion + + #region Disposal Tests + + [SkippableFact] + public async Task Dispose_AfterDisposal_ThrowsObjectDisposedExceptionAsync() + { + this.SkipIfEmulatorNotAvailable(); + + // Arrange + var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); + + // Act + store.Dispose(); + + // Assert + await Assert.ThrowsAsync(() => + store.CreateCheckpointAsync("test-run", checkpointValue).AsTask()); + } + + [SkippableFact] + public void Dispose_MultipleCalls_DoesNotThrow() + { + this.SkipIfEmulatorNotAvailable(); + + // Arrange + var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); + + // Act & Assert (should not throw) + store.Dispose(); + store.Dispose(); + store.Dispose(); + } + + #endregion + + public void Dispose() + { + this.Dispose(true); + GC.SuppressFinalize(this); + } + + protected virtual void Dispose(bool disposing) + { + if (disposing) + { + this._cosmosClient?.Dispose(); + } + } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosDBCollectionFixture.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosDBCollectionFixture.cs new file mode 100644 index 0000000000..195c433de5 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosDBCollectionFixture.cs @@ -0,0 +1,18 @@ +// Copyright (c) Microsoft. All rights reserved. + +using Xunit; + +namespace Microsoft.Agents.AI.CosmosNoSql.UnitTests; + +/// +/// Defines a collection fixture for Cosmos DB tests to ensure they run sequentially. +/// This prevents race conditions and resource conflicts when tests create and delete +/// databases in the Cosmos DB Emulator. +/// +[CollectionDefinition("CosmosDB", DisableParallelization = true)] +public sealed class CosmosDBCollectionFixture +{ + // This class has no code, and is never created. Its purpose is simply + // to be the place to apply [CollectionDefinition] and all the + // ICollectionFixture<> interfaces. +} diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj new file mode 100644 index 0000000000..d60418ee2c --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj @@ -0,0 +1,24 @@ + + + + net10.0;net9.0 + $(NoWarn);MEAI001 + + + + false + + + + + + + + + + + + + + + diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_client.py b/python/packages/azure-ai/agent_framework_azure_ai/_client.py index d27d1d3edb..67bd0bae6a 100644 --- a/python/packages/azure-ai/agent_framework_azure_ai/_client.py +++ b/python/packages/azure-ai/agent_framework_azure_ai/_client.py @@ -155,7 +155,11 @@ class AzureAIClient(OpenAIBaseResponsesClient): self.credential = async_credential self.model_id = azure_ai_settings.model_deployment_name self.conversation_id = conversation_id - self._should_close_client = should_close_client # Track whether we should close client connection + + # Track whether the application endpoint is used + self._is_application_endpoint = "/applications/" in project_client._config.endpoint # type: ignore + # Track whether we should close client connection + self._should_close_client = should_close_client async def setup_azure_ai_observability(self, enable_sensitive_data: bool | None = None) -> None: """Use this method to setup tracing in your Azure AI Project. @@ -308,15 +312,19 @@ class AzureAIClient(OpenAIBaseResponsesClient): return result, instructions async def prepare_options( - self, messages: MutableSequence[ChatMessage], chat_options: ChatOptions + self, + messages: MutableSequence[ChatMessage], + chat_options: ChatOptions, + **kwargs: Any, ) -> dict[str, Any]: """Take ChatOptions and create the specific options for Azure AI.""" - chat_options.store = bool(chat_options.store or chat_options.store is None) prepared_messages, instructions = self._prepare_input(messages) - run_options = await super().prepare_options(prepared_messages, chat_options) - agent_reference = await self._get_agent_reference_or_create(run_options, instructions) + run_options = await super().prepare_options(prepared_messages, chat_options, **kwargs) - run_options["extra_body"] = {"agent": agent_reference} + if not self._is_application_endpoint: + # Application-scoped response APIs do not support "agent" property. + agent_reference = await self._get_agent_reference_or_create(run_options, instructions) + run_options["extra_body"] = {"agent": agent_reference} conversation_id = chat_options.conversation_id or self.conversation_id @@ -378,12 +386,12 @@ class AzureAIClient(OpenAIBaseResponsesClient): self, response: OpenAIResponse | ParsedResponse[BaseModel], store: bool | None ) -> str | None: """Get the conversation ID from the response if store is True.""" - if store: - # If conversation ID exists, it means that we operate with conversation - # so we use conversation ID as input and output. - if response.conversation and response.conversation.id: - return response.conversation.id - # If conversation ID doesn't exist, we operate with responses - # so we use response ID as input and output. - return response.id - return None + if store is False: + return None + # If conversation ID exists, it means that we operate with conversation + # so we use conversation ID as input and output. + if response.conversation and response.conversation.id: + return response.conversation.id + # If conversation ID doesn't exist, we operate with responses + # so we use response ID as input and output. + return response.id diff --git a/python/packages/azure-ai/tests/test_azure_ai_client.py b/python/packages/azure-ai/tests/test_azure_ai_client.py index 5d48e7b2be..8a31abcf35 100644 --- a/python/packages/azure-ai/tests/test_azure_ai_client.py +++ b/python/packages/azure-ai/tests/test_azure_ai_client.py @@ -87,6 +87,7 @@ def create_test_azure_ai_client( client.use_latest_version = use_latest_version client.model_id = azure_ai_settings.model_deployment_name client.conversation_id = conversation_id + client._is_application_endpoint = False # type: ignore client._should_close_client = should_close_client # type: ignore client.additional_properties = {} client.middleware = None @@ -305,6 +306,84 @@ async def test_azure_ai_client_prepare_options_basic(mock_project_client: MagicM assert run_options["extra_body"]["agent"]["name"] == "test-agent" +@pytest.mark.parametrize( + "endpoint,expects_agent", + [ + ("https://example.com/api/projects/my-project/applications/my-application/protocols", False), + ("https://example.com/api/projects/my-project", True), + ], +) +async def test_azure_ai_client_prepare_options_with_application_endpoint( + mock_azure_credential: MagicMock, endpoint: str, expects_agent: bool +) -> None: + client = AzureAIClient( + project_endpoint=endpoint, + model_deployment_name="test-model", + async_credential=mock_azure_credential, + agent_name="test-agent", + agent_version="1", + ) + + messages = [ChatMessage(role=Role.USER, contents=[TextContent(text="Hello")])] + chat_options = ChatOptions() + + with ( + patch.object(client.__class__.__bases__[0], "prepare_options", return_value={"model": "test-model"}), + patch.object( + client, + "_get_agent_reference_or_create", + return_value={"name": "test-agent", "version": "1", "type": "agent_reference"}, + ), + ): + run_options = await client.prepare_options(messages, chat_options) + + if expects_agent: + assert "extra_body" in run_options + assert run_options["extra_body"]["agent"]["name"] == "test-agent" + else: + assert "extra_body" not in run_options + + +@pytest.mark.parametrize( + "endpoint,expects_agent", + [ + ("https://example.com/api/projects/my-project/applications/my-application/protocols", False), + ("https://example.com/api/projects/my-project", True), + ], +) +async def test_azure_ai_client_prepare_options_with_application_project_client( + mock_project_client: MagicMock, endpoint: str, expects_agent: bool +) -> None: + mock_project_client._config = MagicMock() + mock_project_client._config.endpoint = endpoint + + client = AzureAIClient( + project_client=mock_project_client, + model_deployment_name="test-model", + agent_name="test-agent", + agent_version="1", + ) + + messages = [ChatMessage(role=Role.USER, contents=[TextContent(text="Hello")])] + chat_options = ChatOptions() + + with ( + patch.object(client.__class__.__bases__[0], "prepare_options", return_value={"model": "test-model"}), + patch.object( + client, + "_get_agent_reference_or_create", + return_value={"name": "test-agent", "version": "1", "type": "agent_reference"}, + ), + ): + run_options = await client.prepare_options(messages, chat_options) + + if expects_agent: + assert "extra_body" in run_options + assert run_options["extra_body"]["agent"]["name"] == "test-agent" + else: + assert "extra_body" not in run_options + + async def test_azure_ai_client_initialize_client(mock_project_client: MagicMock) -> None: """Test initialize_client method.""" client = create_test_azure_ai_client(mock_project_client) diff --git a/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py b/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py index e0bc3ba51a..e83a9244a7 100644 --- a/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py +++ b/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py @@ -9,6 +9,7 @@ with Azure Durable Entities, enabling stateful and durable AI agent execution. import json import re from collections.abc import Callable, Mapping +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, TypeVar, cast import azure.durable_functions as df @@ -39,6 +40,22 @@ logger = get_logger("agent_framework.azurefunctions") EntityHandler = Callable[[df.DurableEntityContext], None] HandlerT = TypeVar("HandlerT", bound=Callable[..., Any]) + +@dataclass +class AgentMetadata: + """Metadata for a registered agent. + + Attributes: + agent: The agent instance implementing AgentProtocol + http_endpoint_enabled: Whether HTTP endpoint is enabled for this agent + mcp_tool_enabled: Whether MCP tool endpoint is enabled for this agent + """ + + agent: AgentProtocol + http_endpoint_enabled: bool + mcp_tool_enabled: bool + + if TYPE_CHECKING: class DFAppBase: @@ -56,6 +73,15 @@ if TYPE_CHECKING: def activity_trigger(self, input_name: str) -> Callable[[HandlerT], HandlerT]: ... + def mcp_tool_trigger( + self, + arg_name: str, + tool_name: str, + description: str, + tool_properties: str, + data_type: func.DataType, + ) -> Callable[[HandlerT], HandlerT]: ... + else: DFAppBase = df.DFApp # type: ignore[assignment] @@ -117,14 +143,15 @@ class AgentFunctionApp(DFAppBase): agents: Dictionary of agent name to AgentProtocol instance enable_health_check: Whether health check endpoint is enabled enable_http_endpoints: Whether HTTP endpoints are created for agents + enable_mcp_tool_trigger: Whether MCP tool triggers are created for agents max_poll_retries: Maximum polling attempts when waiting for responses poll_interval_seconds: Delay (seconds) between polling attempts """ - agents: dict[str, AgentProtocol] + _agent_metadata: dict[str, AgentMetadata] enable_health_check: bool enable_http_endpoints: bool - agent_http_endpoint_flags: dict[str, bool] + enable_mcp_tool_trigger: bool def __init__( self, @@ -134,6 +161,7 @@ class AgentFunctionApp(DFAppBase): enable_http_endpoints: bool = True, max_poll_retries: int = DEFAULT_MAX_POLL_RETRIES, poll_interval_seconds: float = DEFAULT_POLL_INTERVAL_SECONDS, + enable_mcp_tool_trigger: bool = False, default_callback: AgentResponseCallbackProtocol | None = None, ): """Initialize the AgentFunctionApp. @@ -142,6 +170,8 @@ class AgentFunctionApp(DFAppBase): :param http_auth_level: HTTP authentication level (default: ``func.AuthLevel.FUNCTION``). :param enable_health_check: Enable the built-in health check endpoint (default: ``True``). :param enable_http_endpoints: Enable HTTP endpoints for agents (default: ``True``). + :param enable_mcp_tool_trigger: Enable MCP tool triggers for agents (default: ``False``). + When enabled, agents will be exposed as MCP tools that can be invoked by MCP-compatible clients. :param max_poll_retries: Maximum polling attempts when waiting for a response. Defaults to ``DEFAULT_MAX_POLL_RETRIES``. :param poll_interval_seconds: Delay in seconds between polling attempts. @@ -155,11 +185,11 @@ class AgentFunctionApp(DFAppBase): # Initialize parent DFApp super().__init__(http_auth_level=http_auth_level) - # Initialize agents dictionary - self.agents = {} - self.agent_http_endpoint_flags = {} + # Initialize agent metadata dictionary + self._agent_metadata = {} self.enable_health_check = enable_health_check self.enable_http_endpoints = enable_http_endpoints + self.enable_mcp_tool_trigger = enable_mcp_tool_trigger self.default_callback = default_callback try: @@ -186,11 +216,21 @@ class AgentFunctionApp(DFAppBase): logger.debug("[AgentFunctionApp] Initialization complete") + @property + def agents(self) -> dict[str, AgentProtocol]: + """Returns dict of agent names to agent instances. + + Returns: + Dictionary mapping agent names to their AgentProtocol instances. + """ + return {name: metadata.agent for name, metadata in self._agent_metadata.items()} + def add_agent( self, agent: AgentProtocol, callback: AgentResponseCallbackProtocol | None = None, enable_http_endpoint: bool | None = None, + enable_mcp_tool_trigger: bool | None = None, ) -> None: """Add an agent to the function app after initialization. @@ -198,8 +238,10 @@ class AgentFunctionApp(DFAppBase): agent: The Microsoft Agent Framework agent instance (must implement AgentProtocol) The agent must have a 'name' attribute. callback: Optional callback invoked during agent execution - enable_http_endpoint: Optional flag that overrides the app-level - HTTP endpoint setting for this agent + enable_http_endpoint: Optional flag to enable/disable HTTP endpoint for this agent. + The app level enable_http_endpoints setting will override this setting. + enable_mcp_tool_trigger: Optional flag to enable/disable MCP tool trigger for this agent. + The app level enable_mcp_tool_trigger setting will override this setting. Raises: ValueError: If the agent doesn't have a 'name' attribute or if an agent @@ -210,12 +252,17 @@ class AgentFunctionApp(DFAppBase): if name is None: raise ValueError("Agent does not have a 'name' attribute. All agents must have a 'name' attribute.") - if name in self.agents: + if name in self._agent_metadata: raise ValueError(f"Agent with name '{name}' is already registered. Each agent must have a unique name.") effective_enable_http_endpoint = ( self.enable_http_endpoints if enable_http_endpoint is None else self._coerce_to_bool(enable_http_endpoint) ) + effective_enable_mcp_endpoint = ( + self.enable_mcp_tool_trigger + if enable_mcp_tool_trigger is None + else self._coerce_to_bool(enable_mcp_tool_trigger) + ) logger.debug(f"[AgentFunctionApp] Adding agent: {name}") logger.debug(f"[AgentFunctionApp] Route: /api/agents/{name}") @@ -224,17 +271,21 @@ class AgentFunctionApp(DFAppBase): "enabled" if effective_enable_http_endpoint else "disabled", name, ) + logger.debug( + f"[AgentFunctionApp] MCP tool trigger: {'enabled' if effective_enable_mcp_endpoint else 'disabled'}" + ) - self.agents[name] = agent - self.agent_http_endpoint_flags[name] = effective_enable_http_endpoint + # Store agent metadata + self._agent_metadata[name] = AgentMetadata( + agent=agent, + http_endpoint_enabled=effective_enable_http_endpoint, + mcp_tool_enabled=effective_enable_mcp_endpoint, + ) effective_callback = callback or self.default_callback self._setup_agent_functions( - agent, - name, - effective_callback, - effective_enable_http_endpoint, + agent, name, effective_callback, effective_enable_http_endpoint, effective_enable_mcp_endpoint ) logger.debug(f"[AgentFunctionApp] Agent '{name}' added successfully") @@ -258,7 +309,7 @@ class AgentFunctionApp(DFAppBase): """ normalized_name = str(agent_name) - if normalized_name not in self.agents: + if normalized_name not in self._agent_metadata: raise ValueError(f"Agent '{normalized_name}' is not registered with this app.") return DurableAIAgent(context, normalized_name) @@ -269,15 +320,16 @@ class AgentFunctionApp(DFAppBase): agent_name: str, callback: AgentResponseCallbackProtocol | None, enable_http_endpoint: bool, + enable_mcp_tool_trigger: bool, ) -> None: - """Set up the HTTP trigger and entity for a specific agent. + """Set up the HTTP trigger, entity, and MCP tool trigger for a specific agent. Args: agent: The agent instance agent_name: The name to use for routing and entity registration callback: Optional callback to receive response updates - enable_http_endpoint: Whether the HTTP run route is enabled for - this agent + enable_http_endpoint: Whether to create HTTP endpoint + enable_mcp_tool_trigger: Whether to create MCP tool trigger """ logger.debug(f"[AgentFunctionApp] Setting up functions for agent '{agent_name}'...") @@ -290,6 +342,12 @@ class AgentFunctionApp(DFAppBase): ) self._setup_agent_entity(agent, agent_name, callback) + if enable_mcp_tool_trigger: + agent_description = agent.description + self._setup_mcp_tool_trigger(agent_name, agent_description) + else: + logger.debug(f"[AgentFunctionApp] MCP tool trigger disabled for agent '{agent_name}'") + def _setup_http_run_route(self, agent_name: str) -> None: """Register the POST route that triggers agent execution. @@ -448,6 +506,159 @@ class AgentFunctionApp(DFAppBase): entity_function.__name__ = entity_name_with_prefix self.entity_trigger(context_name="context", entity_name=entity_name_with_prefix)(entity_function) + def _setup_mcp_tool_trigger(self, agent_name: str, agent_description: str | None) -> None: + """Register an MCP tool trigger for an agent using Azure Functions native MCP support. + + This creates a native Azure Functions MCP tool trigger that exposes the agent + as an MCP tool, allowing it to be invoked by MCP-compatible clients. + + Args: + agent_name: The agent name (used as the MCP tool name) + agent_description: Optional description for the MCP tool (shown to clients) + """ + mcp_function_name = self._build_function_name(agent_name, "mcptool") + + # Define tool properties as JSON (MCP tool parameters) + tool_properties = json.dumps([ + { + "propertyName": "query", + "propertyType": "string", + "description": "The query to send to the agent.", + "isRequired": True, + "isArray": False, + }, + { + "propertyName": "threadId", + "propertyType": "string", + "description": "Optional thread identifier for conversation continuity.", + "isRequired": False, + "isArray": False, + }, + ]) + + function_name_decorator = self.function_name(mcp_function_name) + mcp_tool_decorator = self.mcp_tool_trigger( + arg_name="context", + tool_name=agent_name, + description=agent_description or f"Interact with {agent_name} agent", + tool_properties=tool_properties, + data_type=func.DataType.UNDEFINED, + ) + durable_client_decorator = self.durable_client_input(client_name="client") + + @function_name_decorator + @mcp_tool_decorator + @durable_client_decorator + async def mcp_tool_handler(context: str, client: df.DurableOrchestrationClient) -> str: + """Handle MCP tool invocation for the agent. + + Args: + context: MCP tool invocation context containing arguments (query, threadId) + client: Durable orchestration client for entity communication + + Returns: + Agent response text + """ + logger.debug("[MCP Tool Trigger] Received invocation for agent: %s", agent_name) + return await self._handle_mcp_tool_invocation(agent_name=agent_name, context=context, client=client) + + logger.debug("[AgentFunctionApp] Registered MCP tool trigger for agent: %s", agent_name) + + async def _handle_mcp_tool_invocation( + self, agent_name: str, context: str, client: df.DurableOrchestrationClient + ) -> str: + """Handle an MCP tool invocation. + + This method processes MCP tool requests and delegates to the agent entity. + + Args: + agent_name: Name of the agent being invoked + context: MCP tool invocation context as a JSON string + client: Durable orchestration client + + Returns: + Agent response text + + Raises: + ValueError: If required arguments are missing or context is invalid JSON + RuntimeError: If agent execution fails + """ + logger.debug("[MCP Tool Handler] Processing invocation for agent '%s'", agent_name) + + # Parse JSON context string + try: + parsed_context = json.loads(context) + except json.JSONDecodeError as e: + raise ValueError(f"Invalid MCP context format: {e}") from e + + # Extract arguments from MCP context + arguments = parsed_context.get("arguments", {}) if isinstance(parsed_context, dict) else {} + + # Validate required 'query' argument + query = arguments.get("query") + if not query or not isinstance(query, str): + raise ValueError("MCP Tool invocation is missing required 'query' argument of type string.") + + # Extract optional threadId + thread_id = arguments.get("threadId") + + # Create or parse session ID + if thread_id and isinstance(thread_id, str) and thread_id.strip(): + try: + session_id = AgentSessionId.parse(thread_id) + except ValueError as e: + logger.warning( + "Failed to parse AgentSessionId from thread_id '%s': %s. Falling back to new session ID.", + thread_id, + e, + ) + session_id = AgentSessionId(name=agent_name, key=thread_id) + else: + # Generate new session ID + session_id = AgentSessionId.with_random_key(agent_name) + + # Build entity instance ID + entity_instance_id = session_id.to_entity_id() + + # Create run request + correlation_id = self._generate_unique_id() + run_request = self._build_request_data( + req_body={"message": query, "role": "user"}, + message=query, + thread_id=str(session_id), + correlation_id=correlation_id, + request_response_format=REQUEST_RESPONSE_FORMAT_TEXT, + ) + + query_preview = query[:50] + "..." if len(query) > 50 else query + logger.info("[MCP Tool] Invoking agent '%s' with query: %s", agent_name, query_preview) + + # Signal entity to run agent + await client.signal_entity(entity_instance_id, "run_agent", run_request) + + # Poll for response (similar to HTTP handler) + try: + result = await self._get_response_from_entity( + client=client, + entity_instance_id=entity_instance_id, + correlation_id=correlation_id, + message=query, + thread_id=str(session_id), + ) + + # Extract and return response text + if result.get("status") == "success": + response_text = str(result.get("response", "No response")) + logger.info("[MCP Tool] Agent '%s' responded successfully", agent_name) + return response_text + error_msg = result.get("error", "Unknown error") + logger.error("[MCP Tool] Agent '%s' execution failed: %s", agent_name, error_msg) + raise RuntimeError(f"Agent execution failed: {error_msg}") + + except Exception as exc: + logger.error("[MCP Tool] Error invoking agent '%s': %s", agent_name, exc, exc_info=True) + raise + def _setup_health_route(self) -> None: """Register the optional health check route.""" health_route = self.route(route="health", methods=["GET"]) @@ -458,16 +669,14 @@ class AgentFunctionApp(DFAppBase): agent_info = [ { "name": name, - "type": type(agent).__name__, - "http_endpoint_enabled": self.agent_http_endpoint_flags.get( - name, - self.enable_http_endpoints, - ), + "type": type(metadata.agent).__name__, + "http_endpoint_enabled": metadata.http_endpoint_enabled, + "mcp_tool_enabled": metadata.mcp_tool_enabled, } - for name, agent in self.agents.items() + for name, metadata in self._agent_metadata.items() ] return func.HttpResponse( - json.dumps({"status": "healthy", "agents": agent_info, "agent_count": len(self.agents)}), + json.dumps({"status": "healthy", "agents": agent_info, "agent_count": len(self._agent_metadata)}), status_code=200, mimetype=MIMETYPE_APPLICATION_JSON, ) diff --git a/python/packages/azurefunctions/tests/test_app.py b/python/packages/azurefunctions/tests/test_app.py index ebf6eef3e6..c65368e79e 100644 --- a/python/packages/azurefunctions/tests/test_app.py +++ b/python/packages/azurefunctions/tests/test_app.py @@ -2,6 +2,7 @@ """Unit tests for AgentFunctionApp.""" +import json from collections.abc import Awaitable, Callable from typing import Any, TypeVar from unittest.mock import ANY, AsyncMock, Mock, patch @@ -87,7 +88,7 @@ class TestAgentFunctionAppInit: app.add_agent(mock_agent, callback=specific_callback) setup_mock.assert_called_once() - _, _, passed_callback, enable_http_endpoint = setup_mock.call_args[0] + _, _, passed_callback, enable_http_endpoint, enable_mcp_tool_trigger = setup_mock.call_args[0] assert passed_callback is specific_callback assert enable_http_endpoint is True @@ -103,7 +104,7 @@ class TestAgentFunctionAppInit: app.add_agent(mock_agent) setup_mock.assert_called_once() - _, _, passed_callback, enable_http_endpoint = setup_mock.call_args[0] + _, _, passed_callback, enable_http_endpoint, enable_mcp_tool_trigger = setup_mock.call_args[0] assert passed_callback is default_callback assert enable_http_endpoint is True @@ -118,7 +119,7 @@ class TestAgentFunctionAppInit: AgentFunctionApp(agents=[mock_agent], default_callback=default_callback) setup_mock.assert_called_once() - _, _, passed_callback, enable_http_endpoint = setup_mock.call_args[0] + _, _, passed_callback, enable_http_endpoint, enable_mcp_tool_trigger = setup_mock.call_args[0] assert passed_callback is default_callback assert enable_http_endpoint is True @@ -239,7 +240,7 @@ class TestAgentFunctionAppSetup: http_route_mock.assert_called_once_with("OverrideAgent") agent_entity_mock.assert_called_once_with(mock_agent, "OverrideAgent", ANY) - assert app.agent_http_endpoint_flags["OverrideAgent"] is True + assert app._agent_metadata["OverrideAgent"].http_endpoint_enabled is True def test_agent_override_disables_http_route_when_app_enabled(self) -> None: """Agent-level override should disable HTTP route even when app enables it.""" @@ -256,7 +257,7 @@ class TestAgentFunctionAppSetup: http_route_mock.assert_not_called() agent_entity_mock.assert_called_once_with(mock_agent, "DisabledOverride", ANY) - assert app.agent_http_endpoint_flags["DisabledOverride"] is False + assert app._agent_metadata["DisabledOverride"].http_endpoint_enabled is False def test_multiple_apps_independent(self) -> None: """Test that multiple AgentFunctionApp instances are independent.""" @@ -797,5 +798,271 @@ class TestHttpRunRoute: client.signal_entity.assert_not_called() +class TestMCPToolEndpoint: + """Test suite for MCP tool endpoint functionality.""" + + def test_init_with_mcp_tool_endpoint_enabled(self) -> None: + """Test initialization with MCP tool endpoint enabled.""" + mock_agent = Mock() + mock_agent.name = "TestAgent" + + app = AgentFunctionApp(agents=[mock_agent], enable_mcp_tool_trigger=True) + + assert app.enable_mcp_tool_trigger is True + + def test_init_with_mcp_tool_endpoint_disabled(self) -> None: + """Test initialization with MCP tool endpoint disabled (default).""" + mock_agent = Mock() + mock_agent.name = "TestAgent" + + app = AgentFunctionApp(agents=[mock_agent]) + + assert app.enable_mcp_tool_trigger is False + + def test_add_agent_with_mcp_tool_trigger_enabled(self) -> None: + """Test adding an agent with MCP tool trigger explicitly enabled.""" + mock_agent = Mock() + mock_agent.name = "MCPAgent" + mock_agent.description = "Test MCP Agent" + + with patch.object(AgentFunctionApp, "_setup_agent_functions") as setup_mock: + app = AgentFunctionApp() + app.add_agent(mock_agent, enable_mcp_tool_trigger=True) + + setup_mock.assert_called_once() + _, _, _, _, enable_mcp = setup_mock.call_args[0] + assert enable_mcp is True + + def test_add_agent_with_mcp_tool_trigger_disabled(self) -> None: + """Test adding an agent with MCP tool trigger explicitly disabled.""" + mock_agent = Mock() + mock_agent.name = "NoMCPAgent" + + with patch.object(AgentFunctionApp, "_setup_agent_functions") as setup_mock: + app = AgentFunctionApp(enable_mcp_tool_trigger=True) + app.add_agent(mock_agent, enable_mcp_tool_trigger=False) + + setup_mock.assert_called_once() + _, _, _, _, enable_mcp = setup_mock.call_args[0] + assert enable_mcp is False + + def test_agent_override_enables_mcp_when_app_disabled(self) -> None: + """Test that per-agent override can enable MCP when app-level is disabled.""" + mock_agent = Mock() + mock_agent.name = "OverrideAgent" + + with patch.object(AgentFunctionApp, "_setup_mcp_tool_trigger") as mcp_setup_mock: + app = AgentFunctionApp(enable_mcp_tool_trigger=False) + app.add_agent(mock_agent, enable_mcp_tool_trigger=True) + + mcp_setup_mock.assert_called_once() + + def test_agent_override_disables_mcp_when_app_enabled(self) -> None: + """Test that per-agent override can disable MCP when app-level is enabled.""" + mock_agent = Mock() + mock_agent.name = "NoOverrideAgent" + + with patch.object(AgentFunctionApp, "_setup_mcp_tool_trigger") as mcp_setup_mock: + app = AgentFunctionApp(enable_mcp_tool_trigger=True) + app.add_agent(mock_agent, enable_mcp_tool_trigger=False) + + mcp_setup_mock.assert_not_called() + + def test_setup_mcp_tool_trigger_registers_decorators(self) -> None: + """Test that _setup_mcp_tool_trigger registers the correct decorators.""" + mock_agent = Mock() + mock_agent.name = "MCPToolAgent" + mock_agent.description = "Test MCP Tool" + + app = AgentFunctionApp() + + # Mock the decorators + with ( + patch.object(app, "function_name") as func_name_mock, + patch.object(app, "mcp_tool_trigger") as mcp_trigger_mock, + patch.object(app, "durable_client_input") as client_mock, + ): + # Setup mock decorator chain + func_name_mock.return_value = lambda f: f + mcp_trigger_mock.return_value = lambda f: f + client_mock.return_value = lambda f: f + + app._setup_mcp_tool_trigger(mock_agent.name, mock_agent.description) + + # Verify decorators were called with correct parameters + func_name_mock.assert_called_once() + mcp_trigger_mock.assert_called_once_with( + arg_name="context", + tool_name=mock_agent.name, + description=mock_agent.description, + tool_properties=ANY, + data_type=func.DataType.UNDEFINED, + ) + client_mock.assert_called_once_with(client_name="client") + + def test_setup_mcp_tool_trigger_uses_default_description(self) -> None: + """Test that _setup_mcp_tool_trigger uses default description when none provided.""" + mock_agent = Mock() + mock_agent.name = "NoDescAgent" + + app = AgentFunctionApp() + + with ( + patch.object(app, "function_name", return_value=lambda f: f), + patch.object(app, "mcp_tool_trigger") as mcp_trigger_mock, + patch.object(app, "durable_client_input", return_value=lambda f: f), + ): + mcp_trigger_mock.return_value = lambda f: f + + app._setup_mcp_tool_trigger(mock_agent.name, None) + + # Verify default description was used + call_args = mcp_trigger_mock.call_args + assert call_args[1]["description"] == f"Interact with {mock_agent.name} agent" + + async def test_handle_mcp_tool_invocation_with_json_string(self) -> None: + """Test _handle_mcp_tool_invocation with JSON string context.""" + mock_agent = Mock() + mock_agent.name = "TestAgent" + + app = AgentFunctionApp(agents=[mock_agent]) + client = AsyncMock() + + # Mock the entity response + mock_state = Mock() + mock_state.entity_state = { + "schemaVersion": "1.0.0", + "data": {"conversationHistory": []}, + } + client.read_entity_state.return_value = mock_state + + # Create JSON string context + context = '{"arguments": {"query": "test query", "threadId": "test-thread"}}' + + with patch.object(app, "_get_response_from_entity") as get_response_mock: + get_response_mock.return_value = {"status": "success", "response": "Test response"} + + result = await app._handle_mcp_tool_invocation("TestAgent", context, client) + + assert result == "Test response" + get_response_mock.assert_called_once() + + async def test_handle_mcp_tool_invocation_with_json_context(self) -> None: + """Test _handle_mcp_tool_invocation with JSON string context.""" + mock_agent = Mock() + mock_agent.name = "TestAgent" + + app = AgentFunctionApp(agents=[mock_agent]) + client = AsyncMock() + + # Mock the entity response + mock_state = Mock() + mock_state.entity_state = { + "schemaVersion": "1.0.0", + "data": {"conversationHistory": []}, + } + client.read_entity_state.return_value = mock_state + + # Create JSON string context + context = json.dumps({"arguments": {"query": "test query", "threadId": "test-thread"}}) + + with patch.object(app, "_get_response_from_entity") as get_response_mock: + get_response_mock.return_value = {"status": "success", "response": "Test response"} + + result = await app._handle_mcp_tool_invocation("TestAgent", context, client) + + assert result == "Test response" + get_response_mock.assert_called_once() + + async def test_handle_mcp_tool_invocation_missing_query(self) -> None: + """Test _handle_mcp_tool_invocation raises ValueError when query is missing.""" + mock_agent = Mock() + mock_agent.name = "TestAgent" + + app = AgentFunctionApp(agents=[mock_agent]) + client = AsyncMock() + + # Context missing query (as JSON string) + context = json.dumps({"arguments": {}}) + + with pytest.raises(ValueError, match="missing required 'query' argument"): + await app._handle_mcp_tool_invocation("TestAgent", context, client) + + async def test_handle_mcp_tool_invocation_invalid_json(self) -> None: + """Test _handle_mcp_tool_invocation raises ValueError for invalid JSON.""" + mock_agent = Mock() + mock_agent.name = "TestAgent" + + app = AgentFunctionApp(agents=[mock_agent]) + client = AsyncMock() + + # Invalid JSON string + context = "not valid json" + + with pytest.raises(ValueError, match="Invalid MCP context format"): + await app._handle_mcp_tool_invocation("TestAgent", context, client) + + async def test_handle_mcp_tool_invocation_runtime_error(self) -> None: + """Test _handle_mcp_tool_invocation raises RuntimeError when agent fails.""" + mock_agent = Mock() + mock_agent.name = "TestAgent" + + app = AgentFunctionApp(agents=[mock_agent]) + client = AsyncMock() + + # Mock the entity response + mock_state = Mock() + mock_state.entity_state = { + "schemaVersion": "1.0.0", + "data": {"conversationHistory": []}, + } + client.read_entity_state.return_value = mock_state + + context = '{"arguments": {"query": "test query"}}' + + with patch.object(app, "_get_response_from_entity") as get_response_mock: + get_response_mock.return_value = {"status": "failed", "error": "Agent error"} + + with pytest.raises(RuntimeError, match="Agent execution failed"): + await app._handle_mcp_tool_invocation("TestAgent", context, client) + + def test_health_check_includes_mcp_tool_enabled(self) -> None: + """Test that health check endpoint includes mcp_tool_enabled field.""" + mock_agent = Mock() + mock_agent.name = "HealthAgent" + + app = AgentFunctionApp(agents=[mock_agent], enable_mcp_tool_trigger=True) + + # Capture the health check handler function + captured_handler = None + + def capture_decorator(*args, **kwargs): + def decorator(func): + nonlocal captured_handler + captured_handler = func + return func + + return decorator + + with patch.object(app, "route", side_effect=capture_decorator): + app._setup_health_route() + + # Verify we captured the handler + assert captured_handler is not None + + # Call the health handler + request = Mock() + response = captured_handler(request) + + # Verify response includes mcp_tool_enabled + import json + + body = json.loads(response.get_body().decode("utf-8")) + assert "agents" in body + assert len(body["agents"]) == 1 + assert "mcp_tool_enabled" in body["agents"][0] + assert body["agents"][0]["mcp_tool_enabled"] is True + + if __name__ == "__main__": pytest.main([__file__, "-v", "--tb=short"]) diff --git a/python/packages/chatkit/agent_framework_chatkit/_converter.py b/python/packages/chatkit/agent_framework_chatkit/_converter.py index 4c911f5604..0adf040156 100644 --- a/python/packages/chatkit/agent_framework_chatkit/_converter.py +++ b/python/packages/chatkit/agent_framework_chatkit/_converter.py @@ -27,6 +27,7 @@ from chatkit.types import ( EndOfTurnItem, HiddenContextItem, ImageAttachment, + SDKHiddenContextItem, TaskItem, ThreadItem, UserMessageItem, @@ -180,8 +181,10 @@ class ThreadItemConverter: # Subclasses can override this method to provide custom handling return None - def hidden_context_to_input(self, item: HiddenContextItem) -> ChatMessage | list[ChatMessage] | None: - """Convert a ChatKit HiddenContextItem to Agent Framework ChatMessage(s). + def hidden_context_to_input( + self, item: HiddenContextItem | SDKHiddenContextItem + ) -> ChatMessage | list[ChatMessage] | None: + """Convert a ChatKit HiddenContextItem or SDKHiddenContextItem to Agent Framework ChatMessage(s). This method is called internally by `to_agent_input()`. Override this method to customize how hidden context is converted. @@ -522,6 +525,9 @@ class ThreadItemConverter: case HiddenContextItem(): out = self.hidden_context_to_input(item) or [] return out if isinstance(out, list) else [out] + case SDKHiddenContextItem(): + out = self.hidden_context_to_input(item) or [] + return out if isinstance(out, list) else [out] case _: assert_never(item) diff --git a/python/packages/chatkit/pyproject.toml b/python/packages/chatkit/pyproject.toml index 1e2e7bdbd8..a987974a44 100644 --- a/python/packages/chatkit/pyproject.toml +++ b/python/packages/chatkit/pyproject.toml @@ -23,7 +23,7 @@ classifiers = [ ] dependencies = [ "agent-framework-core", - "openai-chatkit>=1.1.0,<2.0.0", + "openai-chatkit>=1.4.0,<2.0.0", ] [tool.uv] diff --git a/python/packages/core/agent_framework/_clients.py b/python/packages/core/agent_framework/_clients.py index 40c13a2037..4d91492822 100644 --- a/python/packages/core/agent_framework/_clients.py +++ b/python/packages/core/agent_framework/_clients.py @@ -568,10 +568,6 @@ class BaseChatClient(SerializationMixin, ABC): additional_properties=additional_properties, ) - # Validate that store is True when conversation_id is set - if chat_options.conversation_id is not None and chat_options.store is not True: - chat_options.store = True - if chat_options.instructions: system_msg = ChatMessage(role="system", text=chat_options.instructions) prepped_messages = [system_msg, *prepare_messages(messages)] @@ -666,10 +662,6 @@ class BaseChatClient(SerializationMixin, ABC): additional_properties=additional_properties, ) - # Validate that store is True when conversation_id is set - if chat_options.conversation_id is not None and chat_options.store is not True: - chat_options.store = True - if chat_options.instructions: system_msg = ChatMessage(role="system", text=chat_options.instructions) prepped_messages = [system_msg, *prepare_messages(messages)] diff --git a/python/packages/core/agent_framework/_workflows/__init__.py b/python/packages/core/agent_framework/_workflows/__init__.py index 35c9f11501..990264df41 100644 --- a/python/packages/core/agent_framework/_workflows/__init__.py +++ b/python/packages/core/agent_framework/_workflows/__init__.py @@ -61,6 +61,8 @@ from ._group_chat import ( GroupChatDirective, GroupChatStateSnapshot, ManagerDirectiveModel, + ManagerSelectionRequest, + ManagerSelectionResponse, ) from ._handoff import HandoffBuilder, HandoffUserInputRequest from ._magentic import ( @@ -147,6 +149,8 @@ __all__ = [ "MagenticPlanReviewReply", "MagenticPlanReviewRequest", "ManagerDirectiveModel", + "ManagerSelectionRequest", + "ManagerSelectionResponse", "Message", "OrchestrationState", "RequestInfoEvent", diff --git a/python/packages/core/agent_framework/_workflows/_group_chat.py b/python/packages/core/agent_framework/_workflows/_group_chat.py index 84859a4f0c..78ddb5c2eb 100644 --- a/python/packages/core/agent_framework/_workflows/_group_chat.py +++ b/python/packages/core/agent_framework/_workflows/_group_chat.py @@ -24,13 +24,12 @@ import logging from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass, field from types import MappingProxyType -from typing import Any, TypeAlias +from typing import Any, TypeAlias, cast from uuid import uuid4 from pydantic import BaseModel, Field -from .._agents import AgentProtocol -from .._clients import ChatClientProtocol +from .._agents import AgentProtocol, ChatAgent from .._types import ChatMessage, Role from ._agent_executor import AgentExecutorRequest, AgentExecutorResponse from ._base_group_chat_orchestrator import BaseGroupChatOrchestrator @@ -87,6 +86,75 @@ class GroupChatDirective: final_message: ChatMessage | None = None +@dataclass +class ManagerSelectionRequest: + """Request sent to manager agent for next speaker selection. + + This dataclass packages the full conversation state and task context + for the manager agent to analyze and make a speaker selection decision. + + Attributes: + task: Original user task message + participants: Mapping of participant names to their descriptions + conversation: Full conversation history including all messages + round_index: Number of manager selection rounds completed so far + metadata: Optional metadata for extensibility + """ + + task: ChatMessage + participants: dict[str, str] # type: ignore + conversation: list[ChatMessage] # type: ignore + round_index: int + metadata: dict[str, Any] | None = None + + def to_dict(self) -> dict[str, Any]: + """Convert to dictionary for serialization.""" + return { + "task": self.task.to_dict(), + "participants": dict(self.participants), + "conversation": [msg.to_dict() for msg in self.conversation], + "round_index": self.round_index, + "metadata": self.metadata, + } + + +class ManagerSelectionResponse(BaseModel): + """Response from manager agent with speaker selection decision. + + The manager agent must produce this structure (or compatible dict/JSON) + to communicate its decision back to the orchestrator. + + Attributes: + selected_participant: Name of participant to speak next (None = finish conversation) + instruction: Optional instruction to provide to the selected participant + finish: Whether the conversation should be completed + final_message: Optional final message string when finishing conversation (will be converted to ChatMessage) + """ + + model_config = {"extra": "forbid"} + + selected_participant: str | None = None + instruction: str | None = None + finish: bool = False + final_message: str | None = Field(default=None, description="Optional text content for final message") + + @staticmethod + def from_dict(data: dict[str, Any]) -> "ManagerSelectionResponse": + """Create from dictionary representation.""" + return ManagerSelectionResponse( + selected_participant=data.get("selected_participant"), + instruction=data.get("instruction"), + finish=data.get("finish", False), + final_message=data.get("final_message"), + ) + + def get_final_message_as_chat_message(self) -> ChatMessage | None: + """Convert final_message string to ChatMessage if present.""" + if self.final_message: + return ChatMessage(role=Role.ASSISTANT, text=self.final_message) + return None + + # endregion @@ -112,17 +180,23 @@ class _GroupChatConfig: """Internal: Configuration passed to factories during workflow assembly. Attributes: - manager: Manager instance responsible for orchestration decisions (None when custom factory handles it) + manager: Manager callable for orchestration decisions (used by set_select_speakers_func) + manager_participant: Manager agent/executor instance (used by set_manager) manager_name: Display name for the manager in conversation history participants: Mapping of participant names to their specifications max_rounds: Optional limit on manager selection rounds to prevent infinite loops + termination_condition: Optional callable that halts the conversation when it returns True orchestrator: Orchestrator executor instance (populated during build) + participant_aliases: Mapping of aliases to executor IDs + participant_executors: Mapping of participant names to their executor instances """ manager: _GroupChatManagerFn | None + manager_participant: AgentProtocol | Executor | None manager_name: str participants: Mapping[str, GroupChatParticipantSpec] max_rounds: int | None = None + termination_condition: Callable[[list[ChatMessage]], bool | Awaitable[bool]] | None = None orchestrator: Executor | None = None participant_aliases: dict[str, str] = field(default_factory=dict) # type: ignore[type-arg] participant_executors: dict[str, Executor] = field(default_factory=dict) # type: ignore[type-arg] @@ -220,6 +294,7 @@ class GroupChatOrchestratorExecutor(BaseGroupChatOrchestrator): participants: Mapping of participant names to descriptions (for manager context) manager_name: Display name for manager in conversation history max_rounds: Optional limit on manager selection rounds (None = unlimited) + termination_condition: Optional callable that halts the conversation when it returns True executor_id: Optional custom ID for observability (auto-generated if not provided) """ @@ -230,6 +305,7 @@ class GroupChatOrchestratorExecutor(BaseGroupChatOrchestrator): participants: Mapping[str, str], manager_name: str, max_rounds: int | None = None, + termination_condition: Callable[[list[ChatMessage]], bool | Awaitable[bool]] | None = None, executor_id: str | None = None, ) -> None: super().__init__(executor_id or f"groupchat_orchestrator_{uuid4().hex[:8]}") @@ -237,9 +313,11 @@ class GroupChatOrchestratorExecutor(BaseGroupChatOrchestrator): self._participants = dict(participants) self._manager_name = manager_name self._max_rounds = max_rounds + self._termination_condition = termination_condition self._history: list[_GroupChatTurn] = [] self._task_message: ChatMessage | None = None self._pending_agent: str | None = None + self._pending_finalization: bool = False # Stashes the initial conversation list until _handle_task_message normalizes it into _conversation. self._pending_initial_conversation: list[ChatMessage] | None = None @@ -317,10 +395,75 @@ class GroupChatOrchestratorExecutor(BaseGroupChatOrchestrator): for turn in metadata["history"] ] + async def _complete_on_termination( + self, + ctx: WorkflowContext[AgentExecutorRequest | _GroupChatRequestMessage, list[ChatMessage]], + ) -> bool: + """Finish the conversation early when the termination condition is met.""" + if not await self._check_termination(): + return False + + if self._is_manager_agent(): + if self._pending_finalization: + return True + + self._pending_finalization = True + termination_prompt = ChatMessage( + role=Role.SYSTEM, + text="Termination condition met. Provide a final manager summary and finish the conversation.", + ) + manager_conversation = [ + self._build_manager_context_message(), + termination_prompt, + *list(self._conversation), + ] + self._pending_agent = self._manager_name + await self._route_to_participant( + participant_name=self._manager_name, + conversation=manager_conversation, + ctx=ctx, + instruction="", + task=self._task_message, + metadata={"termination_condition": True}, + ) + return True + + final_message: ChatMessage | None = None + if self._manager is not None and not self._is_manager_agent(): + try: + directive = await self._manager(self._build_state()) + except Exception: + logger.warning("Manager finalization failed during termination; using default termination message.") + else: + if directive.final_message is not None: + final_message = ensure_author(directive.final_message, self._manager_name) + elif directive.finish: + final_message = ensure_author( + self._create_completion_message( + text="Conversation completed.", + reason="termination_condition_manager_finish", + ), + self._manager_name, + ) + + if final_message is None: + final_message = ensure_author( + self._create_completion_message( + text="Conversation halted after termination condition was met.", + reason="termination_condition", + ), + self._manager_name, + ) + self._conversation.append(final_message) + self._history.append(_GroupChatTurn(self._manager_name, "manager", final_message)) + self._pending_agent = None + await ctx.yield_output(list(self._conversation)) + return True + async def _apply_directive( self, directive: GroupChatDirective, - ctx: WorkflowContext[AgentExecutorRequest | _GroupChatRequestMessage, ChatMessage], + ctx: WorkflowContext[AgentExecutorRequest | _GroupChatRequestMessage, list[ChatMessage]], ) -> None: """Execute a manager directive by either finishing the workflow or routing to a participant. @@ -366,7 +509,7 @@ class GroupChatOrchestratorExecutor(BaseGroupChatOrchestrator): self._conversation.extend((final_message,)) self._history.append(_GroupChatTurn(self._manager_name, "manager", final_message)) self._pending_agent = None - await ctx.yield_output(final_message) + await ctx.yield_output(list(self._conversation)) return agent_name = directive.agent_name @@ -386,6 +529,9 @@ class GroupChatOrchestratorExecutor(BaseGroupChatOrchestrator): self._conversation.extend((manager_message,)) self._history.append(_GroupChatTurn(self._manager_name, "manager", manager_message)) + if await self._complete_on_termination(ctx): + return + self._pending_agent = agent_name self._increment_round() @@ -415,7 +561,7 @@ class GroupChatOrchestratorExecutor(BaseGroupChatOrchestrator): self, participant_name: str, message: ChatMessage, - ctx: WorkflowContext[AgentExecutorRequest | _GroupChatRequestMessage, ChatMessage], + ctx: WorkflowContext[AgentExecutorRequest | _GroupChatRequestMessage, list[ChatMessage]], ) -> None: """Common response ingestion logic shared by agent and custom participants.""" if participant_name not in self._participants: @@ -426,17 +572,213 @@ class GroupChatOrchestratorExecutor(BaseGroupChatOrchestrator): self._history.append(_GroupChatTurn(participant_name, "agent", message)) self._pending_agent = None - if self._check_round_limit(): - await ctx.yield_output( - self._create_completion_message( - text="Conversation halted after reaching manager round limit.", - reason="max_rounds reached after response", - ) - ) + if await self._complete_on_termination(ctx): return - directive = await self._manager(self._build_state()) - await self._apply_directive(directive, ctx) + if self._check_round_limit(): + final_message = self._create_completion_message( + text="Conversation halted after reaching manager round limit.", + reason="max_rounds reached after response", + ) + self._conversation.extend((final_message,)) + self._history.append(_GroupChatTurn(self._manager_name, "manager", final_message)) + await ctx.yield_output(list(self._conversation)) + return + + # Query manager for next speaker selection + if self._is_manager_agent(): + # Agent-based manager: route request through workflow graph + # Prepend system message with participant context + manager_conversation = [self._build_manager_context_message(), *list(self._conversation)] + await self._route_to_participant( + participant_name=self._manager_name, + conversation=manager_conversation, + ctx=ctx, + instruction="", + task=self._task_message, + metadata=None, + ) + else: + # Callable manager: invoke directly + directive = await self._manager(self._build_state()) + await self._apply_directive(directive, ctx) + + def _is_manager_agent(self) -> bool: + """Check if orchestrator is using an agent-based manager (vs callable manager).""" + return self._registry.is_participant_registered(self._manager_name) + + def _build_manager_context_message(self) -> ChatMessage: + """Build system message with participant context for manager agent. + + This message is prepended to the conversation when querying the manager + to provide up-to-date participant information for selection decisions. + + Returns: + System message with participant names and descriptions + """ + participant_list = "\n".join(f"- {name}: {desc}" for name, desc in self._participants.items()) + context_text = ( + "Available participants:\n" + f"{participant_list}\n\n" + "IMPORTANT: Choose only from these exact participant names (case-sensitive)." + ) + return ChatMessage(role=Role.SYSTEM, text=context_text) + + def _parse_manager_selection(self, response: AgentExecutorResponse) -> ManagerSelectionResponse: + """Extract manager selection decision from agent response. + + Attempts to parse structured output from the manager agent using multiple strategies: + 1. response.value (structured output from response_format) + 2. JSON parsing from message text + 3. Fallback error handling + + Args: + response: AgentExecutor response from manager agent + + Returns: + Parsed ManagerSelectionResponse with speaker selection + + Raises: + RuntimeError: If manager response cannot be parsed into valid selection + """ + import json + + # Strategy 1: agent_run_response.value (structured output) + agent_value = response.agent_run_response.value + if agent_value is not None: + if isinstance(agent_value, ManagerSelectionResponse): + return agent_value + if isinstance(agent_value, dict): + return ManagerSelectionResponse.from_dict(cast(dict[str, Any], agent_value)) + if isinstance(agent_value, str): + try: + data = json.loads(agent_value) + return ManagerSelectionResponse.from_dict(data) + except (json.JSONDecodeError, TypeError, KeyError) as e: + raise RuntimeError(f"Manager response.value contains invalid JSON: {e}") from e + + # Strategy 2: Parse from message text + messages = response.agent_run_response.messages or [] + if messages: + last_msg = messages[-1] + text = last_msg.text or "" + try: + return ManagerSelectionResponse.model_validate_json(text) + except (json.JSONDecodeError, TypeError, KeyError): + pass + + # Fallback: Cannot parse manager decision + raise RuntimeError( + "Manager response did not contain valid selection data. " + "Ensure manager agent uses response_format=ManagerSelectionResponse " + "or returns compatible JSON structure." + ) + + async def _handle_manager_response( + self, + response: AgentExecutorResponse, + ctx: WorkflowContext[AgentExecutorRequest | _GroupChatRequestMessage, list[ChatMessage]], + ) -> None: + """Process manager agent's speaker selection decision. + + Parses the manager's response and either finishes the conversation or routes + to the selected participant. This method implements the core orchestration + logic for agent-based managers. + + Args: + response: AgentExecutor response from manager agent + ctx: Workflow context for routing and output + + Behavior: + - Parses manager selection from response + - If finish=True: yields final message and completes workflow + - If participant selected: routes request to that participant + - Validates selected participant exists + - Enforces round limits if configured + + Raises: + ValueError: If manager selects invalid/unknown participant + RuntimeError: If manager response cannot be parsed + """ + selection = self._parse_manager_selection(response) + + if self._pending_finalization: + self._pending_finalization = False + final_message_obj = selection.get_final_message_as_chat_message() + if final_message_obj is None: + final_message_obj = self._create_completion_message( + text="Conversation halted after termination condition was met.", + reason="termination_condition_manager", + ) + final_message_obj = ensure_author(final_message_obj, self._manager_name) + + self._conversation.append(final_message_obj) + self._history.append(_GroupChatTurn(self._manager_name, "manager", final_message_obj)) + self._pending_agent = None + await ctx.yield_output(list(self._conversation)) + return + + if selection.finish: + # Manager decided to complete conversation + final_message_obj = selection.get_final_message_as_chat_message() + if final_message_obj is None: + final_message_obj = self._create_completion_message( + text="Conversation completed.", + reason="manager_finish", + ) + final_message_obj = ensure_author(final_message_obj, self._manager_name) + + self._conversation.append(final_message_obj) + self._history.append(_GroupChatTurn(self._manager_name, "manager", final_message_obj)) + self._pending_agent = None + await ctx.yield_output(list(self._conversation)) + return + + # Manager selected next participant + selected = selection.selected_participant + if not selected: + raise ValueError("Manager selection missing selected_participant when finish=False.") + if selected not in self._participants: + raise ValueError(f"Manager selected unknown participant: '{selected}'") + + # Route to selected participant + instruction = selection.instruction or "" + conversation = list(self._conversation) + if instruction: + manager_message = ensure_author( + self._create_completion_message(text=instruction, reason="manager_instruction"), + self._manager_name, + ) + conversation.append(manager_message) + self._conversation.append(manager_message) + self._history.append(_GroupChatTurn(self._manager_name, "manager", manager_message)) + + if await self._complete_on_termination(ctx): + return + + self._pending_agent = selected + self._increment_round() + + await self._route_to_participant( + participant_name=selected, + conversation=conversation, + ctx=ctx, + instruction=instruction, + task=self._task_message, + metadata=None, + ) + + if self._check_round_limit(): + await self._apply_directive( + GroupChatDirective( + finish=True, + final_message=self._create_completion_message( + text="Conversation halted after reaching manager round limit.", + reason="max_rounds reached after manager selection", + ), + ), + ctx, + ) @staticmethod def _extract_agent_message(response: AgentExecutorResponse, participant_name: str) -> ChatMessage: @@ -469,7 +811,7 @@ class GroupChatOrchestratorExecutor(BaseGroupChatOrchestrator): async def _handle_task_message( self, task_message: ChatMessage, - ctx: WorkflowContext[AgentExecutorRequest | _GroupChatRequestMessage, ChatMessage], + ctx: WorkflowContext[AgentExecutorRequest | _GroupChatRequestMessage, list[ChatMessage]], ) -> None: """Initialize orchestrator state and start the manager-directed conversation loop. @@ -519,14 +861,33 @@ class GroupChatOrchestratorExecutor(BaseGroupChatOrchestrator): self._history = [_GroupChatTurn("user", "user", task_message)] self._pending_agent = None self._round_index = 0 - directive = await self._manager(self._build_state()) - await self._apply_directive(directive, ctx) + + if await self._complete_on_termination(ctx): + return + + # Query manager for first speaker selection + if self._is_manager_agent(): + # Agent-based manager: route request through workflow graph + # Prepend system message with participant context + manager_conversation = [self._build_manager_context_message(), *list(self._conversation)] + await self._route_to_participant( + participant_name=self._manager_name, + conversation=manager_conversation, + ctx=ctx, + instruction="", + task=self._task_message, + metadata=None, + ) + else: + # Callable manager: invoke directly + directive = await self._manager(self._build_state()) + await self._apply_directive(directive, ctx) @handler async def handle_str( self, task: str, - ctx: WorkflowContext[AgentExecutorRequest | _GroupChatRequestMessage, ChatMessage], + ctx: WorkflowContext[AgentExecutorRequest | _GroupChatRequestMessage, list[ChatMessage]], ) -> None: """Handler for string input as workflow entry point. @@ -545,7 +906,7 @@ class GroupChatOrchestratorExecutor(BaseGroupChatOrchestrator): async def handle_chat_message( self, task_message: ChatMessage, - ctx: WorkflowContext[AgentExecutorRequest | _GroupChatRequestMessage, ChatMessage], + ctx: WorkflowContext[AgentExecutorRequest | _GroupChatRequestMessage, list[ChatMessage]], ) -> None: """Handler for ChatMessage input as workflow entry point. @@ -564,7 +925,7 @@ class GroupChatOrchestratorExecutor(BaseGroupChatOrchestrator): async def handle_conversation( self, conversation: list[ChatMessage], - ctx: WorkflowContext[AgentExecutorRequest | _GroupChatRequestMessage, ChatMessage], + ctx: WorkflowContext[AgentExecutorRequest | _GroupChatRequestMessage, list[ChatMessage]], ) -> None: """Handler for conversation history as workflow entry point. @@ -602,7 +963,7 @@ class GroupChatOrchestratorExecutor(BaseGroupChatOrchestrator): async def handle_agent_response( self, response: _GroupChatResponseMessage, - ctx: WorkflowContext[AgentExecutorRequest | _GroupChatRequestMessage, ChatMessage], + ctx: WorkflowContext[AgentExecutorRequest | _GroupChatRequestMessage, list[ChatMessage]], ) -> None: """Handle responses from custom participant executors.""" await self._ingest_participant_message(response.agent_name, response.message, ctx) @@ -611,9 +972,14 @@ class GroupChatOrchestratorExecutor(BaseGroupChatOrchestrator): async def handle_agent_executor_response( self, response: AgentExecutorResponse, - ctx: WorkflowContext[AgentExecutorRequest | _GroupChatRequestMessage, ChatMessage], + ctx: WorkflowContext[AgentExecutorRequest | _GroupChatRequestMessage, list[ChatMessage]], ) -> None: - """Handle direct AgentExecutor responses.""" + """Handle responses from both manager agent and regular participants. + + Routes responses based on whether they come from the manager or a participant: + - Manager responses: parsed for speaker selection decisions + - Participant responses: ingested as conversation messages + """ participant_name = self._registry.get_participant_name(response.executor_id) if participant_name is None: logger.debug( @@ -621,8 +987,14 @@ class GroupChatOrchestratorExecutor(BaseGroupChatOrchestrator): response.executor_id, ) return - message = self._extract_agent_message(response, participant_name) - await self._ingest_participant_message(participant_name, message, ctx) + + # Check if response is from manager agent + if participant_name == self._manager_name and self._is_manager_agent(): + await self._handle_manager_response(response, ctx) + else: + # Regular participant response + message = self._extract_agent_message(response, participant_name) + await self._ingest_participant_message(participant_name, message, ctx) def _default_orchestrator_factory(wiring: _GroupChatConfig) -> Executor: @@ -640,8 +1012,9 @@ def _default_orchestrator_factory(wiring: _GroupChatConfig) -> Executor: Behavior: - Extracts participant names and descriptions for manager context - - Forwards manager instance, manager name, and max_rounds settings + - Forwards manager instance, manager name, max_rounds, and termination_condition settings - Allows orchestrator to auto-generate its executor ID + - Supports both callable managers (set_select_speakers_func) and agent-based managers (set_manager) Why descriptions are extracted: The manager needs participant descriptions (not full specs) to make informed @@ -649,16 +1022,30 @@ def _default_orchestrator_factory(wiring: _GroupChatConfig) -> Executor: since routing is handled by the workflow graph. Raises: - RuntimeError: If manager is None (should not happen when using default factory) + RuntimeError: If neither manager nor manager_participant is configured """ - if wiring.manager is None: - raise RuntimeError("Default orchestrator factory requires a manager to be set") + if wiring.manager is None and wiring.manager_participant is None: + raise RuntimeError( + "Default orchestrator factory requires a manager to be configured. " + "Call set_manager(...) or set_select_speakers_func(...) before build()." + ) + + manager_callable = wiring.manager + if manager_callable is None: + # Keep orchestrator signature satisfied; agent managers are routed via the workflow graph + async def _agent_manager_placeholder(_: GroupChatStateSnapshot) -> GroupChatDirective: # noqa: RUF029 + raise RuntimeError( + "Manager callable invoked unexpectedly. Agent-based managers should route through the workflow graph." + ) + + manager_callable = _agent_manager_placeholder return GroupChatOrchestratorExecutor( - manager=wiring.manager, + manager=manager_callable, participants={name: spec.description for name, spec in wiring.participants.items()}, manager_name=wiring.manager_name, max_rounds=wiring.max_rounds, + termination_condition=wiring.termination_condition, ) @@ -684,8 +1071,41 @@ def assemble_group_chat_workflow( wiring.orchestrator = orchestrator workflow_builder = builder or WorkflowBuilder() - workflow_builder = workflow_builder.set_start_executor(orchestrator) + start_executor = getattr(workflow_builder, "_start_executor", None) + if start_executor is None: + workflow_builder = workflow_builder.set_start_executor(orchestrator) + # Wire manager as participant if agent-based manager is configured + if wiring.manager_participant is not None: + manager_spec = GroupChatParticipantSpec( + name=wiring.manager_name, + participant=wiring.manager_participant, + description="Coordination manager", + ) + manager_pipeline = list(participant_factory(manager_spec, wiring)) + if not manager_pipeline: + raise ValueError("Participant factory returned empty pipeline for manager.") + + manager_entry = manager_pipeline[0] + manager_exit = manager_pipeline[-1] + + # Register manager with orchestrator + register_entry = getattr(orchestrator, "register_participant_entry", None) + if callable(register_entry): + register_entry( + wiring.manager_name, + entry_id=manager_entry.id, + is_agent=not isinstance(wiring.manager_participant, Executor), + ) + + # Wire manager edges: Orchestrator ↔ Manager + workflow_builder = workflow_builder.add_edge(orchestrator, manager_entry) + for upstream, downstream in itertools.pairwise(manager_pipeline): + workflow_builder = workflow_builder.add_edge(upstream, downstream) + if manager_exit is not orchestrator: + workflow_builder = workflow_builder.add_edge(manager_exit, orchestrator) + + # Wire regular participants for name, spec in wiring.participants.items(): pipeline = list(participant_factory(spec, wiring)) if not pipeline: @@ -733,12 +1153,14 @@ class GroupChatBuilder: r"""High-level builder for manager-directed group chat workflows with dynamic orchestration. GroupChat coordinates multi-agent conversations using a manager that selects which participant - speaks next. The manager can be a simple Python function (select_speakers) or an LLM-based - selector (set_prompt_based_manager). These two approaches are mutually exclusive. + speaks next. The manager can be a simple Python function (:py:meth:`GroupChatBuilder.set_select_speakers_func`) + or an agent-based selector via :py:meth:`GroupChatBuilder.set_manager`. These two approaches are + mutually exclusive. **Core Workflow:** 1. Define participants: list of agents (uses their .name) or dict mapping names to agents - 2. Configure speaker selection: select_speakers() OR set_prompt_based_manager() (not both) + 2. Configure speaker selection: :py:meth:`GroupChatBuilder.set_select_speakers_func` OR + :py:meth:`GroupChatBuilder.set_manager` (not both) 3. Optional: set round limits, checkpointing, termination conditions 4. Build and run the workflow @@ -748,6 +1170,9 @@ class GroupChatBuilder: .. code-block:: python + from agent_framework import GroupChatBuilder, GroupChatStateSnapshot + + def select_next_speaker(state: GroupChatStateSnapshot) -> str | None: # state contains: task, participants, conversation, history, round_index if state["round_index"] >= 5: @@ -760,7 +1185,7 @@ class GroupChatBuilder: workflow = ( GroupChatBuilder() - .select_speakers(select_next_speaker) + .set_select_speakers_func(select_next_speaker) .participants([researcher_agent, writer_agent]) # Uses agent.name .build() ) @@ -769,11 +1194,20 @@ class GroupChatBuilder: .. code-block:: python + from agent_framework import ChatAgent from agent_framework.azure import AzureOpenAIChatClient + manager_agent = AzureOpenAIChatClient().create_agent( + instructions="Coordinate the conversation and pick the next speaker.", + name="Coordinator", + temperature=0.3, + seed=42, + max_tokens=500, + ) + workflow = ( GroupChatBuilder() - .set_prompt_based_manager(chat_client=AzureOpenAIChatClient(), display_name="Coordinator") + .set_manager(manager_agent, display_name="Coordinator") .participants([researcher, writer]) # Or use dict: researcher=r, writer=w .with_max_rounds(10) .build() @@ -782,24 +1216,24 @@ class GroupChatBuilder: **Participant Specification:** Two ways to specify participants: - - List form: ``[agent1, agent2]`` - uses ``agent.name`` attribute for participant names - - Dict form: ``{name1: agent1, name2: agent2}`` - explicit name control - - Keyword form: ``participants(name1=agent1, name2=agent2)`` - explicit name control + - List form: `[agent1, agent2]` - uses `agent.name` attribute for participant names + - Dict form: `{name1: agent1, name2: agent2}` - explicit name control + - Keyword form: `participants(name1=agent1, name2=agent2)` - explicit name control **State Snapshot Structure:** - The GroupChatStateSnapshot passed to select_speakers contains: - - ``task``: ChatMessage - Original user task - - ``participants``: dict[str, str] - Mapping of participant names to descriptions - - ``conversation``: tuple[ChatMessage, ...] - Full conversation history - - ``history``: tuple[GroupChatTurn, ...] - Turn-by-turn record with speaker attribution - - ``round_index``: int - Number of manager selection rounds so far - - ``pending_agent``: str | None - Name of agent currently processing (if any) + The GroupChatStateSnapshot passed to set_select_speakers_func contains: + - `task`: ChatMessage - Original user task + - `participants`: dict[str, str] - Mapping of participant names to descriptions + - `conversation`: tuple[ChatMessage, ...] - Full conversation history + - `history`: tuple[GroupChatTurn, ...] - Turn-by-turn record with speaker attribution + - `round_index`: int - Number of manager selection rounds so far + - `pending_agent`: str | None - Name of agent currently processing (if any) **Important Constraints:** - - Cannot combine select_speakers() and set_prompt_based_manager() - choose one + - Cannot combine :py:meth:`GroupChatBuilder.set_select_speakers_func` and :py:meth:`GroupChatBuilder.set_manager` - Participant names must be unique - - When using list form, agents must have a non-empty ``name`` attribute + - When using list form, agents must have a non-empty `name` attribute """ def __init__( @@ -820,9 +1254,11 @@ class GroupChatBuilder: self._participants: dict[str, AgentProtocol | Executor] = {} self._participant_metadata: dict[str, Any] | None = None self._manager: _GroupChatManagerFn | None = None + self._manager_participant: AgentProtocol | Executor | None = None self._manager_name: str = "manager" self._checkpoint_storage: CheckpointStorage | None = None self._max_rounds: int | None = None + self._termination_condition: Callable[[list[ChatMessage]], bool | Awaitable[bool]] | None = None self._interceptors: list[_InterceptorSpec] = [] self._orchestrator_factory = group_chat_orchestrator(_orchestrator_factory) self._participant_factory = _participant_factory or _default_participant_factory @@ -832,68 +1268,107 @@ class GroupChatBuilder: manager: _GroupChatManagerFn, display_name: str | None, ) -> "GroupChatBuilder": - if self._manager is not None: + if self._manager is not None or self._manager_participant is not None: raise ValueError( "GroupChatBuilder already has a manager configured. " - "Call select_speakers(...) or set_prompt_based_manager(...) at most once." + "Call set_select_speakers_func(...) or set_manager(...) at most once." ) resolved_name = display_name or getattr(manager, "name", None) or "manager" self._manager = manager self._manager_name = resolved_name return self - def set_prompt_based_manager( + def set_manager( self, - chat_client: ChatClientProtocol, + manager: AgentProtocol | Executor, *, - instructions: str | None = None, display_name: str | None = None, ) -> "GroupChatBuilder": - r"""Configure the default prompt-based manager driven by an LLM chat client. + """Configure the manager/coordinator agent for group chat orchestration. - The manager coordinates participants by making selection decisions based on the conversation - state, task, and participant descriptions. It uses structured output (ManagerDirectiveModel) - to ensure reliable parsing of decisions. + The manager coordinates participants by selecting who speaks next based on + conversation state and task requirements. The manager is a full workflow + participant with access to all agent infrastructure (tools, context, observability). + + The manager agent must produce structured output compatible with ManagerSelectionResponse + to communicate its speaker selection decisions. Use response_format for reliable parsing. + GroupChatBuilder enforces this when the manager is a ChatAgent and rejects incompatible + response formats. Args: - chat_client: Chat completion client used to run the coordinator LLM. - instructions: System instructions to steer the coordinator's decision-making. - If not provided, uses DEFAULT_MANAGER_INSTRUCTIONS. These instructions are combined - with the task description, participant list, and structured output format to guide - the LLM in selecting the next speaker or completing the conversation. - display_name: Optional conversational display name for manager messages. + manager: Agent or executor responsible for speaker selection and coordination. + Must return ManagerSelectionResponse or compatible dict/JSON structure. + display_name: Optional name for manager messages in conversation history. + If not provided, uses manager.name for AgentProtocol or manager.id for Executor. Returns: Self for fluent chaining. - Note: - Calling this method and :meth:`set_speaker_selector` together is not allowed; choose one. + Raises: + ValueError: If manager is already configured via :py:meth:`GroupChatBuilder.set_select_speakers_func` + TypeError: If manager is not AgentProtocol or Executor instance Example: .. code-block:: python - from agent_framework import GroupChatBuilder, DEFAULT_MANAGER_INSTRUCTIONS + from agent_framework import GroupChatBuilder, ChatAgent + from agent_framework.openai import OpenAIChatClient - custom_instructions = ( - DEFAULT_MANAGER_INSTRUCTIONS + "\\n\\nPrioritize the researcher for data analysis tasks." + # Coordinator agent - response_format is enforced to ManagerSelectionResponse + coordinator = ChatAgent( + name="Coordinator", + description="Coordinates multi-agent collaboration", + instructions=''' + You coordinate a team conversation. Review the conversation history + and select the next participant to speak. + + When ready to finish, set finish=True and provide a summary in final_message. + ''', + chat_client=OpenAIChatClient(), ) workflow = ( GroupChatBuilder() - .set_prompt_based_manager(chat_client, instructions=custom_instructions, display_name="Coordinator") - .participants(researcher=researcher, writer=writer) + .set_manager(coordinator, display_name="Orchestrator") + .participants([researcher, writer]) .build() ) - """ - manager = _PromptBasedGroupChatManager( - chat_client, - instructions=instructions, - name=display_name, - ) - return self._set_manager_function(manager, display_name) - def select_speakers( + Note: + The manager agent's response_format must be ManagerSelectionResponse for structured output. + Custom response formats raise ValueError instead of being overridden. + """ + if self._manager is not None or self._manager_participant is not None: + raise ValueError( + "GroupChatBuilder already has a manager configured. " + "Call set_select_speakers_func(...) or set_manager(...) at most once." + ) + + if not isinstance(manager, (AgentProtocol, Executor)): + raise TypeError(f"Manager must be AgentProtocol or Executor instance. Got {type(manager).__name__}.") + + # Infer display name from manager if not provided + if display_name is None: + display_name = manager.id if isinstance(manager, Executor) else manager.name or "manager" + + # Enforce ManagerSelectionResponse for ChatAgent managers + if isinstance(manager, ChatAgent): + configured_format = manager.chat_options.response_format + if configured_format is None: + manager.chat_options.response_format = ManagerSelectionResponse + elif configured_format is not ManagerSelectionResponse: + configured_format_name = getattr(configured_format, "__name__", str(configured_format)) + raise ValueError( + "Manager ChatAgent response_format must be ManagerSelectionResponse. " + f"Received '{configured_format_name}' for manager '{display_name}'." + ) + + self._manager_participant = manager + self._manager_name = display_name + return self + + def set_select_speakers_func( self, selector: ( Callable[[GroupChatStateSnapshot], Awaitable[str | None]] | Callable[[GroupChatStateSnapshot], str | None] @@ -908,6 +1383,15 @@ class GroupChatBuilder: function receives an immutable snapshot of the current conversation state and returns the name of the next participant to speak, or None to finish the conversation. + The selector function can implement any logic including: + - Simple round-robin or rule-based selection + - LLM-based decision making with custom prompts + - Conversation summarization before routing to the next agent + - Custom metadata or context passing + + For advanced scenarios, return a GroupChatDirective instead of a string to include + custom instructions or metadata for the next participant. + The selector function signature: def select_next_speaker(state: GroupChatStateSnapshot) -> str | None: # state contains: task, participants, conversation, history, round_index @@ -917,6 +1401,7 @@ class GroupChatBuilder: Args: selector: Function that takes GroupChatStateSnapshot and returns the next speaker's name (str) to continue the conversation, or None to finish. May be sync or async. + Can also return GroupChatDirective for advanced control (instruction, metadata). display_name: Optional name shown in conversation history for orchestrator messages (defaults to "manager"). final_message: Optional final message (or factory) emitted when selector returns None @@ -925,7 +1410,7 @@ class GroupChatBuilder: Returns: Self for fluent chaining. - Example: + Example (simple): .. code-block:: python @@ -940,13 +1425,37 @@ class GroupChatBuilder: workflow = ( GroupChatBuilder() - .select_speakers(select_next_speaker) + .set_select_speakers_func(select_next_speaker) .participants(researcher=researcher_agent, writer=writer_agent) .build() ) + Example (with LLM and custom instructions): + + .. code-block:: python + + from agent_framework import GroupChatDirective + + + async def llm_based_selector(state: GroupChatStateSnapshot) -> GroupChatDirective | None: + if state["round_index"] >= 5: + return GroupChatDirective(finish=True) + + # Use LLM to decide next speaker and summarize conversation + conversation_summary = await summarize_with_llm(state["conversation"]) + next_agent = await pick_agent_with_llm(state["participants"], state["task"]) + + # Pass custom instruction to the selected agent + return GroupChatDirective( + agent_name=next_agent, + instruction=f"Context summary: {conversation_summary}", + ) + + + workflow = GroupChatBuilder().set_select_speakers_func(llm_based_selector).participants(...).build() + Note: - Cannot be combined with set_prompt_based_manager(). Choose one orchestration strategy. + Cannot be combined with :py:meth:`GroupChatBuilder.set_manager`. Choose one orchestration strategy. """ manager_name = display_name or "manager" adapter = _SpeakerSelectorAdapter( @@ -985,10 +1494,7 @@ class GroupChatBuilder: from agent_framework import GroupChatBuilder workflow = ( - GroupChatBuilder() - .set_prompt_based_manager(chat_client) - .participants([writer_agent, reviewer_agent]) - .build() + GroupChatBuilder().set_manager(manager_agent).participants([writer_agent, reviewer_agent]).build() ) """ combined: dict[str, AgentProtocol | Executor] = {} @@ -998,6 +1504,11 @@ class GroupChatBuilder: raise ValueError("participant names must be non-empty strings") if name in combined or name in self._participants: raise ValueError(f"Duplicate participant name '{name}' supplied.") + if name == self._manager_name: + raise ValueError( + f"Participant name '{name}' conflicts with manager name. " + "Manager is automatically registered as a participant." + ) combined[name] = participant if participants: @@ -1050,7 +1561,7 @@ class GroupChatBuilder: storage = MemoryCheckpointStorage() workflow = ( GroupChatBuilder() - .set_prompt_based_manager(chat_client) + .set_manager(manager_agent) .participants(agent1=agent1, agent2=agent2) .with_checkpointing(storage) .build() @@ -1088,6 +1599,40 @@ class GroupChatBuilder: self._interceptors.append((factory, condition)) return self + def with_termination_condition( + self, + condition: Callable[[list[ChatMessage]], bool | Awaitable[bool]], + ) -> "GroupChatBuilder": + """Define a custom termination condition for the group chat workflow. + + The condition receives the full conversation (including manager and agent messages) and may be async. + When it returns True, the orchestrator halts the conversation and emits a completion message authored + by the manager. + + Example: + + .. code-block:: python + + from agent_framework import ChatMessage, GroupChatBuilder, Role + + + def stop_after_two_calls(conversation: list[ChatMessage]) -> bool: + calls = sum(1 for msg in conversation if msg.role == Role.ASSISTANT and msg.author_name == "specialist") + return calls >= 2 + + + specialist_agent = ... + workflow = ( + GroupChatBuilder() + .set_select_speakers_func(lambda _: "specialist") + .participants(specialist=specialist_agent) + .with_termination_condition(stop_after_two_calls) + .build() + ) + """ + self._termination_condition = condition + return self + def with_max_rounds(self, max_rounds: int | None) -> "GroupChatBuilder": """Set a maximum number of manager rounds to prevent infinite conversations. @@ -1109,7 +1654,7 @@ class GroupChatBuilder: # Limit to 15 rounds workflow = ( GroupChatBuilder() - .set_prompt_based_manager(chat_client) + .set_manager(manager_agent) .participants(agent1=agent1, agent2=agent2) .with_max_rounds(15) .build() @@ -1117,11 +1662,7 @@ class GroupChatBuilder: # Unlimited rounds workflow = ( - GroupChatBuilder() - .set_prompt_based_manager(chat_client) - .participants(agent1=agent1) - .with_max_rounds(None) - .build() + GroupChatBuilder().set_manager(manager_agent).participants(agent1=agent1).with_max_rounds(None).build() ) """ self._max_rounds = max_rounds @@ -1182,19 +1723,21 @@ class GroupChatBuilder: from agent_framework import GroupChatBuilder # Execute the workflow - workflow = ( - GroupChatBuilder() - .set_prompt_based_manager(chat_client) - .participants(agent1=agent1, agent2=agent2) - .build() - ) + workflow = GroupChatBuilder().set_manager(manager_agent).participants(agent1=agent1, agent2=agent2).build() async for message in workflow.run("Solve this problem collaboratively"): print(message.text) """ # Manager is only required when using the default orchestrator factory # Custom factories (e.g., MagenticBuilder) provide their own orchestrator with embedded manager - if self._manager is None and self._orchestrator_factory == _default_orchestrator_factory: - raise ValueError("manager must be configured before build() when using default orchestrator") + if ( + self._manager is None + and self._manager_participant is None + and self._orchestrator_factory == _default_orchestrator_factory + ): + raise ValueError( + "manager must be configured before build() when using default orchestrator. " + "Call set_manager(...) or set_select_speakers_func(...) before build()." + ) if not self._participants: raise ValueError("participants must be configured before build()") @@ -1202,9 +1745,11 @@ class GroupChatBuilder: participant_specs = self._build_participant_specs() wiring = _GroupChatConfig( manager=self._manager, + manager_participant=self._manager_participant, manager_name=self._manager_name, participants=participant_specs, max_rounds=self._max_rounds, + termination_condition=self._termination_condition, participant_aliases=metadata["aliases"], participant_executors=metadata["executors"], ) @@ -1262,117 +1807,6 @@ class ManagerDirectiveModel(BaseModel): ) -class _PromptBasedGroupChatManager: - """LLM-backed manager that produces directives via structured output. - - This is the default manager implementation for group chat workflows. It uses an LLM - to make speaker selection decisions based on conversation state, participant - descriptions, and custom instructions. - - Coordination strategy: - - Receives immutable state snapshot with full conversation history - - Formats system prompt with instructions, task, and participant descriptions - - Appends conversation context and uses structured output (Pydantic model) for reliable parsing - - Converts LLM response to GroupChatDirective - - Flexibility: - - Custom instructions allow domain-specific coordination strategies - - Participant descriptions guide the LLM's selection logic - - Structured output ensures reliable parsing (no regex or brittle prompts) - - Example coordination patterns: - - Round-robin: "Rotate between participants in order" - - Task-based: "Select the participant best suited for the current sub-task" - - Dependency-aware: "Only call analyst after researcher provides data" - - Args: - chat_client: ChatClientProtocol implementation for LLM inference - instructions: Custom system instructions (defaults to DEFAULT_MANAGER_INSTRUCTIONS). - These instructions are combined with the task, participant list, and - structured output format (ManagerDirectiveModel) to coordinate the conversation. - name: Display name for the manager in conversation history - - Raises: - RuntimeError: If LLM response cannot be parsed into the directive payload - If directive is missing next_agent when finish=False - If selected agent is not in participants - """ - - def __init__( - self, - chat_client: ChatClientProtocol, - *, - instructions: str | None = None, - name: str | None = None, - ) -> None: - self._chat_client = chat_client - self._instructions = instructions or DEFAULT_MANAGER_INSTRUCTIONS - self._name = name or "GroupChatManager" - - @property - def name(self) -> str: - return self._name - - async def __call__(self, state: GroupChatStateSnapshot) -> GroupChatDirective: - participants = state["participants"] - task_message = state["task"] - conversation = state["conversation"] - - participants_section = "\n".join(f"- {agent}: {description}" for agent, description in participants.items()) - - system_message = ChatMessage( - role=Role.SYSTEM, - text=( - f"{self._instructions}\n\n" - f"Task:\n{task_message.text}\n\n" - f"Participants:\n{participants_section}\n\n" - f"{DEFAULT_MANAGER_STRUCTURED_OUTPUT_PROMPT}" - ), - ) - - messages: list[ChatMessage] = [system_message, *conversation] - - response = await self._chat_client.get_response(messages, response_format=ManagerDirectiveModel) - - directive_model: ManagerDirectiveModel - if response.value is not None: - if isinstance(response.value, ManagerDirectiveModel): - directive_model = response.value - elif isinstance(response.value, str): - directive_model = ManagerDirectiveModel.model_validate_json(response.value) - elif isinstance(response.value, dict): - directive_model = ManagerDirectiveModel.model_validate(response.value) # type: ignore[arg-type] - else: - raise RuntimeError(f"Unexpected response.value type: {type(response.value)}") - elif response.messages: - text = response.messages[-1].text or "{}" - directive_model = ManagerDirectiveModel.model_validate_json(text) - else: - raise RuntimeError("LLM response did not contain structured output.") - - if directive_model.finish: - final_text = directive_model.final_response or "" - return GroupChatDirective( - finish=True, - final_message=ChatMessage( - role=Role.ASSISTANT, - text=final_text, - author_name=self._name, - ), - ) - - next_agent = directive_model.next_agent - if not next_agent: - raise RuntimeError("Manager directive missing next_agent while finish is False.") - if next_agent not in participants: - raise RuntimeError(f"Manager selected unknown participant '{next_agent}'.") - - return GroupChatDirective( - agent_name=next_agent, - instruction=directive_model.message or "", - ) - - class _SpeakerSelectorAdapter: """Adapter that turns a simple speaker selector into a full manager directive.""" diff --git a/python/packages/core/agent_framework/_workflows/_handoff.py b/python/packages/core/agent_framework/_workflows/_handoff.py index d18bc59562..054c53f6e3 100644 --- a/python/packages/core/agent_framework/_workflows/_handoff.py +++ b/python/packages/core/agent_framework/_workflows/_handoff.py @@ -1424,6 +1424,7 @@ class HandoffBuilder: prompt=self._request_prompt, id="handoff-user-input", ) + builder = WorkflowBuilder(name=self._name, description=self._description).set_start_executor(input_node) specialist_aliases = {alias: exec_id for alias, exec_id in self._aliases.items() if exec_id in specialists} @@ -1440,6 +1441,7 @@ class HandoffBuilder: wiring = _GroupChatConfig( manager=None, + manager_participant=None, manager_name=self._starting_agent_id, participants=participant_specs, max_rounds=None, @@ -1453,14 +1455,13 @@ class HandoffBuilder: orchestrator_factory=_handoff_orchestrator_factory, interceptors=(), checkpoint_storage=self._checkpoint_storage, - builder=WorkflowBuilder(name=self._name, description=self._description), + builder=builder, return_builder=True, ) if not isinstance(result, tuple): raise TypeError("Expected tuple from assemble_group_chat_workflow with return_builder=True") builder, coordinator = result - builder = builder.set_start_executor(input_node) builder = builder.add_edge(input_node, starting_executor) builder = builder.add_edge(coordinator, user_gateway) builder = builder.add_edge(user_gateway, coordinator) diff --git a/python/packages/core/agent_framework/_workflows/_magentic.py b/python/packages/core/agent_framework/_workflows/_magentic.py index 624c6f50ae..1a6aaf2999 100644 --- a/python/packages/core/agent_framework/_workflows/_magentic.py +++ b/python/packages/core/agent_framework/_workflows/_magentic.py @@ -961,7 +961,7 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): async def _emit_orchestrator_message( self, - ctx: WorkflowContext[Any, ChatMessage], + ctx: WorkflowContext[Any, list[ChatMessage]], message: ChatMessage, kind: str, ) -> None: @@ -1110,7 +1110,7 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): self, message: _MagenticStartMessage, context: WorkflowContext[ - _MagenticResponseMessage | _MagenticRequestMessage | _MagenticPlanReviewRequest, ChatMessage + _MagenticResponseMessage | _MagenticRequestMessage | _MagenticPlanReviewRequest, list[ChatMessage] ], ) -> None: """Handle the initial start message to begin orchestration.""" @@ -1145,7 +1145,7 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): # Start the inner loop ctx2 = cast( - WorkflowContext[_MagenticResponseMessage | _MagenticRequestMessage, ChatMessage], + WorkflowContext[_MagenticResponseMessage | _MagenticRequestMessage, list[ChatMessage]], context, ) await self._run_inner_loop(ctx2) @@ -1155,7 +1155,7 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): self, task_text: str, context: WorkflowContext[ - _MagenticResponseMessage | _MagenticRequestMessage | _MagenticPlanReviewRequest, ChatMessage + _MagenticResponseMessage | _MagenticRequestMessage | _MagenticPlanReviewRequest, list[ChatMessage] ], ) -> None: await self.handle_start_message(_MagenticStartMessage.from_string(task_text), context) @@ -1165,7 +1165,7 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): self, task_message: ChatMessage, context: WorkflowContext[ - _MagenticResponseMessage | _MagenticRequestMessage | _MagenticPlanReviewRequest, ChatMessage + _MagenticResponseMessage | _MagenticRequestMessage | _MagenticPlanReviewRequest, list[ChatMessage] ], ) -> None: await self.handle_start_message(_MagenticStartMessage(task_message), context) @@ -1175,7 +1175,7 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): self, conversation: list[ChatMessage], context: WorkflowContext[ - _MagenticResponseMessage | _MagenticRequestMessage | _MagenticPlanReviewRequest, ChatMessage + _MagenticResponseMessage | _MagenticRequestMessage | _MagenticPlanReviewRequest, list[ChatMessage] ], ) -> None: await self.handle_start_message(_MagenticStartMessage(conversation), context) @@ -1184,7 +1184,7 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): async def handle_response_message( self, message: _MagenticResponseMessage, - context: WorkflowContext[_MagenticResponseMessage | _MagenticRequestMessage, ChatMessage], + context: WorkflowContext[_MagenticResponseMessage | _MagenticRequestMessage, list[ChatMessage]], ) -> None: """Handle responses from agents.""" if getattr(self, "_terminated", False): @@ -1216,7 +1216,7 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): response: _MagenticPlanReviewReply, context: WorkflowContext[ # may broadcast ledger next, or ask for another round of review - _MagenticResponseMessage | _MagenticRequestMessage | _MagenticPlanReviewRequest, ChatMessage + _MagenticResponseMessage | _MagenticRequestMessage | _MagenticPlanReviewRequest, list[ChatMessage] ], ) -> None: if getattr(self, "_terminated", False): @@ -1262,7 +1262,7 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): # Enter the normal coordination loop ctx2 = cast( - WorkflowContext[_MagenticResponseMessage | _MagenticRequestMessage, ChatMessage], + WorkflowContext[_MagenticResponseMessage | _MagenticRequestMessage, list[ChatMessage]], context, ) await self._run_inner_loop(ctx2) @@ -1289,7 +1289,7 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): self._context.chat_history.append(self._task_ledger) # No further review requests; proceed directly into coordination ctx2 = cast( - WorkflowContext[_MagenticResponseMessage | _MagenticRequestMessage, ChatMessage], + WorkflowContext[_MagenticResponseMessage | _MagenticRequestMessage, list[ChatMessage]], context, ) await self._run_inner_loop(ctx2) @@ -1324,7 +1324,7 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): async def _run_outer_loop( self, - context: WorkflowContext[_MagenticResponseMessage | _MagenticRequestMessage, ChatMessage], + context: WorkflowContext[_MagenticResponseMessage | _MagenticRequestMessage, list[ChatMessage]], ) -> None: """Run the outer orchestration loop - planning phase.""" if self._context is None: @@ -1347,7 +1347,7 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): async def _run_inner_loop( self, - context: WorkflowContext[_MagenticResponseMessage | _MagenticRequestMessage, ChatMessage], + context: WorkflowContext[_MagenticResponseMessage | _MagenticRequestMessage, list[ChatMessage]], ) -> None: """Run the inner orchestration loop. Coordination phase. Serialized with a lock.""" if self._context is None or self._task_ledger is None: @@ -1357,7 +1357,7 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): async def _run_inner_loop_helper( self, - context: WorkflowContext[_MagenticResponseMessage | _MagenticRequestMessage, ChatMessage], + context: WorkflowContext[_MagenticResponseMessage | _MagenticRequestMessage, list[ChatMessage]], ) -> None: """Run inner loop with exclusive access.""" # Narrow optional context for the remainder of this method @@ -1442,7 +1442,7 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): async def _reset_and_replan( self, - context: WorkflowContext[_MagenticResponseMessage | _MagenticRequestMessage, ChatMessage], + context: WorkflowContext[_MagenticResponseMessage | _MagenticRequestMessage, list[ChatMessage]], ) -> None: """Reset context and replan.""" if self._context is None: @@ -1468,7 +1468,7 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): async def _prepare_final_answer( self, - context: WorkflowContext[_MagenticResponseMessage | _MagenticRequestMessage, ChatMessage], + context: WorkflowContext[_MagenticResponseMessage | _MagenticRequestMessage, list[ChatMessage]], ) -> None: """Prepare the final answer using the manager.""" if self._context is None: @@ -1478,11 +1478,11 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): final_answer = await self._manager.prepare_final_answer(self._context.clone(deep=True)) # Emit a completed event for the workflow - await context.yield_output(final_answer) + await context.yield_output([final_answer]) async def _check_within_limits_or_complete( self, - context: WorkflowContext[_MagenticResponseMessage | _MagenticRequestMessage, ChatMessage], + context: WorkflowContext[_MagenticResponseMessage | _MagenticRequestMessage, list[ChatMessage]], ) -> bool: """Check if orchestrator is within operational limits.""" if self._context is None: @@ -1509,7 +1509,7 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): ) # Yield the partial result and signal completion - await context.yield_output(partial_result) + await context.yield_output([partial_result]) return False return True @@ -2283,21 +2283,22 @@ class MagenticWorkflow: return # At this point, checkpoint is guaranteed to be WorkflowCheckpoint - executor_states: dict[str, Any] = checkpoint.shared_state.get(EXECUTOR_STATE_KEY, {}) + executor_states = cast(dict[str, Any], checkpoint.shared_state.get(EXECUTOR_STATE_KEY, {})) orchestrator_id = getattr(orchestrator, "id", "") - orchestrator_state = executor_states.get(orchestrator_id) + orchestrator_state = cast(Any, executor_states.get(orchestrator_id)) if orchestrator_state is None: - orchestrator_state = executor_states.get("magentic_orchestrator") + orchestrator_state = cast(Any, executor_states.get("magentic_orchestrator")) if not isinstance(orchestrator_state, dict): return - context_payload = orchestrator_state.get("magentic_context") + orchestrator_state_dict = cast(dict[str, Any], orchestrator_state) + context_payload = cast(Any, orchestrator_state_dict.get("magentic_context")) if not isinstance(context_payload, dict): return context_dict = cast(dict[str, Any], context_payload) - restored_participants = context_dict.get("participant_descriptions") + restored_participants = cast(Any, context_dict.get("participant_descriptions")) if not isinstance(restored_participants, dict): return diff --git a/python/packages/core/agent_framework/_workflows/_orchestrator_helpers.py b/python/packages/core/agent_framework/_workflows/_orchestrator_helpers.py index 4b17dda414..9da726faf4 100644 --- a/python/packages/core/agent_framework/_workflows/_orchestrator_helpers.py +++ b/python/packages/core/agent_framework/_workflows/_orchestrator_helpers.py @@ -186,6 +186,10 @@ class ParticipantRegistry: """Check if a participant is registered.""" return name in self._participant_entry_ids + def is_participant_registered(self, name: str) -> bool: + """Check if a participant is registered (alias for is_registered for compatibility).""" + return self.is_registered(name) + def all_participants(self) -> set[str]: """Get all registered participant names.""" return set(self._participant_entry_ids.keys()) diff --git a/python/packages/core/agent_framework/openai/_responses_client.py b/python/packages/core/agent_framework/openai/_responses_client.py index c3f84aab93..6d4fce7bb2 100644 --- a/python/packages/core/agent_framework/openai/_responses_client.py +++ b/python/packages/core/agent_framework/openai/_responses_client.py @@ -90,7 +90,7 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient): **kwargs: Any, ) -> ChatResponse: client = await self.ensure_client() - run_options = await self.prepare_options(messages, chat_options) + run_options = await self.prepare_options(messages, chat_options, **kwargs) response_format = run_options.pop("response_format", None) text_config = run_options.pop("text", None) text_format, text_config = self._prepare_text_config(response_format=response_format, text_config=text_config) @@ -135,7 +135,7 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient): **kwargs: Any, ) -> AsyncIterable[ChatResponseUpdate]: client = await self.ensure_client() - run_options = await self.prepare_options(messages, chat_options) + run_options = await self.prepare_options(messages, chat_options, **kwargs) function_call_ids: dict[int, tuple[str, str]] = {} # output_index: (call_id, name) response_format = run_options.pop("response_format", None) text_config = run_options.pop("text", None) @@ -248,7 +248,7 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient): self, response: OpenAIResponse | ParsedResponse[BaseModel], store: bool | None ) -> str | None: """Get the conversation ID from the response if store is True.""" - return response.id if store else None + return None if store is False else response.id # region Prep methods @@ -386,9 +386,17 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient): return mcp async def prepare_options( - self, messages: MutableSequence[ChatMessage], chat_options: ChatOptions + self, + messages: MutableSequence[ChatMessage], + chat_options: ChatOptions, + **kwargs: Any, ) -> dict[str, Any]: """Take ChatOptions and create the specific options for Responses API.""" + conversation_id = kwargs.pop("conversation_id", None) + + if conversation_id: + chat_options.conversation_id = conversation_id + run_options: dict[str, Any] = chat_options.to_dict( exclude={ "type", @@ -437,8 +445,6 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient): for key, value in additional_properties.items(): if value is not None: run_options[key] = value - if "store" not in run_options: - run_options["store"] = False if (tool_choice := run_options.get("tool_choice")) and len(tool_choice.keys()) == 1: run_options["tool_choice"] = tool_choice["mode"] return run_options @@ -815,8 +821,11 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient): "additional_properties": metadata, "raw_representation": response, } - if chat_options.store: - args["conversation_id"] = self.get_conversation_id(response, chat_options.store) + + conversation_id = self.get_conversation_id(response, chat_options.store) + + if conversation_id: + args["conversation_id"] = conversation_id if response.usage and (usage_details := self._usage_details_from_openai(response.usage)): args["usage_details"] = usage_details if structured_response: diff --git a/python/packages/core/tests/openai/test_openai_responses_client.py b/python/packages/core/tests/openai/test_openai_responses_client.py index 4700950439..c4d824d31d 100644 --- a/python/packages/core/tests/openai/test_openai_responses_client.py +++ b/python/packages/core/tests/openai/test_openai_responses_client.py @@ -1423,12 +1423,12 @@ async def test_prepare_options_store_parameter_handling() -> None: chat_options = ChatOptions(store=None, conversation_id=None) options = await client.prepare_options(messages, chat_options) - assert options["store"] is False + assert "store" not in options assert "previous_response_id" not in options chat_options = ChatOptions() options = await client.prepare_options(messages, chat_options) - assert options["store"] is False + assert "store" not in options assert "previous_response_id" not in options diff --git a/python/packages/core/tests/workflow/test_group_chat.py b/python/packages/core/tests/workflow/test_group_chat.py index 8d1798a7a4..5d11e64c79 100644 --- a/python/packages/core/tests/workflow/test_group_chat.py +++ b/python/packages/core/tests/workflow/test_group_chat.py @@ -1,9 +1,10 @@ # Copyright (c) Microsoft. All rights reserved. from collections.abc import AsyncIterable, Callable -from typing import Any +from typing import Any, cast import pytest +from pydantic import BaseModel from agent_framework import ( MAGENTIC_EVENT_TYPE_AGENT_DELTA, @@ -14,6 +15,7 @@ from agent_framework import ( AgentThread, BaseAgent, ChatMessage, + Executor, GroupChatBuilder, GroupChatDirective, GroupChatStateSnapshot, @@ -23,21 +25,27 @@ from agent_framework import ( Role, TextContent, Workflow, + WorkflowContext, WorkflowOutputEvent, + handler, ) from agent_framework._workflows._checkpoint import InMemoryCheckpointStorage from agent_framework._workflows._group_chat import ( GroupChatOrchestratorExecutor, + ManagerSelectionResponse, _default_orchestrator_factory, # type: ignore + _default_participant_factory, # type: ignore _GroupChatConfig, # type: ignore - _PromptBasedGroupChatManager, # type: ignore _SpeakerSelectorAdapter, # type: ignore + assemble_group_chat_workflow, ) from agent_framework._workflows._magentic import ( _MagenticProgressLedger, # type: ignore _MagenticProgressLedgerItem, # type: ignore _MagenticStartMessage, # type: ignore ) +from agent_framework._workflows._participant_utils import GroupChatParticipantSpec +from agent_framework._workflows._workflow_builder import WorkflowBuilder class StubAgent(BaseAgent): @@ -70,6 +78,73 @@ class StubAgent(BaseAgent): return _stream() +class StubManagerAgent(BaseAgent): + def __init__(self) -> None: + super().__init__(name="manager_agent", description="Stub manager") + self._call_count = 0 + + async def run( + self, + messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None, + *, + thread: AgentThread | None = None, + **kwargs: Any, + ) -> AgentRunResponse: # type: ignore[override] + if self._call_count == 0: + self._call_count += 1 + payload = {"selected_participant": "agent", "finish": False, "final_message": None} + return AgentRunResponse( + messages=[ + ChatMessage( + role=Role.ASSISTANT, + text='{"selected_participant": "agent", "finish": false}', + author_name=self.name, + ) + ], + value=payload, + ) + + payload = {"selected_participant": None, "finish": True, "final_message": "agent manager final"} + return AgentRunResponse( + messages=[ + ChatMessage( + role=Role.ASSISTANT, + text='{"finish": true, "final_message": "agent manager final"}', + author_name=self.name, + ) + ], + value=payload, + ) + + def run_stream( + self, + messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None, + *, + thread: AgentThread | None = None, + **kwargs: Any, + ) -> AsyncIterable[AgentRunResponseUpdate]: # type: ignore[override] + if self._call_count == 0: + self._call_count += 1 + + async def _stream_initial() -> AsyncIterable[AgentRunResponseUpdate]: + yield AgentRunResponseUpdate( + contents=[TextContent(text='{"selected_participant": "agent", "finish": false}')], + role=Role.ASSISTANT, + author_name=self.name, + ) + + return _stream_initial() + + async def _stream_final() -> AsyncIterable[AgentRunResponseUpdate]: + yield AgentRunResponseUpdate( + contents=[TextContent(text='{"finish": true, "final_message": "agent manager final"}')], + role=Role.ASSISTANT, + author_name=self.name, + ) + + return _stream_final() + + def make_sequence_selector() -> Callable[[GroupChatStateSnapshot], Any]: state_counter = {"value": 0} @@ -123,6 +198,22 @@ class StubMagenticManager(MagenticManagerBase): return ChatMessage(role=Role.ASSISTANT, text="final", author_name="magentic_manager") +class PassthroughExecutor(Executor): + @handler + async def forward(self, message: Any, ctx: WorkflowContext[Any]) -> None: + await ctx.send_message(message) + + +class CountingWorkflowBuilder(WorkflowBuilder): + def __init__(self) -> None: + super().__init__() + self.start_calls = 0 + + def set_start_executor(self, executor: Any) -> "CountingWorkflowBuilder": + self.start_calls += 1 + return cast("CountingWorkflowBuilder", super().set_start_executor(executor)) + + async def test_group_chat_builder_basic_flow() -> None: selector = make_sequence_selector() alpha = StubAgent("alpha", "ack from alpha") @@ -130,21 +221,23 @@ async def test_group_chat_builder_basic_flow() -> None: workflow = ( GroupChatBuilder() - .select_speakers(selector, display_name="manager", final_message="done") + .set_select_speakers_func(selector, display_name="manager", final_message="done") .participants(alpha=alpha, beta=beta) .build() ) - outputs: list[ChatMessage] = [] + outputs: list[list[ChatMessage]] = [] async for event in workflow.run_stream("coordinate task"): if isinstance(event, WorkflowOutputEvent): data = event.data - if isinstance(data, ChatMessage): - outputs.append(data) + if isinstance(data, list): + outputs.append(cast(list[ChatMessage], data)) assert len(outputs) == 1 - assert outputs[0].text == "done" - assert outputs[0].author_name == "manager" + assert len(outputs[0]) >= 1 + # The final message should be "done" from the manager + assert outputs[0][-1].text == "done" + assert outputs[0][-1].author_name == "manager" async def test_magentic_builder_returns_workflow_and_runs() -> None: @@ -169,11 +262,13 @@ async def test_magentic_builder_returns_workflow_and_runs() -> None: agent_event_count += 1 if isinstance(event, WorkflowOutputEvent): msg = event.data - if isinstance(msg, ChatMessage): - outputs.append(msg) + if isinstance(msg, list): + outputs.append(cast(list[ChatMessage], msg)) assert outputs, "Expected a final output message" - final = outputs[-1] + conversation = outputs[-1] + assert len(conversation) >= 1 + final = conversation[-1] assert final.text == "final" assert final.author_name == "magentic_manager" assert orchestrator_event_count > 0, "Expected orchestrator events to be emitted" @@ -187,7 +282,7 @@ async def test_group_chat_as_agent_accepts_conversation() -> None: workflow = ( GroupChatBuilder() - .select_speakers(selector, display_name="manager", final_message="done") + .set_select_speakers_func(selector, display_name="manager", final_message="done") .participants(alpha=alpha, beta=beta) .build() ) @@ -239,7 +334,7 @@ class TestGroupChatBuilder: def selector(state: GroupChatStateSnapshot) -> str | None: return None - builder = GroupChatBuilder().select_speakers(selector) + builder = GroupChatBuilder().set_select_speakers_func(selector) with pytest.raises(ValueError, match="participants must be configured before build"): builder.build() @@ -250,10 +345,10 @@ class TestGroupChatBuilder: def selector(state: GroupChatStateSnapshot) -> str | None: return None - builder = GroupChatBuilder().select_speakers(selector) + builder = GroupChatBuilder().set_select_speakers_func(selector) with pytest.raises(ValueError, match="already has a manager configured"): - builder.select_speakers(selector) + builder.set_select_speakers_func(selector) def test_empty_participants_raises_error(self) -> None: """Test that empty participants list raises ValueError.""" @@ -261,7 +356,7 @@ class TestGroupChatBuilder: def selector(state: GroupChatStateSnapshot) -> str | None: return None - builder = GroupChatBuilder().select_speakers(selector) + builder = GroupChatBuilder().set_select_speakers_func(selector) with pytest.raises(ValueError, match="participants cannot be empty"): builder.participants([]) @@ -274,7 +369,7 @@ class TestGroupChatBuilder: def selector(state: GroupChatStateSnapshot) -> str | None: return None - builder = GroupChatBuilder().select_speakers(selector) + builder = GroupChatBuilder().set_select_speakers_func(selector) with pytest.raises(ValueError, match="Duplicate participant name 'test'"): builder.participants([agent1, agent2]) @@ -302,7 +397,7 @@ class TestGroupChatBuilder: def selector(state: GroupChatStateSnapshot) -> str | None: return None - builder = GroupChatBuilder().select_speakers(selector) + builder = GroupChatBuilder().set_select_speakers_func(selector) with pytest.raises(ValueError, match="must define a non-empty 'name' attribute"): builder.participants([agent]) @@ -314,11 +409,53 @@ class TestGroupChatBuilder: def selector(state: GroupChatStateSnapshot) -> str | None: return None - builder = GroupChatBuilder().select_speakers(selector) + builder = GroupChatBuilder().set_select_speakers_func(selector) with pytest.raises(ValueError, match="participant names must be non-empty strings"): builder.participants({"": agent}) + def test_assemble_group_chat_respects_existing_start_executor(self) -> None: + """Ensure assemble_group_chat_workflow does not override preconfigured start executor.""" + + async def manager(_: GroupChatStateSnapshot) -> GroupChatDirective: + return GroupChatDirective(finish=True) + + builder = CountingWorkflowBuilder() + entry = PassthroughExecutor(id="entry") + builder = builder.set_start_executor(entry) + + participant = PassthroughExecutor(id="participant") + participant_spec = GroupChatParticipantSpec( + name="participant", + participant=participant, + description="participant", + ) + + wiring = _GroupChatConfig( + manager=manager, + manager_participant=None, + manager_name="manager", + participants={"participant": participant_spec}, + max_rounds=None, + termination_condition=None, + participant_aliases={}, + participant_executors={"participant": participant}, + ) + + result = assemble_group_chat_workflow( + wiring=wiring, + participant_factory=_default_participant_factory, + orchestrator_factory=_default_orchestrator_factory, + builder=builder, + return_builder=True, + ) + + assert isinstance(result, tuple) + assembled_builder, _ = result + assert assembled_builder is builder + assert builder.start_calls == 1 + assert assembled_builder._start_executor is entry # type: ignore + class TestGroupChatOrchestrator: """Tests for GroupChatOrchestratorExecutor core functionality.""" @@ -336,25 +473,116 @@ class TestGroupChatOrchestrator: workflow = ( GroupChatBuilder() - .select_speakers(selector) + .set_select_speakers_func(selector) .participants([agent]) .with_max_rounds(2) # Limit to 2 rounds .build() ) - outputs: list[ChatMessage] = [] + outputs: list[list[ChatMessage]] = [] async for event in workflow.run_stream("test task"): if isinstance(event, WorkflowOutputEvent): data = event.data - if isinstance(data, ChatMessage): - outputs.append(data) + if isinstance(data, list): + outputs.append(cast(list[ChatMessage], data)) # Should have terminated due to max_rounds, expect at least one output assert len(outputs) >= 1 - # The final message should be about round limit - final_output = outputs[-1] + # The final message in the conversation should be about round limit + conversation = outputs[-1] + assert len(conversation) >= 1 + final_output = conversation[-1] assert "round limit" in final_output.text.lower() + async def test_termination_condition_halts_conversation(self) -> None: + """Test that a custom termination condition stops the workflow.""" + + def selector(state: GroupChatStateSnapshot) -> str | None: + return "agent" + + def termination_condition(conversation: list[ChatMessage]) -> bool: + replies = [msg for msg in conversation if msg.role == Role.ASSISTANT and msg.author_name == "agent"] + return len(replies) >= 2 + + agent = StubAgent("agent", "response") + + workflow = ( + GroupChatBuilder() + .set_select_speakers_func(selector) + .participants([agent]) + .with_termination_condition(termination_condition) + .build() + ) + + outputs: list[list[ChatMessage]] = [] + async for event in workflow.run_stream("test task"): + if isinstance(event, WorkflowOutputEvent): + data = event.data + if isinstance(data, list): + outputs.append(cast(list[ChatMessage], data)) + + assert outputs, "Expected termination to yield output" + conversation = outputs[-1] + agent_replies = [msg for msg in conversation if msg.author_name == "agent" and msg.role == Role.ASSISTANT] + assert len(agent_replies) == 2 + final_output = conversation[-1] + assert final_output.author_name == "manager" + assert "termination condition" in final_output.text.lower() + + async def test_termination_condition_uses_manager_final_message(self) -> None: + """Test that manager-provided final message is used on termination.""" + + async def selector(state: GroupChatStateSnapshot) -> str | None: + return None + + agent = StubAgent("agent", "response") + final_text = "manager summary on termination" + + workflow = ( + GroupChatBuilder() + .set_select_speakers_func(selector, final_message=final_text) + .participants([agent]) + .with_termination_condition(lambda _: True) + .build() + ) + + outputs: list[list[ChatMessage]] = [] + async for event in workflow.run_stream("test task"): + if isinstance(event, WorkflowOutputEvent): + data = event.data + if isinstance(data, list): + outputs.append(cast(list[ChatMessage], data)) + + assert outputs, "Expected termination to yield output" + conversation = outputs[-1] + assert conversation[-1].text == final_text + assert conversation[-1].author_name == "manager" + + async def test_termination_condition_agent_manager_finalizes(self) -> None: + """Test that agent-based manager can provide final message on termination.""" + manager = StubManagerAgent() + worker = StubAgent("agent", "response") + + workflow = ( + GroupChatBuilder() + .set_manager(manager, display_name="Manager") + .participants([worker]) + .with_termination_condition(lambda conv: any(msg.author_name == "agent" for msg in conv)) + .build() + ) + + outputs: list[list[ChatMessage]] = [] + async for event in workflow.run_stream("test task"): + if isinstance(event, WorkflowOutputEvent): + data = event.data + if isinstance(data, list): + outputs.append(cast(list[ChatMessage], data)) + + assert outputs, "Expected termination to yield output" + conversation = outputs[-1] + assert conversation[-1].text == "agent manager final" + assert conversation[-1].author_name == "Manager" + async def test_unknown_participant_error(self) -> None: """Test that _apply_directive raises error for unknown participants.""" @@ -363,7 +591,7 @@ class TestGroupChatOrchestrator: agent = StubAgent("agent", "response") - workflow = GroupChatBuilder().select_speakers(selector).participants([agent]).build() + workflow = GroupChatBuilder().set_select_speakers_func(selector).participants([agent]).build() with pytest.raises(ValueError, match="Manager selected unknown participant 'unknown_agent'"): async for _ in workflow.run_stream("test task"): @@ -379,7 +607,7 @@ class TestGroupChatOrchestrator: agent = StubAgent("agent", "response") # The _SpeakerSelectorAdapter will catch this and raise TypeError - workflow = GroupChatBuilder().select_speakers(bad_selector).participants([agent]).build() # type: ignore + workflow = GroupChatBuilder().set_select_speakers_func(bad_selector).participants([agent]).build() # type: ignore # This should raise a TypeError because selector doesn't return str or None with pytest.raises(TypeError, match="must return a participant name \\(str\\) or None"): @@ -394,7 +622,7 @@ class TestGroupChatOrchestrator: agent = StubAgent("agent", "response") - workflow = GroupChatBuilder().select_speakers(selector).participants([agent]).build() + workflow = GroupChatBuilder().set_select_speakers_func(selector).participants([agent]).build() with pytest.raises(ValueError, match="requires at least one chat message"): async for _ in workflow.run_stream([]): @@ -529,69 +757,76 @@ class TestCheckpointing: storage = InMemoryCheckpointStorage() workflow = ( - GroupChatBuilder().select_speakers(selector).participants([agent]).with_checkpointing(storage).build() + GroupChatBuilder() + .set_select_speakers_func(selector) + .participants([agent]) + .with_checkpointing(storage) + .build() ) - outputs: list[ChatMessage] = [] + outputs: list[list[ChatMessage]] = [] async for event in workflow.run_stream("test task"): if isinstance(event, WorkflowOutputEvent): data = event.data - if isinstance(data, ChatMessage): - outputs.append(data) + if isinstance(data, list): + outputs.append(cast(list[ChatMessage], data)) assert len(outputs) == 1 # Should complete normally -class TestPromptBasedManager: - """Tests for _PromptBasedGroupChatManager.""" +class TestAgentManagerConfiguration: + """Tests for agent-based manager configuration.""" - async def test_manager_with_missing_next_agent_raises_error(self) -> None: - """Test that manager directive without next_agent raises RuntimeError.""" + async def test_set_manager_configures_response_format(self) -> None: + """Ensure ChatAgent managers receive default ManagerSelectionResponse formatting.""" + from unittest.mock import MagicMock - class MockChatClient: - async def get_response(self, messages: Any, response_format: Any = None) -> Any: - # Return response that has finish=False but no next_agent - class MockResponse: - def __init__(self) -> None: - self.value = {"finish": False, "next_agent": None} - self.messages: list[Any] = [] + from agent_framework import ChatAgent - return MockResponse() + chat_client = MagicMock() + manager_agent = ChatAgent(chat_client=chat_client, name="Coordinator") + assert manager_agent.chat_options.response_format is None - manager = _PromptBasedGroupChatManager(MockChatClient()) # type: ignore + worker = StubAgent("worker", "response") - state = { - "participants": {"agent": "desc"}, - "task": ChatMessage(role=Role.USER, text="test"), - "conversation": (), - } + builder = GroupChatBuilder().set_manager(manager_agent).participants([worker]) - with pytest.raises(RuntimeError, match="missing next_agent while finish is False"): - await manager(state) + assert manager_agent.chat_options.response_format is ManagerSelectionResponse + assert builder._manager_participant is manager_agent # type: ignore[attr-defined] - async def test_manager_with_unknown_participant_raises_error(self) -> None: - """Test that manager selecting unknown participant raises RuntimeError.""" + async def test_set_manager_accepts_agent_manager(self) -> None: + """Verify agent-based manager can be set and workflow builds.""" + from unittest.mock import MagicMock - class MockChatClient: - async def get_response(self, messages: Any, response_format: Any = None) -> Any: - # Return response selecting unknown participant - class MockResponse: - def __init__(self) -> None: - self.value = {"finish": False, "next_agent": "unknown"} - self.messages: list[Any] = [] + from agent_framework import ChatAgent - return MockResponse() + chat_client = MagicMock() + manager_agent = ChatAgent(chat_client=chat_client, name="Coordinator") + worker = StubAgent("worker", "response") - manager = _PromptBasedGroupChatManager(MockChatClient()) # type: ignore + builder = GroupChatBuilder().set_manager(manager_agent, display_name="Orchestrator") + builder = builder.participants([worker]).with_max_rounds(1) - state = { - "participants": {"agent": "desc"}, - "task": ChatMessage(role=Role.USER, text="test"), - "conversation": (), - } + assert builder._manager_participant is manager_agent # type: ignore[attr-defined] + assert "worker" in builder._participants # type: ignore[attr-defined] - with pytest.raises(RuntimeError, match="Manager selected unknown participant 'unknown'"): - await manager(state) + async def test_set_manager_rejects_custom_response_format(self) -> None: + """Reject custom response_format on ChatAgent managers.""" + from unittest.mock import MagicMock + + from agent_framework import ChatAgent + + class CustomResponse(BaseModel): + value: str + + chat_client = MagicMock() + manager_agent = ChatAgent(chat_client=chat_client, name="Coordinator", response_format=CustomResponse) + worker = StubAgent("worker", "response") + + with pytest.raises(ValueError, match="response_format must be ManagerSelectionResponse"): + GroupChatBuilder().set_manager(manager_agent).participants([worker]) + + assert manager_agent.chat_options.response_format is CustomResponse class TestFactoryFunctions: @@ -599,9 +834,9 @@ class TestFactoryFunctions: def test_default_orchestrator_factory_without_manager_raises_error(self) -> None: """Test that default factory requires manager to be set.""" - config = _GroupChatConfig(manager=None, manager_name="test", participants={}) + config = _GroupChatConfig(manager=None, manager_participant=None, manager_name="test", participants={}) - with pytest.raises(RuntimeError, match="requires a manager to be set"): + with pytest.raises(RuntimeError, match="requires a manager to be configured"): _default_orchestrator_factory(config) @@ -619,14 +854,14 @@ class TestConversationHandling: agent = StubAgent("agent", "response") - workflow = GroupChatBuilder().select_speakers(selector).participants([agent]).build() + workflow = GroupChatBuilder().set_select_speakers_func(selector).participants([agent]).build() - outputs: list[ChatMessage] = [] + outputs: list[list[ChatMessage]] = [] async for event in workflow.run_stream("test string"): if isinstance(event, WorkflowOutputEvent): data = event.data - if isinstance(data, ChatMessage): - outputs.append(data) + if isinstance(data, list): + outputs.append(cast(list[ChatMessage], data)) assert len(outputs) == 1 @@ -641,14 +876,14 @@ class TestConversationHandling: agent = StubAgent("agent", "response") - workflow = GroupChatBuilder().select_speakers(selector).participants([agent]).build() + workflow = GroupChatBuilder().set_select_speakers_func(selector).participants([agent]).build() - outputs: list[ChatMessage] = [] + outputs: list[list[ChatMessage]] = [] async for event in workflow.run_stream(task_message): if isinstance(event, WorkflowOutputEvent): data = event.data - if isinstance(data, ChatMessage): - outputs.append(data) + if isinstance(data, list): + outputs.append(cast(list[ChatMessage], data)) assert len(outputs) == 1 @@ -667,14 +902,14 @@ class TestConversationHandling: agent = StubAgent("agent", "response") - workflow = GroupChatBuilder().select_speakers(selector).participants([agent]).build() + workflow = GroupChatBuilder().set_select_speakers_func(selector).participants([agent]).build() - outputs: list[ChatMessage] = [] + outputs: list[list[ChatMessage]] = [] async for event in workflow.run_stream(conversation): if isinstance(event, WorkflowOutputEvent): data = event.data - if isinstance(data, ChatMessage): - outputs.append(data) + if isinstance(data, list): + outputs.append(cast(list[ChatMessage], data)) assert len(outputs) == 1 @@ -695,23 +930,25 @@ class TestRoundLimitEnforcement: workflow = ( GroupChatBuilder() - .select_speakers(selector) + .set_select_speakers_func(selector) .participants([agent]) .with_max_rounds(1) # Very low limit .build() ) - outputs: list[ChatMessage] = [] + outputs: list[list[ChatMessage]] = [] async for event in workflow.run_stream("test"): if isinstance(event, WorkflowOutputEvent): data = event.data - if isinstance(data, ChatMessage): - outputs.append(data) + if isinstance(data, list): + outputs.append(cast(list[ChatMessage], data)) # Should have at least one output (the round limit message) assert len(outputs) >= 1 - # The last message should be about round limit - final_output = outputs[-1] + # The last message in the conversation should be about round limit + conversation = outputs[-1] + assert len(conversation) >= 1 + final_output = conversation[-1] assert "round limit" in final_output.text.lower() async def test_round_limit_in_ingest_participant_message(self) -> None: @@ -728,23 +965,25 @@ class TestRoundLimitEnforcement: workflow = ( GroupChatBuilder() - .select_speakers(selector) + .set_select_speakers_func(selector) .participants([agent]) .with_max_rounds(1) # Hit limit after first response .build() ) - outputs: list[ChatMessage] = [] + outputs: list[list[ChatMessage]] = [] async for event in workflow.run_stream("test"): if isinstance(event, WorkflowOutputEvent): data = event.data - if isinstance(data, ChatMessage): - outputs.append(data) + if isinstance(data, list): + outputs.append(cast(list[ChatMessage], data)) # Should have at least one output (the round limit message) assert len(outputs) >= 1 - # The last message should be about round limit - final_output = outputs[-1] + # The last message in the conversation should be about round limit + conversation = outputs[-1] + assert len(conversation) >= 1 + final_output = conversation[-1] assert "round limit" in final_output.text.lower() @@ -758,12 +997,12 @@ async def test_group_chat_checkpoint_runtime_only() -> None: agent_b = StubAgent("agentB", "Reply from B") selector = make_sequence_selector() - wf = GroupChatBuilder().participants([agent_a, agent_b]).select_speakers(selector).build() + wf = GroupChatBuilder().participants([agent_a, agent_b]).set_select_speakers_func(selector).build() baseline_output: list[ChatMessage] | None = None async for ev in wf.run_stream("runtime checkpoint test", checkpoint_storage=storage): if isinstance(ev, WorkflowOutputEvent): - baseline_output = ev.data # type: ignore[assignment] + baseline_output = cast(list[ChatMessage], ev.data) if isinstance(ev.data, list) else None if isinstance(ev, WorkflowStatusEvent) and ev.state in ( WorkflowRunState.IDLE, WorkflowRunState.IDLE_WITH_PENDING_REQUESTS, @@ -794,7 +1033,7 @@ async def test_group_chat_checkpoint_runtime_overrides_buildtime() -> None: wf = ( GroupChatBuilder() .participants([agent_a, agent_b]) - .select_speakers(selector) + .set_select_speakers_func(selector) .with_checkpointing(buildtime_storage) .build() ) @@ -802,7 +1041,7 @@ async def test_group_chat_checkpoint_runtime_overrides_buildtime() -> None: baseline_output: list[ChatMessage] | None = None async for ev in wf.run_stream("override test", checkpoint_storage=runtime_storage): if isinstance(ev, WorkflowOutputEvent): - baseline_output = ev.data # type: ignore[assignment] + baseline_output = cast(list[ChatMessage], ev.data) if isinstance(ev.data, list) else None if isinstance(ev, WorkflowStatusEvent) and ev.state in ( WorkflowRunState.IDLE, WorkflowRunState.IDLE_WITH_PENDING_REQUESTS, @@ -816,3 +1055,30 @@ async def test_group_chat_checkpoint_runtime_overrides_buildtime() -> None: assert len(runtime_checkpoints) > 0, "Runtime storage should have checkpoints" assert len(buildtime_checkpoints) == 0, "Build-time storage should have no checkpoints when overridden" + + +class _StubExecutor(Executor): + """Minimal executor used to satisfy workflow wiring in tests.""" + + def __init__(self, id: str) -> None: + super().__init__(id=id) + + @handler + async def handle(self, message: object, ctx: WorkflowContext[ChatMessage]) -> None: + await ctx.yield_output(message) + + +def test_set_manager_builds_with_agent_manager() -> None: + """GroupChatBuilder should build when using an agent-based manager.""" + + manager = _StubExecutor("manager_executor") + participant = _StubExecutor("participant_executor") + + workflow = ( + GroupChatBuilder().set_manager(manager, display_name="Moderator").participants({"worker": participant}).build() + ) + + orchestrator = workflow.get_start_executor() + + assert isinstance(orchestrator, GroupChatOrchestratorExecutor) + assert orchestrator._is_manager_agent() diff --git a/python/packages/core/tests/workflow/test_handoff.py b/python/packages/core/tests/workflow/test_handoff.py index 5dfd7522df..3bbed7681e 100644 --- a/python/packages/core/tests/workflow/test_handoff.py +++ b/python/packages/core/tests/workflow/test_handoff.py @@ -23,7 +23,22 @@ from agent_framework import ( WorkflowOutputEvent, ) from agent_framework._mcp import MCPTool +from agent_framework._workflows import _handoff as handoff_module # type: ignore from agent_framework._workflows._handoff import _clone_chat_agent # type: ignore[reportPrivateUsage] +from agent_framework._workflows._workflow_builder import WorkflowBuilder + + +class _CountingWorkflowBuilder(WorkflowBuilder): + created: list["_CountingWorkflowBuilder"] = [] + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.start_calls = 0 + _CountingWorkflowBuilder.created.append(self) + + def set_start_executor(self, executor: Any) -> "_CountingWorkflowBuilder": # type: ignore[override] + self.start_calls += 1 + return cast("_CountingWorkflowBuilder", super().set_start_executor(executor)) @dataclass @@ -478,6 +493,27 @@ async def test_return_to_previous_enabled(): assert len(specialist_a.calls) == 2, "Specialist A should handle follow-up with return_to_previous enabled" +def test_handoff_builder_sets_start_executor_once(monkeypatch: pytest.MonkeyPatch) -> None: + """Ensure HandoffBuilder.build sets the start executor only once when assembling the workflow.""" + _CountingWorkflowBuilder.created.clear() + monkeypatch.setattr(handoff_module, "WorkflowBuilder", _CountingWorkflowBuilder) + + coordinator = _RecordingAgent(name="coordinator") + specialist = _RecordingAgent(name="specialist") + + workflow = ( + HandoffBuilder(participants=[coordinator, specialist]) + .set_coordinator("coordinator") + .with_termination_condition(lambda conv: len(conv) > 0) + .build() + ) + + assert workflow is not None + assert _CountingWorkflowBuilder.created, "Expected CountingWorkflowBuilder to be instantiated" + builder = _CountingWorkflowBuilder.created[-1] + assert builder.start_calls == 1, "set_start_executor should be invoked exactly once" + + async def test_tool_choice_preserved_from_agent_config(): """Verify that agent-level tool_choice configuration is preserved and not overridden.""" from unittest.mock import AsyncMock diff --git a/python/packages/core/tests/workflow/test_magentic.py b/python/packages/core/tests/workflow/test_magentic.py index 7dc8a2c471..b41d243a3e 100644 --- a/python/packages/core/tests/workflow/test_magentic.py +++ b/python/packages/core/tests/workflow/test_magentic.py @@ -33,6 +33,7 @@ from agent_framework import ( WorkflowStatusEvent, handler, ) +from agent_framework._workflows import _group_chat as group_chat_module # type: ignore from agent_framework._workflows._checkpoint import InMemoryCheckpointStorage from agent_framework._workflows._magentic import ( # type: ignore[reportPrivateUsage] MagenticAgentExecutor, @@ -42,6 +43,7 @@ from agent_framework._workflows._magentic import ( # type: ignore[reportPrivate _MagenticProgressLedgerItem, # type: ignore _MagenticStartMessage, # type: ignore ) +from agent_framework._workflows._workflow_builder import WorkflowBuilder if sys.version_info >= (3, 12): from typing import override @@ -162,6 +164,19 @@ class FakeManager(MagenticManagerBase): return ChatMessage(role=Role.ASSISTANT, text="FINAL", author_name="magentic_manager") +class _CountingWorkflowBuilder(WorkflowBuilder): + created: list["_CountingWorkflowBuilder"] = [] + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.start_calls = 0 + _CountingWorkflowBuilder.created.append(self) + + def set_start_executor(self, executor: Any) -> "_CountingWorkflowBuilder": # type: ignore[override] + self.start_calls += 1 + return cast("_CountingWorkflowBuilder", super().set_start_executor(executor)) + + async def test_standard_manager_plan_and_replan_combined_ledger(): manager = FakeManager(max_round_count=10, max_stall_count=3, max_reset_count=2) ctx = MagenticContext( @@ -210,7 +225,7 @@ async def test_magentic_workflow_plan_review_approval_to_completion(): assert req_event is not None completed = False - output: ChatMessage | None = None + output: list[ChatMessage] | None = None async for ev in wf.send_responses_streaming( responses={req_event.request_id: MagenticPlanReviewReply(decision=MagenticPlanReviewDecision.APPROVE)} ): @@ -222,7 +237,8 @@ async def test_magentic_workflow_plan_review_approval_to_completion(): break assert completed assert output is not None - assert isinstance(output, ChatMessage) + assert isinstance(output, list) + assert all(isinstance(msg, ChatMessage) for msg in output) async def test_magentic_plan_review_approve_with_comments_replans_and_proceeds(): @@ -300,8 +316,10 @@ async def test_magentic_orchestrator_round_limit_produces_partial_result(): output_event = next((e for e in events if isinstance(e, WorkflowOutputEvent)), None) assert output_event is not None data = output_event.data - assert isinstance(data, ChatMessage) - assert data.role == Role.ASSISTANT + assert isinstance(data, list) + assert all(isinstance(msg, ChatMessage) for msg in data) + assert len(data) > 0 + assert data[-1].role == Role.ASSISTANT async def test_magentic_checkpoint_resume_round_trip(): @@ -374,6 +392,23 @@ class _DummyExec(Executor): pass +def test_magentic_builder_sets_start_executor_once(monkeypatch: pytest.MonkeyPatch) -> None: + """Ensure MagenticBuilder wiring sets the start executor only once.""" + _CountingWorkflowBuilder.created.clear() + monkeypatch.setattr(group_chat_module, "WorkflowBuilder", _CountingWorkflowBuilder) + + manager = FakeManager() + + workflow = ( + MagenticBuilder().participants(agentA=_DummyExec("agentA")).with_standard_manager(manager=manager).build() + ) + + assert workflow is not None + assert _CountingWorkflowBuilder.created, "Expected CountingWorkflowBuilder to be instantiated" + builder = _CountingWorkflowBuilder.created[-1] + assert builder.start_calls == 1, "set_start_executor should be called exactly once" + + async def test_magentic_agent_executor_on_checkpoint_save_and_restore_roundtrip(): backing_executor = _DummyExec("backing") agent_exec = MagenticAgentExecutor(backing_executor, "agentA") @@ -746,9 +781,11 @@ async def test_magentic_stall_and_reset_successfully(): assert idle_status is not None output_event = next((e for e in events if isinstance(e, WorkflowOutputEvent)), None) assert output_event is not None - assert isinstance(output_event.data, ChatMessage) - assert output_event.data.text is not None - assert output_event.data.text == "re-ledger" + assert isinstance(output_event.data, list) + assert all(isinstance(msg, ChatMessage) for msg in output_event.data) + assert len(output_event.data) > 0 + assert output_event.data[-1].text is not None + assert output_event.data[-1].text == "re-ledger" async def test_magentic_checkpoint_runtime_only() -> None: diff --git a/python/samples/README.md b/python/samples/README.md index bb7d01527a..3434fa639d 100644 --- a/python/samples/README.md +++ b/python/samples/README.md @@ -268,97 +268,7 @@ This directory contains samples demonstrating the capabilities of Microsoft Agen ## Workflows -### Start Here - -| File | Description | -|------|-------------| -| [`getting_started/workflows/_start-here/step1_executors_and_edges.py`](./getting_started/workflows/_start-here/step1_executors_and_edges.py) | Step 1: Foundational patterns: Executors and edges | -| [`getting_started/workflows/_start-here/step2_agents_in_a_workflow.py`](./getting_started/workflows/_start-here/step2_agents_in_a_workflow.py) | Step 2: Agents in a Workflow non-streaming | -| [`getting_started/workflows/_start-here/step3_streaming.py`](./getting_started/workflows/_start-here/step3_streaming.py) | Step 3: Agents in a workflow with streaming | - -### Agents in Workflows - -| File | Description | -|------|-------------| -| [`getting_started/workflows/agents/azure_ai_agents_streaming.py`](./getting_started/workflows/agents/azure_ai_agents_streaming.py) | Sample: Agents in a workflow with streaming | -| [`getting_started/workflows/agents/azure_chat_agents_function_bridge.py`](./getting_started/workflows/agents/azure_chat_agents_function_bridge.py) | Sample: Two agents connected by a function executor bridge | -| [`getting_started/workflows/agents/azure_chat_agents_streaming.py`](./getting_started/workflows/agents/azure_chat_agents_streaming.py) | Sample: Agents in a workflow with streaming | -| [`getting_started/workflows/agents/azure_chat_agents_tool_calls_with_feedback.py`](./getting_started/workflows/agents/azure_chat_agents_tool_calls_with_feedback.py) | Sample: Tool-enabled agents with human feedback | -| [`getting_started/workflows/agents/custom_agent_executors.py`](./getting_started/workflows/agents/custom_agent_executors.py) | Step 2: Agents in a Workflow non-streaming | -| [`getting_started/workflows/agents/workflow_as_agent_human_in_the_loop.py`](./getting_started/workflows/agents/workflow_as_agent_human_in_the_loop.py) | Sample: Workflow Agent with Human-in-the-Loop | -| [`getting_started/workflows/agents/workflow_as_agent_reflection_pattern.py`](./getting_started/workflows/agents/workflow_as_agent_reflection_pattern.py) | Sample: Workflow as Agent with Reflection and Retry Pattern | - -### Checkpoint - -| File | Description | -|------|-------------| -| [`getting_started/workflows/checkpoint/checkpoint_with_human_in_the_loop.py`](./getting_started/workflows/checkpoint/checkpoint_with_human_in_the_loop.py) | Sample: Checkpoint + human-in-the-loop quickstart | -| [`getting_started/workflows/checkpoint/checkpoint_with_resume.py`](./getting_started/workflows/checkpoint/checkpoint_with_resume.py) | Sample: Checkpointing and Resuming a Workflow (with an Agent stage) | -| [`getting_started/workflows/checkpoint/sub_workflow_checkpoint.py`](./getting_started/workflows/checkpoint/sub_workflow_checkpoint.py) | Sample: Checkpointing for workflows that embed sub-workflows | - -### Composition - -| File | Description | -|------|-------------| -| [`getting_started/workflows/composition/sub_workflow_basics.py`](./getting_started/workflows/composition/sub_workflow_basics.py) | Sample: Sub-Workflows (Basics) | -| [`getting_started/workflows/composition/sub_workflow_parallel_requests.py`](./getting_started/workflows/composition/sub_workflow_parallel_requests.py) | Sample: Sub-workflow with parallel request handling by specialized interceptors | -| [`getting_started/workflows/composition/sub_workflow_request_interception.py`](./getting_started/workflows/composition/sub_workflow_request_interception.py) | Sample: Sub-Workflows with Request Interception | - -### Control Flow - -| File | Description | -|------|-------------| -| [`getting_started/workflows/control-flow/edge_condition.py`](./getting_started/workflows/control-flow/edge_condition.py) | Sample: Conditional routing with structured outputs | -| [`getting_started/workflows/control-flow/multi_selection_edge_group.py`](./getting_started/workflows/control-flow/multi_selection_edge_group.py) | Step 06b — Multi-Selection Edge Group sample | -| [`getting_started/workflows/control-flow/sequential_executors.py`](./getting_started/workflows/control-flow/sequential_executors.py) | Sample: Sequential workflow with streaming | -| [`getting_started/workflows/control-flow/sequential_streaming.py`](./getting_started/workflows/control-flow/sequential_streaming.py) | Sample: Foundational sequential workflow with streaming using function-style executors | -| [`getting_started/workflows/control-flow/simple_loop.py`](./getting_started/workflows/control-flow/simple_loop.py) | Sample: Simple Loop (with an Agent Judge) | -| [`getting_started/workflows/control-flow/switch_case_edge_group.py`](./getting_started/workflows/control-flow/switch_case_edge_group.py) | Sample: Switch-Case Edge Group with an explicit Uncertain branch | - -### Human-in-the-Loop - -| File | Description | -|------|-------------| -| [`getting_started/workflows/human-in-the-loop/guessing_game_with_human_input.py`](./getting_started/workflows/human-in-the-loop/guessing_game_with_human_input.py) | Sample: Human in the loop guessing game | -| [`getting_started/workflows/human-in-the-loop/agents_with_approval_requests.py`](./getting_started/workflows/human-in-the-loop/agents_with_approval_requests.py) | Sample: Agents with Approval Requests in Workflows | - -### Orchestration - -| File | Description | -|------|-------------| -| [`getting_started/workflows/orchestration/concurrent_agents.py`](./getting_started/workflows/orchestration/concurrent_agents.py) | Sample: Concurrent fan-out/fan-in (agent-only API) with default aggregator | -| [`getting_started/workflows/orchestration/concurrent_custom_agent_executors.py`](./getting_started/workflows/orchestration/concurrent_custom_agent_executors.py) | Sample: Concurrent Orchestration with Custom Agent Executors | -| [`getting_started/workflows/orchestration/concurrent_custom_aggregator.py`](./getting_started/workflows/orchestration/concurrent_custom_aggregator.py) | Sample: Concurrent Orchestration with Custom Aggregator | -| [`getting_started/workflows/orchestration/group_chat_prompt_based_manager.py`](./getting_started/workflows/orchestration/group_chat_prompt_based_manager.py) | Sample: Group Chat Orchestration with LLM-based manager | -| [`getting_started/workflows/orchestration/group_chat_simple_selector.py`](./getting_started/workflows/orchestration/group_chat_simple_selector.py) | Sample: Group Chat Orchestration with function-based speaker selector | -| [`getting_started/workflows/orchestration/handoff_simple.py`](./getting_started/workflows/orchestration/handoff_simple.py) | Sample: Handoff Orchestration with simple agent handoff pattern | -| [`getting_started/workflows/orchestration/handoff_specialist_to_specialist.py`](./getting_started/workflows/orchestration/handoff_specialist_to_specialist.py) | Sample: Handoff Orchestration with specialist-to-specialist routing | -| [`getting_started/workflows/orchestration/handoff_return_to_previous`](./getting_started/workflows/orchestration/handoff_return_to_previous.py) | Return-to-previous routing: after user input, routes back to the previous specialist instead of coordinator using `.enable_return_to_previous()` | -| [`getting_started/workflows/orchestration/magentic.py`](./getting_started/workflows/orchestration/magentic.py) | Sample: Magentic Orchestration (agentic task planning with multi-agent execution) | -| [`getting_started/workflows/orchestration/magentic_checkpoint.py`](./getting_started/workflows/orchestration/magentic_checkpoint.py) | Sample: Magentic Orchestration with Checkpointing | -| [`getting_started/workflows/orchestration/magentic_human_plan_update.py`](./getting_started/workflows/orchestration/magentic_human_plan_update.py) | Sample: Magentic Orchestration with Human Plan Review | -| [`getting_started/workflows/orchestration/sequential_agents.py`](./getting_started/workflows/orchestration/sequential_agents.py) | Sample: Sequential workflow (agent-focused API) with shared conversation context | -| [`getting_started/workflows/orchestration/sequential_custom_executors.py`](./getting_started/workflows/orchestration/sequential_custom_executors.py) | Sample: Sequential workflow mixing agents and a custom summarizer executor | - -### Parallelism - -| File | Description | -|------|-------------| -| [`getting_started/workflows/parallelism/aggregate_results_of_different_types.py`](./getting_started/workflows/parallelism/aggregate_results_of_different_types.py) | Sample: Concurrent fan out and fan in with two different tasks that output results of different types | -| [`getting_started/workflows/parallelism/fan_out_fan_in_edges.py`](./getting_started/workflows/parallelism/fan_out_fan_in_edges.py) | Sample: Concurrent fan out and fan in with three domain agents | -| [`getting_started/workflows/parallelism/map_reduce_and_visualization.py`](./getting_started/workflows/parallelism/map_reduce_and_visualization.py) | Sample: Map reduce word count with fan out and fan in over file backed intermediate results | - -### State Management - -| File | Description | -|------|-------------| -| [`getting_started/workflows/state-management/shared_states_with_agents.py`](./getting_started/workflows/state-management/shared_states_with_agents.py) | Sample: Shared state with agents and conditional routing | - -### Visualization - -| File | Description | -|------|-------------| -| [`getting_started/workflows/visualization/concurrent_with_visualization.py`](./getting_started/workflows/visualization/concurrent_with_visualization.py) | Sample: Concurrent (Fan-out/Fan-in) with Agents + Visualization | +View the list of Workflows samples [here](./getting_started/workflows/README.md). ## Sample Guidelines diff --git a/python/samples/getting_started/agents/azure_ai/README.md b/python/samples/getting_started/agents/azure_ai/README.md index 0513dab1df..17a944524a 100644 --- a/python/samples/getting_started/agents/azure_ai/README.md +++ b/python/samples/getting_started/agents/azure_ai/README.md @@ -16,6 +16,7 @@ This folder contains examples demonstrating different ways to create and use age | [`azure_ai_with_code_interpreter.py`](azure_ai_with_code_interpreter.py) | Shows how to use the `HostedCodeInterpreterTool` with Azure AI agents to write and execute Python code for mathematical problem solving and data analysis. | | [`azure_ai_with_existing_agent.py`](azure_ai_with_existing_agent.py) | Shows how to work with a pre-existing agent by providing the agent name and version to the Azure AI client. Demonstrates agent reuse patterns for production scenarios. | | [`azure_ai_with_existing_conversation.py`](azure_ai_with_existing_conversation.py) | Demonstrates how to use an existing conversation created on the service side with Azure AI agents. Shows two approaches: specifying conversation ID at the client level and using AgentThread with an existing conversation ID. | +| [`azure_ai_with_application_endpoint.py`](azure_ai_with_application_endpoint.py) | Demonstrates calling the Azure AI application-scoped endpoint. | | [`azure_ai_with_explicit_settings.py`](azure_ai_with_explicit_settings.py) | Shows how to create an agent with explicitly configured `AzureAIClient` settings, including project endpoint, model deployment, and credentials rather than relying on environment variable defaults. | | [`azure_ai_with_file_search.py`](azure_ai_with_file_search.py) | Shows how to use the `HostedFileSearchTool` with Azure AI agents to upload files, create vector stores, and enable agents to search through uploaded documents to answer user questions. | | [`azure_ai_with_hosted_mcp.py`](azure_ai_with_hosted_mcp.py) | Shows how to integrate hosted Model Context Protocol (MCP) tools with Azure AI Agent. | diff --git a/python/samples/getting_started/agents/azure_ai/azure_ai_with_application_endpoint.py b/python/samples/getting_started/agents/azure_ai/azure_ai_with_application_endpoint.py new file mode 100644 index 0000000000..89bb77af11 --- /dev/null +++ b/python/samples/getting_started/agents/azure_ai/azure_ai_with_application_endpoint.py @@ -0,0 +1,39 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import os + +from agent_framework import ChatAgent +from agent_framework.azure import AzureAIClient +from azure.ai.projects.aio import AIProjectClient +from azure.identity.aio import AzureCliCredential + +""" +Azure AI Agent with Application Endpoint Example + +This sample demonstrates working with pre-existing Azure AI Agents by providing +application endpoint instead of project endpoint. +""" + + +async def main() -> None: + # Create the client + async with ( + AzureCliCredential() as credential, + # Endpoint here should be application endpoint with format: + # /api/projects//applications//protocols + AIProjectClient(endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], credential=credential) as project_client, + ChatAgent( + chat_client=AzureAIClient( + project_client=project_client, + ), + ) as agent, + ): + query = "How are you?" + print(f"User: {query}") + result = await agent.run(query) + print(f"Agent: {result}\n") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/python/samples/getting_started/agents/openai/openai_responses_client_with_thread.py b/python/samples/getting_started/agents/openai/openai_responses_client_with_thread.py index 41192ca977..ca52b4f074 100644 --- a/python/samples/getting_started/agents/openai/openai_responses_client_with_thread.py +++ b/python/samples/getting_started/agents/openai/openai_responses_client_with_thread.py @@ -67,19 +67,19 @@ async def example_with_thread_persistence_in_memory() -> None: # First conversation query1 = "What's the weather like in Tokyo?" print(f"User: {query1}") - result1 = await agent.run(query1, thread=thread) + result1 = await agent.run(query1, thread=thread, store=False) print(f"Agent: {result1.text}") # Second conversation using the same thread - maintains context query2 = "How about London?" print(f"\nUser: {query2}") - result2 = await agent.run(query2, thread=thread) + result2 = await agent.run(query2, thread=thread, store=False) print(f"Agent: {result2.text}") # Third conversation - agent should remember both previous cities query3 = "Which of the cities I asked about has better weather?" print(f"\nUser: {query3}") - result3 = await agent.run(query3, thread=thread) + result3 = await agent.run(query3, thread=thread, store=False) print(f"Agent: {result3.text}") print("Note: The agent remembers context from previous messages in the same thread.\n") @@ -105,8 +105,7 @@ async def example_with_existing_thread_id() -> None: query1 = "What's the weather in Paris?" print(f"User: {query1}") - # Enable OpenAI conversation state by setting `store` parameter to True - result1 = await agent.run(query1, thread=thread, store=True) + result1 = await agent.run(query1, thread=thread) print(f"Agent: {result1.text}") # The thread ID is set after the first response @@ -127,7 +126,7 @@ async def example_with_existing_thread_id() -> None: query2 = "What was the last city I asked about?" print(f"User: {query2}") - result2 = await agent.run(query2, thread=thread, store=True) + result2 = await agent.run(query2, thread=thread) print(f"Agent: {result2.text}") print("Note: The agent continues the conversation from the previous thread by using thread ID.\n") diff --git a/python/samples/getting_started/azure_functions/08_mcp_server/README.md b/python/samples/getting_started/azure_functions/08_mcp_server/README.md new file mode 100644 index 0000000000..ed8ecfb1e9 --- /dev/null +++ b/python/samples/getting_started/azure_functions/08_mcp_server/README.md @@ -0,0 +1,187 @@ +# Agent as MCP Tool Sample + +This sample demonstrates how to configure AI agents to be accessible as both HTTP endpoints and [Model Context Protocol (MCP)](https://modelcontextprotocol.io/) tools, enabling flexible integration patterns for AI agent consumption. + +## Key Concepts Demonstrated + +- **Multi-trigger Agent Configuration**: Configure agents to support HTTP triggers, MCP tool triggers, or both +- **Microsoft Agent Framework Integration**: Use the framework to define AI agents with specific roles and capabilities +- **Flexible Agent Registration**: Register agents with customizable trigger configurations +- **MCP Server Hosting**: Expose agents as MCP tools for consumption by MCP-compatible clients + +## Sample Architecture + +This sample creates three agents with different trigger configurations: + +| Agent | Role | HTTP Trigger | MCP Tool Trigger | Description | +|-------|------|--------------|------------------|-------------| +| **Joker** | Comedy specialist | ✅ Enabled | ❌ Disabled | Accessible only via HTTP requests | +| **StockAdvisor** | Financial data | ❌ Disabled | ✅ Enabled | Accessible only as MCP tool | +| **PlantAdvisor** | Indoor plant recommendations | ✅ Enabled | ✅ Enabled | Accessible via both HTTP and MCP | + +## Environment Setup + +See the [README.md](../README.md) file in the parent directory for complete setup instructions, including: + +- Prerequisites installation +- Azure OpenAI configuration +- Durable Task Scheduler setup +- Storage emulator configuration + +## Configuration + +Update your `local.settings.json` with your Azure OpenAI credentials: + +```json +{ + "Values": { + "AZURE_OPENAI_ENDPOINT": "https://your-resource.openai.azure.com/", + "AZURE_OPENAI_CHAT_DEPLOYMENT_NAME": "your-deployment-name", + "AZURE_OPENAI_KEY": "your-api-key-if-not-using-rbac" + } +} +``` + +## Running the Sample + +1. **Start the Function App**: + ```bash + cd python/samples/getting_started/azure_functions/08_mcp_server + func start + ``` + +2. **Note the MCP Server Endpoint**: When the app starts, you'll see the MCP server endpoint in the terminal output. It will look like: + ``` + MCP server endpoint: http://localhost:7071/runtime/webhooks/mcp + ``` + +## Testing MCP Tool Integration + +### Using MCP Inspector + +1. Install the [MCP Inspector](https://modelcontextprotocol.io/docs/tools/inspector) +2. Connect using the MCP server endpoint from your terminal output +3. Select **"Streamable HTTP"** as the transport method +4. Test the available MCP tools: + - `StockAdvisor` - Available only as MCP tool + - `PlantAdvisor` - Available as both HTTP and MCP tool + +### Using Other MCP Clients + +Any MCP-compatible client can connect to the server endpoint and utilize the exposed agent tools. The agents will appear as callable tools within the MCP protocol. + +## Testing HTTP Endpoints + +For agents with HTTP triggers enabled (Joker and PlantAdvisor), you can test them using curl: + +```bash +# Test Joker agent (HTTP only) +curl -X POST http://localhost:7071/api/agents/Joker/run \ + -H "Content-Type: application/json" \ + -d '{"message": "Tell me a joke"}' + +# Test PlantAdvisor agent (HTTP and MCP) +curl -X POST http://localhost:7071/api/agents/PlantAdvisor/run \ + -H "Content-Type: application/json" \ + -d '{"message": "Recommend an indoor plant"}' +``` + +Note: StockAdvisor does not have HTTP endpoints and is only accessible via MCP tool triggers. + +## Expected Output + +**HTTP Responses** will be returned directly to your HTTP client. + +**MCP Tool Responses** will be visible in: +- The terminal where `func start` is running +- Your MCP client interface +- The DTS dashboard at `http://localhost:8080` (if using Durable Task Scheduler) + +## Health Check + +Check the health endpoint to see which agents have which triggers enabled: + +```bash +curl http://localhost:7071/api/health +``` + +Expected response: + +```json +{ + "status": "healthy", + "agents": [ + { + "name": "Joker", + "type": "Agent", + "http_endpoint_enabled": true, + "mcp_tool_enabled": false + }, + { + "name": "StockAdvisor", + "type": "Agent", + "http_endpoint_enabled": false, + "mcp_tool_enabled": true + }, + { + "name": "PlantAdvisor", + "type": "Agent", + "http_endpoint_enabled": true, + "mcp_tool_enabled": true + } + ], + "agent_count": 3 +} +``` + +## Code Structure + +The sample shows how to enable MCP tool triggers with flexible agent configuration: + +```python +from agent_framework.azure import AgentFunctionApp, AzureOpenAIChatClient + +# Create Azure OpenAI Chat Client +chat_client = AzureOpenAIChatClient() + +# Define agents with different roles +joker_agent = chat_client.create_agent( + name="Joker", + instructions="You are good at telling jokes.", +) + +stock_agent = chat_client.create_agent( + name="StockAdvisor", + instructions="Check stock prices.", +) + +plant_agent = chat_client.create_agent( + name="PlantAdvisor", + instructions="Recommend plants.", + description="Get plant recommendations.", +) + +# Create the AgentFunctionApp +app = AgentFunctionApp(enable_health_check=True) + +# Configure agents with different trigger combinations: +# HTTP trigger only (default) +app.add_agent(joker_agent) + +# MCP tool trigger only (HTTP disabled) +app.add_agent(stock_agent, enable_http_endpoint=False, enable_mcp_tool_trigger=True) + +# Both HTTP and MCP tool triggers enabled +app.add_agent(plant_agent, enable_http_endpoint=True, enable_mcp_tool_trigger=True) +``` + +This automatically creates the following endpoints based on agent configuration: +- `POST /api/agents/{AgentName}/run` - HTTP endpoint (when `enable_http_endpoint=True`) +- MCP tool triggers for agents with `enable_mcp_tool_trigger=True` +- `GET /api/health` - Health check endpoint showing agent configurations + +## Learn More + +- [Model Context Protocol Documentation](https://modelcontextprotocol.io/) +- [Microsoft Agent Framework Documentation](https://github.com/microsoft/agent-framework) +- [Azure Functions Documentation](https://learn.microsoft.com/azure/azure-functions/) diff --git a/python/samples/getting_started/azure_functions/08_mcp_server/function_app.py b/python/samples/getting_started/azure_functions/08_mcp_server/function_app.py new file mode 100644 index 0000000000..14bd230b2a --- /dev/null +++ b/python/samples/getting_started/azure_functions/08_mcp_server/function_app.py @@ -0,0 +1,63 @@ +""" +Example showing how to configure AI agents with different trigger configurations. + +This sample demonstrates how to configure agents to be accessible as both HTTP endpoints +and Model Context Protocol (MCP) tools, enabling flexible integration patterns for AI agent +consumption. + +Key concepts demonstrated: +- Multi-trigger Agent Configuration: Configure agents to support HTTP triggers, MCP tool triggers, or both +- Microsoft Agent Framework Integration: Use the framework to define AI agents with specific roles +- Flexible Agent Registration: Register agents with customizable trigger configurations + +This sample creates three agents with different trigger configurations: +- Joker: HTTP trigger only (default) +- StockAdvisor: MCP tool trigger only (HTTP disabled) +- PlantAdvisor: Both HTTP and MCP tool triggers enabled + +Required environment variables: +- AZURE_OPENAI_ENDPOINT: Your Azure OpenAI endpoint +- AZURE_OPENAI_CHAT_DEPLOYMENT_NAME: Your Azure OpenAI deployment name + +Authentication uses AzureCliCredential (Azure Identity). +""" + +from agent_framework.azure import AgentFunctionApp, AzureOpenAIChatClient + +# Create Azure OpenAI Chat Client +# This uses AzureCliCredential for authentication (requires 'az login') +chat_client = AzureOpenAIChatClient() + +# Define three AI agents with different roles +# Agent 1: Joker - HTTP trigger only (default) +agent1 = chat_client.create_agent( + name="Joker", + instructions="You are good at telling jokes.", +) + +# Agent 2: StockAdvisor - MCP tool trigger only +agent2 = chat_client.create_agent( + name="StockAdvisor", + instructions="Check stock prices.", +) + +# Agent 3: PlantAdvisor - Both HTTP and MCP tool triggers +agent3 = chat_client.create_agent( + name="PlantAdvisor", + instructions="Recommend plants.", + description="Get plant recommendations.", +) + +# Create the AgentFunctionApp with selective trigger configuration +app = AgentFunctionApp( + enable_health_check=True, +) + +# Agent 1: HTTP trigger only (default) +app.add_agent(agent1) + +# Agent 2: Disable HTTP trigger, enable MCP tool trigger only +app.add_agent(agent2, enable_http_endpoint=False, enable_mcp_tool_trigger=True) + +# Agent 3: Enable both HTTP and MCP tool triggers +app.add_agent(agent3, enable_http_endpoint=True, enable_mcp_tool_trigger=True) diff --git a/python/samples/getting_started/azure_functions/08_mcp_server/host.json b/python/samples/getting_started/azure_functions/08_mcp_server/host.json new file mode 100644 index 0000000000..b7e5ad1c0b --- /dev/null +++ b/python/samples/getting_started/azure_functions/08_mcp_server/host.json @@ -0,0 +1,7 @@ +{ + "version": "2.0", + "extensionBundle": { + "id": "Microsoft.Azure.Functions.ExtensionBundle", + "version": "[4.*, 5.0.0)" + } +} diff --git a/python/samples/getting_started/azure_functions/08_mcp_server/local.settings.json.template b/python/samples/getting_started/azure_functions/08_mcp_server/local.settings.json.template new file mode 100644 index 0000000000..6c98a7d1cb --- /dev/null +++ b/python/samples/getting_started/azure_functions/08_mcp_server/local.settings.json.template @@ -0,0 +1,10 @@ +{ + "IsEncrypted": false, + "Values": { + "FUNCTIONS_WORKER_RUNTIME": "python", + "AzureWebJobsStorage": "UseDevelopmentStorage=true", + "DURABLE_TASK_SCHEDULER_CONNECTION_STRING": "Endpoint=http://localhost:8080;TaskHub=default;Authentication=None", + "AZURE_OPENAI_ENDPOINT": "", + "AZURE_OPENAI_CHAT_DEPLOYMENT_NAME": "" + } +} diff --git a/python/samples/getting_started/azure_functions/08_mcp_server/requirements.txt b/python/samples/getting_started/azure_functions/08_mcp_server/requirements.txt new file mode 100644 index 0000000000..39ad8a124f --- /dev/null +++ b/python/samples/getting_started/azure_functions/08_mcp_server/requirements.txt @@ -0,0 +1,2 @@ +agent-framework-azurefunctions +azure-identity diff --git a/python/samples/getting_started/workflows/README.md b/python/samples/getting_started/workflows/README.md index c7d1622577..4dbeeb6071 100644 --- a/python/samples/getting_started/workflows/README.md +++ b/python/samples/getting_started/workflows/README.md @@ -92,7 +92,8 @@ For observability samples in Agent Framework, see the [observability getting sta | Concurrent Orchestration (Default Aggregator) | [orchestration/concurrent_agents.py](./orchestration/concurrent_agents.py) | Fan-out to multiple agents; fan-in with default aggregator returning combined ChatMessages | | Concurrent Orchestration (Custom Aggregator) | [orchestration/concurrent_custom_aggregator.py](./orchestration/concurrent_custom_aggregator.py) | Override aggregator via callback; summarize results with an LLM | | Concurrent Orchestration (Custom Agent Executors) | [orchestration/concurrent_custom_agent_executors.py](./orchestration/concurrent_custom_agent_executors.py) | Child executors own ChatAgents; concurrent fan-out/fan-in via ConcurrentBuilder | -| Group Chat Orchestration with Prompt Based Manager | [orchestration/group_chat_prompt_based_manager.py](./orchestration/group_chat_prompt_based_manager.py) | LLM Manager-directed conversation using GroupChatBuilder | +| Group Chat with Agent Manager | [orchestration/group_chat_agent_manager.py](./orchestration/group_chat_agent_manager.py) | Agent-based manager using `set_manager()` to select next speaker | +| Group Chat Philosophical Debate | [orchestration/group_chat_philosophical_debate.py](./orchestration/group_chat_philosophical_debate.py) | Agent manager moderates long-form, multi-round debate across diverse participants | | Group Chat with Simple Function Selector | [orchestration/group_chat_simple_selector.py](./orchestration/group_chat_simple_selector.py) | Group chat with a simple function selector for next speaker | | Handoff (Simple) | [orchestration/handoff_simple.py](./orchestration/handoff_simple.py) | Single-tier routing: triage agent routes to specialists, control returns to user after each specialist response | | Handoff (Specialist-to-Specialist) | [orchestration/handoff_specialist_to_specialist.py](./orchestration/handoff_specialist_to_specialist.py) | Multi-tier routing: specialists can hand off to other specialists using `.add_handoff()` fluent API | diff --git a/python/samples/getting_started/workflows/orchestration/group_chat_agent_manager.py b/python/samples/getting_started/workflows/orchestration/group_chat_agent_manager.py new file mode 100644 index 0000000000..3bc79fcddc --- /dev/null +++ b/python/samples/getting_started/workflows/orchestration/group_chat_agent_manager.py @@ -0,0 +1,112 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import logging +from typing import cast + +from agent_framework import ( + AgentRunUpdateEvent, + ChatAgent, + ChatMessage, + GroupChatBuilder, + Role, + WorkflowOutputEvent, +) +from agent_framework.azure import AzureOpenAIChatClient +from azure.identity import AzureCliCredential + +logging.basicConfig(level=logging.INFO) + +""" +Sample: Group Chat with Agent-Based Manager + +What it does: +- Demonstrates the new set_manager() API for agent-based coordination +- Manager is a full ChatAgent with access to tools, context, and observability +- Coordinates a researcher and writer agent to solve tasks collaboratively + +Prerequisites: +- OpenAI environment variables configured for OpenAIChatClient +""" + + +def _get_chat_client() -> AzureOpenAIChatClient: + return AzureOpenAIChatClient(credential=AzureCliCredential()) + + +async def main() -> None: + # Create coordinator agent with structured output for speaker selection + # Note: response_format is enforced to ManagerSelectionResponse by set_manager() + coordinator = ChatAgent( + name="Coordinator", + description="Coordinates multi-agent collaboration by selecting speakers", + instructions=""" +You coordinate a team conversation to solve the user's task. + +Review the conversation history and select the next participant to speak. + +Guidelines: +- Start with Researcher to gather information +- Then have Writer synthesize the final answer +- Only finish after both have contributed meaningfully +- Allow for multiple rounds of information gathering if needed +""", + chat_client=_get_chat_client(), + ) + + researcher = ChatAgent( + name="Researcher", + description="Collects relevant background information", + instructions="Gather concise facts that help a teammate answer the question.", + chat_client=_get_chat_client(), + ) + + writer = ChatAgent( + name="Writer", + description="Synthesizes polished answers from gathered information", + instructions="Compose clear and structured answers using any notes provided.", + chat_client=_get_chat_client(), + ) + + workflow = ( + GroupChatBuilder() + .set_manager(coordinator, display_name="Orchestrator") + .with_termination_condition(lambda messages: sum(1 for msg in messages if msg.role == Role.ASSISTANT) >= 2) + .participants([researcher, writer]) + .build() + ) + + task = "What are the key benefits of using async/await in Python? Provide a concise summary." + + print("\nStarting Group Chat with Agent-Based Manager...\n") + print(f"TASK: {task}\n") + print("=" * 80) + + final_conversation: list[ChatMessage] = [] + last_executor_id: str | None = None + async for event in workflow.run_stream(task): + if isinstance(event, AgentRunUpdateEvent): + eid = event.executor_id + if eid != last_executor_id: + if last_executor_id is not None: + print() + print(f"{eid}:", end=" ", flush=True) + last_executor_id = eid + print(event.data, end="", flush=True) + elif isinstance(event, WorkflowOutputEvent): + final_conversation = cast(list[ChatMessage], event.data) + + if final_conversation and isinstance(final_conversation, list): + print("\n\n" + "=" * 80) + print("FINAL CONVERSATION") + print("=" * 80) + for msg in final_conversation: + author = getattr(msg, "author_name", "Unknown") + text = getattr(msg, "text", str(msg)) + print(f"\n[{author}]") + print(text) + print("-" * 80) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/python/samples/getting_started/workflows/orchestration/group_chat_philosophical_debate.py b/python/samples/getting_started/workflows/orchestration/group_chat_philosophical_debate.py new file mode 100644 index 0000000000..7059a84e32 --- /dev/null +++ b/python/samples/getting_started/workflows/orchestration/group_chat_philosophical_debate.py @@ -0,0 +1,408 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import logging +from typing import cast + +from agent_framework import ( + AgentRunUpdateEvent, + ChatAgent, + ChatMessage, + GroupChatBuilder, + Role, + WorkflowOutputEvent, +) +from agent_framework.azure import AzureOpenAIChatClient +from azure.identity import AzureCliCredential + +logging.basicConfig(level=logging.WARNING) + +""" +Sample: Philosophical Debate with Agent-Based Manager + +What it does: +- Creates a diverse group of agents representing different global perspectives +- Uses an agent-based manager to guide a philosophical discussion +- Demonstrates longer, multi-round discourse with natural conversation flow +- Manager decides when discussion has reached meaningful conclusion + +Topic: "What does a good life mean to you personally?" + +Participants represent: +- Farmer from Southeast Asia (tradition, sustainability, land connection) +- Software Developer from United States (innovation, technology, work-life balance) +- History Teacher from Eastern Europe (legacy, learning, cultural continuity) +- Activist from South America (social justice, environmental rights) +- Spiritual Leader from Middle East (morality, community service) +- Artist from Africa (creative expression, storytelling) +- Immigrant Entrepreneur from Asia in Canada (tradition + adaptation) +- Doctor from Scandinavia (public health, equity, societal support) + +Prerequisites: +- OpenAI environment variables configured for OpenAIChatClient +""" + + +def _get_chat_client() -> AzureOpenAIChatClient: + return AzureOpenAIChatClient(credential=AzureCliCredential()) + + +async def main() -> None: + # Create debate moderator with structured output for speaker selection + # Note: Participant names and descriptions are automatically injected by the orchestrator + moderator = ChatAgent( + name="Moderator", + description="Guides philosophical discussion by selecting next speaker", + instructions=""" +You are a thoughtful moderator guiding a philosophical discussion on the topic handed to you by the user. + +Your participants bring diverse global perspectives. Select speakers strategically to: +- Create natural conversation flow and responses to previous points +- Ensure all voices are heard throughout the discussion +- Build on themes and contrasts that emerge +- Allow for respectful challenges and counterpoints +- Guide toward meaningful conclusions + +Select speakers who can: +1. Respond directly to points just made +2. Introduce fresh perspectives when needed +3. Bridge or contrast different viewpoints +4. Deepen the philosophical exploration + +Finish when: +- Multiple rounds have occurred (at least 6-8 exchanges) +- Key themes have been explored from different angles +- Natural conclusion or synthesis has emerged +- Diminishing returns in new insights + +In your final_message, provide a brief synthesis highlighting key themes that emerged. +""", + chat_client=_get_chat_client(), + ) + + farmer = ChatAgent( + name="Farmer", + description="A rural farmer from Southeast Asia", + instructions=""" +You're a farmer from Southeast Asia. Your life is deeply connected to land and family. +You value tradition and sustainability. You are in a philosophical debate. + +Share your perspective authentically. Feel free to: +- Challenge other participants respectfully +- Build on points others have made +- Use concrete examples from your experience +- Keep responses thoughtful but concise (2-4 sentences) +""", + chat_client=_get_chat_client(), + ) + + developer = ChatAgent( + name="Developer", + description="An urban software developer from the United States", + instructions=""" +You're a software developer from the United States. Your life is fast-paced and technology-driven. +You value innovation, freedom, and work-life balance. You are in a philosophical debate. + +Share your perspective authentically. Feel free to: +- Challenge other participants respectfully +- Build on points others have made +- Use concrete examples from your experience +- Keep responses thoughtful but concise (2-4 sentences) +""", + chat_client=_get_chat_client(), + ) + + teacher = ChatAgent( + name="Teacher", + description="A retired history teacher from Eastern Europe", + instructions=""" +You're a retired history teacher from Eastern Europe. You bring historical and philosophical +perspectives to discussions. You value legacy, learning, and cultural continuity. +You are in a philosophical debate. + +Share your perspective authentically. Feel free to: +- Challenge other participants respectfully +- Build on points others have made +- Use concrete examples from history or your teaching experience +- Keep responses thoughtful but concise (2-4 sentences) +""", + chat_client=_get_chat_client(), + ) + + activist = ChatAgent( + name="Activist", + description="A young activist from South America", + instructions=""" +You're a young activist from South America. You focus on social justice, environmental rights, +and generational change. You are in a philosophical debate. + +Share your perspective authentically. Feel free to: +- Challenge other participants respectfully +- Build on points others have made +- Use concrete examples from your activism +- Keep responses thoughtful but concise (2-4 sentences) +""", + chat_client=_get_chat_client(), + ) + + spiritual_leader = ChatAgent( + name="SpiritualLeader", + description="A spiritual leader from the Middle East", + instructions=""" +You're a spiritual leader from the Middle East. You provide insights grounded in religion, +morality, and community service. You are in a philosophical debate. + +Share your perspective authentically. Feel free to: +- Challenge other participants respectfully +- Build on points others have made +- Use examples from spiritual teachings or community work +- Keep responses thoughtful but concise (2-4 sentences) +""", + chat_client=_get_chat_client(), + ) + + artist = ChatAgent( + name="Artist", + description="An artist from Africa", + instructions=""" +You're an artist from Africa. You view life through creative expression, storytelling, +and collective memory. You are in a philosophical debate. + +Share your perspective authentically. Feel free to: +- Challenge other participants respectfully +- Build on points others have made +- Use examples from your art or cultural traditions +- Keep responses thoughtful but concise (2-4 sentences) +""", + chat_client=_get_chat_client(), + ) + + immigrant = ChatAgent( + name="Immigrant", + description="An immigrant entrepreneur from Asia living in Canada", + instructions=""" +You're an immigrant entrepreneur from Asia living in Canada. You balance tradition with adaptation. +You focus on family success, risk, and opportunity. You are in a philosophical debate. + +Share your perspective authentically. Feel free to: +- Challenge other participants respectfully +- Build on points others have made +- Use examples from your immigrant and entrepreneurial journey +- Keep responses thoughtful but concise (2-4 sentences) +""", + chat_client=_get_chat_client(), + ) + + doctor = ChatAgent( + name="Doctor", + description="A doctor from Scandinavia", + instructions=""" +You're a doctor from Scandinavia. Your perspective is shaped by public health, equity, +and structured societal support. You are in a philosophical debate. + +Share your perspective authentically. Feel free to: +- Challenge other participants respectfully +- Build on points others have made +- Use examples from healthcare and societal systems +- Keep responses thoughtful but concise (2-4 sentences) +""", + chat_client=_get_chat_client(), + ) + + workflow = ( + GroupChatBuilder() + .set_manager(moderator, display_name="Moderator") + .participants([farmer, developer, teacher, activist, spiritual_leader, artist, immigrant, doctor]) + .with_termination_condition(lambda messages: sum(1 for msg in messages if msg.role == Role.ASSISTANT) >= 10) + .build() + ) + + topic = "What does a good life mean to you personally?" + + print("\n" + "=" * 80) + print("PHILOSOPHICAL DEBATE: Perspectives on a Good Life") + print("=" * 80) + print(f"\nTopic: {topic}") + print("\nParticipants:") + print(" - Farmer (Southeast Asia)") + print(" - Developer (United States)") + print(" - Teacher (Eastern Europe)") + print(" - Activist (South America)") + print(" - SpiritualLeader (Middle East)") + print(" - Artist (Africa)") + print(" - Immigrant (Asia → Canada)") + print(" - Doctor (Scandinavia)") + print("\n" + "=" * 80) + print("DISCUSSION BEGINS") + print("=" * 80 + "\n") + + final_conversation: list[ChatMessage] = [] + current_speaker: str | None = None + + async for event in workflow.run_stream(f"Please begin the discussion on: {topic}"): + if isinstance(event, AgentRunUpdateEvent): + speaker_id = event.executor_id.replace("groupchat_agent:", "") + + if speaker_id != current_speaker: + if current_speaker is not None: + print("\n") + print(f"[{speaker_id}]", flush=True) + current_speaker = speaker_id + + print(event.data, end="", flush=True) + + elif isinstance(event, WorkflowOutputEvent): + final_conversation = cast(list[ChatMessage], event.data) + + print("\n\n" + "=" * 80) + print("DISCUSSION SUMMARY") + print("=" * 80) + + if final_conversation and isinstance(final_conversation, list) and final_conversation: + final_msg = final_conversation[-1] + if hasattr(final_msg, "author_name") and final_msg.author_name == "Moderator": + print(f"\n{final_msg.text}") + + """ + Sample Output: + + ================================================================================ + PHILOSOPHICAL DEBATE: Perspectives on a Good Life + ================================================================================ + + Topic: What does a good life mean to you personally? + + Participants: + - Farmer (Southeast Asia) + - Developer (United States) + - Teacher (Eastern Europe) + - Activist (South America) + - SpiritualLeader (Middle East) + - Artist (Africa) + - Immigrant (Asia → Canada) + - Doctor (Scandinavia) + + ================================================================================ + DISCUSSION BEGINS + ================================================================================ + + [Moderator] + {"selected_participant":"Farmer","instruction":"Please start by sharing what living a good life means to you, + especially from your perspective living in a rural area in Southeast Asia.","finish":false,"final_message":null} + + [Farmer] + To me, a good life is deeply intertwined with the rhythm of the land and the nurturing of relationships with my + family and community. It means cultivating crops that respect our environment, ensuring sustainability for future + generations, and sharing meals made from our harvests around the dinner table. The joy found in everyday + tasks—planting rice or tending to our livestock—creates a sense of fulfillment that cannot be measured by material + wealth. It's the simple moments, like sharing stories with my children under the stars, that truly define a good + life. What good is progress if it isolates us from those we love and the land that sustains us? + + [Moderator] + {"selected_participant":"Developer","instruction":"Given the insights shared by the Farmer, please discuss what a + good life means to you as a software developer in an urban setting in the United States and how it might contrast + with or complement the Farmer's view.","finish":false,"final_message":null} + + [Developer] + As a software developer in an urban environment, a good life for me hinges on the intersection of innovation, + creativity, and balance. It's about having the freedom to explore new technologies that can solve real-world + problems while ensuring that my work doesn't encroach on my personal life. For instance, I value remote work + flexibility, which allows me to maintain connections with family and friends, similar to how the Farmer values + community. While our lifestyles may differ markedly, both of us seek fulfillment—whether through meaningful work or + rich personal experiences. The challenge is finding harmony between technological progress and preserving the + intimate human connections that truly enrich our lives. + + [Moderator] + {"selected_participant":"SpiritualLeader","instruction":"Reflect on both the Farmer's and Developer's perspectives + and share your view of what constitutes a good life, particularly from your spiritual and cultural standpoint in + the Middle East.","finish":false,"final_message":null} + + [SpiritualLeader] + From my spiritual perspective, a good life embodies a balance between personal fulfillment and service to others, + rooted in compassion and community. In our teachings, we emphasize that true happiness comes from helping those in + need and fostering strong connections with our families and neighbors. Whether it's the Farmer nurturing the earth + or the Developer creating tools to enhance lives, both contribute to the greater good. The essence of a good life + lies in our intentions and actions—finding ways to serve our communities, spread kindness, and live harmoniously + with those around us. Ultimately, as we align our personal beliefs with our communal responsibilities, we cultivate + a richness that transcends material wealth. + + [Moderator] + {"selected_participant":"Activist","instruction":"Add to the discussion by sharing your perspective on what a good + life entails, particularly from your background as a young activist in South America.","finish":false, + "final_message":null} + + [Activist] + As a young activist in South America, a good life for me is about advocating for social justice and environmental + sustainability. It means living in a society where everyone's rights are respected and where marginalized voices, + particularly those of Indigenous communities, are amplified. I see a good life as one where we work collectively to + dismantle oppressive systems—such as deforestation and inequality—while nurturing our planet. For instance, through + my activism, I've witnessed the transformative power of community organizing, where collective efforts lead to real + change, like resisting destructive mining practices that threaten our rivers and lands. A good life, therefore, is + not just lived for oneself but is deeply tied to the well-being of our communities and the health of our + environment. How can we, regardless of our backgrounds, collaborate to foster these essential changes? + + [Moderator] + {"selected_participant":"Teacher","instruction":"Considering the views shared so far, tell us how your experience + as a retired history teacher from Eastern Europe shapes your understanding of a good life, perhaps reflecting on + lessons from the past and their impact on present-day life choices.","finish":false,"final_message":null} + + [Teacher] + As a retired history teacher from Eastern Europe, my understanding of a good life is deeply rooted in the lessons + drawn from history and the struggle for freedom and dignity. Historical events, such as the fall of the Iron + Curtain, remind us of the profound importance of liberty and collective resilience. A good life, therefore, is about + cherishing our freedoms and working towards a society where everyone has a voice, much as my students and I + discussed the impacts of totalitarian regimes. Additionally, I believe it involves fostering cultural continuity, + where we honor our heritage while embracing progressive values. We must learn from the past—especially the + consequences of neglecting empathy and solidarity—so that we can cultivate a future that values every individual's + contributions to the rich tapestry of our shared humanity. How can we ensure that the lessons of history inform a + more compassionate and just society moving forward? + + [Moderator] + {"selected_participant":"Artist","instruction":"Expound on the themes and perspectives discussed so far by sharing + how, as an artist from Africa, you define a good life and how art plays a role in that vision.","finish":false, + "final_message":null} + + [Artist] + As an artist from Africa, I define a good life as one steeped in cultural expression, storytelling, and the + celebration of our collective memories. Art is a powerful medium through which we capture our histories, struggles, + and triumphs, creating a tapestry that connects generations. For instance, in my work, I often draw from folktales + and traditional music, weaving narratives that reflect the human experience, much like how the retired teacher + emphasizes learning from history. A good life involves not only personal fulfillment but also the responsibility to + share our narratives and use our creativity to inspire change, whether addressing social injustices or environmental + issues. It's in this interplay of art and activism that we can transcend individual existence and contribute to a + collective good, fostering empathy and understanding among diverse communities. How can we harness art to bridge + differences and amplify marginalized voices in our pursuit of a good life? + + [Moderator] + {"selected_participant":null,"instruction":null,"finish":true,"final_message":"As our discussion unfolds, several + key themes have gracefully emerged, reflecting the richness of diverse perspectives on what constitutes a good life. + From the rural farmer's integration with the land to the developer's search for balance between technology and + personal connection, each viewpoint validates that fulfillment, at its core, transcends material wealth. The + spiritual leader and the activist highlight the importance of community and social justice, while the history + teacher and the artist remind us of the lessons and narratives that shape our cultural and personal identities. + + Ultimately, the good life seems to revolve around meaningful relationships, honoring our legacies while striving for + progress, and nurturing both our inner selves and external communities. This dialogue demonstrates that despite our + varied backgrounds and experiences, the quest for a good life binds us together, urging cooperation and empathy in + our shared human journey."} + + ================================================================================ + DISCUSSION SUMMARY + ================================================================================ + + As our discussion unfolds, several key themes have gracefully emerged, reflecting the richness of diverse + perspectives on what constitutes a good life. From the rural farmer's integration with the land to the developer's + search for balance between technology and personal connection, each viewpoint validates that fulfillment, at its + core, transcends material wealth. The spiritual leader and the activist highlight the importance of community and + social justice, while the history teacher and the artist remind us of the lessons and narratives that shape our + cultural and personal identities. + + Ultimately, the good life seems to revolve around meaningful relationships, honoring our legacies while striving for + progress, and nurturing both our inner selves and external communities. This dialogue demonstrates that despite our + varied backgrounds and experiences, the quest for a good life binds us together, urging cooperation and empathy in + our shared human journey. + """ + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/python/samples/getting_started/workflows/orchestration/group_chat_prompt_based_manager.py b/python/samples/getting_started/workflows/orchestration/group_chat_prompt_based_manager.py deleted file mode 100644 index 6a6d3a5e22..0000000000 --- a/python/samples/getting_started/workflows/orchestration/group_chat_prompt_based_manager.py +++ /dev/null @@ -1,75 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -import asyncio -import logging - -from agent_framework import AgentRunUpdateEvent, ChatAgent, GroupChatBuilder, WorkflowOutputEvent -from agent_framework.openai import OpenAIChatClient, OpenAIResponsesClient - -logging.basicConfig(level=logging.INFO) - -""" -Sample: Group Chat Orchestration (manager-directed) - -What it does: -- Demonstrates the generic GroupChatBuilder with a language-model manager directing two agents. -- The manager coordinates a researcher (chat completions) and a writer (responses API) to solve a task. -- Uses the default group chat orchestration pipeline shared with Magentic. - -Prerequisites: -- OpenAI environment variables configured for `OpenAIChatClient` and `OpenAIResponsesClient`. -""" - - -async def main() -> None: - researcher = ChatAgent( - name="Researcher", - description="Collects relevant background information.", - instructions="Gather concise facts that help a teammate answer the question.", - chat_client=OpenAIChatClient(model_id="gpt-4o-mini"), - ) - - writer = ChatAgent( - name="Writer", - description="Synthesizes a polished answer using the gathered notes.", - instructions="Compose clear and structured answers using any notes provided.", - chat_client=OpenAIResponsesClient(), - ) - - workflow = ( - GroupChatBuilder() - .set_prompt_based_manager(chat_client=OpenAIChatClient(), display_name="Coordinator") - .participants(researcher=researcher, writer=writer) - .build() - ) - - task = "Outline the core considerations for planning a community hackathon, and finish with a concise action plan." - - print("\nStarting Group Chat Workflow...\n") - print(f"TASK: {task}\n") - - final_response = None - last_executor_id: str | None = None - async for event in workflow.run_stream(task): - if isinstance(event, AgentRunUpdateEvent): - # Handle the streaming agent update as it's produced - eid = event.executor_id - if eid != last_executor_id: - if last_executor_id is not None: - print() - print(f"{eid}:", end=" ", flush=True) - last_executor_id = eid - print(event.data, end="", flush=True) - elif isinstance(event, WorkflowOutputEvent): - final_response = getattr(event.data, "text", str(event.data)) - - if final_response: - print("=" * 60) - print("FINAL RESPONSE") - print("=" * 60) - print(final_response) - print("=" * 60) - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/python/samples/getting_started/workflows/orchestration/group_chat_simple_selector.py b/python/samples/getting_started/workflows/orchestration/group_chat_simple_selector.py index ba4d16accb..1fd074ca4d 100644 --- a/python/samples/getting_started/workflows/orchestration/group_chat_simple_selector.py +++ b/python/samples/getting_started/workflows/orchestration/group_chat_simple_selector.py @@ -2,8 +2,9 @@ import asyncio import logging +from typing import cast -from agent_framework import ChatAgent, GroupChatBuilder, GroupChatStateSnapshot, WorkflowOutputEvent +from agent_framework import ChatAgent, ChatMessage, GroupChatBuilder, GroupChatStateSnapshot, WorkflowOutputEvent from agent_framework.openai import OpenAIChatClient logging.basicConfig(level=logging.INFO) @@ -12,7 +13,7 @@ logging.basicConfig(level=logging.INFO) Sample: Group Chat with Simple Speaker Selector Function What it does: -- Demonstrates the select_speakers() API for GroupChat orchestration +- Demonstrates the set_select_speakers_func() API for GroupChat orchestration - Uses a pure Python function to control speaker selection based on conversation state - Alternates between researcher and writer agents in a simple round-robin pattern - Shows how to access conversation history, round index, and participant metadata @@ -84,7 +85,7 @@ async def main() -> None: # 2. Dict form - explicit names: .participants(researcher=researcher, writer=writer) workflow = ( GroupChatBuilder() - .select_speakers(select_next_speaker, display_name="Orchestrator") + .set_select_speakers_func(select_next_speaker, display_name="Orchestrator") .participants([researcher, writer]) # Uses agent.name for participant names .build() ) @@ -97,11 +98,14 @@ async def main() -> None: async for event in workflow.run_stream(task): if isinstance(event, WorkflowOutputEvent): - final_message = event.data - author = getattr(final_message, "author_name", "Unknown") - text = getattr(final_message, "text", str(final_message)) - print(f"\n[{author}]\n{text}\n") - print("-" * 80) + conversation = cast(list[ChatMessage], event.data) + if isinstance(conversation, list): + print("\n===== Final Conversation =====\n") + for msg in conversation: + author = getattr(msg, "author_name", "Unknown") + text = getattr(msg, "text", str(msg)) + print(f"[{author}]\n{text}\n") + print("-" * 80) print("\nWorkflow completed.") diff --git a/python/samples/getting_started/workflows/orchestration/magentic.py b/python/samples/getting_started/workflows/orchestration/magentic.py index 5010172e2b..0e265cb931 100644 --- a/python/samples/getting_started/workflows/orchestration/magentic.py +++ b/python/samples/getting_started/workflows/orchestration/magentic.py @@ -2,20 +2,21 @@ import asyncio import logging +from typing import cast from agent_framework import ( + MAGENTIC_EVENT_TYPE_AGENT_DELTA, + MAGENTIC_EVENT_TYPE_ORCHESTRATOR, + AgentRunUpdateEvent, ChatAgent, + ChatMessage, HostedCodeInterpreterTool, - MagenticAgentDeltaEvent, - MagenticAgentMessageEvent, MagenticBuilder, - MagenticFinalResultEvent, - MagenticOrchestratorMessageEvent, WorkflowOutputEvent, ) from agent_framework.openai import OpenAIChatClient, OpenAIResponsesClient -logging.basicConfig(level=logging.DEBUG) +logging.basicConfig(level=logging.WARNING) logger = logging.getLogger(__name__) """ @@ -97,35 +98,30 @@ async def main() -> None: try: output: str | None = None async for event in workflow.run_stream(task): - if isinstance(event, MagenticOrchestratorMessageEvent): - print(f"\n[ORCH:{event.kind}]\n\n{getattr(event.message, 'text', '')}\n{'-' * 26}") - elif isinstance(event, MagenticAgentDeltaEvent): - if last_stream_agent_id != event.agent_id or not stream_line_open: - if stream_line_open: - print() - print(f"\n[STREAM:{event.agent_id}]: ", end="", flush=True) - last_stream_agent_id = event.agent_id - stream_line_open = True - if event.text: - print(event.text, end="", flush=True) - elif isinstance(event, MagenticAgentMessageEvent): - if stream_line_open: - print(" (final)") - stream_line_open = False - print() - msg = event.message - if msg is not None: - response_text = (msg.text or "").replace("\n", " ") - print(f"\n[AGENT:{event.agent_id}] {msg.role.value}\n\n{response_text}\n{'-' * 26}") - elif isinstance(event, MagenticFinalResultEvent): - print("\n" + "=" * 50) - print("FINAL RESULT:") - print("=" * 50) - if event.message is not None: - print(event.message.text) - print("=" * 50) + if isinstance(event, AgentRunUpdateEvent): + props = event.data.additional_properties if event.data else None + event_type = props.get("magentic_event_type") if props else None + + if event_type == MAGENTIC_EVENT_TYPE_ORCHESTRATOR: + kind = props.get("orchestrator_message_kind", "") if props else "" + text = event.data.text if event.data else "" + print(f"\n[ORCH:{kind}]\n\n{text}\n{'-' * 26}") + elif event_type == MAGENTIC_EVENT_TYPE_AGENT_DELTA: + agent_id = props.get("agent_id", event.executor_id) if props else event.executor_id + if last_stream_agent_id != agent_id or not stream_line_open: + if stream_line_open: + print() + print(f"\n[STREAM:{agent_id}]: ", end="", flush=True) + last_stream_agent_id = agent_id + stream_line_open = True + if event.data and event.data.text: + print(event.data.text, end="", flush=True) + elif event.data and event.data.text: + print(event.data.text, end="", flush=True) elif isinstance(event, WorkflowOutputEvent): - output = str(event.data) if event.data is not None else None + output_messages = cast(list[ChatMessage], event.data) + if output_messages: + output = output_messages[-1].text if stream_line_open: print() diff --git a/python/samples/semantic-kernel-migration/orchestrations/group_chat.py b/python/samples/semantic-kernel-migration/orchestrations/group_chat.py index 42142b5363..8ea37e8f5e 100644 --- a/python/samples/semantic-kernel-migration/orchestrations/group_chat.py +++ b/python/samples/semantic-kernel-migration/orchestrations/group_chat.py @@ -233,8 +233,8 @@ async def run_agent_framework_example(task: str) -> str: workflow = ( GroupChatBuilder() - .set_prompt_based_manager( - chat_client=AzureOpenAIChatClient(credential=credential), + .set_manager( + manager=AzureOpenAIChatClient(credential=credential).create_agent(), display_name="Coordinator", ) .participants(researcher=researcher, planner=planner) @@ -245,7 +245,12 @@ async def run_agent_framework_example(task: str) -> str: async for event in workflow.run_stream(task): if isinstance(event, WorkflowOutputEvent): data = event.data - final_response = data.text or "" if isinstance(data, ChatMessage) else str(data) + if isinstance(data, list) and len(data) > 0: + # Get the final message from the conversation + final_message = data[-1] + final_response = final_message.text or "" if isinstance(final_message, ChatMessage) else str(data) + else: + final_response = str(data) return final_response diff --git a/python/uv.lock b/python/uv.lock index 767d0b70a2..87081e2bc5 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -280,7 +280,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "agent-framework-core", editable = "packages/core" }, - { name = "openai-chatkit", specifier = ">=1.1.0,<2.0.0" }, + { name = "openai-chatkit", specifier = ">=1.4.0,<2.0.0" }, ] [[package]] @@ -798,7 +798,7 @@ wheels = [ [[package]] name = "anthropic" -version = "0.74.1" +version = "0.75.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -810,9 +810,9 @@ dependencies = [ { name = "sniffio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/d7/7b/609eea5c54ae69b1a4a94169d4b0c86dc5c41b43509989913f6cdc61b81d/anthropic-0.74.1.tar.gz", hash = "sha256:04c087b2751385c524f6d332d066a913870e4de8b3e335fb0a0c595f1f88dc6e", size = 428981, upload-time = "2025-11-19T22:17:31.533Z" } +sdist = { url = "https://files.pythonhosted.org/packages/04/1f/08e95f4b7e2d35205ae5dcbb4ae97e7d477fc521c275c02609e2931ece2d/anthropic-0.75.0.tar.gz", hash = "sha256:e8607422f4ab616db2ea5baacc215dd5f028da99ce2f022e33c7c535b29f3dfb", size = 439565, upload-time = "2025-11-24T20:41:45.28Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/dd/45/6b18d0692302b8cbc01a10c35b43953d3c4172fbd4f83337b8ed21a8eaa4/anthropic-0.74.1-py3-none-any.whl", hash = "sha256:b07b998d1cee7f41d9f02530597d7411672b362cc2417760a40c0167b81c6e65", size = 371473, upload-time = "2025-11-19T22:17:29.998Z" }, + { url = "https://files.pythonhosted.org/packages/60/1c/1cd02b7ae64302a6e06724bf80a96401d5313708651d277b1458504a1730/anthropic-0.75.0-py3-none-any.whl", hash = "sha256:ea8317271b6c15d80225a9f3c670152746e88805a7a61e14d4a374577164965b", size = 388164, upload-time = "2025-11-24T20:41:43.587Z" }, ] [[package]] @@ -1799,7 +1799,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "(python_full_version < '3.13' and sys_platform == 'darwin') or (python_full_version < '3.13' and sys_platform == 'linux') or (python_full_version < '3.13' and sys_platform == 'win32')" }, + { name = "typing-extensions", marker = "(python_full_version < '3.11' and sys_platform == 'darwin') or (python_full_version < '3.11' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [ @@ -1817,7 +1817,7 @@ wheels = [ [[package]] name = "fastapi" -version = "0.121.3" +version = "0.122.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "annotated-doc", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -1825,9 +1825,9 @@ dependencies = [ { name = "starlette", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/80/f0/086c442c6516195786131b8ca70488c6ef11d2f2e33c9a893576b2b0d3f7/fastapi-0.121.3.tar.gz", hash = "sha256:0055bc24fe53e56a40e9e0ad1ae2baa81622c406e548e501e717634e2dfbc40b", size = 344501, upload-time = "2025-11-19T16:53:39.243Z" } +sdist = { url = "https://files.pythonhosted.org/packages/b2/de/3ee97a4f6ffef1fb70bf20561e4f88531633bb5045dc6cebc0f8471f764d/fastapi-0.122.0.tar.gz", hash = "sha256:cd9b5352031f93773228af8b4c443eedc2ac2aa74b27780387b853c3726fb94b", size = 346436, upload-time = "2025-11-24T19:17:47.95Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/98/b6/4f620d7720fc0a754c8c1b7501d73777f6ba43b57c8ab99671f4d7441eb8/fastapi-0.121.3-py3-none-any.whl", hash = "sha256:0c78fc87587fcd910ca1bbf5bc8ba37b80e119b388a7206b39f0ecc95ebf53e9", size = 109801, upload-time = "2025-11-19T16:53:37.918Z" }, + { url = "https://files.pythonhosted.org/packages/7a/93/aa8072af4ff37b795f6bbf43dcaf61115f40f49935c7dbb180c9afc3f421/fastapi-0.122.0-py3-none-any.whl", hash = "sha256:a456e8915dfc6c8914a50d9651133bd47ec96d331c5b44600baa635538a30d67", size = 110671, upload-time = "2025-11-24T19:17:45.96Z" }, ] [[package]] @@ -3730,17 +3730,18 @@ wheels = [ [[package]] name = "openai-chatkit" -version = "1.3.1" +version = "1.4.0" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "openai", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "openai-agents", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "pydantic", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "uvicorn", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/f7/81/906844005976afb3c415e27f6506ed0f2b8b1040f27b4bc9ef118a256986/openai_chatkit-1.3.1.tar.gz", hash = "sha256:91f39b04584f969642a6c3b4099fdad74c2a357e25d8f746f9709046304a06cf", size = 50730, upload-time = "2025-11-21T21:22:11.62Z" } +sdist = { url = "https://files.pythonhosted.org/packages/c5/89/bf2f094997c8e5cad5334e8a02e05fc458823e65fb2675f45b56b6d1ab73/openai_chatkit-1.4.0.tar.gz", hash = "sha256:e2527dffc3794a05596ad75efa66bdc4efb4ded5a77a013a55496cc989bcf2e6", size = 55269, upload-time = "2025-11-25T21:02:58.503Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/fb/61/235e5f49bd068bbe5f7278bc7d7c4bd92226858698fcac53ab92090baf04/openai_chatkit-1.3.1-py3-none-any.whl", hash = "sha256:5626492e5752879e66b2b6d4fbac51994407d46429de99b91515a77c2e0c6148", size = 35899, upload-time = "2025-11-21T21:22:10.37Z" }, + { url = "https://files.pythonhosted.org/packages/90/bf/68d42561dd8a674b6f8541d879dd165b5ac4d81fcf1027462e154de66a4f/openai_chatkit-1.4.0-py3-none-any.whl", hash = "sha256:35d00ca8398908bd70d63e2284adcd836641cc11746f68d7cfa91d276e3dad3d", size = 39077, upload-time = "2025-11-25T21:02:57.288Z" }, ] [[package]] @@ -5075,7 +5076,7 @@ wheels = [ [[package]] name = "qdrant-client" -version = "1.16.0" +version = "1.16.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "grpcio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -5087,9 +5088,9 @@ dependencies = [ { name = "pydantic", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "urllib3", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/fa/16/366541897d270ee3f9c3f87da145baa8a5c9cc5190e0e53e8bbec1267cff/qdrant_client-1.16.0.tar.gz", hash = "sha256:0716aa0b7cca39745829c2e8ea0beb275fe2990e743ad803eabd6218e4b35c1b", size = 284128, upload-time = "2025-11-17T13:19:52.726Z" } +sdist = { url = "https://files.pythonhosted.org/packages/d9/68/fec3816a223c0b73b0e0036460be45c61ce2770ffb9197ac371e4f615ddc/qdrant_client-1.16.1.tar.gz", hash = "sha256:676c7c10fd4d4cb2981b8fcb32fd764f5f661b04b7334d024034d07212f971fd", size = 332130, upload-time = "2025-11-25T04:31:54.212Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a5/ff/3a69bb56835c4b2e9fa780655790937011ac389b0408b9a1147eaa2cee22/qdrant_client-1.16.0-py3-none-any.whl", hash = "sha256:6b932393e84e4c0233e5b2eb96b0918e968725855adae4d9c541761f4c50cf11", size = 328579, upload-time = "2025-11-17T13:19:51.092Z" }, + { url = "https://files.pythonhosted.org/packages/60/e2/60a20d04b0595c641516463168909c5bbcc192d3d6eacb637c1677109c6a/qdrant_client-1.16.1-py3-none-any.whl", hash = "sha256:1eefe89f66e8a468ba0de1680e28b441e69825cfb62e8fb2e457c15e24ce5e3b", size = 378481, upload-time = "2025-11-25T04:31:52.629Z" }, ] [[package]]