# Streaming input and output using WebSockets

## **TL;DR**  
- Learn how to build an agent chat application using [WebSockets](https://fastapi.tiangolo.com/advanced/websockets/) and [`IOStream`](https://docs.ag2.ai/0.8.5/docs/api-reference/autogen/io/IOStream)  
- Explore a hands-on example of connecting a web application to a responsive chat with agents over [WebSockets](https://fastapi.tiangolo.com/advanced/websockets/).  
- **Streamlined Real-Time Interactions**: [WebSockets](https://fastapi.tiangolo.com/advanced/websockets/) offer a low-latency, persistent connection for sending and receiving data in real time.

---

## **Real-Time Applications: Why WebSockets?**  
[WebSockets](https://fastapi.tiangolo.com/advanced/websockets/) provide a powerful framework for real-time communication between a client and server. Unlike traditional HTTP requests, which require polling for updates, [WebSockets](https://fastapi.tiangolo.com/advanced/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](https://fastapi.tiangolo.com/advanced/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](https://fastapi.tiangolo.com/advanced/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](https://fastapi.tiangolo.com/advanced/websockets/) to the browser where they can be displayed.

## **Example: Creating a Weather Chat App**  
Let’s walk through an example that integrates [WebSockets](https://fastapi.tiangolo.com/advanced/websockets/) with a weather-focused chat.

### **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/](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`](https://github.com/ag2ai/agentchat-over-websockets/blob/main/agentchat-over-websockets/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.
```python
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`**](https://docs.ag2.ai/0.8.5/docs/api-reference/autogen/ConversableAgent#autogen.ConversableAgent.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`.
```python
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.
```python
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("",