├── .gitattributes ├── .env.example ├── requirements.txt ├── .gitignore ├── Dockerfile ├── README.md ├── frontend.py └── main.py /.gitattributes: -------------------------------------------------------------------------------- 1 | # Auto detect text files and perform LF normalization 2 | * text=auto 3 | -------------------------------------------------------------------------------- /.env.example: -------------------------------------------------------------------------------- 1 | OPENAI_API_KEY=sk-proj-1234 2 | 3 | API_KEY=1234 4 | 5 | NOTION_API_KEY=secret_1234 6 | NOTION_TASKS_DB_ID=123456789 -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | pydantic 2 | asyncio 3 | dotenv 4 | fastapi 5 | uvicorn 6 | openai 7 | requests 8 | websockets 9 | streamlit 10 | openai-agents -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .env 2 | .env.local 3 | .env.development.local 4 | .env.test.local 5 | .env.production.local 6 | .env.development 7 | .env.test 8 | .env.production 9 | 10 | .venv 11 | venv 12 | env/ 13 | ENV/ 14 | 15 | __pycache__ 16 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | FROM python:3.13 2 | 3 | WORKDIR /app 4 | 5 | COPY requirements.txt . 6 | RUN pip install --no-cache-dir -r requirements.txt 7 | 8 | COPY . . 9 | 10 | EXPOSE 8000 11 | 12 | CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"] -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 | # Task Manager Agent 3 | 4 | A WebSocket-powered AI agent that helps users organize, manage, and track their tasks efficiently. This project demonstrates an easy use case of tools and function calling with an LLM agent, making it a great starting point for developers interested in building AI-powered assistants with external tool integration. 5 | 6 | **Demo Purpose:** 7 | 8 | This project was created to demonstrate an easy use case of tools and function calling with LLM agents. It is ideal for learning, prototyping, or as a foundation for more advanced agent-based applications. 9 | 10 | ## 🚀 Features 11 | 12 | - **Real-time WebSocket Communication**: Interactive chat interface with instant responses 13 | - **Task Management**: Create, update, delete, and list tasks with priorities and due dates 14 | - **Notion Integration**: Syncs tasks with a Notion database (optional) 15 | - **Session Management**: Maintains conversation context across multiple interactions 16 | - **API Security**: Protected endpoints with API key authentication 17 | 18 | 19 | ## 🏗️ Architecture 20 | 21 | The system consists of two main components: 22 | 23 | 1. **WebSocket Agent** ([main.py](main.py)): Core AI agent with FastAPI WebSocket server 24 | 2. **Streamlit Frontend** ([frontend.py](frontend.py)): Optional web interface for testing and demo 25 | 26 | ### WebSocket Agent 27 | 28 | The agent is built using the `openai-agents` library and provides: 29 | 30 | - **Task Management Tools**: Create, update, delete, and list tasks 31 | - **Notion Integration**: Syncs tasks with Notion (if configured) 32 | - **Conversation Memory**: Maintains context throughout the chat session 33 | - **Error Handling**: Robust error management and user feedback 34 | 35 | 36 | ## 🛠️ Setup 37 | 38 | ### Prerequisites 39 | 40 | - Python 3.13+ 41 | - OpenAI API key 42 | - Notion API access (optional, for Notion sync) 43 | 44 | ### Installation 45 | 46 | 1. Clone the repository: 47 | ```bash 48 | git clone 49 | cd task-manager-agent 50 | ``` 51 | 52 | 2. Install dependencies: 53 | ```bash 54 | pip install -r requirements.txt 55 | ``` 56 | 57 | 3. Configure environment variables in `.env`: 58 | ```env 59 | OPENAI_API_KEY=your_openai_api_key 60 | API_KEY=your_api_security_key 61 | NOTION_API_KEY=your_notion_api_key 62 | NOTION_TASKS_DB_ID=your_notion_tasks_db_id 63 | ``` 64 | 65 | 4. Run the WebSocket server: 66 | ```bash 67 | uvicorn main:app --host 0.0.0.0 --port 8000 68 | ``` 69 | 70 | 71 | ## 🔌 WebSocket API 72 | 73 | ### Connection 74 | 75 | Connect to the WebSocket endpoint: 76 | ``` 77 | ws://localhost:8000/ws/chat 78 | ``` 79 | 80 | ### Message Format 81 | 82 | Send messages as JSON with the following structure: 83 | 84 | ```json 85 | { 86 | "session_id": "unique-session-id", 87 | "message": "Create a task to review code for the new feature", 88 | "user_id": "user123", 89 | "timestamp": "2025-01-01T10:00:00Z", 90 | "message_type": "chat" 91 | } 92 | ``` 93 | 94 | ### Response Format 95 | 96 | Receive responses in the following format: 97 | 98 | ```json 99 | { 100 | "session_id": "unique-session-id", 101 | "message": "Task 'Review code for the new feature' created successfully.", 102 | "sender": "agent", 103 | "timestamp": "2025-01-01T10:00:01Z", 104 | "message_type": "response", 105 | "success": true, 106 | "error": null 107 | } 108 | ``` 109 | 110 | ### Special Commands 111 | 112 | - Send `"/clear"` or set `message_type: "clear"` to reset conversation history 113 | 114 | 115 | ## 🛠️ Available Tools 116 | 117 | The agent has access to several function tools for task management: 118 | 119 | ### 1. Create Task 120 | - **Function**: [`create_task`](main.py) 121 | - **Purpose**: Create a new task with a title and optional due date 122 | 123 | ### 2. List Tasks 124 | - **Function**: [`list_tasks`](main.py) 125 | - **Purpose**: List all current tasks with their status and priority 126 | 127 | ### 3. Update Task Status 128 | - **Function**: [`update_task_status`](main.py) 129 | - **Purpose**: Update the status of a task (e.g., mark as done) 130 | 131 | ### 4. Delete Task 132 | - **Function**: [`delete_task`](main.py) 133 | - **Purpose**: Delete a task by name 134 | 135 | ### 5. Count Tasks 136 | - **Function**: [`count_tasks`](main.py) 137 | - **Purpose**: Get the total number of tasks 138 | 139 | 140 | ## 📋 REST API Endpoints 141 | 142 | ### Health Check 143 | ``` 144 | GET / 145 | ``` 146 | Returns API status (requires API key authentication) 147 | 148 | ### Clear Chat History 149 | ``` 150 | POST /clear-chat 151 | ``` 152 | Clears the global conversation history 153 | 154 | 155 | ## 🐳 Docker Deployment 156 | 157 | The project includes a [Dockerfile](Dockerfile) for containerized deployment: 158 | 159 | ```bash 160 | # Build the image 161 | docker build -t task-manager-agent . 162 | 163 | # Run the container 164 | docker run -p 8000:8000 --env-file .env task-manager-agent 165 | ``` 166 | 167 | 168 | ## 💡 Usage Examples 169 | 170 | ### Create a Task 171 | ```json 172 | { 173 | "session_id": "session123", 174 | "message": "Create a task to review code for the new feature", 175 | "message_type": "chat" 176 | } 177 | ``` 178 | 179 | ### List All Tasks 180 | ```json 181 | { 182 | "session_id": "session123", 183 | "message": "List all my tasks", 184 | "message_type": "chat" 185 | } 186 | ``` 187 | 188 | ### Mark a Task as Done 189 | ```json 190 | { 191 | "session_id": "session123", 192 | "message": "Mark task 'Review code for the new feature' as done", 193 | "message_type": "chat" 194 | } 195 | ``` 196 | 197 | 198 | ## 🧪 Testing with Frontend 199 | 200 | An optional Streamlit frontend is provided for testing and demo: 201 | 202 | ```bash 203 | streamlit run frontend.py 204 | ``` 205 | 206 | Access the web interface at `http://localhost:8501` 207 | 208 | 209 | ## 📊 Task Data Structure 210 | 211 | The agent works with task data stored in Notion (if configured) and can manage tasks with the following properties: 212 | 213 | - **Title** 214 | - **Due Date** (optional) 215 | - **Status** (Not Started, In Progress, Done) 216 | - **Priority** (Low, Medium, High) 217 | 218 | 219 | ## 🔧 Configuration 220 | 221 | Key configuration options in [main.py](main.py): 222 | 223 | - **Model**: GPT-4o for optimal reasoning 224 | - **Session Management**: UUID-based session tracking 225 | - **Error Handling**: Comprehensive exception management 226 | - **API Security**: Header-based authentication 227 | 228 | 229 | ## 🤝 Contributing 230 | 231 | 1. Fork the repository 232 | 2. Create a feature branch 233 | 3. Make your changes 234 | 4. Test with the WebSocket API or Streamlit frontend 235 | 5. Submit a pull request 236 | 237 | 238 | ## 📝 License 239 | 240 | This project is licensed under the MIT License. 241 | 242 | --- 243 | 244 | Built with ❤️ by [Tom Shaw](https://tomshaw.dev) -------------------------------------------------------------------------------- /frontend.py: -------------------------------------------------------------------------------- 1 | import streamlit as st 2 | import json 3 | from datetime import datetime 4 | import uuid 5 | import asyncio 6 | import websockets 7 | from typing import Optional 8 | import threading 9 | from concurrent.futures import ThreadPoolExecutor 10 | 11 | # Page configuration 12 | st.set_page_config( 13 | page_title="AI Task Manager Agent", 14 | page_icon="📋", 15 | layout="wide", 16 | initial_sidebar_state="expanded" 17 | ) 18 | 19 | # Initialize session state 20 | if "messages" not in st.session_state: 21 | st.session_state.messages = [] 22 | if "session_id" not in st.session_state: 23 | st.session_state.session_id = str(uuid.uuid4()) 24 | 25 | # WebSocket communication function 26 | def send_websocket_message(websocket_url: str, message_payload: dict, timeout: int = 30) -> Optional[str]: 27 | """Send message via WebSocket and return response""" 28 | async def _send_message(): 29 | try: 30 | async with websockets.connect(websocket_url) as websocket: 31 | # Send message 32 | await websocket.send(json.dumps(message_payload)) 33 | 34 | # Wait for response with timeout 35 | response = await asyncio.wait_for( 36 | websocket.recv(), 37 | timeout=timeout 38 | ) 39 | return response 40 | except Exception as e: 41 | raise e 42 | 43 | # Run the async function in a new event loop 44 | try: 45 | loop = asyncio.new_event_loop() 46 | asyncio.set_event_loop(loop) 47 | return loop.run_until_complete(_send_message()) 48 | except Exception as e: 49 | raise e 50 | finally: 51 | loop.close() 52 | 53 | # Configuration 54 | WEBSOCKET_URL = st.sidebar.text_input( 55 | "WebSocket URL", 56 | value="ws://localhost:8000/ws/chat", 57 | help="Enter the WebSocket URL of your AI agent" 58 | ) 59 | 60 | REQUEST_TIMEOUT = st.sidebar.slider( 61 | "Request Timeout (seconds)", 62 | min_value=5, 63 | max_value=60, 64 | value=30, 65 | help="How long to wait for agent response" 66 | ) 67 | 68 | # Sidebar info 69 | st.sidebar.markdown("---") 70 | st.sidebar.markdown(f"**Session ID:** `{st.session_state.session_id[:8]}...`") 71 | st.sidebar.markdown(f"**Messages:** {len(st.session_state.messages)}") 72 | 73 | # Clear chat button 74 | if st.sidebar.button("Clear Chat", type="secondary"): 75 | st.session_state.messages = [] 76 | 77 | # Also clear the backend conversation 78 | clear_payload = { 79 | "session_id": st.session_state.session_id, 80 | "message": "/clear", 81 | "timestamp": datetime.now().isoformat(), 82 | "user_id": "streamlit_user", 83 | "message_type": "clear" 84 | } 85 | 86 | try: 87 | # Send clear command via WebSocket 88 | clear_response = send_websocket_message(WEBSOCKET_URL, clear_payload, REQUEST_TIMEOUT) 89 | if clear_response: 90 | response_data = json.loads(clear_response) 91 | st.sidebar.success("Chat cleared on backend") 92 | except Exception as e: 93 | st.sidebar.warning(f"Failed to clear backend: {str(e)}") 94 | 95 | st.rerun() 96 | 97 | # Main chat interface 98 | st.title("📋 AI Task Manager Agent") 99 | st.markdown("Manage your tasks efficiently with your AI assistant!") 100 | st.markdown(""" 101 | You can ask me to create, update, delete, or organize your tasks. Just tell me what you need! 102 | 103 | **Examples:** 104 | - "Create a task to review code for the new feature" 105 | - "List all my high priority tasks" 106 | - "Mark task 3 as completed" 107 | - "Show me a summary of all my tasks" 108 | """) 109 | 110 | # Display chat messages 111 | chat_container = st.container() 112 | with chat_container: 113 | for message in st.session_state.messages: 114 | with st.chat_message(message["role"]): 115 | st.markdown(message["content"]) 116 | if "timestamp" in message: 117 | st.caption(f"_{message['timestamp']}_") 118 | 119 | # Chat input 120 | if prompt := st.chat_input("Type your message here..."): 121 | # Add user message to chat 122 | timestamp = datetime.now().strftime("%H:%M:%S") 123 | user_message = { 124 | "role": "user", 125 | "content": prompt, 126 | "timestamp": timestamp 127 | } 128 | st.session_state.messages.append(user_message) 129 | 130 | # Display user message immediately 131 | with st.chat_message("user"): 132 | st.markdown(prompt) 133 | st.caption(f"_{timestamp}_") 134 | 135 | # Prepare WebSocket message payload 136 | message_payload = { 137 | "session_id": st.session_state.session_id, 138 | "message": prompt, 139 | "timestamp": datetime.now().isoformat(), 140 | "user_id": "streamlit_user", 141 | "message_type": "chat" 142 | } 143 | 144 | # Show loading spinner and send WebSocket message 145 | with st.chat_message("assistant"): 146 | with st.spinner("Agent is thinking..."): 147 | try: 148 | # Send message via WebSocket 149 | agent_response = send_websocket_message(WEBSOCKET_URL, message_payload, REQUEST_TIMEOUT) 150 | 151 | if agent_response: 152 | # Parse response 153 | response_data = json.loads(agent_response) 154 | agent_message = response_data.get("message", "No response from agent") 155 | 156 | # Display agent response 157 | st.markdown(agent_message) 158 | response_timestamp = datetime.now().strftime("%H:%M:%S") 159 | st.caption(f"_{response_timestamp}_") 160 | 161 | # Add agent message to session state 162 | assistant_message = { 163 | "role": "assistant", 164 | "content": agent_message, 165 | "timestamp": response_timestamp 166 | } 167 | st.session_state.messages.append(assistant_message) 168 | 169 | else: 170 | error_msg = "No response received from agent" 171 | st.error(error_msg) 172 | st.session_state.messages.append({ 173 | "role": "assistant", 174 | "content": error_msg, 175 | "timestamp": datetime.now().strftime("%H:%M:%S") 176 | }) 177 | 178 | except Exception as e: 179 | error_msg = f"WebSocket error: {str(e)}" 180 | st.error(error_msg) 181 | st.session_state.messages.append({ 182 | "role": "assistant", 183 | "content": error_msg, 184 | "timestamp": datetime.now().strftime("%H:%M:%S") 185 | }) 186 | 187 | # Rerun to update the chat display 188 | st.rerun() 189 | 190 | # Footer 191 | st.markdown("---") 192 | st.markdown( 193 | """ 194 |
195 | AI Agent Chat Interface | Built by Tom Shaw 196 |
197 | """, 198 | unsafe_allow_html=True 199 | ) 200 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import uuid 4 | from datetime import datetime, date 5 | from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, Depends 6 | from fastapi.security import APIKeyHeader 7 | from pydantic import BaseModel 8 | from starlette.status import HTTP_403_FORBIDDEN 9 | from dotenv import load_dotenv 10 | from agents import Agent, Runner, TResponseInputItem, function_tool 11 | from typing import List, Dict, Optional 12 | import requests 13 | import traceback 14 | 15 | # Load environment variables 16 | load_dotenv() 17 | 18 | API_KEY = os.getenv("API_KEY") 19 | API_KEY_NAME = "x-api-key" 20 | 21 | notion_api_key = os.environ.get("NOTION_API_KEY") 22 | notion_tasks_db_id = os.environ.get("NOTION_TASKS_DB_ID") 23 | notion_headers = { 24 | "Authorization": f"Bearer {notion_api_key}", 25 | "Content-Type": "application/json", 26 | "Notion-Version": "2022-06-28" 27 | } 28 | 29 | # Agent Tools 30 | 31 | 32 | # Internal Notion API request functions 33 | def _notion_create_task(title: str, due_date: str = None) -> str: 34 | print("Creating task:", title) 35 | if not title: 36 | return "Please provide a title for the task." 37 | data = { 38 | "parent": {"database_id": notion_tasks_db_id}, 39 | "properties": { 40 | "Name": {"title": [{"text": {"content": title}}]} 41 | } 42 | } 43 | if due_date: 44 | data["properties"]["Due"] = {"date": {"start": due_date}} 45 | response = requests.post( 46 | "https://api.notion.com/v1/pages", 47 | headers=notion_headers, 48 | data=json.dumps(data) 49 | ) 50 | if response.status_code == 200: 51 | return f"Task '{title}' created successfully{f' with due date {due_date}' if due_date else ''}." 52 | return f"Failed to create task: {response.text}" 53 | 54 | def _notion_list_tasks() -> list: 55 | print("Listing tasks...") 56 | query = { 57 | "filter": { 58 | "or": [ 59 | {"property": "Status", "status": {"equals": "Not Started"}}, 60 | {"property": "Status", "status": {"equals": "In Progress"}} 61 | ] 62 | }, 63 | "sorts": [{"timestamp": "last_edited_time", "direction": "descending"}] 64 | } 65 | response = requests.post( 66 | f"https://api.notion.com/v1/databases/{notion_tasks_db_id}/query", 67 | headers=notion_headers, 68 | data=json.dumps(query) 69 | ) 70 | if response.status_code != 200: 71 | return [] 72 | results = response.json().get("results", []) 73 | formatted_tasks = [] 74 | for task in results: 75 | props = task["properties"] 76 | formatted_tasks.append({ 77 | "id": task["id"], 78 | "title": props["Name"]["title"][0]["plain_text"] if props["Name"]["title"] else "", 79 | "due": props["Due"]["date"]["start"] if props.get("Due") and props["Due"]["date"] else None, 80 | "status": props["Status"]["status"]["name"] if props.get("Status") and props["Status"]["status"] else "", 81 | "priority": props["Priority"]["select"]["name"] if props.get("Priority") and props["Priority"]["select"] else "", 82 | "event_id": props["event_id"]["rich_text"][0]["plain_text"] if props.get("event_id") and props["event_id"]["rich_text"] else None, 83 | }) 84 | return formatted_tasks 85 | 86 | def _notion_count_tasks() -> int: 87 | print("Counting tasks...") 88 | tasks = _notion_list_tasks() 89 | return len(tasks) 90 | 91 | def _notion_update_task_status(task_name: str, status: str) -> str: 92 | try: 93 | print('updating task status: ', task_name, status) 94 | tasks = _notion_list_tasks() 95 | task = next((t for t in tasks if t["title"].lower() == task_name.lower()), None) 96 | print("Found task:", task) 97 | if not task: 98 | return f"Task '{task_name}' not found." 99 | valid_statuses = ["Done", "In Progress", "Not Started"] 100 | if status not in valid_statuses: 101 | return f"Invalid status. Choose from: {', '.join(valid_statuses)}." 102 | data = { 103 | "properties": { 104 | "Status": {"status": {"name": status}} 105 | } 106 | } 107 | response = requests.patch( 108 | f"https://api.notion.com/v1/pages/{task['id']}", 109 | headers=notion_headers, 110 | data=json.dumps(data) 111 | ) 112 | print("Response status code:", response.status_code) 113 | if response.status_code == 200: 114 | return f"Task '{task_name}' status updated to {status}." 115 | return f"Failed to update task status: {response.text}" 116 | except Exception as e: 117 | print("Error updating task status:", e) 118 | return "An error occurred while updating the task status." 119 | 120 | def _notion_delete_task(task_name: str) -> str: 121 | if not task_name: 122 | return "Please provide the name of the task to delete." 123 | tasks = _notion_list_tasks() 124 | task = next((t for t in tasks if t["title"].lower() == task_name.lower()), None) 125 | if not task: 126 | return f"Task '{task_name}' not found." 127 | # Notion does not support hard delete, so we set the archived property to True 128 | data = {"archived": True} 129 | response = requests.patch( 130 | f"https://api.notion.com/v1/pages/{task['id']}", 131 | headers=notion_headers, 132 | data=json.dumps(data) 133 | ) 134 | if response.status_code == 200: 135 | return f"Task '{task_name}' deleted successfully." 136 | return f"Failed to delete task: {response.text}" 137 | 138 | # Agent Tool functions (decorated) 139 | @function_tool 140 | def create_task(title: str, due_date: str = None) -> str: 141 | return _notion_create_task(title, due_date) 142 | 143 | @function_tool 144 | def list_tasks() -> list: 145 | return _notion_list_tasks() 146 | 147 | @function_tool 148 | def count_tasks() -> int: 149 | return _notion_count_tasks() 150 | 151 | @function_tool 152 | def update_task_status(task_name: str, status: str) -> str: 153 | return _notion_update_task_status(task_name, status) 154 | 155 | @function_tool 156 | def delete_task(task_name: str) -> str: 157 | return _notion_delete_task(task_name) 158 | 159 | # Agent Function 160 | 161 | conversation: list[TResponseInputItem] = [] 162 | 163 | async def run_task_manager_agent(message: str): 164 | 165 | agent = Agent( 166 | name="Task Manager Agent", 167 | instructions=f""" 168 | 169 | You are a helpful Task Manager agent that assists users in organizing and managing their tasks efficiently. You can create, update, delete, and organize tasks with different priorities and statuses. 170 | 171 | 172 | Help users manage their tasks effectively by providing a comprehensive task management system. You can handle task creation, status updates, priority management, and provide summaries and insights about their task list. 173 | 174 | 175 | You have access to the following task management functions: 176 | - Create new tasks with titles, descriptions, priorities (low/medium/high), and due dates 177 | - List tasks with filtering by status (todo/in_progress/done) and priority 178 | - Update task status and priority 179 | - Get detailed information about specific tasks 180 | - Delete tasks when no longer needed 181 | - Provide task summaries and statistics 182 | 183 | 184 | If you have been asked to make a change or remove a specific task, but you do not have a list of tasks stored in your memory, you should retreive the list of tasks first before trying to make the change. If a due date has not been specified, leave it blank. 185 | - Always be aware of today's date which is {datetime.now().strftime('%Y-%m-%d')} when working with due dates and scheduling tasks 186 | 188 | - Be friendly and encouraging when helping with task management 189 | - Provide clear confirmations when tasks are created, updated, or deleted 190 | - Offer helpful suggestions for task organization and prioritization 191 | - When listing tasks, present them in a clear, organized format 192 | - Ask clarifying questions if task details are unclear 193 | - Celebrate completed tasks and encourage productivity 194 | 195 | 196 | - Use clear, structured responses 197 | - Format task lists in an easy-to-read manner 198 | - Include relevant task IDs for reference 199 | - Provide actionable next steps when appropriate 200 | 201 | """, 202 | tools=[ 203 | create_task, 204 | list_tasks, 205 | update_task_status, 206 | delete_task, 207 | count_tasks, 208 | ], 209 | model="gpt-4o", 210 | ) 211 | 212 | # Add the message to the conversation 213 | conversation.append({ 214 | "role": "user", 215 | "content": message 216 | }) 217 | 218 | result = await Runner.run(agent, conversation) 219 | 220 | if result and result.final_output: 221 | conversation.append({ 222 | "role": "assistant", 223 | "content": result.final_output 224 | }) 225 | else: 226 | conversation.append({ 227 | "role": "assistant", 228 | "content": "No response from agent." 229 | }) 230 | 231 | return result 232 | 233 | # FastAPI Setup 234 | 235 | api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False) 236 | 237 | app = FastAPI() 238 | 239 | def verify_api_key(api_key: str = Depends(api_key_header)): 240 | if api_key != API_KEY: 241 | raise HTTPException(status_code=HTTP_403_FORBIDDEN, detail="Invalid API Key") 242 | return api_key 243 | 244 | @app.get("/") 245 | def read_root(api_key: str = Depends(verify_api_key)): 246 | return {"message": "API is working"} 247 | 248 | @app.post("/clear-chat") 249 | def clear_chat(api_key: str = Depends(verify_api_key)): 250 | """Clear the conversation history""" 251 | global conversation 252 | conversation.clear() 253 | return {"message": "Chat history cleared successfully", "success": True} 254 | 255 | class ConnectionManager: 256 | def __init__(self): 257 | self.active_connections: list[WebSocket] = [] 258 | 259 | async def connect(self, websocket: WebSocket): 260 | await websocket.accept() 261 | self.active_connections.append(websocket) 262 | 263 | def disconnect(self, websocket: WebSocket): 264 | self.active_connections.remove(websocket) 265 | 266 | async def send_message(self, message: str, websocket: WebSocket): 267 | await websocket.send_text(message) 268 | 269 | async def broadcast(self, message: str): 270 | for connection in self.active_connections: 271 | await self.send_message(message, connection) 272 | 273 | manager = ConnectionManager() 274 | 275 | class ChatMessage(BaseModel): 276 | session_id: str 277 | message: str 278 | user_id: str = "user" 279 | timestamp: str | None = None 280 | message_type: str = "chat" 281 | 282 | class ChatResponse(BaseModel): 283 | session_id: str 284 | message: str 285 | sender: str = "agent" 286 | timestamp: str 287 | message_type: str = "response" 288 | success: bool = True 289 | error: str | None = None 290 | 291 | @app.websocket("/ws/chat") 292 | async def websocket_endpoint(websocket: WebSocket): 293 | await manager.connect(websocket) 294 | session_id = str(uuid.uuid4()) 295 | print(f"New WebSocket connection established with session ID: {session_id}") 296 | 297 | try: 298 | while True: 299 | # Receive message from client 300 | data = await websocket.receive_text() 301 | 302 | try: 303 | # Parse incoming JSON message 304 | message_data = json.loads(data) 305 | chat_message = ChatMessage(**message_data) 306 | 307 | # Check if this is a clear command 308 | if chat_message.message_type == "clear" or chat_message.message.lower().strip() == "/clear": 309 | global conversation 310 | conversation.clear() 311 | response = ChatResponse( 312 | session_id=chat_message.session_id, 313 | message="Chat history has been cleared.", 314 | timestamp=datetime.now().isoformat() 315 | ) 316 | else: 317 | # Process the message with the agent 318 | agent_response = await process_chat_message(chat_message.message) 319 | 320 | # Create response 321 | response = ChatResponse( 322 | session_id=chat_message.session_id, 323 | message=agent_response, 324 | timestamp=datetime.now().isoformat() 325 | ) 326 | 327 | # Send response back to client 328 | await manager.send_message(response.model_dump_json(), websocket) 329 | 330 | except json.JSONDecodeError: 331 | # Handle plain text messages for backward compatibility 332 | agent_response = await process_chat_message(data) 333 | response = ChatResponse( 334 | session_id=session_id, 335 | message=agent_response, 336 | timestamp=datetime.now().isoformat() 337 | ) 338 | await manager.send_message(response.model_dump_json(), websocket) 339 | 340 | except Exception as e: 341 | # Handle errors 342 | error_response = ChatResponse( 343 | session_id=session_id, 344 | message="Sorry, I encountered an error processing your message.", 345 | timestamp=datetime.now().isoformat(), 346 | success=False, 347 | error=str(e) 348 | ) 349 | await manager.send_message(error_response.model_dump_json(), websocket) 350 | 351 | except WebSocketDisconnect: 352 | manager.disconnect(websocket) 353 | print(f"WebSocket connection closed for session: {session_id}") 354 | traceback.print_exc() 355 | 356 | async def process_chat_message(message: str) -> str: 357 | """Process chat message with the task manager agent""" 358 | try: 359 | # Run the task manager agent with the user's message 360 | result = await run_task_manager_agent(message) 361 | 362 | # Return the agent's final output 363 | if result and result.final_output: 364 | return result.final_output 365 | else: 366 | return "I couldn't process your request at the moment. Please try again." 367 | 368 | except Exception as e: 369 | print(f"Error processing message: {e}") 370 | return "I encountered an error while processing your message. Please try again." --------------------------------------------------------------------------------