import gradio as gr
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_groq import ChatGroq
from langchain.prompts import PromptTemplate
import os
GROQ_API = os.getenv("GROQ_API") 
# Define embeddings and text splitter
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)

# Initialize the LLM (ChatGroq)
llm = ChatGroq(
    model="mixtral-8x7b-32768",
    temperature=1.25,
    max_tokens=512,
    timeout=None,
    api_key=GROQ_API
)

# Define the custom prompt template
custom_prompt = PromptTemplate.from_template("""
You are a helpful AI assistant. 
Answer the question using only the following context:

{context}

Question: {input}

Provide a detailed and accurate response.
""")

# Create the document chain
document_chain = create_stuff_documents_chain(
    llm=llm,
    prompt=custom_prompt,
)

# Function to process the uploaded PDF and create the retriever
def process_pdf(file):
    # Load the PDF file
    loader = PyPDFLoader(file.name)
    pages = loader.load_and_split(text_splitter)
    
    # Create the vector store
    vector_store = Chroma.from_documents(pages, embeddings)
    
    # Create the retriever
    retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})
    
    # Create the QA chain
    qa_chain = create_retrieval_chain(
        retriever=retriever,
        combine_docs_chain=document_chain
    )
    
    return qa_chain

# Function to get the answer from the QA chain
def get_answer(file, user_prompt):
    # Process the PDF and create the QA chain
    qa_chain = process_pdf(file)
    
    # Get the answer
    response = qa_chain.invoke({"input": user_prompt})
    return response["answer"]

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# PDF Q&A with ChatGroq and LangChain")
    
    with gr.Row():
        pdf_input = gr.File(label="Upload PDF", type="filepath")
        question_input = gr.Textbox(label="Ask a question", placeholder="Type your question here...")
    
    output = gr.Textbox(label="Answer", interactive=False)
    
    submit_button = gr.Button("Submit")
    submit_button.click(fn=get_answer, inputs=[pdf_input, question_input], outputs=output)

# Launch the Gradio app
demo.launch()