Authors: Mark Sze, Tvrtko Sternak, Davor Runje

TL;DR

• Learn how to build an agent chat application using WebSockets and IOStream

• Explore a hands-on example of connecting a web application to a responsive chat with agents over WebSockets.

• Streamlined Real-Time Interactions: WebSockets offer a low-latency, persistent connection for sending and receiving data in real time.

​Real-Time Applications: Why WebSockets?

WebSockets provide a powerful framework for real-time communication between a client and server. Unlike traditional HTTP requests, which require polling for updates, WebSockets establish a persistent, full-duplex connection that allows for continuous data exchange.

This capability is critical for applications that use AG2, where seamless interaction is essential.

​Key Benefits of WebSockets

1. Low Latency: WebSockets reduce latency by maintaining a direct, open connection, avoiding the overhead of repeated HTTP handshakes.

2. Efficient Data Streaming: Continuous, two-way data streams enable smooth user experiences in real-time applications.

3. Event-Driven Communication: With WebSocket protocols, the server can “push” updates to the client as events occur.

4. Simplified Architecture: WebSockets eliminate the need for separate polling mechanisms, reducing server load and complexity.

​Building a Chat System

This example demonstrates how to create a WebSocket-based chat system that streams real-time input and output from AG2 Agents.

​How It Works

1. WebSocket Connection: The client establishes a persistent WebSocket connection to the server.

2. Real-Time Data Flow: Events in the conversation are streamed over WebSockets to the browser where they can be displayed.

​Example: Creating a Weather chat app

Let’s walk through an example that integrates WebSockets with a weather-focused chat.

You can explore the full example code here.

​1. Clone the Repository

git clone https://github.com/ag2ai/agentchat-over-websockets.git
cd agentchat-over-websockets

​2. Set Up Environment Variables

Create a OAI_CONFIG_LIST file based on the provided OAI_CONFIG_LIST_sample:

cp OAI_CONFIG_LIST_sample OAI_CONFIG_LIST

In the OAI_CONFIG_LIST file, update the api_key to your OpenAI API key.

​(Optional) Create and use a virtual environment

To reduce cluttering your global Python environment on your machine, you can create a virtual environment. On your command line, enter:

python3 -m venv env
source env/bin/activate

​3. Install Dependencies

Install the required Python packages using pip:

pip install -r requirements.txt

​4. Start the Server

Run the main.py file:

python agentchat-over-websockets/main.py

​Test the App

With the server running, open the client application in your browser by navigating to http://localhost:8001/. And send a message to the chat and watch the conversation between agents roll out in your browser.

​Code review

​Backend Code: main.py

The backend is responsible for serving the frontend, managing WebSocket connections, and hosting the AI-powered conversational agent. Below is a step-by-step breakdown.

​Setting Up the WebSocket Server

The IOWebsockets.run_server_in_thread utility is used to run a WebSocket server. The on_connect function handles new client connections and initializes the chatbot.

from autogen.io.websockets import IOWebsockets
from datetime import datetime

def on_connect(iostream: IOWebsockets) -> None:
    print(f"Connected to client: {iostream}")
    initial_msg = iostream.input()  # Receive the first message from the client.
    print(f"Initial message: {initial_msg}")

    # Define the agent
    agent = autogen.ConversableAgent(
        name="chatbot",
        system_message="Complete tasks and reply TERMINATE when done. Use the 'weather_forecast' tool for weather-related queries.",
        llm_config={"stream": False},
    )

    # Define the user proxy
    user_proxy = autogen.UserProxyAgent(
        name="user_proxy",
        system_message="A proxy for the user.",
        is_termination_msg=lambda msg: msg.get("content", "").endswith("TERMINATE"),
        human_input_mode="NEVER",
    )

    # Register tool functions
    def weather_forecast(city: str) -> str:
        return f"The weather forecast for {city} is sunny as of {datetime.now()}."

    autogen.register_function(
        weather_forecast,
        caller=agent,
        executor=user_proxy,
        description="Provides a mock weather forecast.",
    )

    # Initiate conversation
    user_proxy.initiate_chat(agent, message=initial_msg)

Explanation:

1. on_connect: Handles client connections and manages the interaction between the ConversableAgent and the client.

2. Tool Registration: The weather_forecast function provides a mock weather report and is linked to the agent for handling weather-related queries.​Serving the Frontend

The SimpleHTTPRequestHandler is used to serve HTML files. A custom handler class overrides the behavior for the root path to serve chat.html.

class MyRequestHandler(SimpleHTTPRequestHandler):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, directory=Path(__file__).parent / "website_files" / "templates", **kwargs)

    def do_GET(self):
        if self.path == "/":
            self.path = "/chat.html"
        return super().do_GET()

Explanation:

• The MyRequestHandler class ensures that the default page served is chat.html.

• Files are served from the website_files/templates directory.

​Running the Servers

Finally, both the WebSocket and HTTP servers are started.

from http.server import HTTPServer

PORT = 8001

handler = MyRequestHandler

# Start WebSocket server
with IOWebsockets.run_server_in_thread(on_connect=on_connect, port=8080) as uri:
    print(f"WebSocket server started at {uri}")

    # Start HTTP server
    with HTTPServer(("", PORT), handler) as httpd:
        print(f"HTTP server started at http://localhost:{PORT}")
        try:
            httpd.serve_forever()
        except KeyboardInterrupt:
            print("HTTP server stopped.")

Explanation:

• The WebSocket server listens on port 8080, while the HTTP server listens on port 8001.

• The WebSocket server handles real-time communication, while the HTTP server serves static files.

​Frontend Code: chat.html

The frontend provides a simple interface for users to interact with the chatbot.

HTML Structure

The HTML structure defines an input form for sending messages and a list for displaying them.

    
    


    
AI Chat Interface
     onsubmit="sendMessage(event)">
         type="text" id="messageText" autocomplete="off" />
        Send
    
     id="messages">

​JavaScript Logic

The JavaScript code establishes a WebSocket connection, handles incoming messages, and sends user input to the backend.

var ws = new WebSocket("ws://localhost:8080");

ws.onmessage = function(event) {
    var messages = document.getElementById('messages');
    var message = document.createElement('li');
    message.textContent = event.data;  // Display the message content.
    messages.appendChild(message);
};

function sendMessage(event) {
    var input = document.getElementById("messageText");
    ws.send(input.value);  // Send the input value to the backend.
    input.value = '';  // Clear the input field.
    event.preventDefault();  // Prevent form submission.
}

Explanation:

1. WebSocket Initialization: Connects to the WebSocket server at ws://localhost:8080.

2. Message Display: Appends incoming messages to the #messages list.

3. Sending Messages: Captures user input, sends it to the server, and clears the input field.

​Conclusion

Building an AgentChat system with WebSockets unlocks the potential for real-time, interactive applications. By maintaining a persistent connection, WebSockets enable seamless communication, enhancing user experience with minimal latency.

The example of a weather chatbot demonstrates the ease of integrating AG2 with WebSockets to create dynamic conversational agents. Whether for customer support, virtual assistants, or personalized services, this architecture provides a robust foundation for building next-generation applications.

Ready to start building? Explore the full example code here.

Finding this useful?

The AG2 team is working hard to create content like this, not to mention building a powerful, open-source, end-to-end platform for multi-agent automation.

The easiest way to show your support is just to star AG2 repo, but also take a look at it for contributions or simply to give it a try.

Also, let us know if you have any interesting use cases for AgentChat with WebSockets? Or maybe you would like to see more features or improvements? Do join our Discord server for discussion.