hello chat gpt this is my code with i can chat to an ai and the ai can do action can you help me to imrpove it animation on this pls :
import os
import datetime
import json
import asyncio
from dateutil.parser import parse as parse_date
import dateparser
from google.oauth2.credentials import Credentials
from google_auth_oauthlib.flow import InstalledAppFlow
from googleapiclient.discovery import build
from email.mime.text import MIMEText
import base64

from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
from starlette.websockets import WebSocketState
from fastapi.templating import Jinja2Templates
from fastapi.responses import HTMLResponse
from fastapi.middleware.cors import CORSMiddleware

# Import the Groq client
from groq import Groq

# Import dotenv to manage environment variables
from dotenv import load_dotenv

load_dotenv()

# Import GoogleAuthRequest to avoid conflicts with FastAPI's Request
from google.auth.transport.requests import Request as GoogleAuthRequest

# Set the scopes for Google API permissions
SCOPES = [
    'https://www.googleapis.com/auth/documents',
    'https://www.googleapis.com/auth/calendar',
    'https://www.googleapis.com/auth/drive',
    'https://www.googleapis.com/auth/gmail.send',
    'https://www.googleapis.com/auth/gmail.readonly'
]

app = FastAPI()

templates = Jinja2Templates(directory="templates")

# Add CORS middleware (if necessary)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Adjust according to your needs
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

def authenticate_google_services():
    creds = None
    # Check if credentials already exist
    if os.path.exists('token.json'):
        creds = Credentials.from_authorized_user_file('token.json', SCOPES)
    # If no valid credentials, initiate the authentication flow
    if not creds or not creds.valid:
        if creds and creds.expired and creds.refresh_token:
            creds.refresh(GoogleAuthRequest())
        else:
            if not os.path.exists('credentials.json'):
                print("The 'credentials.json' file is missing.")
                return None, None, None, None
            flow = InstalledAppFlow.from_client_secrets_file(
                'credentials.json', SCOPES)
            creds = flow.run_local_server(port=0)
        # Save the credentials for the next run
        with open('token.json', 'w') as token:
            token.write(creds.to_json())
    # Create Google Docs, Calendar, Drive, and Gmail services
    try:
        docs_service = build('docs', 'v1', credentials=creds)
        calendar_service = build('calendar', 'v3', credentials=creds)
        drive_service = build('drive', 'v3', credentials=creds)
        gmail_service = build('gmail', 'v1', credentials=creds)
        return docs_service, calendar_service, drive_service, gmail_service
    except Exception as e:
        print(f"An error occurred while creating Google services: {e}")
        return None, None, None, None

# Initialize Google services
docs_service, calendar_service, drive_service, gmail_service = authenticate_google_services()
if not docs_service or not calendar_service or not drive_service or not gmail_service:
    print("Authentication failed.")

# Get the Groq API key from environment variables
groq_api_key = os.getenv("GROQ_API_KEY")
if not groq_api_key:
    raise Exception("Please set the 'GROQ_API_KEY' environment variable with your Groq API key.")

# Initialize the Groq client
groq_client = Groq(api_key=groq_api_key)

@app.get("/", response_class=HTMLResponse)
async def get_index(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        # Initialize the conversation history
        conversation_history = []
        while True:
            # Check if WebSocket is still connected
            if websocket.client_state != WebSocketState.CONNECTED:
                break

            # Receive the user's message
            data = await websocket.receive_text()
            request_data = json.loads(data)
            user_message = request_data.get('message')

            if not user_message:
                await websocket.send_text("The 'message' field is required.")
                continue  # Continue to listen for the next messages

            # Add the user's message to the history
            conversation_history.append({"role": "user", "content": user_message})

            # Process the user's request in a separate thread to avoid blocking
            assistant_response = await process_user_request(user_message, conversation_history, websocket)

            # Add the assistant's response to the history
            conversation_history.append({"role": "assistant", "content": assistant_response})

            # Send the response to the user
            await websocket.send_text(assistant_response)

    except WebSocketDisconnect:
        print("Client disconnected")
    except Exception as e:
        await websocket.send_text(f"Server error: {e}")

def generate_with_groq(conversation_history, temperature=0.7, max_tokens=512):
    response = groq_client.chat.completions.create(
        messages=conversation_history,
        model="llama-3.1-70b-versatile",
        temperature=temperature,
        max_tokens=max_tokens
    )
    return response.choices[0].message.content.strip()

def generate_simple_response(prompt, temperature=0.0, max_tokens=150):
    response = groq_client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model="llama-3.2-90b-vision-preview",
        temperature=temperature,
        max_tokens=max_tokens
    )
    return response.choices[0].message.content.strip()

