Introduction

You have probably heard about Gemini AI already. Gemini is a family of generative AI models that can generate content and solve problems. There are different models with their own set of capabilities. We will be focusing on the Gemini 1.0 Pro model. Throughout this tutorial, we will be using the name gemini-pro. It is an alias for Gemini 1.0 Pro. To know more about gemini models, Refer Here.

We will be creating an AI Chat Bot powered by Gemini AI, with JWT authentication using python.

You can find the complete code on my Github repo: Converse3.0

Tech Stack

• Python

• Flask

• PostgreSQL

• Docker

Let's Start

Obtain API Key to use Gemini

1. Go to Google AI Studio

2. Login with google account

3. Create API Key

Install PostgreSQL

Download and install postgreSQL, from here. Go with the latest version.

If you are using linux, make sure that you install libpq-dev

• Linux (with apt package manager)

    sudo apt-get install libpq-dev
  

Install Docker (Optional)

Installing docker is completely optional. Having the Dockerfile and docker compose setup will help in faster deployment.

You can also spin up a docker container for postgreSQL instead of installing it in your system.

docker run -d \
  --name gemini_postgres \
  -p 5432:5432 \
  -e POSTGRES_USER=postgres \
  -e POSTGRES_PASSWORD=password\
  -e POSTGRES_DB=postgres\
  postgres:latest

Postgres now accessible on port 5432.

Setup Environment

Create Virtual Environment

python -m venv venv

Activate Virtual Environment

• Windows (CMD)

    venv\Scripts\activate
  

• Mac/Linux

    source venv/bin/activate
  

Install Python Packages

With the virtual environment activated, you can install Python packages using pip.

pip install Flask
pip install Flask-JWT-Extended
pip install Flask-SQLAlchemy
pip install Flask-Cors

• Flask is a lightweight and flexible web framework for Python

• Flask-JWT-Extended - provides support for JSON Web Tokens (JWTs). We will be using JWTs for implementing authentication and authorization in this web application.

• Flask-SQLAlchemy - integrates SQLAlchemy with Flask. We will be using SQLAlchemy to interact with PostgreSQL.

• Flask-Cors - enables Cross-Origin Resource Sharing (CORS)

pip install psycopg2

• Psycopg2 is a PostgreSQL database adapter

pip install python-decouple

• Decouple helps in loading values from .env files

pip install jsonpickle

• jsonpickle helps in the conversion of complex python object to JSON and vice-versa. We will use this library to convert and store the chat history.

Install Python SDK for Gemini API

pip install google-generativeai

Start Coding

Configure Python Decouple

Python decouple automatically picks up .env file. But, since we are also planning to deploy the application into production, we need it to pick up a different .env when in production environment. So, rather than using the default config provided by decouple, we will create a method to return a decouple config which includes our additional logic.

import os
import decouple
from decouple import RepositoryEnv
import pathlib


class DecoupleConfigUtil:

    @classmethod
    def get_env_config(cls) -> decouple.Config:
        """
        Creates and returns a Config object based on the environment setting.
        It uses .env for development and .prod.env for production.
        """

        ENVIRONMENT = os.getenv("ENVIRONMENT", default="DEVELOPMENT")

        env_files = {
            "DEVELOPMENT": ".env",
            "PRODUCTION": ".env.prod",
        }

        app_dir_path = pathlib.Path(__file__).resolve().parent.parent
        env_file_name = env_files.get(ENVIRONMENT, ".env")
        file_path = app_dir_path / env_file_name

        if not file_path.is_file():
            raise FileNotFoundError(f"Environment file not found: {file_path}")

        return decouple.Config(RepositoryEnv(file_path))

We create this class inside utils package. This would check for an environment variable called "ENVIRONMENT" in the system. According to the variable's value, it would decide from which .env, it should take the configurations from.

In the production environment, it would take values from .env.prod. So, make sure that you have this file when deploying the application to production.

You can add these values to the .env file:

APP_SECRET_KEY=this_is_my_secret_key
JWT_SECRET_KEY=this_is_my_secret_jwt_secret_key

CORS_ORIGINS=http://localhost:5173,http://localhost:4173

DATABASE_URI=postgresql://postgres:password@localhost:5432/postgres

HOST=0.0.0.0
PORT=8000

GOOGLE_API_KEY=AIzaSyDTzAF3jNsbktskJLC_EIBz0_QKPFdnHds

Create file for storing static messages

We create a messages.py file to store all the static messages we need.

