app.py

import os
import gradio as gr
from config import SanatanConfig
from drive_downloader import ZipDownloader
from graph_helper import generate_graph
import uuid
import logging
from dotenv import load_dotenv

logging.basicConfig()
logger = logging.getLogger()
logger.setLevel(logging.INFO)

graph = generate_graph()

def init():
    load_dotenv(override=True)

    downloader = ZipDownloader(
        service_account_json=os.getenv("GOOGLE_SERVICE_ACCOUNT_JSON")
    )
    # 1. Download zip file
    zip_path = downloader.download_zip_from_drive(
        file_id=os.getenv("CHROMADB_FILE_ID"),
        output_path=SanatanConfig.dbStorePath,
    )

    # 2. Extract it
    downloader.unzip(zip_path, extract_to="./")

def init_session():
    # Generate a unique ID per browser session
    return str(uuid.uuid4())


def chat(message, history, thread_id):
    print("received chat message for thread:", thread_id)
    config = {"configurable": {"thread_id": thread_id}}
    response = graph.invoke(
        {"messages": [{"role": "user", "content": message}]}, config=config
    )
    return response["messages"][-1].content


thread_id = gr.State(init_session)
supported_scriptures = "\n - ".join(
    [scripture["title"] for scripture in SanatanConfig.scriptures]
)
intro_messages = [
    {"role" : "assistant","content" : f"Namaskaram! I can assist with the following scriptures: \n - {supported_scriptures}"},
    {"role" : "assistant","content" : """
    #### You can ask me questions like:
      - How do all scriptures describe the form of Vishnu?
      - What did Krishna teach in Gita?
      - How did Arjun react on seeing the Vishwarupa form of Krishna?
      - Give some names of Vishnu from Sahasranamam.
      """}
]
chatbot = gr.Chatbot(
    elem_id="chatbot",
    avatar_images=("assets/avatar_user.png", "assets/adiyen_bot.png"),
    value=intro_messages,
    label="Sanatan-Bot",
    show_copy_button=True,
    show_copy_all_button=True,
    type="messages",
    height=800
)
chatInterface = gr.ChatInterface(title="Sanatan-AI",
    fn=chat, additional_inputs=[thread_id], chatbot=chatbot
)
# initializze database
init()
chatInterface.launch()