PDF_QA

[sauragrkatonic]

`import streamlit as st
from langchain_groq import ChatGroq
from tools.pdf_agent import pdf_qa
from dotenv import load_dotenv
from tools.research_paper import arxiv_tool

load_dotenv()

Initialize Groq LLM

llm = ChatGroq(
model="llama-3.1-8b-instant",
temperature=0,
max_tokens=None,
timeout=None,
max_retries=2,
)

st.title("ChatBot powered by Groq")

Initialize session messages

if "messages" not in st.session_state:
st.session_state["messages"] = [{"role": "assistant", "content": "How can I help you?"}]

Display previous messages

for msg in st.session_state.messages:
st.chat_message(msg["role"]).write(msg["content"])

Chat input + file upload

if prompt := st.chat_input(
"Say something and/or attach an image",
accept_file=True,
file_type=["jpg", "jpeg", "png", "pdf"]
):
st.session_state.messages.append({"role": "user", "content": prompt.text})
st.chat_message("user").write(prompt.text)

if prompt.text.lower().startswith("search arxiv:"):
    query = prompt.text[len("search arxiv:"):].strip()
    with st.spinner(f"Searching arXiv for: {query}"):
        arxiv_results = search_arxiv(query)
    st.chat_message("assistant").markdown(arxiv_results)
    st.session_state.messages.append({"role": "assistant", "content": arxiv_results})

elif prompt.files:
    # Process PDF
    retriever, images_by_page = pdf_qa(prompt.files[0])  # Assuming one PDF for now

    if prompt.text:
        with st.spinner("Generating an answer..."):
            # Retrieve relevant docs
            relevant_docs = retriever.get_relevant_documents(prompt.text)
            matched_pages = set()

            max_chars = 6000
            context = ""
            for doc in relevant_docs:
                if len(context) + len(doc.page_content) <= max_chars:
                    context += doc.page_content + "\n\n"
                    matched_pages.add(doc.metadata.get("page", -1))
                else:
                    break

            # Ask Groq model
            question = f"""Answer the following question based on the provided context:\n\n{context}\n\nQuestion: {prompt.text}"""
            response = llm.invoke(question)

        # Display assistant's answer
        st.session_state.messages.append({"role": "assistant", "content": response.content})
        st.chat_message("assistant").write(response.content)

        # Display matched images from the PDF
        if matched_pages:
            st.subheader("Relevant Images from PDF:")
            for page in matched_pages:
                for img_bytes in images_by_page.get(page, []):
                    st.image(img_bytes, caption=f"Page {page + 1}")

else:
    # No file attached, just LLM reply
    with st.spinner("Thinking..."):
        response = llm.invoke(prompt.text)
    st.session_state.messages.append({"role": "assistant", "content": response.content})
    st.chat_message("assistant").write(response.content)`