USERNAME_REQUIRED = "Username is required"
PASSWORD_REQUIRED = "Password is required"
USERNAME_EXISTS = "Username already exists"
USER_REGISTRATION_SUCCESSFUL = "User has been successfully registered"

INVALID_USERNAME_PASSWORD = "Username or Password is invalid"

INVALID_PROMPT = "Provided prompt is invalid"

CHAT_HISTORY_LIST_RETRIEVED = "Chat history list has been successfully retrieved"
CHAT_HISTORY_RETRIEVED = "Chat history has been successfully retrieved"
CHAT_HISTORY_UNAVAILABLE = "Chat history is not available"

CHAT_RESPONSE_NOT_SAFE = "Chat response is not safe to be viewed publicly"
CHAT_NOT_FOUND = "Chat not found. Please provide correct details"

Create Models

We create models using SQLAlchemy and place these models inside the models package.

models.py

from flask_sqlalchemy import SQLAlchemy


db = SQLAlchemy()


def bulk_save_objects(objects):
    db.session.bulk_save_objects(objects)
    db.session.commit()


def delete_objects(objects):
    for obj in objects:
        db.session.delete(obj)

    db.session.commit()

Here, we initialize an instance of SQLAlchemy. This instance will be used to create further models. We also define two methods to bulk save objects and bulk delete objects.

user.py

from models.models import db
from models.document import Document
from werkzeug.security import generate_password_hash, check_password_hash


class User(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(80), unique=True, nullable=False)
    password_hash = db.Column(db.String(256), nullable=False)

    documents = db.relationship(Document, back_populates="user")

    def __repr__(self):
        return "<User %r>" % self.username

    def set_password(self, password):
        self.password_hash = generate_password_hash(password)

    def check_password(self, password):
        return check_password_hash(self.password_hash, password)

    def save(self):
        db.session.add(self)
        db.session.commit()

We specify relationship between User model and Document model.

document.py

from models.models import db


