Implementing chat functionality with WebSockets

Real-time chat functionality is a common feature in many modern web applications, enabling users to communicate instantly with each other. In this recipe, we’ll explore how to implement chat functionality using WebSockets in FastAPI applications.

By leveraging WebSockets, we will create a bidirectional communication channel between the server and multiple clients, allowing messages to be sent and received in real time.

Getting ready

To follow the recipe, you need to have a good understanding of WebSockets and know how to build a WebSocket endpoint using FastAPI.

Additionally, having some basic knowledge of HTML and JavaScript can help create simple web pages for the application. The recipe we’ll be using is the foundation of our chat application.

Also, we will use the jinja2 package to apply basic templating for the HTML page. Make sure to have it in your environment. If you didn’t install packages with requirements.txt, install jinja2 with pip:

$ pip install jinja2

Once the installation is complete, we are ready to start with the recipe.

How to do it…

To build the application, we will need to build three core pieces – the WebSocket connections manager, the WebSocket endpoint, and the chat HTML page:

1. Let’s start by building the connection manager. The role of the connection manager is to keep track of open WebSocket connections and broadcast messages to active ones. Let’s define the ConnectionManager class in a dedicated ws_manager.py module under the app folder:

   import asyncio
   from fastapi import WebSocket
   class ConnectionManager:
       def __init__(self):
           self.active_connections: list[WebSocket] = []
       async def connect(self, websocket: WebSocket):
           await websocket.accept()
           self.active_connections.append(websocket)
       def disconnect(self, websocket: WebSocket):
           self.active_connections.remove(websocket)
       async def send_personal_message(
           self, message: dict, websocket: WebSocket
       ):
           await websocket.send_json(message)
       async def broadcast(
           self, message: json, exclude: WebSocket = None
       ):
           tasks = [
               connection.send_json(message)
               for connection in self.active_connections
               if connection != exclude
           ]
           await asyncio.gather(*tasks)

   The async def connect method will be responsible for the handshake and adding the WebSocket to the list of active ones. The def disconnect method will remove the WebSocket from the list of active connections. The async def send_personal_message method will send a message to a specific WebSocket. Finally, async def broadcast will send the message to all the active connections except one, if specified.

   The connection manager will then be used in the chat WebSocket endpoint.

2. Let’s create the WebSocket endpoint in a separate module called chat.py. Let’s initialize the connection manager:

   from app.ws_manager import ConnectionManager
   conn_manager = ConnectionManager()

   Then we define the router:

   from fastapi import APIRouter
   router = APIRouter()

   And finally, we can define the WebSocket endpoint:

   from fastapi import WebSocket, WebSocketDisconnect
   @router.websocket("/chatroom/{username}")
   async def chatroom_endpoint(
       websocket: WebSocket, username: str
   ):
       await conn_manager.connect(websocket)
       await conn_manager.broadcast(
           f"{username} joined the chat",
           exclude=websocket,
       )
       try:
           while True:
               data = await websocket.receive_text()
               await conn_manager.broadcast(
                   {"sender": username, "message": data},
                   exclude=websocket,
               )
               await conn_manager.send_personal_message(
                   {"sender": "You", "message": data},
                   websocket,
               )
       except WebSocketDisconnect:
           conn_manager.disconnect(websocket)
           await connection_manager.broadcast(
               {
                   "sender": "system",
                   "message": f"Client #{username} "
                   "left the chat",
               }
           )

3. After a new client joins a chat, the connection manager sends a message to all chat participants to notify them of the new arrival. The endpoint uses the username path parameter to retrieve the client’s name. Don’t forget to add the router to the FastAPI object in the main.py file:

   from app.chat import router as chat_router
   # rest of the code
   app = FastAPI()
   app.include_router(chat_router)

   Once the WebSocket endpoint is ready, we can create the endpoint to return the HTML chat page.

4. The page endpoint will return an HTML page rendered with Jinja2.

   The HTML chat page named chatroom.html should be stored in a templates folder in the project root. We will keep the page simple with the JavaScript tag embedded.

   The HTML part will look like this:

   <!doctype html>
   <html>
     <head>
       <title>Chat</title>
     </head>
     <body>
       <h1>WebSocket Chat</h1>
       <h2>Your ID: <span id="ws-id"></span></h2>
       <form action="" onsubmit="sendMessage(event)">
         <input
           type="text"
           id="messageText"
           autocomplete="off"
         />
         <button>Send</button>
       </form>
       <ul id="messages"></ul>
       <script>
           <!—content of js script -->
       <script/>
   </body>
   </html>

   The <script> tag will contain the Javascript code that will connect to the WebSocket /chatroom/{username} endpoint with the client name as a parameter, send the message from the client page, receive messages from the server, and render the message text on the page in the messages list section.

   You can find an example in the GitHub repository, in the templates/chatroom.html file. Feel free to make your own version or download it.

5. To conclude, we need to build the endpoint that returns the HTML page. We can build it in the same chat.py module:

   from fastapi.responses import HTMLResponse
   from fastapi.templating import Jinja2Templates
   from app.ws_manager import ConnectionManager
   conn_manager = ConnectionManager()
   templates = Jinja2Templates(directory="templates")
   @router.get("/chatroom/{username}")
   async def chatroom_page_endpoint(
       request: Request, username: str
   ) -> HTMLResponse:
       return templates.TemplateResponse(
           request=request,
           name="chatroom.html",
           context={"username": username},
       )

The endpoint will take as a path parameter the username of the client that will show in the chat conversation.

You have set up a basic chat room within your FastAPI application with the WebSockets protocol. You only have to spin up the server with uvicorn app.main:app and connect to http://localhost:8000/chatroom/your-username from your browser. Then, from another page, connect to the same address with a different username and start exchanging messages between the two browsers.

How it works…

When connecting to the GET /chatroom/{username} endpoint address (http://localhost:8000/chatroom/{username}), the server will use the username to render the HTML page customized to the username.

The HTML will contain the code to make the connection to the /chatroom WebSocket endpoint and create a new WebSocket connection for each user.

The endpoint will then use the ConnectionManager() connection manager object to exchange messages between all clients through the HTML page.

See also

We have used a basic feature of the Jinja2 templating library. However, you can free your creativity and discover the potential of this package by looking at the documentation:

• Jinja2 Documentation: https://jinja.palletsprojects.com/en/3.1.x/