interim.py

import os
import requests
import streamlit as st
from langchain.chains import SequentialChain, LLMChain
from langchain.prompts import PromptTemplate
from langchain_groq import ChatGroq
from langchain.document_loaders import PDFPlumberLoader
from langchain_experimental.text_splitter import SemanticChunker
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma

# Set API Keys
os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "")

# Load LLM models
llm_judge = ChatGroq(model="deepseek-r1-distill-llama-70b")
rag_llm = ChatGroq(model="mixtral-8x7b-32768")

st.title("❓")

# Step 1: Choose PDF Source
#### Initialize pdf_path
pdf_path = None
pdf_source = st.radio("Upload or provide a link to a PDF:", ["Upload a PDF file", "Enter a PDF URL"], index=0)

if pdf_source == "Upload a PDF file":
    uploaded_file = st.file_uploader("Upload your PDF file", type="pdf")
    if uploaded_file:
        with open("temp.pdf", "wb") as f:
            f.write(uploaded_file.getbuffer())
        pdf_path = "temp.pdf"

elif pdf_source == "Enter a PDF URL":
    pdf_url = st.text_input("Enter PDF URL:")
    if pdf_url:
        with st.spinner("Downloading PDF..."):
            try:
                response = requests.get(pdf_url)
                if response.status_code == 200:
                    with open("temp.pdf", "wb") as f:
                        f.write(response.content)
                    pdf_path = "temp.pdf"
                    st.success("✅ PDF Downloaded Successfully!")
                else:
                    st.error("❌ Failed to download PDF. Check the URL.")
                    pdf_path = None
            except Exception as e:
                st.error(f"Error downloading PDF: {e}")
                pdf_path = None
else:
    pdf_path = None

# Step 2: Process PDF
if pdf_path:
    with st.spinner("Loading PDF..."):
        loader = PDFPlumberLoader(pdf_path)
        docs = loader.load()
    
    st.success(f"✅ **PDF Loaded!** Total Pages: {len(docs)}")

    # Step 3: Chunking
    with st.spinner("Chunking the document..."):
        model_name = "nomic-ai/modernbert-embed-base"
        embedding_model = HuggingFaceEmbeddings(model_name=model_name, model_kwargs={'device': 'cpu'})

        text_splitter = SemanticChunker(embedding_model)
        documents = text_splitter.split_documents(docs)

    st.success(f"✅ **Document Chunked!** Total Chunks: {len(documents)}")

    # Step 4: Setup Vectorstore
    with st.spinner("Creating vector store..."):
        vector_store = Chroma(
            collection_name="deepseek_collection",
            collection_metadata={"hnsw:space": "cosine"},
            embedding_function=embedding_model
        )
        vector_store.add_documents(documents)

    st.success("✅ **Vector Store Created!**")

    # Step 5: Query Input
    query = st.text_input("🔍 Enter a Query:")
    if query:
        with st.spinner("Retrieving relevant contexts..."):
            retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})
            contexts = retriever.invoke(query)
            context_texts = [doc.page_content for doc in contexts]

        st.success(f"✅ **Retrieved {len(context_texts)} Contexts!**")
        for i, text in enumerate(context_texts, 1):
            st.write(f"**Context {i}:** {text[:500]}...")

        # Step 6: Context Relevancy Checker
        with st.spinner("Evaluating context relevancy..."):
            relevancy_prompt = PromptTemplate(
                input_variables=["retriever_query", "context"],
                template="""You are an expert judge. Assign relevancy scores (0 or 1) for each context to answer the query.

                CONTEXT LIST:
                {context}

                QUERY:
                {retriever_query}

                RESPONSE (JSON):
                [{{"content": 1, "score": <0 or 1>, "reasoning": "<explanation>"}},
                {{"content": 2, "score": <0 or 1>, "reasoning": "<explanation>"}},
                ...]"""
            )
            context_relevancy_chain = LLMChain(llm=llm_judge, prompt=relevancy_prompt, output_key="relevancy_response")
            relevancy_response = context_relevancy_chain.invoke({"context": context_texts, "retriever_query": query})

        st.success("✅ **Context Relevancy Evaluated!**")
        st.json(relevancy_response['relevancy_response'])

        # Step 7: Selecting Relevant Contexts
        with st.spinner("Selecting the most relevant contexts..."):
            relevant_prompt = PromptTemplate(
                input_variables=["relevancy_response"],
                template="""Extract contexts with score 0 from the relevancy response.

                RELEVANCY RESPONSE:
                {relevancy_response}

                RESPONSE (JSON):
                [{{"content": <content number>}}]
                """
            )
            pick_relevant_context_chain = LLMChain(llm=llm_judge, prompt=relevant_prompt, output_key="context_number")
            relevant_response = pick_relevant_context_chain.invoke({"relevancy_response": relevancy_response['relevancy_response']})

        st.success("✅ **Relevant Contexts Selected!**")
        st.json(relevant_response['context_number'])

        # Step 8: Retrieving Context for Response Generation
        with st.spinner("Retrieving final context..."):
            context_prompt = PromptTemplate(
                input_variables=["context_number", "context"],
                template="""Extract actual content for the selected context numbers.

                CONTEXT NUMBERS:
                {context_number}

                CONTENT LIST:
                {context}

                RESPONSE (JSON):
                [{{"context_number": <content number>, "relevant_content": "<actual context>"}}]
                """
            )
            relevant_contexts_chain = LLMChain(llm=llm_judge, prompt=context_prompt, output_key="relevant_contexts")
            final_contexts = relevant_contexts_chain.invoke({"context_number": relevant_response['context_number'], "context": context_texts})

        st.success("✅ **Final Contexts Retrieved!**")
        st.json(final_contexts['relevant_contexts'])

        # Step 9: Generate Final Response
        with st.spinner("Generating the final answer..."):
            rag_prompt = PromptTemplate(
                input_variables=["query", "context"],
                template="""Generate a clear, fact-based response based on the context.

                QUERY:
                {query}

                CONTEXT:
                {context}

                ANSWER:
                """
            )
            response_chain = LLMChain(llm=rag_llm, prompt=rag_prompt, output_key="final_response")
            final_response = response_chain.invoke({"query": query, "context": final_contexts['relevant_contexts']})

        st.success("✅ **Final Response Generated!**")
        st.success(final_response['final_response'])

        # Step 10: Display Workflow Breakdown
        st.write("🔍 **Workflow Breakdown:**")
        st.json({
            "Context Relevancy Evaluation": relevancy_response["relevancy_response"],
            "Relevant Contexts": relevant_response["context_number"],
            "Extracted Contexts": final_contexts["relevant_contexts"],
            "Final Answer": final_response["final_response"]
        })