Migrate from REST API to SSR for Receiving new message

[farhoud]

Issue Title

Migrate from REST API to WebSocket for Receiving LLM Responses

Description

Currently, our application uses a single POST API endpoint to send messages and receive responses from the LLM (Language Learning Model). This synchronous approach can lead to performance bottlenecks, especially when handling multiple requests or waiting for long responses from the LLM. To improve efficiency and responsiveness, we should transition to using WebSockets for receiving messages.

Current Implementation

• Technology Stack: Next.js (Frontend), FastAPI (Backend)
• Communication Method: HTTP POST
• Endpoint: /api/message
• Process:
  1. The frontend sends a POST request with the message.
  2. The backend processes the request and sends the response back via the same POST request.

Proposed Change

• New Communication Method: WebSocket
• WebSocket Endpoint: /ws/message
• Process:
  1. The frontend establishes a WebSocket connection to the backend.
  2. The frontend sends messages through the WebSocket connection.
  3. The backend processes the message and sends the response back through the WebSocket connection.

Benefits

• Asynchronous Communication: Allows the backend to send responses as soon as they are ready without keeping the HTTP connection open.
• Improved Performance: Reduces the load on the server by maintaining a single connection for multiple messages.
• Better User Experience: More responsive application with faster message handling.

Tasks

1. Backend (FastAPI)

   • Set up a WebSocket endpoint.
   • Implement message handling through WebSocket.
   • Ensure proper connection management (e.g., handling disconnections).

2. Frontend (Next.js)

   • Establish a WebSocket connection to the backend.
   • Modify the message sending logic to use WebSocket instead of HTTP POST.
   • Update the UI to handle asynchronous responses from the WebSocket.

3. Testing

   • Ensure that the WebSocket connection is stable and handles reconnections.
   • Test message sending and receiving to ensure proper functionality.
   • Performance testing to compare the new WebSocket implementation with the existing POST API.

Example Code

Backend (FastAPI)

from fastapi import FastAPI, WebSocket, WebSocketDisconnect

app = FastAPI()

@app.websocket("/ws/message")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            # Process the data and generate response
            response = process_message(data)
            await websocket.send_text(response)
    except WebSocketDisconnect:
        print("Client disconnected")

def process_message(message: str) -> str:
    # LLM processing logic
    return "Processed response"

Frontend (Next.js)

import React, { useEffect, useRef, useState } from 'react';

const MessageComponent = () => {
  const [message, setMessage] = useState('');
  const [response, setResponse] = useState('');
  const ws = useRef(null);

  useEffect(() => {
    ws.current = new WebSocket('ws://localhost:8000/ws/message');
    ws.current.onmessage = (event) => {
      setResponse(event.data);
    };
    ws.current.onclose = () => {
      console.log('WebSocket connection closed');
    };
    return () => {
      ws.current.close();
    };
  }, []);

  const sendMessage = () => {
    if (ws.current.readyState === WebSocket.OPEN) {
      ws.current.send(message);
    }
  };

  return (
    <div>
      <input
        type="text"
        value={message}
        onChange={(e) => setMessage(e.target.value)}
      />
      <button onClick={sendMessage}>Send Message</button>
      <div>Response: {response}</div>
    </div>
  );
};

export default MessageComponent;

Additional Notes

• Ensure security measures are in place for WebSocket communication.
• Monitor the WebSocket server for any potential issues such as memory leaks or connection limits.
• Update the documentation to reflect the changes in the communication method.

References

• FastAPI WebSocket Documentation
• Next.js WebSocket Integration