diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 1129930a64f4..ee8bffb7b5a5 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -33,7 +33,7 @@ jobs: [ # For main use the workflow target { ref: "${{github.ref}}", dest-dir: dev, uv-version: "0.5.13", sphinx-release-override: "dev" }, - { ref: "python-v0.4.5", dest-dir: stable, uv-version: "0.5.13", sphinx-release-override: "stable" }, + { ref: "python-v0.4.6", dest-dir: stable, uv-version: "0.5.13", sphinx-release-override: "stable" }, { ref: "v0.4.0.dev0", dest-dir: "0.4.0.dev0", uv-version: "0.5.11", sphinx-release-override: "" }, { ref: "v0.4.0.dev1", dest-dir: "0.4.0.dev1", uv-version: "0.5.11", sphinx-release-override: "" }, { ref: "v0.4.0.dev2", dest-dir: "0.4.0.dev2", uv-version: "0.5.11", sphinx-release-override: "" }, @@ -54,6 +54,7 @@ jobs: { ref: "v0.4.3", dest-dir: "0.4.3", uv-version: "0.5.13", sphinx-release-override: "" }, { ref: "v0.4.4", dest-dir: "0.4.4", uv-version: "0.5.13", sphinx-release-override: "" }, { ref: "python-v0.4.5", dest-dir: "0.4.5", uv-version: "0.5.13", sphinx-release-override: "" }, + { ref: "python-v0.4.6", dest-dir: "0.4.6", uv-version: "0.5.13", sphinx-release-override: "" }, ] steps: - name: Checkout diff --git a/docs/switcher.json b/docs/switcher.json index a31f2ca58fec..341ee23ffa11 100644 --- a/docs/switcher.json +++ b/docs/switcher.json @@ -1,6 +1,6 @@ [ { - "name": "0.4.5 (stable)", + "name": "0.4.6 (stable)", "version": "stable", "url": "/autogen/stable/", "preferred": true diff --git a/dotnet/src/Microsoft.AutoGen/Core.Grpc/GrpcAgentRuntime.cs b/dotnet/src/Microsoft.AutoGen/Core.Grpc/GrpcAgentRuntime.cs index 324874a0005b..83c77c5ef770 100644 --- a/dotnet/src/Microsoft.AutoGen/Core.Grpc/GrpcAgentRuntime.cs +++ b/dotnet/src/Microsoft.AutoGen/Core.Grpc/GrpcAgentRuntime.cs @@ -253,12 +253,33 @@ private async ValueTask HandlePublish(CloudEvent evt, CancellationToken cancella } } - public ValueTask StartAsync(CancellationToken cancellationToken) + public async ValueTask StartAsync(CancellationToken cancellationToken) { - return this._messageRouter.StartAsync(cancellationToken); + await this._messageRouter.StartAsync(cancellationToken); + if (this._agentsContainer.RegisteredAgentTypes.Count > 0) + { + foreach (var type in this._agentsContainer.RegisteredAgentTypes) + { + await this._client.RegisterAgentAsync(new RegisterAgentTypeRequest + { + Type = type + }, this.CallOptions); + } + } + + if (this._agentsContainer.Subscriptions.Count > 0) + { + foreach (var subscription in this._agentsContainer.Subscriptions.Values) + { + await this._client.AddSubscriptionAsync(new AddSubscriptionRequest + { + Subscription = subscription.ToProtobuf() + }, this.CallOptions); + } + } } - Task IHostedService.StartAsync(CancellationToken cancellationToken) => this._messageRouter.StartAsync(cancellationToken).AsTask(); + Task IHostedService.StartAsync(CancellationToken cancellationToken) => this.StartAsync(cancellationToken).AsTask(); public Task StopAsync(CancellationToken cancellationToken) { @@ -344,30 +365,39 @@ public async ValueTask AddSubscriptionAsync(ISubscriptionDefinition subscription { this._agentsContainer.AddSubscription(subscription); - var _ = await this._client.AddSubscriptionAsync(new AddSubscriptionRequest + if (this._messageRouter.IsChannelOpen) { - Subscription = subscription.ToProtobuf() - }, this.CallOptions); + var _ = await this._client.AddSubscriptionAsync(new AddSubscriptionRequest + { + Subscription = subscription.ToProtobuf() + }, this.CallOptions); + } } public async ValueTask RemoveSubscriptionAsync(string subscriptionId) { this._agentsContainer.RemoveSubscriptionAsync(subscriptionId); - await this._client.RemoveSubscriptionAsync(new RemoveSubscriptionRequest + if (this._messageRouter.IsChannelOpen) { - Id = subscriptionId - }, this.CallOptions); + await this._client.RemoveSubscriptionAsync(new RemoveSubscriptionRequest + { + Id = subscriptionId + }, this.CallOptions); + } } public async ValueTask RegisterAgentFactoryAsync(AgentType type, Func> factoryFunc) { this._agentsContainer.RegisterAgentFactory(type, factoryFunc); - await this._client.RegisterAgentAsync(new RegisterAgentTypeRequest + if (this._messageRouter.IsChannelOpen) { - Type = type, - }, this.CallOptions); + await this._client.RegisterAgentAsync(new RegisterAgentTypeRequest + { + Type = type, + }, this.CallOptions); + } return type; } diff --git a/dotnet/src/Microsoft.AutoGen/Core.Grpc/GrpcMessageRouter.cs b/dotnet/src/Microsoft.AutoGen/Core.Grpc/GrpcMessageRouter.cs index e46b392c708f..32317cbc7a30 100644 --- a/dotnet/src/Microsoft.AutoGen/Core.Grpc/GrpcMessageRouter.cs +++ b/dotnet/src/Microsoft.AutoGen/Core.Grpc/GrpcMessageRouter.cs @@ -34,6 +34,8 @@ public AutoRestartChannel(AgentRpc.AgentRpcClient client, _shutdownCts = CancellationTokenSource.CreateLinkedTokenSource(shutdownCancellation); } + public bool Connected { get => _channel is not null; } + public void EnsureConnected() { _logger.LogInformation("Connecting to gRPC endpoint " + Environment.GetEnvironmentVariable("AGENT_HOST")); @@ -287,6 +289,8 @@ public async Task StopAsync() this._incomingMessageChannel.Dispose(); } + public bool IsChannelOpen => this._incomingMessageChannel.Connected; + public void Dispose() { _outboundMessagesChannel.Writer.TryComplete(); diff --git a/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/AgentGrpcTests.cs b/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/AgentGrpcTests.cs index a2a970bebb41..dc74446ec582 100644 --- a/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/AgentGrpcTests.cs +++ b/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/AgentGrpcTests.cs @@ -10,13 +10,13 @@ namespace Microsoft.AutoGen.Core.Grpc.Tests; [Trait("Category", "GRPC")] -public class AgentGrpcTests +public class AgentGrpcTests : TestBase { [Fact] public async Task AgentShouldNotReceiveMessagesWhenNotSubscribedTest() { var fixture = new GrpcAgentRuntimeFixture(); - var runtime = (GrpcAgentRuntime)await fixture.Start(); + var runtime = (GrpcAgentRuntime)await fixture.StartAsync(); Logger logger = new(new LoggerFactory()); TestProtobufAgent agent = null!; @@ -45,7 +45,7 @@ await runtime.RegisterAgentFactoryAsync("MyAgent", async (id, runtime) => public async Task AgentShouldReceiveMessagesWhenSubscribedTest() { var fixture = new GrpcAgentRuntimeFixture(); - var runtime = (GrpcAgentRuntime)await fixture.Start(); + var runtime = (GrpcAgentRuntime)await fixture.StartAsync(); Logger logger = new(new LoggerFactory()); SubscribedProtobufAgent agent = null!; @@ -80,7 +80,7 @@ public async Task SendMessageAsyncShouldReturnResponseTest() { // Arrange var fixture = new GrpcAgentRuntimeFixture(); - var runtime = (GrpcAgentRuntime)await fixture.Start(); + var runtime = (GrpcAgentRuntime)await fixture.StartAsync(); Logger logger = new(new LoggerFactory()); await runtime.RegisterAgentFactoryAsync("MyAgent", async (id, runtime) => await ValueTask.FromResult(new TestProtobufAgent(id, runtime, logger))); @@ -114,7 +114,7 @@ public ValueTask HandleAsync(TextMessage item, MessageContext messageContext) public async Task SubscribeAsyncRemoveSubscriptionAsyncAndGetSubscriptionsTest() { var fixture = new GrpcAgentRuntimeFixture(); - var runtime = (GrpcAgentRuntime)await fixture.Start(); + var runtime = (GrpcAgentRuntime)await fixture.StartAsync(); ReceiverAgent? agent = null; await runtime.RegisterAgentFactoryAsync("MyAgent", async (id, runtime) => { diff --git a/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/FreePortManager.cs b/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/FreePortManager.cs new file mode 100644 index 000000000000..bbf3dbd54a6c --- /dev/null +++ b/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/FreePortManager.cs @@ -0,0 +1,67 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// FreePortManager.cs + +using System.Diagnostics; + +namespace Microsoft.AutoGen.Core.Grpc.Tests; + +internal sealed class FreePortManager +{ + private HashSet takenPorts = new(); + private readonly object mutex = new(); + + [DebuggerDisplay($"{{{nameof(Port)}}}")] + internal sealed class PortTicket(FreePortManager portManager, int port) : IDisposable + { + private FreePortManager? portManager = portManager; + + public int Port { get; } = port; + + public void Dispose() + { + FreePortManager? localPortManager = Interlocked.Exchange(ref this.portManager, null); + localPortManager?.takenPorts.Remove(this.Port); + } + + public override string ToString() + { + return this.Port.ToString(); + } + + public override bool Equals(object? obj) + { + return obj is PortTicket ticket && ticket.Port == this.Port; + } + + public override int GetHashCode() + { + return this.Port.GetHashCode(); + } + + public static implicit operator int(PortTicket ticket) => ticket.Port; + public static implicit operator string(PortTicket ticket) => ticket.ToString(); + } + + public PortTicket GetAvailablePort() + { + lock (mutex) + { + int port; + do + { + using var listener = new System.Net.Sockets.TcpListener(System.Net.IPAddress.Loopback, 0); + listener.Start(); + port = ((System.Net.IPEndPoint)listener.LocalEndpoint).Port; + listener.Stop(); + listener.Dispose(); + Thread.Yield(); // Let the listener actually shut down before we try to use the port + } while (takenPorts.Contains(port)); + + takenPorts.Add(port); + + Console.WriteLine($"FreePortManager: Yielding port {port}"); + Debug.WriteLine($"FreePortManager: Yielding port {port}"); + return new PortTicket(this, port); + } + } +} diff --git a/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/GrpcAgentRuntimeFixture.cs b/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/GrpcAgentRuntimeFixture.cs index bade7f785757..064f22cfa27f 100644 --- a/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/GrpcAgentRuntimeFixture.cs +++ b/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/GrpcAgentRuntimeFixture.cs @@ -7,69 +7,89 @@ using Microsoft.Extensions.Hosting; namespace Microsoft.AutoGen.Core.Grpc.Tests; + /// /// Fixture for setting up the gRPC agent runtime for testing. /// public sealed class GrpcAgentRuntimeFixture : IDisposable { + private FreePortManager.PortTicket? portTicket; + /// the gRPC agent runtime. - public AgentsApp? Client { get; private set; } + public AgentsApp? AgentsApp { get; private set; } + /// mock server for testing. - public WebApplication? Server { get; private set; } + public WebApplication? GatewayServer { get; private set; } + + public GrpcAgentServiceCollector GrpcRequestCollector { get; } public GrpcAgentRuntimeFixture() { + GrpcRequestCollector = new GrpcAgentServiceCollector(); } + /// /// Start - gets a new port and starts fresh instances /// - public async Task Start(bool initialize = true) + public async Task StartAsync(bool startRuntime = true, bool registerDefaultAgent = true) { - int port = GetAvailablePort(); // Get a new port per test run + this.portTicket = GrpcAgentRuntimeFixture.PortManager.GetAvailablePort(); // Get a new port per test run // Update environment variables so each test runs independently - Environment.SetEnvironmentVariable("ASPNETCORE_HTTPS_PORTS", port.ToString()); - Environment.SetEnvironmentVariable("AGENT_HOST", $"https://localhost:{port}"); + Environment.SetEnvironmentVariable("ASPNETCORE_HTTPS_PORTS", portTicket); + Environment.SetEnvironmentVariable("AGENT_HOST", $"https://localhost:{portTicket}"); Environment.SetEnvironmentVariable("ASPNETCORE_ENVIRONMENT", "Development"); - Server = ServerBuilder().Result; - await Server.StartAsync().ConfigureAwait(true); - Client = ClientBuilder().Result; - await Client.StartAsync().ConfigureAwait(true); - var worker = Client.Services.GetRequiredService(); + this.GatewayServer = await this.InitializeGateway(); + this.AgentsApp = await this.InitializeRuntime(startRuntime, registerDefaultAgent); + var runtime = AgentsApp.Services.GetRequiredService(); - return (worker); + return runtime; } - private static async Task ClientBuilder() + + private async Task InitializeRuntime(bool callStartAsync, bool registerDefaultAgent) { var appBuilder = new AgentsAppBuilder(); appBuilder.AddGrpcAgentWorker(); - appBuilder.AddAgent("TestAgent"); - return await appBuilder.BuildAsync(); + + if (registerDefaultAgent) + { + appBuilder.AddAgent("TestAgent"); + } + + AgentsApp result = await appBuilder.BuildAsync(); + + if (callStartAsync) + { + await result.StartAsync().ConfigureAwait(true); + } + + return result; } - private static async Task ServerBuilder() + + private async Task InitializeGateway() { var builder = WebApplication.CreateBuilder(); builder.Services.AddGrpc(); - var app = builder.Build(); + builder.Services.AddSingleton(this.GrpcRequestCollector); + + WebApplication app = builder.Build(); app.MapGrpcService(); + + await app.StartAsync().ConfigureAwait(true); return app; } - private static int GetAvailablePort() - { - using var listener = new System.Net.Sockets.TcpListener(System.Net.IPAddress.Loopback, 0); - listener.Start(); - int port = ((System.Net.IPEndPoint)listener.LocalEndpoint).Port; - listener.Stop(); - return port; - } + + private static readonly FreePortManager PortManager = new(); + /// /// Stop - stops the agent and ensures cleanup /// public void Stop() { - (Client as IHost)?.StopAsync(TimeSpan.FromSeconds(30)).GetAwaiter().GetResult(); - Server?.StopAsync().GetAwaiter().GetResult(); + (AgentsApp as IHost)?.StopAsync(TimeSpan.FromSeconds(30)).GetAwaiter().GetResult(); + GatewayServer?.StopAsync().GetAwaiter().GetResult(); + portTicket?.Dispose(); } /// diff --git a/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/GrpcAgentRuntimeTests.cs b/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/GrpcAgentRuntimeTests.cs new file mode 100644 index 000000000000..b48942f1dcba --- /dev/null +++ b/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/GrpcAgentRuntimeTests.cs @@ -0,0 +1,50 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// GrpcAgentRuntimeTests.cs + +using FluentAssertions; +using Microsoft.AutoGen.Contracts; +using Microsoft.Extensions.Logging; +using Xunit; + +namespace Microsoft.AutoGen.Core.Grpc.Tests; + +[Trait("Category", "GRPC")] +public class GrpcAgentRuntimeTests : TestBase +{ + [Fact] + public async Task GatewayShouldNotReceiveRegistrationsUntilRuntimeStart() + { + var fixture = new GrpcAgentRuntimeFixture(); + var runtime = (GrpcAgentRuntime)await fixture.StartAsync(startRuntime: false, registerDefaultAgent: false); + + Logger logger = new(new LoggerFactory()); + + await runtime.RegisterAgentFactoryAsync("MyAgent", async (id, runtime) => + { + return await ValueTask.FromResult(new SubscribedProtobufAgent(id, runtime, logger)); + }); + await runtime.RegisterImplicitAgentSubscriptionsAsync("MyAgent"); + + fixture.GrpcRequestCollector.RegisterAgentTypeRequests.Should().BeEmpty(); + fixture.GrpcRequestCollector.AddSubscriptionRequests.Should().BeEmpty(); + + await fixture.AgentsApp!.StartAsync().ConfigureAwait(true); + + fixture.GrpcRequestCollector.RegisterAgentTypeRequests.Should().NotBeEmpty(); + fixture.GrpcRequestCollector.RegisterAgentTypeRequests.Single().Type.Should().Be("MyAgent"); + fixture.GrpcRequestCollector.AddSubscriptionRequests.Should().NotBeEmpty(); + + fixture.GrpcRequestCollector.Clear(); + + await runtime.RegisterAgentFactoryAsync("MyAgent2", async (id, runtime) => + { + return await ValueTask.FromResult(new TestProtobufAgent(id, runtime, logger)); + }); + + fixture.GrpcRequestCollector.RegisterAgentTypeRequests.Should().NotBeEmpty(); + fixture.GrpcRequestCollector.RegisterAgentTypeRequests.Single().Type.Should().Be("MyAgent2"); + fixture.GrpcRequestCollector.AddSubscriptionRequests.Should().BeEmpty(); + + fixture.Dispose(); + } +} diff --git a/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/GrpcAgentServiceFixture.cs b/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/GrpcAgentServiceFixture.cs index 1ca37809a57e..af008f2419d2 100644 --- a/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/GrpcAgentServiceFixture.cs +++ b/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/GrpcAgentServiceFixture.cs @@ -1,14 +1,37 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // GrpcAgentServiceFixture.cs + using Grpc.Core; using Microsoft.AutoGen.Protobuf; +using Microsoft.Extensions.DependencyInjection; + namespace Microsoft.AutoGen.Core.Grpc.Tests; +public sealed class GrpcAgentServiceCollector +{ + public List AddSubscriptionRequests { get; } = new(); + public List RemoveSubscriptionRequests { get; } = new(); + public List RegisterAgentTypeRequests { get; } = new(); + + internal void Clear() + { + this.AddSubscriptionRequests.Clear(); + this.RemoveSubscriptionRequests.Clear(); + this.RegisterAgentTypeRequests.Clear(); + } +} + /// /// This fixture is largely just a loopback as we are testing the client side logic of the GrpcAgentRuntime in isolation from the rest of the system. /// -public sealed class GrpcAgentServiceFixture() : AgentRpc.AgentRpcBase +public class GrpcAgentServiceFixture : AgentRpc.AgentRpcBase { + private GrpcAgentServiceCollector requestCollector; + public GrpcAgentServiceFixture(IServiceProvider serviceProvider) + { + this.requestCollector = serviceProvider.GetService() ?? new(); + } + public override async Task OpenChannel(IAsyncStreamReader requestStream, IServerStreamWriter responseStream, ServerCallContext context) { try @@ -25,8 +48,27 @@ public override async Task OpenChannel(IAsyncStreamReader requestStream throw; } } - public override async Task AddSubscription(AddSubscriptionRequest request, ServerCallContext context) => new AddSubscriptionResponse { }; - public override async Task RemoveSubscription(RemoveSubscriptionRequest request, ServerCallContext context) => new RemoveSubscriptionResponse { }; + + public List AddSubscriptionRequests => this.requestCollector.AddSubscriptionRequests; + public override async Task AddSubscription(AddSubscriptionRequest request, ServerCallContext context) + { + this.AddSubscriptionRequests.Add(request); + return new AddSubscriptionResponse(); + } + + public List RemoveSubscriptionRequests => this.requestCollector.RemoveSubscriptionRequests; + public override async Task RemoveSubscription(RemoveSubscriptionRequest request, ServerCallContext context) + { + this.RemoveSubscriptionRequests.Add(request); + return new RemoveSubscriptionResponse(); + } + public override async Task GetSubscriptions(GetSubscriptionsRequest request, ServerCallContext context) => new GetSubscriptionsResponse { }; - public override async Task RegisterAgent(RegisterAgentTypeRequest request, ServerCallContext context) => new RegisterAgentTypeResponse { }; + + public List RegisterAgentTypeRequests => this.requestCollector.RegisterAgentTypeRequests; + public override async Task RegisterAgent(RegisterAgentTypeRequest request, ServerCallContext context) + { + this.RegisterAgentTypeRequests.Add(request); + return new RegisterAgentTypeResponse(); + } } diff --git a/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/TestBase.cs b/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/TestBase.cs new file mode 100644 index 000000000000..77d9d552a1c0 --- /dev/null +++ b/dotnet/test/Microsoft.AutoGen.Core.Grpc.Tests/TestBase.cs @@ -0,0 +1,22 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// TestBase.cs + +namespace Microsoft.AutoGen.Core.Grpc.Tests; + +public class TestBase +{ + public TestBase() + { + try + { + // For some reason the first call to StartAsync() throws when these tests + // run in parallel, even though the port does not actually collide between + // different instances of GrpcAgentRuntimeFixture. This is a workaround. + _ = new GrpcAgentRuntimeFixture().StartAsync().Result; + } + catch (Exception e) + { + Console.WriteLine(e); + } + } +} diff --git a/python/packages/autogen-core/docs/src/conf.py b/python/packages/autogen-core/docs/src/conf.py index 5c4fcbb0ccfb..bfb8123da32b 100644 --- a/python/packages/autogen-core/docs/src/conf.py +++ b/python/packages/autogen-core/docs/src/conf.py @@ -178,6 +178,7 @@ "user-guide/core-user-guide/framework/command-line-code-executors.ipynb": "user-guide/core-user-guide/components/command-line-code-executors.ipynb", "user-guide/core-user-guide/framework/model-clients.ipynb": "user-guide/core-user-guide/components/model-clients.ipynb", "user-guide/core-user-guide/framework/tools.ipynb": "user-guide/core-user-guide/components/tools.ipynb", + "user-guide/agentchat-user-guide/tutorial/custom-agents.ipynb": "user-guide/agentchat-user-guide/custom-agents.ipynb", } diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/custom-agents.ipynb b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/custom-agents.ipynb new file mode 100644 index 000000000000..d738e72b60f8 --- /dev/null +++ b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/custom-agents.ipynb @@ -0,0 +1,736 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Custom Agents\n", + "\n", + "You may have agents with behaviors that do not fall into a preset. \n", + "In such cases, you can build custom agents.\n", + "\n", + "All agents in AgentChat inherit from {py:class}`~autogen_agentchat.agents.BaseChatAgent` \n", + "class and implement the following abstract methods and attributes:\n", + "\n", + "- {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages`: The abstract method that defines the behavior of the agent in response to messages. This method is called when the agent is asked to provide a response in {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run`. It returns a {py:class}`~autogen_agentchat.base.Response` object.\n", + "- {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_reset`: The abstract method that resets the agent to its initial state. This method is called when the agent is asked to reset itself.\n", + "- {py:attr}`~autogen_agentchat.agents.BaseChatAgent.produced_message_types`: The list of possible {py:class}`~autogen_agentchat.messages.ChatMessage` message types the agent can produce in its response.\n", + "\n", + "Optionally, you can implement the the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages_stream` method to stream messages as they are generated by the agent. If this method is not implemented, the agent\n", + "uses the default implementation of {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages_stream`\n", + "that calls the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages` method and\n", + "yields all messages in the response." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## CountDownAgent\n", + "\n", + "In this example, we create a simple agent that counts down from a given number to zero,\n", + "and produces a stream of messages with the current count." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "3...\n", + "2...\n", + "1...\n", + "Done!\n" + ] + } + ], + "source": [ + "from typing import AsyncGenerator, List, Sequence\n", + "\n", + "from autogen_agentchat.agents import BaseChatAgent\n", + "from autogen_agentchat.base import Response\n", + "from autogen_agentchat.messages import AgentEvent, ChatMessage, TextMessage\n", + "from autogen_core import CancellationToken\n", + "\n", + "\n", + "class CountDownAgent(BaseChatAgent):\n", + " def __init__(self, name: str, count: int = 3):\n", + " super().__init__(name, \"A simple agent that counts down.\")\n", + " self._count = count\n", + "\n", + " @property\n", + " def produced_message_types(self) -> Sequence[type[ChatMessage]]:\n", + " return (TextMessage,)\n", + "\n", + " async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:\n", + " # Calls the on_messages_stream.\n", + " response: Response | None = None\n", + " async for message in self.on_messages_stream(messages, cancellation_token):\n", + " if isinstance(message, Response):\n", + " response = message\n", + " assert response is not None\n", + " return response\n", + "\n", + " async def on_messages_stream(\n", + " self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken\n", + " ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]:\n", + " inner_messages: List[AgentEvent | ChatMessage] = []\n", + " for i in range(self._count, 0, -1):\n", + " msg = TextMessage(content=f\"{i}...\", source=self.name)\n", + " inner_messages.append(msg)\n", + " yield msg\n", + " # The response is returned at the end of the stream.\n", + " # It contains the final message and all the inner messages.\n", + " yield Response(chat_message=TextMessage(content=\"Done!\", source=self.name), inner_messages=inner_messages)\n", + "\n", + " async def on_reset(self, cancellation_token: CancellationToken) -> None:\n", + " pass\n", + "\n", + "\n", + "async def run_countdown_agent() -> None:\n", + " # Create a countdown agent.\n", + " countdown_agent = CountDownAgent(\"countdown\")\n", + "\n", + " # Run the agent with a given task and stream the response.\n", + " async for message in countdown_agent.on_messages_stream([], CancellationToken()):\n", + " if isinstance(message, Response):\n", + " print(message.chat_message.content)\n", + " else:\n", + " print(message.content)\n", + "\n", + "\n", + "# Use asyncio.run(run_countdown_agent()) when running in a script.\n", + "await run_countdown_agent()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## ArithmeticAgent\n", + "\n", + "In this example, we create an agent class that can perform simple arithmetic operations\n", + "on a given integer. Then, we will use different instances of this agent class\n", + "in a {py:class}`~autogen_agentchat.teams.SelectorGroupChat`\n", + "to transform a given integer into another integer by applying a sequence of arithmetic operations.\n", + "\n", + "The `ArithmeticAgent` class takes an `operator_func` that takes an integer and returns an integer,\n", + "after applying an arithmetic operation to the integer.\n", + "In its `on_messages` method, it applies the `operator_func` to the integer in the input message,\n", + "and returns a response with the result." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from typing import Callable, Sequence\n", + "\n", + "from autogen_agentchat.agents import BaseChatAgent\n", + "from autogen_agentchat.base import Response\n", + "from autogen_agentchat.conditions import MaxMessageTermination\n", + "from autogen_agentchat.messages import ChatMessage\n", + "from autogen_agentchat.teams import SelectorGroupChat\n", + "from autogen_agentchat.ui import Console\n", + "from autogen_core import CancellationToken\n", + "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", + "\n", + "\n", + "class ArithmeticAgent(BaseChatAgent):\n", + " def __init__(self, name: str, description: str, operator_func: Callable[[int], int]) -> None:\n", + " super().__init__(name, description=description)\n", + " self._operator_func = operator_func\n", + " self._message_history: List[ChatMessage] = []\n", + "\n", + " @property\n", + " def produced_message_types(self) -> Sequence[type[ChatMessage]]:\n", + " return (TextMessage,)\n", + "\n", + " async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:\n", + " # Update the message history.\n", + " # NOTE: it is possible the messages is an empty list, which means the agent was selected previously.\n", + " self._message_history.extend(messages)\n", + " # Parse the number in the last message.\n", + " assert isinstance(self._message_history[-1], TextMessage)\n", + " number = int(self._message_history[-1].content)\n", + " # Apply the operator function to the number.\n", + " result = self._operator_func(number)\n", + " # Create a new message with the result.\n", + " response_message = TextMessage(content=str(result), source=self.name)\n", + " # Update the message history.\n", + " self._message_history.append(response_message)\n", + " # Return the response.\n", + " return Response(chat_message=response_message)\n", + "\n", + " async def on_reset(self, cancellation_token: CancellationToken) -> None:\n", + " pass" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "```{note}\n", + "The `on_messages` method may be called with an empty list of messages, in which\n", + "case it means the agent was called previously and is now being called again,\n", + "without any new messages from the caller. So it is important to keep a history\n", + "of the previous messages received by the agent, and use that history to generate\n", + "the response.\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we can create a {py:class}`~autogen_agentchat.teams.SelectorGroupChat` with 5 instances of `ArithmeticAgent`:\n", + "\n", + "- one that adds 1 to the input integer,\n", + "- one that subtracts 1 from the input integer,\n", + "- one that multiplies the input integer by 2,\n", + "- one that divides the input integer by 2 and rounds down to the nearest integer, and\n", + "- one that returns the input integer unchanged.\n", + "\n", + "We then create a {py:class}`~autogen_agentchat.teams.SelectorGroupChat` with these agents,\n", + "and set the appropriate selector settings:\n", + "\n", + "- allow the same agent to be selected consecutively to allow for repeated operations, and\n", + "- customize the selector prompt to tailor the model's response to the specific task." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- user ----------\n", + "Apply the operations to turn the given number into 25.\n", + "---------- user ----------\n", + "10\n", + "---------- multiply_agent ----------\n", + "20\n", + "---------- add_agent ----------\n", + "21\n", + "---------- multiply_agent ----------\n", + "42\n", + "---------- divide_agent ----------\n", + "21\n", + "---------- add_agent ----------\n", + "22\n", + "---------- add_agent ----------\n", + "23\n", + "---------- add_agent ----------\n", + "24\n", + "---------- add_agent ----------\n", + "25\n", + "---------- Summary ----------\n", + "Number of messages: 10\n", + "Finish reason: Maximum number of messages 10 reached, current message count: 10\n", + "Total prompt tokens: 0\n", + "Total completion tokens: 0\n", + "Duration: 2.40 seconds\n" + ] + } + ], + "source": [ + "async def run_number_agents() -> None:\n", + " # Create agents for number operations.\n", + " add_agent = ArithmeticAgent(\"add_agent\", \"Adds 1 to the number.\", lambda x: x + 1)\n", + " multiply_agent = ArithmeticAgent(\"multiply_agent\", \"Multiplies the number by 2.\", lambda x: x * 2)\n", + " subtract_agent = ArithmeticAgent(\"subtract_agent\", \"Subtracts 1 from the number.\", lambda x: x - 1)\n", + " divide_agent = ArithmeticAgent(\"divide_agent\", \"Divides the number by 2 and rounds down.\", lambda x: x // 2)\n", + " identity_agent = ArithmeticAgent(\"identity_agent\", \"Returns the number as is.\", lambda x: x)\n", + "\n", + " # The termination condition is to stop after 10 messages.\n", + " termination_condition = MaxMessageTermination(10)\n", + "\n", + " # Create a selector group chat.\n", + " selector_group_chat = SelectorGroupChat(\n", + " [add_agent, multiply_agent, subtract_agent, divide_agent, identity_agent],\n", + " model_client=OpenAIChatCompletionClient(model=\"gpt-4o\"),\n", + " termination_condition=termination_condition,\n", + " allow_repeated_speaker=True, # Allow the same agent to speak multiple times, necessary for this task.\n", + " selector_prompt=(\n", + " \"Available roles:\\n{roles}\\nTheir job descriptions:\\n{participants}\\n\"\n", + " \"Current conversation history:\\n{history}\\n\"\n", + " \"Please select the most appropriate role for the next message, and only return the role name.\"\n", + " ),\n", + " )\n", + "\n", + " # Run the selector group chat with a given task and stream the response.\n", + " task: List[ChatMessage] = [\n", + " TextMessage(content=\"Apply the operations to turn the given number into 25.\", source=\"user\"),\n", + " TextMessage(content=\"10\", source=\"user\"),\n", + " ]\n", + " stream = selector_group_chat.run_stream(task=task)\n", + " await Console(stream)\n", + "\n", + "\n", + "# Use asyncio.run(run_number_agents()) when running in a script.\n", + "await run_number_agents()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "From the output, we can see that the agents have successfully transformed the input integer\n", + "from 10 to 25 by choosing appropriate agents that apply the arithmetic operations in sequence." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Custom Model Clients in Custom Agents\n", + "\n", + "One of the key features of the {py:class}`~autogen_agentchat.agents.AssistantAgent` preset in AgentChat is that it takes a `model_client` argument and can use it in responding to messages. However, in some cases, you may want your agent to use a custom model client that is not currently supported (see [supported model clients](https://microsoft.github.io/autogen/dev/user-guide/core-user-guide/components/model-clients.html)) or custom model behaviours. \n", + "\n", + "You can accomplish this with a custom agent that implements *your custom model client*.\n", + "\n", + "In the example below, we will walk through an example of a custom agent that uses the [Google Gemini SDK](https://github.com/googleapis/python-genai) directly to respond to messages.\n", + "\n", + "> **Note:** You will need to install the [Google Gemini SDK](https://github.com/googleapis/python-genai) to run this example. You can install it using the following command: \n", + "\n", + "```bash\n", + "pip install google-genai\n", + "``` " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# !pip install google-genai\n", + "import os\n", + "from typing import AsyncGenerator, Sequence\n", + "\n", + "from autogen_agentchat.agents import BaseChatAgent\n", + "from autogen_agentchat.base import Response\n", + "from autogen_agentchat.messages import AgentEvent, ChatMessage\n", + "from autogen_core import CancellationToken\n", + "from autogen_core.model_context import UnboundedChatCompletionContext\n", + "from autogen_core.models import AssistantMessage, RequestUsage, UserMessage\n", + "from google import genai\n", + "from google.genai import types\n", + "\n", + "\n", + "class GeminiAssistantAgent(BaseChatAgent):\n", + " def __init__(\n", + " self,\n", + " name: str,\n", + " description: str = \"An agent that provides assistance with ability to use tools.\",\n", + " model: str = \"gemini-1.5-flash-002\",\n", + " api_key: str = os.environ[\"GEMINI_API_KEY\"],\n", + " system_message: str\n", + " | None = \"You are a helpful assistant that can respond to messages. Reply with TERMINATE when the task has been completed.\",\n", + " ):\n", + " super().__init__(name=name, description=description)\n", + " self._model_context = UnboundedChatCompletionContext()\n", + " self._model_client = genai.Client(api_key=api_key)\n", + " self._system_message = system_message\n", + " self._model = model\n", + "\n", + " @property\n", + " def produced_message_types(self) -> Sequence[type[ChatMessage]]:\n", + " return (TextMessage,)\n", + "\n", + " async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:\n", + " final_response = None\n", + " async for message in self.on_messages_stream(messages, cancellation_token):\n", + " if isinstance(message, Response):\n", + " final_response = message\n", + "\n", + " if final_response is None:\n", + " raise AssertionError(\"The stream should have returned the final result.\")\n", + "\n", + " return final_response\n", + "\n", + " async def on_messages_stream(\n", + " self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken\n", + " ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]:\n", + " # Add messages to the model context\n", + " for msg in messages:\n", + " await self._model_context.add_message(UserMessage(content=msg.content, source=msg.source))\n", + "\n", + " # Get conversation history\n", + " history = [\n", + " (msg.source if hasattr(msg, \"source\") else \"system\")\n", + " + \": \"\n", + " + (msg.content if isinstance(msg.content, str) else \"\")\n", + " + \"\\n\"\n", + " for msg in await self._model_context.get_messages()\n", + " ]\n", + " # Generate response using Gemini\n", + " response = self._model_client.models.generate_content(\n", + " model=self._model,\n", + " contents=f\"History: {history}\\nGiven the history, please provide a response\",\n", + " config=types.GenerateContentConfig(\n", + " system_instruction=self._system_message,\n", + " temperature=0.3,\n", + " ),\n", + " )\n", + "\n", + " # Create usage metadata\n", + " usage = RequestUsage(\n", + " prompt_tokens=response.usage_metadata.prompt_token_count,\n", + " completion_tokens=response.usage_metadata.candidates_token_count,\n", + " )\n", + "\n", + " # Add response to model context\n", + " await self._model_context.add_message(AssistantMessage(content=response.text, source=self.name))\n", + "\n", + " # Yield the final response\n", + " yield Response(\n", + " chat_message=TextMessage(content=response.text, source=self.name, models_usage=usage),\n", + " inner_messages=[],\n", + " )\n", + "\n", + " async def on_reset(self, cancellation_token: CancellationToken) -> None:\n", + " \"\"\"Reset the assistant by clearing the model context.\"\"\"\n", + " await self._model_context.clear()" + ] + }, + { + "cell_type": "code", + "execution_count": 38, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- user ----------\n", + "What is the capital of New York?\n", + "---------- gemini_assistant ----------\n", + "Albany\n", + "TERMINATE\n", + "\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What is the capital of New York?', type='TextMessage'), TextMessage(source='gemini_assistant', models_usage=RequestUsage(prompt_tokens=46, completion_tokens=5), content='Albany\\nTERMINATE\\n', type='TextMessage')], stop_reason=None)" + ] + }, + "execution_count": 38, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "gemini_assistant = GeminiAssistantAgent(\"gemini_assistant\")\n", + "await Console(gemini_assistant.run_stream(task=\"What is the capital of New York?\"))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In the example above, we have chosen to provide `model`, `api_key` and `system_message` as arguments - you can choose to provide any other arguments that are required by the model client you are using or fits with your application design. \n", + "\n", + "Now, let us explore how to use this custom agent as part of a team in AgentChat." + ] + }, + { + "cell_type": "code", + "execution_count": 39, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "---------- user ----------\n", + "Write a Haiku poem with 4 lines about the fall season.\n", + "---------- primary ----------\n", + "Crimson leaves cascade, \n", + "Whispering winds sing of change, \n", + "Chill wraps the fading, \n", + "Nature's quilt, rich and warm.\n", + "---------- gemini_critic ----------\n", + "The poem is good, but it has four lines instead of three. A haiku must have three lines with a 5-7-5 syllable structure. The content is evocative of autumn, but the form is incorrect. Please revise to adhere to the haiku's syllable structure.\n", + "\n", + "---------- primary ----------\n", + "Thank you for your feedback! Here’s a revised haiku that follows the 5-7-5 syllable structure:\n", + "\n", + "Crimson leaves drift down, \n", + "Chill winds whisper through the gold, \n", + "Autumn’s breath is near.\n", + "---------- gemini_critic ----------\n", + "The revised haiku is much improved. It correctly follows the 5-7-5 syllable structure and maintains the evocative imagery of autumn. APPROVE\n", + "\n" + ] + }, + { + "data": { + "text/plain": [ + "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='Write a Haiku poem with 4 lines about the fall season.', type='TextMessage'), TextMessage(source='primary', models_usage=RequestUsage(prompt_tokens=33, completion_tokens=31), content=\"Crimson leaves cascade, \\nWhispering winds sing of change, \\nChill wraps the fading, \\nNature's quilt, rich and warm.\", type='TextMessage'), TextMessage(source='gemini_critic', models_usage=RequestUsage(prompt_tokens=86, completion_tokens=60), content=\"The poem is good, but it has four lines instead of three. A haiku must have three lines with a 5-7-5 syllable structure. The content is evocative of autumn, but the form is incorrect. Please revise to adhere to the haiku's syllable structure.\\n\", type='TextMessage'), TextMessage(source='primary', models_usage=RequestUsage(prompt_tokens=141, completion_tokens=49), content='Thank you for your feedback! Here’s a revised haiku that follows the 5-7-5 syllable structure:\\n\\nCrimson leaves drift down, \\nChill winds whisper through the gold, \\nAutumn’s breath is near.', type='TextMessage'), TextMessage(source='gemini_critic', models_usage=RequestUsage(prompt_tokens=211, completion_tokens=32), content='The revised haiku is much improved. It correctly follows the 5-7-5 syllable structure and maintains the evocative imagery of autumn. APPROVE\\n', type='TextMessage')], stop_reason=\"Text 'APPROVE' mentioned\")" + ] + }, + "execution_count": 39, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from autogen_agentchat.agents import AssistantAgent\n", + "from autogen_agentchat.conditions import TextMentionTermination\n", + "from autogen_agentchat.teams import RoundRobinGroupChat\n", + "from autogen_agentchat.ui import Console\n", + "\n", + "# Create the primary agent.\n", + "primary_agent = AssistantAgent(\n", + " \"primary\",\n", + " model_client=OpenAIChatCompletionClient(model=\"gpt-4o-mini\"),\n", + " system_message=\"You are a helpful AI assistant.\",\n", + ")\n", + "\n", + "# Create a critic agent based on our new GeminiAssistantAgent.\n", + "gemini_critic_agent = GeminiAssistantAgent(\n", + " \"gemini_critic\",\n", + " system_message=\"Provide constructive feedback. Respond with 'APPROVE' to when your feedbacks are addressed.\",\n", + ")\n", + "\n", + "\n", + "# Define a termination condition that stops the task if the critic approves or after 10 messages.\n", + "termination = TextMentionTermination(\"APPROVE\") | MaxMessageTermination(10)\n", + "\n", + "# Create a team with the primary and critic agents.\n", + "team = RoundRobinGroupChat([primary_agent, gemini_critic_agent], termination_condition=termination)\n", + "\n", + "await Console(team.run_stream(task=\"Write a Haiku poem with 4 lines about the fall season.\"))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In section above, we show several very important concepts:\n", + "- We have developed a custom agent that uses the Google Gemini SDK to respond to messages. \n", + "- We show that this custom agent can be used as part of the broader AgentChat ecosystem - in this case as a participant in a {py:class}`~autogen_agentchat.teams.RoundRobinGroupChat` as long as it inherits from {py:class}`~autogen_agentchat.agents.BaseChatAgent`.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Making the Custom Agent Declarative \n", + "\n", + "Autogen provides a [Component](https://microsoft.github.io/autogen/dev/user-guide/core-user-guide/framework/component-config.html) interface for making the configuration of components serializable to a declarative format. This is useful for saving and loading configurations, and for sharing configurations with others. \n", + "\n", + "We accomplish this by inheriting from the `Component` class and implementing the `_from_config` and `_to_config` methods.\n", + "The declarative class can be serialized to a JSON format using the `dump_component` method, and deserialized from a JSON format using the `load_component` method." + ] + }, + { + "cell_type": "code", + "execution_count": 40, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "from typing import AsyncGenerator, Sequence\n", + "\n", + "from autogen_agentchat.agents import BaseChatAgent\n", + "from autogen_agentchat.base import Response\n", + "from autogen_agentchat.messages import AgentEvent, ChatMessage\n", + "from autogen_core import CancellationToken, Component\n", + "from pydantic import BaseModel\n", + "from typing_extensions import Self\n", + "\n", + "\n", + "class GeminiAssistantAgentConfig(BaseModel):\n", + " name: str\n", + " description: str = \"An agent that provides assistance with ability to use tools.\"\n", + " model: str = \"gemini-1.5-flash-002\"\n", + " system_message: str | None = None\n", + "\n", + "\n", + "class GeminiAssistantAgent(BaseChatAgent, Component[GeminiAssistantAgentConfig]): # type: ignore[no-redef]\n", + " component_config_schema = GeminiAssistantAgentConfig\n", + " # component_provider_override = \"mypackage.agents.GeminiAssistantAgent\"\n", + "\n", + " def __init__(\n", + " self,\n", + " name: str,\n", + " description: str = \"An agent that provides assistance with ability to use tools.\",\n", + " model: str = \"gemini-1.5-flash-002\",\n", + " api_key: str = os.environ[\"GEMINI_API_KEY\"],\n", + " system_message: str\n", + " | None = \"You are a helpful assistant that can respond to messages. Reply with TERMINATE when the task has been completed.\",\n", + " ):\n", + " super().__init__(name=name, description=description)\n", + " self._model_context = UnboundedChatCompletionContext()\n", + " self._model_client = genai.Client(api_key=api_key)\n", + " self._system_message = system_message\n", + " self._model = model\n", + "\n", + " @property\n", + " def produced_message_types(self) -> Sequence[type[ChatMessage]]:\n", + " return (TextMessage,)\n", + "\n", + " async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:\n", + " final_response = None\n", + " async for message in self.on_messages_stream(messages, cancellation_token):\n", + " if isinstance(message, Response):\n", + " final_response = message\n", + "\n", + " if final_response is None:\n", + " raise AssertionError(\"The stream should have returned the final result.\")\n", + "\n", + " return final_response\n", + "\n", + " async def on_messages_stream(\n", + " self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken\n", + " ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]:\n", + " # Add messages to the model context\n", + " for msg in messages:\n", + " await self._model_context.add_message(UserMessage(content=msg.content, source=msg.source))\n", + "\n", + " # Get conversation history\n", + " history = [\n", + " (msg.source if hasattr(msg, \"source\") else \"system\")\n", + " + \": \"\n", + " + (msg.content if isinstance(msg.content, str) else \"\")\n", + " + \"\\n\"\n", + " for msg in await self._model_context.get_messages()\n", + " ]\n", + "\n", + " # Generate response using Gemini\n", + " response = self._model_client.models.generate_content(\n", + " model=self._model,\n", + " contents=f\"History: {history}\\nGiven the history, please provide a response\",\n", + " config=types.GenerateContentConfig(\n", + " system_instruction=self._system_message,\n", + " temperature=0.3,\n", + " ),\n", + " )\n", + "\n", + " # Create usage metadata\n", + " usage = RequestUsage(\n", + " prompt_tokens=response.usage_metadata.prompt_token_count,\n", + " completion_tokens=response.usage_metadata.candidates_token_count,\n", + " )\n", + "\n", + " # Add response to model context\n", + " await self._model_context.add_message(AssistantMessage(content=response.text, source=self.name))\n", + "\n", + " # Yield the final response\n", + " yield Response(\n", + " chat_message=TextMessage(content=response.text, source=self.name, models_usage=usage),\n", + " inner_messages=[],\n", + " )\n", + "\n", + " async def on_reset(self, cancellation_token: CancellationToken) -> None:\n", + " \"\"\"Reset the assistant by clearing the model context.\"\"\"\n", + " await self._model_context.clear()\n", + "\n", + " @classmethod\n", + " def _from_config(cls, config: GeminiAssistantAgentConfig) -> Self:\n", + " return cls(\n", + " name=config.name, description=config.description, model=config.model, system_message=config.system_message\n", + " )\n", + "\n", + " def _to_config(self) -> GeminiAssistantAgentConfig:\n", + " return GeminiAssistantAgentConfig(\n", + " name=self.name,\n", + " description=self.description,\n", + " model=self._model,\n", + " system_message=self._system_message,\n", + " )" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now that we have the required methods implemented, we can now load and dump the custom agent to and from a JSON format, and then load the agent from the JSON format.\n", + " \n", + " > Note: You should set the `component_provider_override` class variable to the full path of the module containing the custom agent class e.g., (`mypackage.agents.GeminiAssistantAgent`). This is used by `load_component` method to determine how to instantiate the class. \n", + " " + ] + }, + { + "cell_type": "code", + "execution_count": 41, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "{\n", + " \"provider\": \"__main__.GeminiAssistantAgent\",\n", + " \"component_type\": \"agent\",\n", + " \"version\": 1,\n", + " \"component_version\": 1,\n", + " \"description\": null,\n", + " \"label\": \"GeminiAssistantAgent\",\n", + " \"config\": {\n", + " \"name\": \"gemini_assistant\",\n", + " \"description\": \"An agent that provides assistance with ability to use tools.\",\n", + " \"model\": \"gemini-1.5-flash-002\",\n", + " \"system_message\": \"You are a helpful assistant that can respond to messages. Reply with TERMINATE when the task has been completed.\"\n", + " }\n", + "}\n", + "<__main__.GeminiAssistantAgent object at 0x11a5c5a90>\n" + ] + } + ], + "source": [ + "gemini_assistant = GeminiAssistantAgent(\"gemini_assistant\")\n", + "config = gemini_assistant.dump_component()\n", + "print(config.model_dump_json(indent=2))\n", + "loaded_agent = GeminiAssistantAgent.load_component(config)\n", + "print(loaded_agent)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Next Steps \n", + "\n", + "So far, we have seen how to create custom agents, add custom model clients to agents, and make custom agents declarative. There are a few ways in which this basic sample can be extended:\n", + "\n", + "- Extend the Gemini model client to handle function calling similar to the {py:class}`~autogen_agentchat.agents.AssistantAgent` class. https://ai.google.dev/gemini-api/docs/function-calling \n", + "- Implement a package wit a custom agent and experiment with using it's declarative format in a tool like [AutoGen Studio](https://microsoft.github.io/autogen/dev/user-guide/autogenstudio-user-guide/index.html)." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.9" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/index.md b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/index.md index e83288338dc7..b7d124f9e67e 100644 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/index.md +++ b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/index.md @@ -101,7 +101,7 @@ tutorial/agents tutorial/teams tutorial/human-in-the-loop tutorial/termination -tutorial/custom-agents + tutorial/state ``` @@ -111,11 +111,13 @@ tutorial/state :hidden: :caption: Advanced +custom-agents selector-group-chat swarm magentic-one memory serialize-components + ``` ```{toctree} diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/migration-guide.md b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/migration-guide.md index fdb3ca2152d2..de9991609258 100644 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/migration-guide.md +++ b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/migration-guide.md @@ -33,29 +33,35 @@ We provide a detailed guide on how to migrate your existing codebase from `v0.2` See each feature below for detailed information on how to migrate. -- [Model Client](#model-client) -- [Model Client for OpenAI-Compatible APIs](#model-client-for-openai-compatible-apis) -- [Model Client Cache](#model-client-cache) -- [Assistant Agent](#assistant-agent) -- [Multi-Modal Agent](#multi-modal-agent) -- [User Proxy](#user-proxy) -- [Conversable Agent and Register Reply](#conversable-agent-and-register-reply) -- [Save and Load Agent State](#save-and-load-agent-state) -- [Two-Agent Chat](#two-agent-chat) -- [Tool Use](#tool-use) -- [Chat Result](#chat-result) -- [Conversion between v0.2 and v0.4 Messages](#conversion-between-v02-and-v04-messages) -- [Group Chat](#group-chat) -- [Group Chat with Resume](#group-chat-with-resume) -- [Save and Load Group Chat State](#save-and-load-group-chat-state) -- [Group Chat with Tool Use](#group-chat-with-tool-use) -- [Group Chat with Custom Selector (Stateflow)](#group-chat-with-custom-selector-stateflow) -- [Nested Chat](#nested-chat) -- [Sequential Chat](#sequential-chat) -- [GPTAssistantAgent](#gptassistantagent) -- [Long-Context Handling](#long-context-handling) -- [Observability and Control](#observability-and-control) -- [Code Executors](#code-executors) +- [Migration Guide for v0.2 to v0.4](#migration-guide-for-v02-to-v04) + - [What is `v0.4`?](#what-is-v04) + - [New to AutoGen?](#new-to-autogen) + - [What's in this guide?](#whats-in-this-guide) + - [Model Client](#model-client) + - [Use component config](#use-component-config) + - [Use model client class directly](#use-model-client-class-directly) + - [Model Client for OpenAI-Compatible APIs](#model-client-for-openai-compatible-apis) + - [Model Client Cache](#model-client-cache) + - [Assistant Agent](#assistant-agent) + - [Multi-Modal Agent](#multi-modal-agent) + - [User Proxy](#user-proxy) + - [Conversable Agent and Register Reply](#conversable-agent-and-register-reply) + - [Save and Load Agent State](#save-and-load-agent-state) + - [Two-Agent Chat](#two-agent-chat) + - [Tool Use](#tool-use) + - [Chat Result](#chat-result) + - [Conversion between v0.2 and v0.4 Messages](#conversion-between-v02-and-v04-messages) + - [Group Chat](#group-chat) + - [Group Chat with Resume](#group-chat-with-resume) + - [Save and Load Group Chat State](#save-and-load-group-chat-state) + - [Group Chat with Tool Use](#group-chat-with-tool-use) + - [Group Chat with Custom Selector (Stateflow)](#group-chat-with-custom-selector-stateflow) + - [Nested Chat](#nested-chat) + - [Sequential Chat](#sequential-chat) + - [GPTAssistantAgent](#gptassistantagent) + - [Long Context Handling](#long-context-handling) + - [Observability and Control](#observability-and-control) + - [Code Executors](#code-executors) The following features currently in `v0.2` will be provided in the future releases of `v0.4.*` versions: @@ -417,7 +423,7 @@ class CustomAgent(BaseChatAgent): ``` You can then use the custom agent in the same way as the {py:class}`~autogen_agentchat.agents.AssistantAgent`. -See [Custom Agent Tutorial](./tutorial/custom-agents.ipynb) +See [Custom Agent Tutorial](custom-agents.ipynb) for more details. ## Save and Load Agent State diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/custom-agents.ipynb b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/custom-agents.ipynb deleted file mode 100644 index 5b8c4e7f24fa..000000000000 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/custom-agents.ipynb +++ /dev/null @@ -1,313 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Custom Agents\n", - "\n", - "You may have agents with behaviors that do not fall into a preset. \n", - "In such cases, you can build custom agents.\n", - "\n", - "All agents in AgentChat inherit from {py:class}`~autogen_agentchat.agents.BaseChatAgent` \n", - "class and implement the following abstract methods and attributes:\n", - "\n", - "- {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages`: The abstract method that defines the behavior of the agent in response to messages. This method is called when the agent is asked to provide a response in {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run`. It returns a {py:class}`~autogen_agentchat.base.Response` object.\n", - "- {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_reset`: The abstract method that resets the agent to its initial state. This method is called when the agent is asked to reset itself.\n", - "- {py:attr}`~autogen_agentchat.agents.BaseChatAgent.produced_message_types`: The list of possible {py:class}`~autogen_agentchat.messages.ChatMessage` message types the agent can produce in its response.\n", - "\n", - "Optionally, you can implement the the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages_stream` method to stream messages as they are generated by the agent. If this method is not implemented, the agent\n", - "uses the default implementation of {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages_stream`\n", - "that calls the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages` method and\n", - "yields all messages in the response." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## CountDownAgent\n", - "\n", - "In this example, we create a simple agent that counts down from a given number to zero,\n", - "and produces a stream of messages with the current count." - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "3...\n", - "2...\n", - "1...\n", - "Done!\n" - ] - } - ], - "source": [ - "from typing import AsyncGenerator, List, Sequence\n", - "\n", - "from autogen_agentchat.agents import BaseChatAgent\n", - "from autogen_agentchat.base import Response\n", - "from autogen_agentchat.messages import AgentEvent, ChatMessage, TextMessage\n", - "from autogen_core import CancellationToken\n", - "\n", - "\n", - "class CountDownAgent(BaseChatAgent):\n", - " def __init__(self, name: str, count: int = 3):\n", - " super().__init__(name, \"A simple agent that counts down.\")\n", - " self._count = count\n", - "\n", - " @property\n", - " def produced_message_types(self) -> Sequence[type[ChatMessage]]:\n", - " return (TextMessage,)\n", - "\n", - " async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:\n", - " # Calls the on_messages_stream.\n", - " response: Response | None = None\n", - " async for message in self.on_messages_stream(messages, cancellation_token):\n", - " if isinstance(message, Response):\n", - " response = message\n", - " assert response is not None\n", - " return response\n", - "\n", - " async def on_messages_stream(\n", - " self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken\n", - " ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]:\n", - " inner_messages: List[AgentEvent | ChatMessage] = []\n", - " for i in range(self._count, 0, -1):\n", - " msg = TextMessage(content=f\"{i}...\", source=self.name)\n", - " inner_messages.append(msg)\n", - " yield msg\n", - " # The response is returned at the end of the stream.\n", - " # It contains the final message and all the inner messages.\n", - " yield Response(chat_message=TextMessage(content=\"Done!\", source=self.name), inner_messages=inner_messages)\n", - "\n", - " async def on_reset(self, cancellation_token: CancellationToken) -> None:\n", - " pass\n", - "\n", - "\n", - "async def run_countdown_agent() -> None:\n", - " # Create a countdown agent.\n", - " countdown_agent = CountDownAgent(\"countdown\")\n", - "\n", - " # Run the agent with a given task and stream the response.\n", - " async for message in countdown_agent.on_messages_stream([], CancellationToken()):\n", - " if isinstance(message, Response):\n", - " print(message.chat_message.content)\n", - " else:\n", - " print(message.content)\n", - "\n", - "\n", - "# Use asyncio.run(run_countdown_agent()) when running in a script.\n", - "await run_countdown_agent()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## ArithmeticAgent\n", - "\n", - "In this example, we create an agent class that can perform simple arithmetic operations\n", - "on a given integer. Then, we will use different instances of this agent class\n", - "in a {py:class}`~autogen_agentchat.teams.SelectorGroupChat`\n", - "to transform a given integer into another integer by applying a sequence of arithmetic operations.\n", - "\n", - "The `ArithmeticAgent` class takes an `operator_func` that takes an integer and returns an integer,\n", - "after applying an arithmetic operation to the integer.\n", - "In its `on_messages` method, it applies the `operator_func` to the integer in the input message,\n", - "and returns a response with the result." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from typing import Callable, Sequence\n", - "\n", - "from autogen_agentchat.agents import BaseChatAgent\n", - "from autogen_agentchat.base import Response\n", - "from autogen_agentchat.conditions import MaxMessageTermination\n", - "from autogen_agentchat.messages import ChatMessage\n", - "from autogen_agentchat.teams import SelectorGroupChat\n", - "from autogen_agentchat.ui import Console\n", - "from autogen_core import CancellationToken\n", - "from autogen_ext.models.openai import OpenAIChatCompletionClient\n", - "\n", - "\n", - "class ArithmeticAgent(BaseChatAgent):\n", - " def __init__(self, name: str, description: str, operator_func: Callable[[int], int]) -> None:\n", - " super().__init__(name, description=description)\n", - " self._operator_func = operator_func\n", - " self._message_history: List[ChatMessage] = []\n", - "\n", - " @property\n", - " def produced_message_types(self) -> Sequence[type[ChatMessage]]:\n", - " return (TextMessage,)\n", - "\n", - " async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:\n", - " # Update the message history.\n", - " # NOTE: it is possible the messages is an empty list, which means the agent was selected previously.\n", - " self._message_history.extend(messages)\n", - " # Parse the number in the last message.\n", - " assert isinstance(self._message_history[-1], TextMessage)\n", - " number = int(self._message_history[-1].content)\n", - " # Apply the operator function to the number.\n", - " result = self._operator_func(number)\n", - " # Create a new message with the result.\n", - " response_message = TextMessage(content=str(result), source=self.name)\n", - " # Update the message history.\n", - " self._message_history.append(response_message)\n", - " # Return the response.\n", - " return Response(chat_message=response_message)\n", - "\n", - " async def on_reset(self, cancellation_token: CancellationToken) -> None:\n", - " pass" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{note}\n", - "The `on_messages` method may be called with an empty list of messages, in which\n", - "case it means the agent was called previously and is now being called again,\n", - "without any new messages from the caller. So it is important to keep a history\n", - "of the previous messages received by the agent, and use that history to generate\n", - "the response.\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now we can create a {py:class}`~autogen_agentchat.teams.SelectorGroupChat` with 5 instances of `ArithmeticAgent`:\n", - "\n", - "- one that adds 1 to the input integer,\n", - "- one that subtracts 1 from the input integer,\n", - "- one that multiplies the input integer by 2,\n", - "- one that divides the input integer by 2 and rounds down to the nearest integer, and\n", - "- one that returns the input integer unchanged.\n", - "\n", - "We then create a {py:class}`~autogen_agentchat.teams.SelectorGroupChat` with these agents,\n", - "and set the appropriate selector settings:\n", - "\n", - "- allow the same agent to be selected consecutively to allow for repeated operations, and\n", - "- customize the selector prompt to tailor the model's response to the specific task." - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "---------- user ----------\n", - "Apply the operations to turn the given number into 25.\n", - "---------- user ----------\n", - "10\n", - "---------- multiply_agent ----------\n", - "20\n", - "---------- add_agent ----------\n", - "21\n", - "---------- multiply_agent ----------\n", - "42\n", - "---------- divide_agent ----------\n", - "21\n", - "---------- add_agent ----------\n", - "22\n", - "---------- add_agent ----------\n", - "23\n", - "---------- add_agent ----------\n", - "24\n", - "---------- add_agent ----------\n", - "25\n", - "---------- Summary ----------\n", - "Number of messages: 10\n", - "Finish reason: Maximum number of messages 10 reached, current message count: 10\n", - "Total prompt tokens: 0\n", - "Total completion tokens: 0\n", - "Duration: 2.40 seconds\n" - ] - } - ], - "source": [ - "async def run_number_agents() -> None:\n", - " # Create agents for number operations.\n", - " add_agent = ArithmeticAgent(\"add_agent\", \"Adds 1 to the number.\", lambda x: x + 1)\n", - " multiply_agent = ArithmeticAgent(\"multiply_agent\", \"Multiplies the number by 2.\", lambda x: x * 2)\n", - " subtract_agent = ArithmeticAgent(\"subtract_agent\", \"Subtracts 1 from the number.\", lambda x: x - 1)\n", - " divide_agent = ArithmeticAgent(\"divide_agent\", \"Divides the number by 2 and rounds down.\", lambda x: x // 2)\n", - " identity_agent = ArithmeticAgent(\"identity_agent\", \"Returns the number as is.\", lambda x: x)\n", - "\n", - " # The termination condition is to stop after 10 messages.\n", - " termination_condition = MaxMessageTermination(10)\n", - "\n", - " # Create a selector group chat.\n", - " selector_group_chat = SelectorGroupChat(\n", - " [add_agent, multiply_agent, subtract_agent, divide_agent, identity_agent],\n", - " model_client=OpenAIChatCompletionClient(model=\"gpt-4o\"),\n", - " termination_condition=termination_condition,\n", - " allow_repeated_speaker=True, # Allow the same agent to speak multiple times, necessary for this task.\n", - " selector_prompt=(\n", - " \"Available roles:\\n{roles}\\nTheir job descriptions:\\n{participants}\\n\"\n", - " \"Current conversation history:\\n{history}\\n\"\n", - " \"Please select the most appropriate role for the next message, and only return the role name.\"\n", - " ),\n", - " )\n", - "\n", - " # Run the selector group chat with a given task and stream the response.\n", - " task: List[ChatMessage] = [\n", - " TextMessage(content=\"Apply the operations to turn the given number into 25.\", source=\"user\"),\n", - " TextMessage(content=\"10\", source=\"user\"),\n", - " ]\n", - " stream = selector_group_chat.run_stream(task=task)\n", - " await Console(stream)\n", - "\n", - "\n", - "# Use asyncio.run(run_number_agents()) when running in a script.\n", - "await run_number_agents()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "From the output, we can see that the agents have successfully transformed the input integer\n", - "from 10 to 25 by choosing appropriate agents that apply the arithmetic operations in sequence." - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": ".venv", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.11.5" - } - }, - "nbformat": 4, - "nbformat_minor": 2 -} diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/messages.ipynb b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/messages.ipynb index 25dc78641980..a3f5a25519dd 100644 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/messages.ipynb +++ b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/messages.ipynb @@ -97,7 +97,7 @@ "\n", "Examples of these include {py:class}`~autogen_agentchat.messages.ToolCallRequestEvent`, which indicates that a request was made to call a tool, and {py:class}`~autogen_agentchat.messages.ToolCallExecutionEvent`, which contains the results of tool calls.\n", "\n", - "Typically, events are created by the agent itself and are contained in the {py:attr}`~autogen_agentchat.base.Response.inner_messages` field of the {py:class}`~autogen_agentchat.base.Response` returned from {py:class}`~autogen_agentchat.base.ChatAgent.on_messages`. If you are building a custom agent and have events that you want to communicate to other entities (e.g., a UI), you can include these in the {py:attr}`~autogen_agentchat.base.Response.inner_messages` field of the {py:class}`~autogen_agentchat.base.Response`. We will show examples of this in [Custom Agents](./custom-agents.ipynb).\n", + "Typically, events are created by the agent itself and are contained in the {py:attr}`~autogen_agentchat.base.Response.inner_messages` field of the {py:class}`~autogen_agentchat.base.Response` returned from {py:class}`~autogen_agentchat.base.ChatAgent.on_messages`. If you are building a custom agent and have events that you want to communicate to other entities (e.g., a UI), you can include these in the {py:attr}`~autogen_agentchat.base.Response.inner_messages` field of the {py:class}`~autogen_agentchat.base.Response`. We will show examples of this in [Custom Agents](../custom-agents.ipynb).\n", "\n", "\n", "You can read about the full set of messages supported in AgentChat in the {py:mod}`~autogen_agentchat.messages` module. \n", @@ -107,7 +107,7 @@ ], "metadata": { "kernelspec": { - "display_name": ".venv", + "display_name": "agnext", "language": "python", "name": "python3" }, @@ -121,7 +121,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.6" + "version": "3.11.9" } }, "nbformat": 4, diff --git a/python/packages/autogen-ext/pyproject.toml b/python/packages/autogen-ext/pyproject.toml index bf86c3e4004e..c5e2b62bb627 100644 --- a/python/packages/autogen-ext/pyproject.toml +++ b/python/packages/autogen-ext/pyproject.toml @@ -70,6 +70,10 @@ semantic-kernel-core = [ "semantic-kernel>=1.17.1", ] +gemini = [ + "google-genai>=1.0.0", +] + semantic-kernel-google = [ "semantic-kernel[google]>=1.17.1", ] diff --git a/python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py b/python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py index 05924e186643..1ac058a9680f 100644 --- a/python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py +++ b/python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py @@ -12,7 +12,7 @@ from hashlib import sha256 from pathlib import Path from types import TracebackType -from typing import Any, Callable, ClassVar, List, Optional, ParamSpec, Type, Union +from typing import Any, Callable, ClassVar, Dict, List, Optional, ParamSpec, Type, Union from autogen_core import CancellationToken from autogen_core.code_executor import ( @@ -88,6 +88,13 @@ class DockerCommandLineCodeExecutor(CodeExecutor): the Python process exits with atext. Defaults to True. functions (List[Union[FunctionWithRequirements[Any, A], Callable[..., Any]]]): A list of functions that are available to the code executor. Default is an empty list. functions_module (str, optional): The name of the module that will be created to store the functions. Defaults to "functions". + extra_volumes (Optional[Dict[str, Dict[str, str]]], optional): A dictionary of extra volumes (beyond the work_dir) to mount to the container; + key is host source path and value 'bind' is the container path. See Defaults to None. + Example: extra_volumes = {'/home/user1/': {'bind': '/mnt/vol2', 'mode': 'rw'}, '/var/www': {'bind': '/mnt/vol1', 'mode': 'ro'}} + extra_hosts (Optional[Dict[str, str]], optional): A dictionary of host mappings to add to the container. (See Docker docs on extra_hosts) Defaults to None. + Example: extra_hosts = {"kubernetes.docker.internal": "host-gateway"} + init_command (Optional[str], optional): A shell command to run before each shell operation execution. Defaults to None. + Example: init_command="kubectl config use-context docker-hub" """ SUPPORTED_LANGUAGES: ClassVar[List[str]] = [ @@ -126,6 +133,9 @@ def __init__( ] ] = [], functions_module: str = "functions", + extra_volumes: Optional[Dict[str, Dict[str, str]]] = None, + extra_hosts: Optional[Dict[str, str]] = None, + init_command: Optional[str] = None, ): if timeout < 1: raise ValueError("Timeout must be greater than or equal to 1.") @@ -157,6 +167,10 @@ def __init__( self._functions_module = functions_module self._functions = functions + self._extra_volumes = extra_volumes if extra_volumes is not None else {} + self._extra_hosts = extra_hosts if extra_hosts is not None else {} + self._init_command = init_command + # Setup could take some time so we intentionally wait for the first code block to do it. if len(functions) > 0: self._setup_functions_complete = False @@ -354,16 +368,22 @@ async def start(self) -> None: # Let the docker exception escape if this fails. await asyncio.to_thread(client.images.pull, self._image) + # Prepare the command (if needed) + shell_command = "/bin/sh" + command = ["-c", f"{(self._init_command)};exec {shell_command}"] if self._init_command else None + self._container = await asyncio.to_thread( client.containers.create, self._image, name=self.container_name, - entrypoint="/bin/sh", + entrypoint=shell_command, + command=command, tty=True, detach=True, auto_remove=self._auto_remove, - volumes={str(self._bind_dir.resolve()): {"bind": "/workspace", "mode": "rw"}}, + volumes={str(self._bind_dir.resolve()): {"bind": "/workspace", "mode": "rw"}, **self._extra_volumes}, working_dir="/workspace", + extra_hosts=self._extra_hosts, ) await asyncio.to_thread(self._container.start) diff --git a/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime_host_servicer.py b/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime_host_servicer.py index daa4ad65101d..621ed9511eb4 100644 --- a/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime_host_servicer.py +++ b/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime_host_servicer.py @@ -7,6 +7,7 @@ from typing import Any, AsyncIterator, Awaitable, Callable, Dict, Generic, Sequence, Set, Tuple, TypeVar from autogen_core import TopicId +from autogen_core._agent_id import AgentId from autogen_core._runtime_impl_helpers import SubscriptionManager from ._constants import GRPC_IMPORT_ERROR_STR @@ -100,6 +101,9 @@ def __init__(self) -> None: self._data_connections: Dict[ ClientConnectionId, ChannelConnection[agent_worker_pb2.Message, agent_worker_pb2.Message] ] = {} + self._control_connections: Dict[ + ClientConnectionId, ChannelConnection[agent_worker_pb2.ControlMessage, agent_worker_pb2.ControlMessage] + ] = {} self._agent_type_to_client_id_lock = asyncio.Lock() self._agent_type_to_client_id: Dict[str, ClientConnectionId] = {} self._pending_responses: Dict[ClientConnectionId, Dict[str, Future[Any]]] = {} @@ -140,7 +144,23 @@ async def OpenControlChannel( # type: ignore request_iterator: AsyncIterator[agent_worker_pb2.ControlMessage], context: grpc.aio.ServicerContext[agent_worker_pb2.ControlMessage, agent_worker_pb2.ControlMessage], ) -> AsyncIterator[agent_worker_pb2.ControlMessage]: - raise NotImplementedError("Method not implemented.") + client_id = await get_client_id_or_abort(context) + + async def handle_callback(message: agent_worker_pb2.ControlMessage) -> None: + await self._receive_control_message(client_id, message) + + connection = CallbackChannelConnection[agent_worker_pb2.ControlMessage, agent_worker_pb2.ControlMessage]( + request_iterator, client_id, handle_callback=handle_callback + ) + self._control_connections[client_id] = connection + logger.info(f"Client {client_id} connected.") + + try: + async for message in connection: + yield message + finally: + # Clean up the client connection. + del self._control_connections[client_id] async def _on_client_disconnect(self, client_id: ClientConnectionId) -> None: async with self._agent_type_to_client_id_lock: @@ -182,6 +202,29 @@ async def _receive_message(self, client_id: ClientConnectionId, message: agent_w case None: logger.warning("Received empty message") + async def _receive_control_message( + self, client_id: ClientConnectionId, message: agent_worker_pb2.ControlMessage + ) -> None: + logger.info(f"Received message from client {client_id}: {message}") + destination = message.destination + if destination.startswith("agentid="): + agent_id = AgentId.from_str(destination[len("agentid=") :]) + target_client_id = self._agent_type_to_client_id.get(agent_id.type) + if target_client_id is None: + logger.error(f"Agent client id not found for agent type {agent_id.type}.") + return + elif destination.startswith("clientid="): + target_client_id = destination[len("clientid=") :] + else: + logger.error(f"Invalid destination {destination}") + return + + target_send_queue = self._control_connections.get(target_client_id) + if target_send_queue is None: + logger.error(f"Client {target_client_id} not found, failed to deliver message.") + return + await target_send_queue.send(message) + async def _process_request(self, request: agent_worker_pb2.RpcRequest, client_id: ClientConnectionId) -> None: # Deliver the message to a client given the target agent type. async with self._agent_type_to_client_id_lock: diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_factory.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_factory.py index fcb29e3c4a99..3eb8634b3698 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_factory.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_factory.py @@ -12,6 +12,14 @@ async def mcp_server_tools( This factory function connects to an MCP server and returns adapters for all available tools. The adapters can be directly assigned to an AutoGen agent's tools list. + .. note:: + + To use this function, you need to install `mcp` extra for the `autogen-ext` package. + + .. code-block:: bash + + pip install -U "autogen-ext[mcp]" + Args: server_params (McpServerParams): Connection parameters for the MCP server. Can be either StdioServerParams for command-line tools or diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_sse.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_sse.py index 7b875649a4f4..252af7ce50da 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_sse.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_sse.py @@ -25,6 +25,14 @@ class SseMcpToolAdapter( with AutoGen agents. Common use cases include integrating with remote MCP services, cloud-based tools, and web APIs that implement the Model Context Protocol (MCP). + .. note:: + + To use this class, you need to install `mcp` extra for the `autogen-ext` package. + + .. code-block:: bash + + pip install -U "autogen-ext[mcp]" + Args: server_params (SseServerParameters): Parameters for the MCP server connection, including URL, headers, and timeouts diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_stdio.py b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_stdio.py index 629b1a46ec1a..4f827785e903 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_stdio.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/mcp/_stdio.py @@ -24,6 +24,15 @@ class StdioMcpToolAdapter( with AutoGen agents. Common use cases include wrapping command-line tools and local services that implement the Model Context Protocol (MCP). + .. note:: + + To use this class, you need to install `mcp` extra for the `autogen-ext` package. + + .. code-block:: bash + + pip install -U "autogen-ext[mcp]" + + Args: server_params (StdioServerParams): Parameters for the MCP server connection, including command to run and its arguments diff --git a/python/packages/autogen-ext/src/autogen_ext/tools/semantic_kernel/_kernel_function_from_tool.py b/python/packages/autogen-ext/src/autogen_ext/tools/semantic_kernel/_kernel_function_from_tool.py index 8125aef9c5ce..0f3f93e2e7fa 100644 --- a/python/packages/autogen-ext/src/autogen_ext/tools/semantic_kernel/_kernel_function_from_tool.py +++ b/python/packages/autogen-ext/src/autogen_ext/tools/semantic_kernel/_kernel_function_from_tool.py @@ -3,8 +3,8 @@ from autogen_core import CancellationToken from autogen_core.tools import BaseTool from pydantic import BaseModel -from semantic_kernel.functions.kernel_parameter_metadata import KernelParameterMetadata from semantic_kernel.functions import KernelFunctionFromMethod, kernel_function +from semantic_kernel.functions.kernel_parameter_metadata import KernelParameterMetadata InputT = TypeVar("InputT", bound=BaseModel) OutputT = TypeVar("OutputT", bound=BaseModel) diff --git a/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py b/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py index ad74236008fa..6c65835d183d 100644 --- a/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py +++ b/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py @@ -164,3 +164,48 @@ async def test_docker_commandline_code_executor_start_stop_context_manager() -> with tempfile.TemporaryDirectory() as temp_dir: async with DockerCommandLineCodeExecutor(work_dir=temp_dir) as _exec: pass + + +@pytest.mark.asyncio +async def test_docker_commandline_code_executor_extra_args() -> None: + if not docker_tests_enabled(): + pytest.skip("Docker tests are disabled") + + with tempfile.TemporaryDirectory() as temp_dir: + # Create a file in temp_dir to mount + host_file_path = Path(temp_dir) / "host_file.txt" + host_file_path.write_text("This is a test file.") + + container_file_path = "/container/host_file.txt" + + extra_volumes = {str(host_file_path): {"bind": container_file_path, "mode": "rw"}} + init_command = "echo 'Initialization command executed' > /workspace/init_command.txt" + extra_hosts = {"example.com": "127.0.0.1"} + + async with DockerCommandLineCodeExecutor( + work_dir=temp_dir, + extra_volumes=extra_volumes, + init_command=init_command, + extra_hosts=extra_hosts, + ) as executor: + cancellation_token = CancellationToken() + + # Verify init_command was executed + init_command_file_path = Path(temp_dir) / "init_command.txt" + assert init_command_file_path.exists() + + # Verify extra_hosts + ns_lookup_code_blocks = [ + CodeBlock(code="import socket; print(socket.gethostbyname('example.com'))", language="python") + ] + ns_lookup_result = await executor.execute_code_blocks(ns_lookup_code_blocks, cancellation_token) + assert ns_lookup_result.exit_code == 0 + assert "127.0.0.1" in ns_lookup_result.output + + # Verify the file is accessible in the volume mounted in extra_volumes + code_blocks = [ + CodeBlock(code=f"with open('{container_file_path}') as f: print(f.read())", language="python") + ] + code_result = await executor.execute_code_blocks(code_blocks, cancellation_token) + assert code_result.exit_code == 0 + assert "This is a test file." in code_result.output diff --git a/python/uv.lock b/python/uv.lock index accecffb7166..8bc22e84a535 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -574,6 +574,9 @@ file-surfer = [ { name = "autogen-agentchat" }, { name = "markitdown" }, ] +gemini = [ + { name = "google-genai" }, +] graphrag = [ { name = "graphrag" }, ] @@ -681,6 +684,7 @@ requires-dist = [ { name = "diskcache", marker = "extra == 'diskcache'", specifier = ">=5.6.3" }, { name = "docker", marker = "extra == 'docker'", specifier = "~=7.0" }, { name = "ffmpeg-python", marker = "extra == 'video-surfer'" }, + { name = "google-genai", marker = "extra == 'gemini'", specifier = ">=1.0.0" }, { name = "graphrag", marker = "extra == 'graphrag'", specifier = ">=1.0.1" }, { name = "grpcio", marker = "extra == 'grpc'", specifier = "~=1.70.0" }, { name = "httpx", marker = "extra == 'http-tool'", specifier = ">=0.27.0" }, @@ -2199,6 +2203,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/14/fb/54deefe679b7d1c1cc81d83396fcf28ad1a66d213bddeb275a8d28665918/google_crc32c-1.6.0-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:18e311c64008f1f1379158158bb3f0c8d72635b9eb4f9545f8cf990c5668e59d", size = 27866 }, ] +[[package]] +name = "google-genai" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "google-auth" }, + { name = "pydantic" }, + { name = "requests" }, + { name = "websockets" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/6b/30/25443d2dec5fe4efd9f05440e13735ae68b19de2dea043b3b83a91a4e14b/google_genai-1.1.0-py3-none-any.whl", hash = "sha256:c48ac44612ad6aadc0bf96b12fa4314756baa16382c890fff793bcb53e9a9cc8", size = 130299 }, +] + [[package]] name = "google-generativeai" version = "0.8.4"