import gradio as gr
from models import EmbeddingModel, LLM
from utils import MistralPrompts
from vector_store import FaissVectorStore
from chat import ChatBot

VECTOR_DATABASE_PATH = 'vector_db'
# Initialize models and vector store
embedding_model = EmbeddingModel(model_name='sentence-transformers/all-MiniLM-L6-v2')
llm = LLM("mistralai/Mistral-7B-Instruct-v0.1")
vector_store = FaissVectorStore.as_retriever(database_path=VECTOR_DATABASE_PATH)

# Create a ChatBot instance
chat_bot = ChatBot(llm, embedding_model, vector_store)

# Function to handle the user's input and generate a response
def chat_bot(input_text):
    response = chat_bot.chat(input_text)

    return response

# Create a Gradio interface
chatbot_interface = gr.Interface(
    fn=chat_bot,
    inputs=gr.inputs.Textbox(prompt="User:"),
    outputs=gr.inputs.Textbox(prompt="Bot:"),
    title="Chatbot Assitant for PAN card related query",
    theme="compact"
)

# Launch the Gradio interface
chatbot_interface.launch()