def get_action_title_and_steps(action):
    action_mapping = {
        'create_document': ("Creating Google Document", ["Analyzing request", "Generating content", "Creating document", "Formatting document", "Saving and sharing"]),
        'read_calendar': ("Reading Calendar Events", ["Fetching events", "Processing data", "Preparing response"]),
        # Add titles and steps for other actions similarly
    }
    return action_mapping.get(action, ("Performing Action", ["Processing"]))

async def process_user_request(user_message, conversation_history, websocket):
    # Use the LLM to detect the action to perform
    detection_prompt = f"""You are an assistant that analyzes the user's message to determine if there is an action to perform among the following:
- create_document: to create a Google Docs document.
- read_calendar: to read calendar events.
- create_calendar_event: to create an event in the calendar.
- share_document: to share a document with someone.
- send_email: to send an email.
- read_emails: to read emails in Gmail.

For each action, extract the necessary parameters as a JSON object. If no action is required, simply respond with: {{"action": "chat"}}

User's message: "{user_message}"

Respond in JSON with the format:
{{
  "action": "name_of_action",
  "parameters": {{
    "parameter1": "value1",
    "parameter2": "value2",
    ...
  }}
}}
"""

    # Get the LLM's response for the action to perform
    action_response = generate_simple_response(detection_prompt)

    try:
        action_data = json.loads(action_response)
        action = action_data.get("action", "chat")
    except json.JSONDecodeError:
        # If the LLM did not return valid JSON, we assume it's a normal chat
        action = "chat"

    if action == "chat":
        # Normal chat response
        assistant_response = generate_with_groq(conversation_history)
        return assistant_response
    else:
        # Send action start event before performing the action
        action_title, action_steps = get_action_title_and_steps(action)
        await websocket.send_text(json.dumps({
            "type": "action_start",
            "title": action_title,
            "steps": action_steps
        }))

        # Perform the action and handle exceptions
        try:
            if action == 'create_document':
                return await create_document_from_parameters(action_data.get("parameters", {}), conversation_history, websocket)
            elif action == 'read_calendar':
                return await read_calendar_from_parameters(action_data.get("parameters", {}), websocket)
            elif action == 'create_calendar_event':
                return await create_calendar_event_from_parameters(action_data.get("parameters", {}), websocket)
            elif action == 'share_document':
                return await share_document_from_parameters(action_data.get("parameters", {}), websocket)
            elif action == 'send_email':
                return await send_email_from_parameters(action_data.get("parameters", {}), websocket)
            elif action == 'read_emails':
                return await read_emails_from_parameters(action_data.get("parameters", {}), websocket)
            else:
                # Unrecognized action, we respond with normal chat
                assistant_response = generate_with_groq(conversation_history)
                return assistant_response
        except Exception as e:
            # Send action update event for error
            await websocket.send_text(json.dumps({"type": "action_error", "message": str(e)}))
            return f"An error occurred: {e}"

# Specific action functions

async def create_document_from_parameters(parameters, conversation_history, websocket):
    if not docs_service:
        await websocket.send_text(json.dumps({"type": "action_error", "message": "Google Docs service is not available."}))
        return "Google Docs service is not available."
    try:

        await asyncio.sleep(1)
        await websocket.send_text(json.dumps({"type": "action_update", "step": 0, "status": "complete"}))

        # Generate the title
        if 'title' not in parameters or not parameters['title']:
            title_prompt = f"""Generate a concise and relevant title for a document based on the following user input:
{parameters.get('content', '')}"""
            title = generate_simple_response(title_prompt)
        else:
            title = parameters['title']

        await asyncio.sleep(1)
        await websocket.send_text(json.dumps({"type": "action_update", "step": 1, "status": "complete"}))

        if 'content' not in parameters or not parameters['content']:
            content_prompt = f"""Generate detailed content for a document titled '{title}' based on the user's input."""
            content = generate_simple_response(content_prompt, temperature=0.7, max_tokens=1024)
        else:
            content = parameters['content']

        if not content.strip():
            return "The document content is empty."

        # Step 2: Creating document
        await asyncio.sleep(1)
        await websocket.send_text(json.dumps({"type": "action_update", "step": 2, "status": "complete"}))

        # Create the document
        document = {'title': title}
        doc = docs_service.documents().create(body=document).execute()
        doc_id = doc.get('documentId')

        # Step 3: Formatting document
        await asyncio.sleep(1)
        await websocket.send_text(json.dumps({"type": "action_update", "step": 3, "status": "complete"}))

        # Insert the content
        requests_body = [
            {
                'insertText': {
                    'location': {'index': 1},
                    'text': content
                }
            }
        ]

        docs_service.documents().batchUpdate(
            documentId=doc_id,
            body={'requests': requests_body}
        ).execute()

        # Step 4: Saving and sharing
        await asyncio.sleep(1)
        await websocket.send_text(json.dumps({"type": "action_update", "step": 4, "status": "complete"}))

        doc_url = f"https://docs.google.com/document/d/{doc_id}"

        # Send final action completion message
        await websocket.send_text(json.dumps({"type": "action_complete", "message": f"Document '{title}' created successfully: {doc_url}"}))

        return f"Document '{title}' created successfully: {doc_url}"
    except Exception as e:
        # Send action error message
        await websocket.send_text(json.dumps({"type": "action_error", "message": str(e)}))
        return f"An error occurred while creating the document: {e}"

