Streaming input and output using WebSockets - AG2

Streaming input and output using WebSockets

TL;DR

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.

Note

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:

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("",