app.py

import os
import shutil
import time
import streamlit as st
from dotenv import load_dotenv

from langchain_groq import ChatGroq
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate

from langchain_community.document_loaders import PyPDFLoader
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings

from langchain_text_splitters import RecursiveCharacterTextSplitter

load_dotenv()

st.set_page_config(page_title="PuppetGPT", page_icon="🤖")

st.title("🤖 PuppetGPT")
st.caption("Chat with your documents using Retrieval-Augmented Generation")

# -------------------------
# Reset Conversation
# -------------------------

if st.button("Reset Conversation"):
    st.session_state.chat_history = []

# -------------------------
# Upload PDFs
# -------------------------

uploaded_files = st.file_uploader(
    "Upload PDFs",
    type="pdf",
    accept_multiple_files=True
)

if uploaded_files:

    st.subheader("Loaded Documents")

    for file in uploaded_files:
        st.write("📄", file.name)

# -------------------------
# Save Uploaded Files
# -------------------------

if uploaded_files:

    if os.path.exists("uploaded_docs"):
        shutil.rmtree("uploaded_docs")

    os.makedirs("uploaded_docs", exist_ok=True)

    for file in uploaded_files:
        pdf_path = os.path.join("uploaded_docs", file.name)

        with open(pdf_path, "wb") as f:
            f.write(file.read())

    st.success(f"{len(uploaded_files)} PDFs uploaded successfully!")

# -------------------------
# Prompt Template
# -------------------------

template = """
You are a document assistant.

Answer the question ONLY using the provided context.

Rules:
- Determine what the question is asking:
  • If it asks for a LIST of items → return ONLY the items.
  • If it asks for a SPECIFIC value → return ONLY that value.
  • Otherwise → give a concise answer.

- Do NOT include extra details unless explicitly asked.
- Do NOT include metadata (dates, tools, links, etc.) unless relevant to the question.
- If multiple items are requested, return a clean bullet list.
- Keep answers minimal and precise.

If the answer is not contained in the document context, say:
"I cannot find this information in the document."

Context:
{context}

Question:
{question}

Answer:
"""
qa_prompt = PromptTemplate(
    template=template,
    input_variables=["context", "question"]
)

# -------------------------
# Build Vectorstore
# -------------------------

def build_vectorstore():

    docs = []

    for file in os.listdir("uploaded_docs"):
        loader = PyPDFLoader(os.path.join("uploaded_docs", file))
        docs.extend(loader.load())

    if len(docs) == 0:
        return None

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200
    )

    chunks = splitter.split_documents(docs)

    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2"
    )

    with st.spinner("Building document index..."):

        vectorstore = Chroma.from_documents(
            documents=chunks,
            embedding=embeddings
        )

    return vectorstore

# -------------------------
# Build QA Chain
# -------------------------

def build_qa_chain(vectorstore):

    retriever = vectorstore.as_retriever(
        search_type="mmr",
        search_kwargs={"k": 8, "lambda_mult": 0.5}
    )

    llm = ChatGroq(
        model_name="llama-3.1-8b-instant",
        groq_api_key=os.getenv("GROQ_API_KEY")
    )

    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        retriever=retriever,
        chain_type_kwargs={"prompt": qa_prompt},
        return_source_documents=True
    )

    return qa_chain

# -------------------------
# Initialize Vectorstore Once
# -------------------------

if uploaded_files:

    st.session_state.vectorstore = build_vectorstore()

    if st.session_state.vectorstore is not None:
        st.session_state.qa_chain = build_qa_chain(
            st.session_state.vectorstore
        )

# -------------------------
# Chat Memory
# -------------------------

if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

for role, message in st.session_state.chat_history:
    with st.chat_message(role):
        st.write(message)

# -------------------------
# Chat Interface
# -------------------------

if uploaded_files and "qa_chain" in st.session_state:

    qa_chain = st.session_state.qa_chain

    user_question = st.chat_input("Ask a question about the PDF...")

    if user_question:

        st.session_state.chat_history.append(("user", user_question))

        with st.chat_message("user"):
            st.write(user_question)

        with st.chat_message("assistant"):

            with st.spinner("Searching document context..."):

                result = qa_chain.invoke({"query": user_question})
                answer = result["result"]

                # Fix bullet formatting
                answer = answer.replace("• ", "\n• ").strip()

            placeholder = st.empty()
            typed_text = ""

            for char in answer:
                typed_text += char
                placeholder.markdown(typed_text)
                time.sleep(0.01)

        st.session_state.chat_history.append(("assistant", answer))

        # -------------------------
        # Show Sources
        # -------------------------

        with st.expander("Sources"):

            for doc in result["source_documents"]:

                source = doc.metadata.get("source", "Unknown document")
                page = doc.metadata.get("page", "Unknown")

                filename = os.path.basename(source)

                st.markdown(f"📄 **{filename}** (Page {page})")
                st.write(doc.page_content[:300] + "...")