import os
import gradio as gr
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

embeddings=HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
llm=ChatGroq(model="llama-3.3-70b-versatile",temperature=0)

prompt=ChatPromptTemplate.from_template("""
Answer only from the supplied context.
Context:
{context}

Question:
{question}
""")
chain=prompt|llm|StrOutputParser()

def ask(files,question,chunk_size,overlap,k):
    docs=[]
    for f in files:
        docs.extend(PyPDFLoader(f.name).load())
    splitter=RecursiveCharacterTextSplitter(chunk_size=int(chunk_size),chunk_overlap=int(overlap))
    chunks=splitter.split_documents(docs)
    db=Chroma.from_documents(chunks,embeddings)
    retrieved=db.as_retriever(search_kwargs={"k":int(k)}).invoke(question)
    context="\n\n".join(d.page_content for d in retrieved)
    answer=chain.invoke({"context":context,"question":question})
    sources="\n".join(f"{d.metadata.get('source')} | Page {d.metadata.get('page')}" for d in retrieved)
    return answer,sources

gr.Interface(
fn=ask,
inputs=[
gr.File(file_count="multiple",file_types=[".pdf"]),
gr.Textbox(label="Question"),
gr.Slider(300,1500,value=800),
gr.Slider(0,300,value=150),
gr.Slider(1,10,value=4,step=1)
],
outputs=[gr.Textbox(label="Answer"),gr.Textbox(label="Sources")],
title="Document QA using Groq").launch()