app.py

# AI assistant with a RAG system to query information from
#  the gwIAS search pipeline
# using Langchain and deployed with Gradio

from rag import RAG, load_docs
from langchain_community.embeddings import HuggingFaceInstructEmbeddings
from langchain_community.chat_models import ChatOpenAI
import gradio as gr

# Load the documentation
docs = load_docs()
print("Pages loaded:", len(docs))

# LLM model
llm = ChatOpenAI(model="gpt-4") # Fixed model name to use a real model

# Embeddings
embed_model = "sentence-transformers/multi-qa-distilbert-cos-v1"
# embed_model = "nvidia/NV-Embed-v2"
# text-embedding-3-small
embeddings = HuggingFaceInstructEmbeddings(model_name=embed_model)

# RAG chain
rag_chain = RAG(llm, docs, embeddings)

# Function to handle prompt and query the RAG chain
def handle_prompt(message, history):
    try:
        # Stream output
        out = ""
        for chunk in rag_chain.stream(message):
            out += chunk
            yield out
    except Exception as e:
        raise gr.Error(f"An error occurred: {str(e)}")


if __name__ == "__main__":
    # Predefined messages and examples
    description = "AI powered assistant to help with [gwIAS](https://github.com/JayWadekar/gwIAS-HM) gravitational wave search pipeline."
    greeting_message = "Hi, I'm the gwIAS Bot, I'm here to assist you with the search pipeline."
    example_questions = [
        "Can you give me the code for calculating coherent score?",
        "Which module in the code is used for collecting coincident triggers?",
        "How are template banks constructed?"
    ]

    # Define customized Gradio chatbot
    chatbot = gr.Chatbot(
        [{"role": "assistant", "content": greeting_message}],
        type="messages",
        avatar_images=["ims/userpic.png", "ims/gwIASlogo.jpg"],
        height="60vh"
    )

    # Define Gradio interface
    demo = gr.ChatInterface(
        fn=handle_prompt,
        chatbot=chatbot,
        title="gwIAS DocBot",
        description=description,
        examples=example_questions,
        theme=gr.themes.Soft(),
        fill_height=True
    )

    # Launch the interface
    demo.launch()