class Document(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    document_name = db.Column(db.String(80), nullable=False)
    user_id = db.Column(db.Integer, db.ForeignKey("user.id"), nullable=False)

    user = db.relationship("User", back_populates="documents")

    def __repr__(self):
        return "<UploadedDocuments %r>" % self.document_name

    def save(self):
        db.session.add(self)
        db.session.commit()

The Document model has a one-one relationship with the User model.

normal_chat_history.py

from models.models import db
from datetime import datetime


class NormalChatHistory(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    chat_history = db.Column(db.JSON)
    started_at = db.Column(db.DateTime, default=datetime.utcnow)
    user_id = db.Column(db.Integer, db.ForeignKey("user.id"), nullable=False)

    def __repr__(self):
        return "<NormalChatHistory %r>" % self.id + "-" + self.user_id

    def save(self):
        db.session.add(self)
        db.session.commit()

The chat history is stored in JSON format. NormalChatHistory model has a one-one relationship with the User model.

Create Auth Controllers

User Registration

We place the controllers inside the controllers package.

register.py

from flask import request

from utils.common_response import CommonResponse
from service.user_registration import UserRegistration


def register():
    username = request.json.get("username", None)
    password = request.json.get("password", None)

    user_registration = UserRegistration()
    message, status = user_registration.register(username, password)

    return CommonResponse(message, status).format_response()

We also create custom response files to structure the response as we want. Create a file called common_response.py in utils package.

common_response.py

from flask import jsonify


class CommonResponse:

    def __init__(self, message, status, data=None) -> None:
        self.message = message
        self.status = status
        self.data = data

    def format_response(self):
        return (
            jsonify({"message": self.message, "data": self.data}),
            self.status,
        )

User Login

We place it in the controllers package.

login.py

from flask import request

from utils.login_response import LoginResponse
from utils.common_response import CommonResponse
from service.user_login import UserLogin


def login():
    username = request.json.get("username", None)
    password = request.json.get("password", None)

    user_login = UserLogin()
    response, status = user_login.login(username, password)
    if status != 200:
        return CommonResponse(response, status)
    else:
        return LoginResponse(response, status)

Create a file called login_response.py in utils package.

login_response.py

from flask import jsonify


class LoginResponse:

    def __init__(self, access_token, status) -> None:
        self.access_token = access_token
        self.status = status

    def format_response(self):
        return (
            jsonify({"access_token": self.access_token}),
            self.status,
        )

Create Auth Services

User Registration

We place the service classes inside the services package. This file will contain all the logic required for registering a new user.

user_registration.py

from models.user import User
from messages import (
    PASSWORD_REQUIRED,
    USER_REGISTRATION_SUCCESSFUL,
    USERNAME_EXISTS,
    USERNAME_REQUIRED,
)


class UserRegistration:

    def register(self, username, password):
        if not username:
            message = USERNAME_REQUIRED
            status = 400
        elif not password:
            message = PASSWORD_REQUIRED
            status = 400
        elif User.query.filter_by(username=username).first():
            message = USERNAME_EXISTS
            status = 400
        else:
            new_user = User(username=username)
            new_user.set_password(password)
            new_user.save()
            message = USER_REGISTRATION_SUCCESSFUL
            status = 201

        return message, status

User Login

This file will contain all the logic required for successfully logging in a valid user.

user_login.py

from flask_jwt_extended import create_access_token

from messages import INVALID_USERNAME_PASSWORD
from models.user import User


class UserLogin:

    def login(self, username, password):
        user = User.query.filter_by(username=username).first()
        if user and user.check_password(password):
            additional_claims = {"user_id": user.id}
            access_token = create_access_token(
                identity=username, additional_claims=additional_claims
            )
            return access_token, 200
        else:
            return INVALID_USERNAME_PASSWORD, 400

Add these lines to main.py:

main.py

from flask import Flask
from flask_cors import CORS
from flask_jwt_extended import JWTManager

from messages import CHAT_NOT_FOUND, CHAT_RESPONSE_NOT_SAFE
from exceptions import ChatNotFoundException, SafetyException
from models.models import db
from utils.decouple_config_util import DecoupleConfigUtil
from utils.common_response import CommonResponse
from controllers.register import register
from controllers.login import login
from controllers.normal_ai_chat import (
    NormalAiChat,
    NormalAiChatHistory,
    NormalAiChatHistoryList,
    deleteNormalAiChatHistory,
)


config = DecoupleConfigUtil.get_env_config()

app = Flask(__name__)
app.secret_key = config("APP_SECRET_KEY")

app.config["SQLALCHEMY_DATABASE_URI"] = config("DATABASE_URI")
app.config["JWT_SECRET_KEY"] = config("JWT_SECRET_KEY")

db.init_app(app)
jwt = JWTManager(app)

CORS(
    app,
    origins=config(
        "CORS_ORIGINS", cast=lambda v: [item.strip() for item in v.split(",")]
    ),
)

with app.app_context():
    db.create_all()


app.route("/register", methods=["POST"])(register)
app.route("/login", methods=["POST"])(login)

if __name__ == "__main__":
    app.run(host=config("HOST"), port=config("PORT"))

• Initialize the Flask application

• Initialize the DB

• Specify the API routes

Chat with AI

Custom Exceptions

Create custom exceptions and place it in exceptions.py.

class SafetyException(Exception):
    pass


class ChatNotFoundException(Exception):
    pass

We create an abstract class in client/ai_models. This abstract class will serve as the parent class for the other AI model classes.

ai_model.py

from abc import ABC, abstractmethod


class AIModel(ABC):

    @abstractmethod
    def chat(chat_history):
        pass

    @abstractmethod
    def get_chat_history():
        pass

Gemini AI Model

Since we are using Gemini as our AI model, we will put all the Gemini related files inside client/ai_models/gemini.

First, lets setup the required configurations in config.py.

config.py

from utils.decouple_config_util import DecoupleConfigUtil


config = DecoupleConfigUtil.get_env_config()


GOOGLE_API_KEY = config("GOOGLE_API_KEY")
GEMINI_MODEL_NAME = config("GEMINI_MODEL_NAME")

gemini_ai_model.py

import google.generativeai as genai
from google.generativeai.types.generation_types import StopCandidateException
import logging

from client.ai_models.ai_model import AIModel
from client.ai_models.gemini.config import GEMINI_MODEL_NAME, GOOGLE_API_KEY
from exceptions import SafetyException


logging.basicConfig(
    filename="app.log",
    level=logging.DEBUG,
    format="%(asctime)s %(levelname)s %(name)s %(threadName)s : %(message)s",
)

logger = logging.getLogger(__name__)


class GeminiAIModel(AIModel):

    def __init__(self, model_name=GEMINI_MODEL_NAME) -> None:
        genai.configure(api_key=GOOGLE_API_KEY, transport="rest")
        self.model = genai.GenerativeModel(model_name)

    def chat(self, prompt, chat_history):
        try:
            self.chat = self.model.start_chat(history=chat_history)
            response = self.chat.send_message(prompt)
            content_parts = response.candidates[0].content.parts[0]
            return content_parts.text
        except StopCandidateException as sce:
            logger.error(sce)
            raise SafetyException()

    def get_chat_history(self):
        if self.chat:
            return self.chat.history
        return []

• Initializes the Google AI python SDK with the necessary values to interact with Gemini

• chat method to chat with the AI model

• Raises the custom exception SafetyException if the StopCandidateException is caught

• get_chat_history method to get the entire history of that chat

Create AI Chat Controllers

In the controllers package, create file called normal_ai_chat.py.

normal_ai_chat.py

from flask import request
from flask_jwt_extended import get_jwt, jwt_required

from client.ai_models.gemini.gemini_ai_model import GeminiAIModel
from utils.common_response import CommonResponse
from messages import (
    CHAT_HISTORY_LIST_RETRIEVED,
    CHAT_HISTORY_RETRIEVED,
    CHAT_HISTORY_UNAVAILABLE,
    INVALID_PROMPT,
)
from service.normal_ai_chatter import NormalAiChatter
from utils.chat_response import ChatResponse


@jwt_required()
def NormalAiChat():
    prompt = request.json.get("prompt", None)
    chat_id = request.json.get("chat_id", None)
    user_id = get_jwt().get("user_id", None)

    if prompt:
        gemini_ai_model = GeminiAIModel()
        normal_chatter = NormalAiChatter(gemini_ai_model)
        ai_response, chat_id = normal_chatter.chat(prompt, chat_id, user_id)
        status = 200
    else:
        ai_response = INVALID_PROMPT
        status = 400

    return ChatResponse(ai_response, chat_id, status).format_response()

• Calls the Gemini model with the user prompt

@jwt_required()
def NormalAiChatHistoryList():
    user_id = get_jwt().get("user_id", None)

    normal_chatter = NormalAiChatter()
    chat_list = normal_chatter.get_chat_history_list(user_id)
    status = 200

    return CommonResponse(CHAT_HISTORY_LIST_RETRIEVED, status, chat_list).format_response()

• Retrieves the list of chats the user has made with the model

@jwt_required()
def NormalAiChatHistory():
    user_id = get_jwt().get("user_id", None)
    chat_id = request.args.get("chat_id", None)

    normal_chatter = NormalAiChatter()
    chat_history = normal_chatter.get_chat_history(user_id, chat_id)
    if chat_history:
        status = 200
        return CommonResponse(CHAT_HISTORY_RETRIEVED, status, chat_history).format_response()

    status = 404
    return CommonResponse(CHAT_HISTORY_UNAVAILABLE, status, None).format_response()

• Retrieves all the contents of a particular chat

@jwt_required()
def deleteNormalAiChatHistory():
    user_id = get_jwt().get("user_id", None)
    chat_id = request.args.get("chat_id", None)

    normal_chatter = NormalAiChatter()
    normal_chatter.delete_chat_history(user_id, chat_id)

    status = 204
    return "", status

• Delete chats based on user_id and chat_id

Create a file called chat_response.py in utils package.

chat_response.py

from flask import jsonify


class ChatResponse:

    def __init__(self, ai_response, chat_id, status) -> None:
        self.ai_response = ai_response
        self.chat_id = chat_id
        self.status = status

    def format_response(self):
        return (
            jsonify(
                {
                    "ai_response": self.ai_response,
                    "chat_id": self.chat_id,
                }
            ),
            self.status,
        )

Create AI Chat Service

We will create a file in the services package that handles chatting with the AI model.

normal_ai_chatter.py

import jsonpickle

from exceptions import ChatNotFoundException
from models.models import delete_objects
from models.normal_chat_history import NormalChatHistory


class NormalAiChatter:

    def __init__(self, model=None) -> None:
        self.model = model

    def chat(self, prompt, chat_id, user_id):
        chat_history = []
        if chat_id:
            normal_chat_history = (
                NormalChatHistory.query.filter_by(user_id=user_id)
                .filter_by(id=chat_id)
                .first()
            )
            if not normal_chat_history:
                raise ChatNotFoundException()
            chat_history = jsonpickle.decode(normal_chat_history.chat_history)
        else:
            normal_chat_history = NormalChatHistory(
                chat_history=chat_history, user_id=user_id
            )
            normal_chat_history.save()
            chat_id = normal_chat_history.id

        ai_response = self.model.chat(prompt, chat_history)
        self.save_chat_history(normal_chat_history)

        return ai_response, chat_id

    def save_chat_history(self, normal_chat_history):
        chat_history = self.model.get_chat_history()
        chat_history_json_string = jsonpickle.encode(chat_history)
        normal_chat_history.chat_history = chat_history_json_string
        normal_chat_history.save()

        return True

• If there is a chat_id, then it will fetch the chat history corresponding to that chat_id. Otherwise, it would create a new chat.

• After getting the response from the model, it would save the entire chat history.

def get_chat_history_list(self, user_id):
        chat_history_list = NormalChatHistory.query.filter_by(user_id=user_id).all()
        chat_list = [
            {"chat_id": chat.id, "chat_name": chat.started_at}
            for chat in chat_history_list
        ]

        return chat_list

• Get the list of all the chats of the particular user

def get_chat_history(self, user_id, chat_id):
        normal_chat_history = (
            NormalChatHistory.query.filter_by(user_id=user_id)
            .filter_by(id=chat_id)
            .first()
        )

        if normal_chat_history:
            chat_history = jsonpickle.decode(normal_chat_history.chat_history)
            chat_history_data = {
                "chat_id": chat_id,
                "chat_history": [
                    {"text": chat.parts[0].text, "role": chat.role}
                    for chat in chat_history
                ],
            }
            return chat_history_data

        return []

• Get all the contents of a particular chat of that user

def delete_chat_history(self, user_id, chat_id):
        if chat_id:
            normal_chat_history = (
                NormalChatHistory.query.filter_by(user_id=user_id)
                .filter_by(id=chat_id)
                .all()
            )
        else:
            normal_chat_history = NormalChatHistory.query.filter_by(
                user_id=user_id
            ).all()

        delete_objects(normal_chat_history)

        return True

• Delete a particular chat of the user if the chat_id is provided

• Delete all the chats of the user if the chat_id is not provided

Add Chat Routes

Update main.py with the following routes.

main.py

from controllers.normal_ai_chat import (
    NormalAiChat,
    NormalAiChatHistory,
    NormalAiChatHistoryList,
    deleteNormalAiChatHistory,
)


app.route("/normal_chat_with_ai", methods=["POST"])(NormalAiChat)
app.route("/normal_chat_history_list", methods=["GET"])(NormalAiChatHistoryList)
app.route("/normal_chat_history", methods=["GET"])(NormalAiChatHistory)
app.route("/normal_chat_history_list", methods=["DELETE"])(deleteNormalAiChatHistory)

Handle Custom Exceptions

Update main.py with the following code.

main.py

@app.errorhandler(SafetyException)
def handle_safety_exception(error):
    status = 403
    return CommonResponse(CHAT_RESPONSE_NOT_SAFE, status, None).format_response()


@app.errorhandler(ChatNotFoundException)
def handle_chat_not_found_exception(error):
    status = 404
    return CommonResponse(CHAT_NOT_FOUND, status, None).format_response()

Run the Application

We can now run the application by executing:

python main.py

• Make sure that PostgreSQL is already up and configured correctly

• Ensure that the correct details are provided in the .env file

• You can access the application on http://localhost:8000

Dockerize the Application (Optional)

Add these lines to the Dockerfile.

Dockerfile

FROM python:3.10-bookworm

ENV PYTHONDONTWRITEBYTECODE 1
ENV PYTHONUNBUFFERED 1
ENV ENVIRONMENT DEVELOPMENT

WORKDIR /app

RUN apt-get update -y 
RUN apt-get install -y git ca-certificates
RUN apt-get update && \
    apt-get install -y \
    python3-dev \
    libpq-dev \
    libmariadb-dev \
    build-essential \
    cmake \
    swig \
    pkgconf

RUN apt-get install curl
RUN curl -k -L https://github.com/jwilder/dockerize/releases/download/v0.6.1/dockerize-linux-amd64-v0.6.1.tar.gz | tar -C /usr/local/bin -xzv

COPY requirements.txt /app/
RUN pip install wheel
RUN pip install --upgrade pip
RUN pip install pyarrow
RUN pip install pyOpenSSL
RUN pip install -r requirements.txt

COPY . /app/

EXPOSE 5000
EXPOSE 8000

RUN chmod +x deploy.sh

CMD ["python", "main.py"]

Run in Production

To run the application in production, use a WSGI server like Gunicorn to serve the Flask application.

Add the configurations in .prod.env instead of .env when in production