async def read_calendar_from_parameters(parameters, websocket):
    if not calendar_service:
        return "Google Calendar service is not available."
    try:
        start_date_str = parameters.get('start_date')
        end_date_str = parameters.get('end_date')

        # Convert dates
        if not start_date_str:
            start_date = datetime.datetime.utcnow()
            start_date_str = start_date.strftime('%Y-%m-%d')
        else:
            start_date = parse_date(start_date_str)

        if not end_date_str:
            end_date = start_date + datetime.timedelta(days=7)
            end_date_str = end_date.strftime('%Y-%m-%d')
        else:
            end_date = parse_date(end_date_str)

        # Format ISO dates
        start_date_iso = start_date.isoformat() + 'Z'
        end_date_iso = end_date.isoformat() + 'Z'

        # Get events
        events_result = calendar_service.events().list(
            calendarId='primary', timeMin=start_date_iso, timeMax=end_date_iso,
            singleEvents=True, orderBy='startTime').execute()
        events = events_result.get('items', [])

        if not events:
            return "No events found for the specified period."

        # Prepare the list of events
        event_list = []
        for event in events:
            start_event = event['start'].get('dateTime', event['start'].get('date'))
            summary = event.get('summary', 'No title')
            event_list.append(f"{start_event} - {summary}")

        # Generate a response
        events_str = "\n".join(event_list)
        return f"Here are your events between {start_date_str} and {end_date_str}:\n{events_str}"

    except Exception as e:
        return f"An error occurred while reading the calendar: {e}"

async def create_calendar_event_from_parameters(parameters, websocket):
    if not calendar_service:
        return "Google Calendar service is not available."
    try:
        summary = parameters.get('summary', 'New Event')
        description = parameters.get('description', '')
        location = parameters.get('location', '')
        start_datetime_str = parameters.get('start_datetime')
        end_datetime_str = parameters.get('end_datetime')
        time_zone = parameters.get('time_zone', 'UTC')

        if not start_datetime_str or not end_datetime_str:
            return "Date and time information is missing."

        # Parse dates
        start_datetime = dateparser.parse(start_datetime_str)
        end_datetime = dateparser.parse(end_datetime_str)

        if not start_datetime or not end_datetime:
            return "Unable to parse the provided dates."

        # Format ISO dates
        start_datetime_iso = start_datetime.isoformat()
        end_datetime_iso = end_datetime.isoformat()

        # Create the event
        event = {
            'summary': summary,
            'location': location,
            'description': description,
            'start': {
                'dateTime': start_datetime_iso,
                'timeZone': time_zone,
            },
            'end': {
                'dateTime': end_datetime_iso,
                'timeZone': time_zone,
            },
        }

        event = calendar_service.events().insert(calendarId='primary', body=event).execute()
        event_link = event.get('htmlLink')
        return f"Event '{summary}' created successfully: {event_link}"
    except Exception as e:
        return f"Unable to create the event: {e}"

async def share_document_from_parameters(parameters, websocket):
    if not drive_service:
        return "Google Drive service is not available."
    try:
        file_name = parameters.get('file_name')
        user_email = parameters.get('user_email')

        if not file_name or not user_email:
            return "The 'file_name' or 'user_email' information is missing."

        # Search for the file by name
        search_query = f"name = '{file_name}'"
        response = drive_service.files().list(q=search_query, spaces='drive').execute()
        files = response.get('files', [])

        if not files:
            return f"No file found with the name '{file_name}'."

        file_id = files[0].get('id')

        # Share the file
        permission = {
            'type': 'user',
            'role': 'writer',
            'emailAddress': user_email
        }
        drive_service.permissions().create(
            fileId=file_id,
            body=permission,
            fields='id',
            sendNotificationEmail=True,
        ).execute()
        return f"The file '{file_name}' has been shared with {user_email}."
    except Exception as e:
        return f"Unable to share the document: {e}"

