Building an AI-Powered Chat Interface Using FastAPI and Gemini

#python #fastapi #google #ai

In this blog, we'll walk through creating a WebSocket endpoint using FastAPI to handle real-time chat messages. WebSockets provide a full-duplex communication channel over a single TCP connection, which is perfect for applications requiring real-time updates like chat applications. For frontend we will create a simple html from with one text field, one send button, and you will be able to see the text you text as you and text from gemini as AI.

here are is demo of that product

Here I am talking to the AI and tell it about the me and my friends play dota 2.

I am creating simple and cool things for beginners.
let's start
First create project in new directory then install following Package

pip install fastapi
pip install google-generativeai 
pip install python-dotenv

Go to this following link to get gemini api key
https://aistudio.google.com/app/apikey

in the above picture you can see that *create API key * click on it and get your api key

Now, create a file named .env in your project directory and add your key like this:
API_KEY = "here you copy paste the API key "

This is blog I will show you how can add LLM chat using fastapi

first import all libraries


import os
import google.generativeai as genai
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi import APIRouter
from dotenv import load_dotenv

then load you API key

load_dotenv()


API_KEY = os.getenv("API_KEY")

Now we are going to create fastapi app which

app = FastAPI(
    title="AI Chat API",
    docs_url='/',
    description="This API allows you to chat with an AI model using WebSocket connections.",
    version="1.0.0"
)

after that you need to add some permisson and in the cast I am open it for all but in the production you cannot do this.

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

creating router and loading the model name gemini-1.5-flash

router = APIRouter(prefix="/chat", tags=["Chat"])
genai.configure(api_key=API_KEY)
model = genai.GenerativeModel('gemini-1.5-flash')

Setting Up the WebSocket Endpoint

We'll start by defining our WebSocket endpoint. This endpoint will allow clients to send messages to our AI model and receive streamed responses.

@router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """
    WebSocket endpoint for handling chat messages.

    This WebSocket endpoint allows clients to send messages to the AI model
    and receive streamed responses.

    To use this endpoint, establish a WebSocket connection to `/ws`.

    - Send a message to the WebSocket.
    - Receive a response from the AI model.
    - If the message "exit" is sent, the chat session will end.
    """
    await websocket.accept()
    chat = model.start_chat(history=[])
    try:
        while True:
            data = await websocket.receive_text()
            if data.lower().startswith("you: "):
                user_message = data[5:]
                if user_message.lower() == "exit":
                    await websocket.send_text("AI: Ending chat session.")
                    break
                response = chat.send_message(user_message, stream=True)
                full_response = ""
                for chunk in response:
                    full_response += chunk.text
                await websocket.send_text("AI: " + full_response)
            else:
                await websocket.send_text("AI: Please start your message with 'You: '")
    except WebSocketDisconnect:
        print("Client disconnected")
    finally:
        await websocket.close()

Setting Up the Chat Interface

We'll start by defining an HTTP endpoint that serves an HTML page. This page contains a form for sending messages and a script for handling WebSocket communication.

Here's the code for the HTTP endpoint:
Note: I am here just showing you how to js is working I have not include the CSS that make chat box like shown above for that check following repo on github and please give it start:
https://github.com/GoAndPyMasters/fastapichatbot/tree/main

@router.get("/")
async def get():
    return HTMLResponse("""
    <html>
        <head>
            <title>Chat</title>
        </head>
        <body>
            <h1>Chat with AI</h1>
            <form action="" onsubmit="sendMessage(event)">
                <input type="text" id="messageText" autocomplete="off"/>
                <button>Send</button>
            </form>
            <ul id="messages">
            </ul>
            <script>
                var ws = new WebSocket("ws://0.0.0.0:8080/ws");
                ws.onmessage = function(event) {
                    var messages = document.getElementById('messages')
                    var message = document.createElement('li')
                    var content = document.createTextNode(event.data)
                    message.appendChild(content)
                    messages.appendChild(message)
                };
                function sendMessage(event) {
                    var input = document.getElementById("messageText")
                    ws.send("You: " + input.value)
                    input.value = ''
                    event.preventDefault()
                }
            </script>
        </body>
    </html>
    """)

Now to run the fastapi APP use following command

uvicorn main:app --host 0.0.0.0 --port 8000

after that use will see

then ctrl+click on this link http://0.0.0.0:8000
then show the docs like this

then go to the url and write /chat/ like as shown in the image

and then you are ready to chat

Conclusion

In this blog, we demonstrated how to create a real-time chat application using FastAPI and WebSockets, powered by a generative AI model. By following the steps outlined, you can set up a fully functional WebSocket endpoint and a basic HTML interface to interact with the AI. This combination allows for seamless, real-time communication, making it a powerful solution for chat applications and other real-time systems. With the provided code and instructions, you're equipped to build and customize your own AI-powered chat interface.

check the code on following repo on github
https://github.com/GoAndPyMasters/fastapichatbot/tree/main

here is my github profile

https://github.com/MuhammadNizamani

If you need any help contact with on linkedin
https://www.linkedin.com/in/muhammad-ishaque-nizamani-109a13194/

Please like and comment it will give me motivation to write new things
Thanks for Reading