Add interactive WebSocket UI and FastAPI backend for Groupchat

notebook/agentchat_fastapi_websocket/README.md

@@ -0,0 +1,202 @@
+# 🔌 AG2AI WebSocket Agentic Workflow UI & Backend
+
+An interactive WebSocket-based UI and FastAPI backend for orchestrating agentic workflows using [AG2AI Autogen](https://github.com/ag2ai/ag2). This project demonstrates real-time, multi-agent collaboration for solving problems through a WebSocket-powered interface — ideal for tasks like data analysis, EDA, and more.
+
+![App UI Screenshot](./docs/screenshot.png)
+
+---
+
+## 📘 How to Use the WebSocket UI
+
+To learn how to interact with the UI step-by-step, check out the full guide:
+
+➡️ [Usage Guide (UI Walkthrough)](./docs/USAGE_GUIDE.md)
+
+
+---
+
+## 📦 Overview
+
+This repository includes:
+
+- A clean **frontend UI** to interact with WebSockets visually — more intuitive than Postman or raw clients.
+- A **FastAPI backend** that manages real-time WebSocket communication and coordinates multiple agents via AG2AI Autogen.
+- A custom **orchestrator agent** (`agent_aligner`) that manages execution flow, ensuring orderly agent coordination.
+
+---
+
+## ✨ Features
+
+### ✅ Frontend
+
+- Interactive WebSocket client with formatted JSON display
+- Message blocks styled for clarity and separation
+- UUID-based client tracking
+- Send/receive messages with live updates
+- Manual message construction and quick templates
+
+### ✅ Backend
+
+- Built with **FastAPI** and **async WebSocket** handling
+- Modular architecture using manager classes
+- Custom `AgentChat` class for group chat orchestration
+- Real-time message streaming to frontend
+- Manual user input integration during live chat
+- Environment-based configuration via `.env`
+
+---
+
+## 🧩 Problem & Solution
+
+### ❌ The Problem We Faced
+
+While working with WebSocket-based agent systems using AG2AI/Autogen, we encountered several major bottlenecks that affected productivity and developer experience:
+
+- **Postman and raw WebSocket clients are not interactive**
+  These tools make it hard to follow multi-agent conversations. They lack formatting, which slows down debugging and understanding the data flow.
+
+- **Reading agent messages is time-consuming**
+  When working with multiple agents, reviewing each step (especially during prompt tuning or alignment) becomes tedious and error-prone.
+
+- **Lack of message formatting**
+  JSON responses from agents are dumped as raw strings, making them hard to read and troubleshoot — especially when nested or streamed.
+
+- **Frontend development was not feasible**
+  Building a fully custom UI in frameworks like React or Vue would add significant overhead and distract from core system development.
+
+- **No streamlined session management**
+  Keeping track of WebSocket sessions and switching between different chats was a manual and error-prone task.
+
+---
+
+### ✅ The Solution We Implemented
+
+To overcome these challenges, we built a minimal yet powerful **interactive WebSocket UI**, paired with a **FastAPI backend**, enabling seamless development and debugging of agent workflows.
+
+Key benefits:
+
+- **Clean, interactive WebSocket communication**
+  Live messages stream directly to the browser with proper formatting and role-based separation.
+
+- **Well-structured message display**
+  All messages are styled in blocks and automatically formatted as JSON, making it easy to inspect agent responses.
+
+- **Faster prompt tuning & agent alignment**
+  Developers can instantly see how agents respond, helping fine-tune prompts with clarity and speed.
+
+- **Quick session switching**
+  Chat sessions can be created and reused easily, improving workflow efficiency during development and testing.
+
+- **Minimal development effort**
+  A lightweight HTML/JS UI replaces the need for building a full-fledged frontend framework — saving time while still improving UX significantly.
+
+---
+
+This setup dramatically reduced the friction in debugging, testing, and managing agentic workflows — allowing us to focus on what matters: building smart and responsive agents.
+
+---
+
+## 🤖 Agents Overview
+
+This system supports the following agents for structured task completion:
+
+- **`planner_agent`**: Produces a step-by-step execution plan (no code).
+- **`code_writer`**: Converts the plan into working `python` code.
+- **`code_executor`**: Executes code in the local runtime environment.
+- **`debugger`**: Detects and resolves runtime errors; retries code.
+- **`process_completion`**: Summarizes results and guides the next steps.
+- `agent_aligner`: Coordinates the overall workflow, ensuring agents operate in the correct sequence. Enforces the execution flow (Plan → Confirm → Write → Confirm → Execute) to maintain structure, avoid loops, and ensure safe progression.
+
+
+---
+
+## 📁 Project Structure
+
+```
+
+.
+├── managers/
+│   ├── cancellation_token.py     # Manages cancellation signals
+│   ├── connection.py             # Handles socket connections
+│   ├── groupchat.py              # AgentChat orchestration logic
+│   ├── prompts.py                # Prompt templates and roles
+│
+├── templates/
+│   └── index.html                # Frontend UI (WebSocket client)
+│
+├── .env                          # API keys and environment variables
+├── dependencies.py               # AG2AI agent setup and configuration
+├── helpers.py                    # Utility functions
+├── main.py                       # FastAPI server entry point
+├── ws.py                         # WebSocket route handler
+├── requirements.txt              # Python dependencies
+└── Readme.md                     # You're reading it!
+
+````
+
+---
+
+## 🚀 Getting Started
+
+### 1. Clone the Repository
+
+```bash
+git clone https://github.com/Suryaaa-Rathore/websocket-ag2ai.git
+cd websocket-ag2ai
+````
+
+### 2. Create and Activate Virtual Environment
+
+```bash
+python -m venv venv
+source venv/bin/activate  # On Windows: venv\Scripts\activate
+```
+
+### 3. Install Dependencies
+
+```bash
+pip install -r requirements.txt
+```
+
+### 4. Set Your OpenAI API Key
+
+Create a `.env` file in the root directory:
+
+```
+OPENAI_API_KEY=your-openai-api-key
+```
+
+Or export the variable directly:
+
+```bash
+export OPENAI_API_KEY=your-openai-api-key
+```
+
+### 5. Run the Server
+
+```bash
+python main.py
+```
+
+### 6. Access the Frontend
+
+Open your browser and go to:
+
+```
+http://localhost:8000
+```
+
+---
+
+## 🔍 Discoverability Tags
+
+* FastAPI WebSocket Manager
+* AG2AI Autogen Orchestrator
+* Real-time agent workflows
+* WebSocket frontend UI
+* AI agent orchestration
+* Agentic problem solving with Python
+* Multi-agent system with streaming responses
+* Custom group chat with Autogen
+
+---

notebook/agentchat_fastapi_websocket/dependencies.py

@@ -0,0 +1,61 @@
+from fastapi import HTTPException, status
+from managers.connection import WebSocketManager
+
+_websocket_manager: WebSocketManager | None = None
+
+# Singleton instance of MongoDB service
+
+
+async def get_websocket_manager() -> WebSocketManager:
+    """Dependency provider for connection manager"""
+    if not _websocket_manager:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Connection manager not initialized",
+        )
+    return _websocket_manager
+
+
+# Manager initialization and cleanup
+async def init_managers() -> None:
+    """Initialize all manager instances"""
+    global _websocket_manager
+
+    try:
+        # Initialize connection manager
+        _websocket_manager = WebSocketManager()
+
+    except Exception:
+        await cleanup_managers()  # Cleanup any partially initialized managers
+        raise
+
+
+async def cleanup_managers() -> None:
+    """Cleanup and shutdown all manager instances"""
+    global _websocket_manager
+
+    # Cleanup connection manager first to ensure all active connections are closed
+    if _websocket_manager:
+        try:
+            await _websocket_manager.cleanup()
+        except Exception as e:
+            print(e)
+            raise
+        finally:
+            _websocket_manager = None
+
+
+# Utility functions for dependency management
+
+
+# Error handling for manager operations
+
+
+class ManagerOperationError(Exception):
+    """Custom exception for manager operation errors"""
+
+    def __init__(self, manager_name: str, operation: str, detail: str):
+        self.manager_name = manager_name
+        self.operation = operation
+        self.detail = detail
+        super().__init__(f"{manager_name} failed during {operation}: {detail}")