async def send_email_from_parameters(parameters, websocket):
    if not gmail_service:
        return "Gmail service is not available."
    try:
        to = parameters.get('to')
        subject = parameters.get('subject', '')
        message_text = parameters.get('message', '')

        if not to or not message_text:
            return "The 'to' or 'message' information is missing."

        # Send the email
        message = MIMEText(message_text)
        message['to'] = to
        message['from'] = 'your_email_address@example.com'  # Replace with your email address
        message['subject'] = subject
        raw_message = base64.urlsafe_b64encode(message.as_bytes()).decode()
        body = {'raw': raw_message}
        gmail_service.users().messages().send(userId='me', body=body).execute()
        return f"Email sent to {to}."
    except Exception as e:
        return f"Unable to send the email: {e}"

async def read_emails_from_parameters(parameters, websocket):
    if not gmail_service:
        return "Gmail service is not available."
    try:
        query = parameters.get('query', 'is:unread')

        # Read emails matching the query
        response = gmail_service.users().messages().list(userId='me', q=query).execute()
        messages = response.get('messages', [])

        if not messages:
            return "No emails found."

        emails = []
        for message in messages[:5]:  # Limit to 5 emails to avoid a long response
            msg = gmail_service.users().messages().get(userId='me', id=message['id'], format='full').execute()
            headers = msg.get('payload', {}).get('headers', [])
            subject = "No subject"
            from_email = "Unknown"
            for header in headers:
                if header['name'] == 'Subject':
                    subject = header['value']
                elif header['name'] == 'From':
                    from_email = header['value']
            snippet = msg.get('snippet', '')
            emails.append(f"From: {from_email}\nSubject: {subject}\nPreview: {snippet}\n")

        email_list = "\n---\n".join(emails)
        return f"Here are the emails found:\n{email_list}"
    except Exception as e:
        return f"Unable to read emails: {e}"

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>AI Assistant with Google Integration</title>

    <!-- Font Awesome for icons -->
    <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css">

    <!-- Google Fonts -->
    <link href="https://fonts.googleapis.com/css2?family=Roboto:wght@400;500;700&display=swap" rel="stylesheet">

    <style>
        /* Base Styles */
        * {
            margin: 0;
            padding: 0;
            box-sizing: border-box;
        }

        body {
            font-family: 'Roboto', Arial, sans-serif;
            background: #f0f2f5;
            color: #333;
            line-height: 1.6;
        }

        .container {
            max-width: 1200px;
            margin: 0 auto;
            padding: 20px;
        }

        /* Chat Interface */
        .chat-container {
            background: white;
            border-radius: 12px;
            box-shadow: 0 2px 10px rgba(0,0,0,0.1);
            height: 80vh;
            display: flex;
            flex-direction: column;
        }

        .chat-messages {
            flex: 1;
            overflow-y: auto;
            padding: 20px;
            display: flex;
            flex-direction: column;
        }

        /* Message Styles */
        .message {
            margin-bottom: 20px;
            max-width: 80%;
            animation: fadeIn 0.5s ease-in-out;
            padding: 12px 16px;
            border-radius: 12px;
            line-height: 1.4;
            word-wrap: break-word;
        }

        .user-message {
            margin-left: auto;
            background: #dcf8c6;
            color: #27ae60;
            border-radius: 12px 12px 0 12px;
        }

        .assistant-message {
            margin-right: auto;
            background: #f1f2f6;
            color: #333;
            border-radius: 12px 12px 12px 0;
        }

        /* Action Indicator Styles */
        .action-indicator {
            background: #ffffff;
            border-radius: 12px;
            padding: 20px;
            box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
            margin: 15px 0;
            border-left: 4px solid #2196f3;
            animation: fadeIn 0.5s ease-in-out;
            overflow: hidden;
        }

        .action-header {
            display: flex;
            align-items: center;
            margin-bottom: 15px;
        }

        .action-title {
            font-size: 16px;
            font-weight: 600;
            color: #333;
            margin-left: 10px;
        }

        .action-steps {
            display: flex;
            flex-direction: column;
            gap: 12px;
        }

        .step {
            display: flex;
            align-items: center;
            padding: 8px;
            border-radius: 8px;
            background: #f1f2f6;
            transition: all 0.3s cubic-bezier(0.25, 0.1, 0.25, 1);
            transform: translateX(-100%);
            opacity: 0;
        }

        .step.active {
            transform: translateX(0);
            opacity: 1;
            background: #e0f7fa;
        }

        .step.complete {
            background: #e8f5e9;
            transform: translateX(0);
            opacity: 1;
        }

        .step.error {
            background: #fff3cd;
            transform: translateX(0);
            opacity: 1;
        }

        .step-icon {
            width: 24px;
            height: 24px;
            margin-right: 12px;
            position: relative;
            display: flex;
            align-items: center;
            justify-content: center;
            transition: all 0.3s ease;
        }

        .step-loading {
            border: 3px solid #e0e0e0;
            border-radius: 50%;
            border-top: 3px solid #2196f3;
            animation: spin 1s linear infinite, pulse 1.5s infinite alternate;
        }

        .step-complete {
            color: #4caf50;
            transform: scale(1.2);
        }

        .step-error {
            color: #f44336;
            transform: scale(1.2) rotate(180deg);
        }

        .step-text {
            font-size: 14px;
            color: #757575;
        }

        /* Input Area */
        .input-area {
            padding: 20px;
            border-top: 1px solid #e0e0e0;
        }

        .input-form {
            display: flex;
            gap: 10px;
        }

        .input-field {
            flex: 1;
            padding: 12px 16px;
            border: 1px solid #e0e0e0;
            border-radius: 24px;
            font-size: 16px;
            transition: all 0.3s ease;
        }

        .input-field:focus {
            outline: none;
            border-color: #2196f3;
            box-shadow: 0 0 0 2px rgba(33, 150, 243, 0.2);
        }

        .send-button {
            background: #2196f3;
            color: white;
            border: none;
            border-radius: 24px;
            padding: 12px 24px;
            font-size: 16px;
            cursor: pointer;
            transition: all 0.3s ease;
        }

        .send-button:hover {
            background: #1976d2;
        }

        /* Animations */
        @keyframes spin {
            0% { transform: rotate(0deg); }
            100% { transform: rotate(360deg); }
        }

        @keyframes pulse {
            from { transform: scale(0.9); }
            to { transform: scale(1.1); }
        }

        @keyframes fadeIn {
            from {
                opacity: 0;
                transform: translateY(10px);
            }
            to {
                opacity: 1;
                transform: translateY(0);
            }
        }

        @keyframes fadeOut {
            from {
                opacity: 1;
                transform: translateY(0);
            }
            to {
                opacity: 0;
                transform: translateY(-10px);
            }
        }

        @keyframes staggerFadeIn {
            0% {
                opacity: 0;
                transform: translateX(-100%);
            }
            100% {
                opacity: 1;
                transform: translateX(0);
            }
        }
    </style>
