import gradio as gr
import os
from langchain.vectorstores import Chroma  # Verwenden des normalen Chroma Moduls
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from pdfplumber import open as open_pdf  # Verwenden von pdfplumber zum Extrahieren von Text aus PDFs

def process_pdf_and_query(pdf_path, question):
    # Lade die PDF und extrahiere den Text
    with open_pdf(pdf_path) as pdf:
        text = ""
        for page in pdf.pages:
            text += page.extract_text()

    # Text als Dokumente in den Chroma Vektor-Datenbank laden
    documents = [{"content": text, "metadata": {"source": pdf_path}}]

    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    vectordb = Chroma.from_documents(documents, embeddings)

    retriever = vectordb.as_retriever()
    prompt_template = "Beantworte die folgende Frage basierend auf dem Dokument: {context}\nFrage: {question}\nAntwort:"
    prompt = PromptTemplate(input_variables=["context", "question"], template=prompt_template)

    qa_chain = RetrievalQA.from_chain_type(llm=None, retriever=retriever, chain_type_kwargs={"prompt": prompt})
    response = qa_chain.run(input_documents=documents, question=question)
    return response

def chatbot_response(pdf, question):
    # Gradio gibt uns die PDF als NamedString, wir extrahieren den Inhalt als Byte-Stream
    pdf_path = "/tmp/uploaded_pdf.pdf"
    
    # Extrahiere den Inhalt der Datei als Bytes
    pdf_content = pdf.read()  # Hier holen wir den Inhalt der PDF als Byte-Stream

    # Speichern des Byte-Streams von der Datei
    with open(pdf_path, "wb") as f:
        f.write(pdf_content)

    # Frage beantworten basierend auf dem Text der PDF
    answer = process_pdf_and_query(pdf_path, question)

    # Temporäre Datei löschen
    os.remove(pdf_path)
    
    return answer

# Gradio Interface
pdf_input = gr.File(label="PDF-Datei hochladen")
question_input = gr.Textbox(label="Frage eingeben")
response_output = gr.Textbox(label="Antwort")

interface = gr.Interface(
    fn=chatbot_response,
    inputs=[pdf_input, question_input],
    outputs=response_output,
    title="RAG Chatbot mit PDF-Unterstützung",
    description="Lade eine PDF-Datei hoch und stelle Fragen zu ihrem Inhalt."
)

if __name__ == "__main__":
    interface.launch()