import streamlit as st
from groq import Groq
import pdfplumber
import os
import gc

# Set your Groq API key (make sure to set this in your environment variables)
os.environ["GROQ_API_KEY"] = "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"  # Set your API key here

# Initialize the Groq client
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

def generate_response_groq(context, query):
    """Generate response using Groq API."""
    prompt = f"Context: {context}\nQuestion: {query}\nAnswer:"
    chat_completion = client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": prompt,
            }
        ],
        model="llama3-8b-8192",
    )
    response = chat_completion.choices[0].message.content
    return response

def extract_text_from_pdf(pdf_file):
    """Extract text from PDF file using pdfplumber."""
    text = ""
    with pdfplumber.open(pdf_file) as pdf:
        for page in pdf.pages:
            text += page.extract_text() or ""
    return text

# Streamlit UI
st.title("PDF Query Application")

uploaded_file = st.file_uploader("Upload a PDF file", type="pdf")

if uploaded_file:
    st.write("PDF uploaded successfully!")
    document_text = extract_text_from_pdf(uploaded_file)
    st.text_area("Extracted Text", document_text, height=300)

    query = st.text_input("Enter your query")

    if st.button("Get Answer"):
        if query:
            with st.spinner("Generating response..."):
                response = generate_response_groq(document_text, query)
                st.write("Response:")
                st.write(response)

            # Clear memory after generating response
            gc.collect()

        else:
            st.error("Please enter a query.")