</head>
<body>
    <div class="container">
        <div class="chat-container">
            <div class="chat-messages" id="chat-messages">
                <!-- Messages will be inserted here -->
            </div>

            <div class="input-area">
                <form class="input-form" id="chat-form">
                    <input type="text"
                           class="input-field"
                           id="user-input"
                           placeholder="Type your message..."
                           autocomplete="off">
                    <button type="submit" class="send-button">
                        <i class="fas fa-paper-plane"></i>
                    </button>
                </form>
            </div>
        </div>
    </div>

    <script>
        const chatMessages = document.getElementById('chat-messages');
        const chatForm = document.getElementById('chat-form');
        const userInput = document.getElementById('user-input');
        let ws = new WebSocket('ws://localhost:8000/ws');

        class ActionIndicator {
            constructor(action) {
                this.element = document.createElement('div');
                this.element.className = 'action-indicator';
                this.steps = [];
                this.currentStep = 0;
                this.render(action);
            }

            render(action) {
                this.element.innerHTML = `
                    <div class="action-header">
                        <i class="fas fa-cog fa-spin"></i>
                        <span class="action-title">${action.title}</span>
                    </div>
                    <div class="action-steps">
                        ${action.steps.map((step, index) => `
                            <div class="step ${index === 0 ? 'active' : ''}">
                                <div class="step-icon">${index === 0 ? '<i class="fas fa-spinner fa-spin"></i>' : ''}</div>
                                <span class="step-text">${step}</span>
                            </div>
                        `).join('')}
                    </div>
                `;

                this.steps = Array.from(this.element.querySelectorAll('.step'));
            }

            updateStep(index, status) {
                if (this.steps[index]) {
                    const step = this.steps[index];
                    const icon = step.querySelector('.step-icon');

                    // Update step classes
                    step.classList.remove('active', 'complete', 'error');
                    step.classList.add(status);

                    // Update icon
                    if (status === 'complete') {
                        icon.innerHTML = '<i class="fas fa-check-circle step-complete"></i>';
                    } else if (status === 'error') {
                        icon.innerHTML = '<i class="fas fa-times-circle step-error"></i>';
                    } else if (status === 'active') {
                        icon.innerHTML = '<i class="fas fa-spinner fa-spin"></i>';
                    }

                    // Activate next step
                    if (status === 'complete' && this.steps[index + 1]) {
                        this.steps[index + 1].classList.add('active');
                        const nextIcon = this.steps[index + 1].querySelector('.step-icon');
                        nextIcon.innerHTML = '<i class="fas fa-spinner fa-spin"></i>';
                    }

                    // Check if all steps are complete
                    if (status === 'complete' && index === this.steps.length - 1) {
                        this.hide();
                    }
                }
            }

            completeAction() {
                this.steps.forEach((step, index) => {
                    step.classList.add('complete');
                    const icon = step.querySelector('.step-icon');
                    icon.innerHTML = '<i class="fas fa-check-circle step-complete"></i>';
                });
            }

            errorAction() {
                this.steps.forEach((step, index) => {
                    step.classList.add('error');
                    const icon = step.querySelector('.step-icon');
                    icon.innerHTML = '<i class="fas fa-times-circle step-error"></i>';
                });
                this.hide();
            }

            hide() {
                this.element.style.animation = 'fadeOut 0.5s ease-in-out forwards';
                setTimeout(() => {
                    this.element.remove();
                }, 500);
            }
        }

        let currentActionIndicator = null;

        // WebSocket event handlers
        ws.onopen = () => {
            console.log('Connected to WebSocket');
            addMessage('Connected to AI Assistant', 'assistant');
        };

        ws.onclose = () => {
            console.log('Disconnected from WebSocket');
            addMessage('Connection closed. Please refresh the page.', 'assistant');
        };

        ws.onmessage = (event) => {
            try {
                const data = JSON.parse(event.data);

                if (data.type === 'action_start') {
                    // Create new action indicator
                    currentActionIndicator = new ActionIndicator({
                        title: data.title,
                        steps: data.steps
                    });
                    addElement(currentActionIndicator.element);
                } else if (data.type === 'action_update') {
                    // Update the current action indicator
                    if (currentActionIndicator) {
                        currentActionIndicator.updateStep(data.step, data.status);
                    }
                } else if (data.type === 'action_complete') {
                    // Action completed successfully
                    if (currentActionIndicator) {
                        currentActionIndicator.completeAction();
                    }
                    addMessage(data.message, 'assistant');
                } else if (data.type === 'action_error') {
                    // Action encountered an error
                    if (currentActionIndicator) {
                        currentActionIndicator.errorAction();
                    }
                    addMessage(`Error: ${data.message}`, 'assistant');
                } else {
                    // Regular message
                    addMessage(event.data, 'assistant');
                }
            } catch (e) {
                // If not JSON, treat as regular message
                addMessage(event.data, 'assistant');
            }
        };

        // Form submission handler
        chatForm.addEventListener('submit', (e) => {
            e.preventDefault();
            const message = userInput.value.trim();
            if (message) {
                // Send message to server
                ws.send(JSON.stringify({
                    message: message
                }));

                // Add user message to chat
                addMessage(message, 'user');

                // Clear input
                userInput.value = '';
            }
        });

        // Helper function to add messages to chat
        function addMessage(content, type) {
            const messageDiv = document.createElement('div');
            messageDiv.className = `message ${type}-message`;
            messageDiv.textContent = content;
            chatMessages.appendChild(messageDiv);
            chatMessages.scrollTop = chatMessages.scrollHeight;
        }

        // Helper function to add any element to chat
        function addElement(element) {
            chatMessages.appendChild(element);
            chatMessages.scrollTop = chatMessages.scrollHeight;
        }

        // Optional: Add error handling for WebSocket
        ws.onerror = (error) => {
            console.error('WebSocket error:', error);
            addMessage('Error: Could not connect to the server. Please try again later.', 'assistant');
        };
    </script>
</body>
</html>