notebook/agentchat_fastapi_websocket/docs/USAGE_GUIDE.md

@@ -0,0 +1,104 @@
+# 🧪 AG2AI WebSocket UI – Step-by-Step Usage Guide
+
+This guide explains how to use the interactive WebSocket UI to send and receive agentic messages through a running AG2AI Autogen backend.
+
+---
+
+## 🖼️ Screenshot Reference
+
+![WebSocket UI Screenshot](./usage_screenshot.png)
+
+---
+
+## 🚶‍♂️ Step-by-Step Instructions
+
+### 1️⃣ Create a New Chat Session
+Click on the **`Create New Chat`** button to generate a unique `chat_id`.
+This will be used to manage your session across messages.
+
+---
+
+### 2️⃣ Use the Chat ID
+After creating a chat, click **`Use`** next to the desired `chat_id`.
+This tells the system which session you're actively working in.
+
+---
+
+### 3️⃣ Connect to WebSocket
+Click the **`Connect`** button to establish a WebSocket connection with the backend server.
+Once connected, messages can be sent and received in real time.
+
+---
+
+### 4️⃣ Load a Query Template
+Click **`User Query`** to auto-fill a JSON template into the input box.
+This template follows the format expected by the backend.
+
+---
+
+### 5️⃣ Enter Your Task
+In the JSON input, fill in your prompt inside the `task` key.
+
+#### Example:
+```json
+{
+  "type": "start",
+  "task": "Analyze a dataset and generate EDA visualizations",
+  "files": []
+}
+````
+
+---
+
+### 6️⃣ Send the Message
+
+Click **`Send`** to dispatch the message.
+Wait for streaming agent responses to appear in the UI.
+
+---
+
+## 🧠 Optional: Respond to User Proxy Input Requests
+
+If your setup includes a **User Proxy agent** (i.e., the agent needs user confirmation/input during the flow), you can respond using the `User Input` feature.
+
+### When a prompt from the User Proxy appears:
+
+1. Click **`User Input`** to load a template into the message input box.
+2. Fill your response inside the `response` key.
+
+#### Example:
+
+```json
+{
+  "type": "input_response",
+  "response": "Yes, proceed with the next step",
+  "files": []
+}
+```
+
+3. Click **`Send`** to reply to the ongoing agent conversation.
+
+---
+
+## ✅ Tips
+
+* You can view each message block (sent/received) formatted with timestamps and roles.
+* JSON inputs are auto-formatted for better readability.
+* You can disconnect anytime using the **`Disconnect`** button.
+
+---
+
+## 🧪 Use Case Example
+
+> Want to try a full agentic workflow?
+> Try entering a task like:
+
+```json
+{
+  "type": "start",
+  "task": "Perform EDA on /path/to/titanic.csv dataset using Python",
+  "files": []
+}
+```
+
+and watch the agents plan, code, execute, debug, and summarize the task live.

notebook/agentchat_fastapi_websocket/helpers.py

@@ -0,0 +1,6 @@
+import asyncio
+
+
+def async_to_sync(awaitable):
+    loop = asyncio.get_event_loop()
+    return loop.run_until_complete(awaitable)