Update app.py

app.py

@@ -1,14 +1,16 @@
 import streamlit as st
 import os
 import requests
+import pdfplumber
 import chromadb
 from langchain.document_loaders import PDFPlumberLoader
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_experimental.text_splitter import SemanticChunker
 from langchain_chroma import Chroma
-from langchain.chains import LLMChain, SequentialChain
+from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
 from langchain_groq import ChatGroq
+import re
 from prompts import rag_prompt, relevancy_prompt, relevant_context_picker_prompt, response_synth
 
 # ----------------- Streamlit UI Setup -----------------
@@ -18,8 +20,9 @@ st.title("Blah-1")
 # ----------------- API Keys -----------------
 os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "")
 
-# ----------------- Clear ChromaDB Cache -----------------
-chromadb.api.client.SharedSystemClient.clear_system_cache()
+# ----------------- ChromaDB Persistent Directory -----------------
+CHROMA_DB_DIR = "/mnt/data/chroma_db"  # Hugging Face Spaces persistent storage
+os.makedirs(CHROMA_DB_DIR, exist_ok=True)
 
 # ----------------- Initialize Session State -----------------
 if "pdf_loaded" not in st.session_state:
@@ -33,22 +36,41 @@ if "processed_chunks" not in st.session_state:
 if "vector_store" not in st.session_state:
     st.session_state.vector_store = None
 
-# ----------------- Load Models -----------------
-llm_judge = ChatGroq(model="deepseek-r1-distill-llama-70b")
-rag_llm = ChatGroq(model="mixtral-8x7b-32768")
+# ----------------- Extract Metadata (Title, Author, Emails, Affiliations) -----------------
+def extract_metadata(pdf_path):
+    """Extract metadata such as Title, Author, Emails, and Affiliations."""
+    with pdfplumber.open(pdf_path) as pdf:
+        metadata = pdf.metadata or {}
 
-# Enable verbose logging for debugging
-llm_judge.verbose = True
-rag_llm.verbose = True
+        # Extract title
+        title = metadata.get("Title", "").strip()
+        if not title and pdf.pages:
+            text = pdf.pages[0].extract_text()
+            title = text.split("\n")[0] if text else "Untitled Document"
+
+        # Extract author
+        author = metadata.get("Author", "").strip()
+        if not author and pdf.pages:
+            author_matches = re.findall(r"By ([A-Za-z\s,]+)", pdf.pages[0].extract_text() or "")
+            author = author_matches[0] if author_matches else "Unknown Author"
+
+        # Extract emails
+        emails = re.findall(r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}", pdf.pages[0].extract_text() or "")
+        email_str = ", ".join(emails) if emails else "No emails found"
+
+        # Extract affiliations
+        affiliations = re.findall(r"(?:Department|Faculty|Institute|University|College|School)\s+[\w\s]+", pdf.pages[0].extract_text() or "")
+        affiliation_str = ", ".join(affiliations) if affiliations else "No affiliations found"
+
+    return title, author, email_str, affiliation_str
 
 # ----------------- PDF Selection -----------------
-#st.subheader("PDF Selection")
 pdf_source = st.radio("Choose a PDF source:", ["Upload a PDF file", "Enter a PDF URL"], index=0, horizontal=True)
 
 if pdf_source == "Upload a PDF file":
     uploaded_file = st.file_uploader("Upload your PDF file", type=["pdf"])
     if uploaded_file:
-        st.session_state.pdf_path = "temp.pdf"
+        st.session_state.pdf_path = "/mnt/data/temp.pdf"
         with open(st.session_state.pdf_path, "wb") as f:
             f.write(uploaded_file.getbuffer())
         st.session_state.pdf_loaded = False
@@ -62,7 +84,7 @@ elif pdf_source == "Enter a PDF URL":
             try:
                 response = requests.get(pdf_url)
                 if response.status_code == 200:
-                    st.session_state.pdf_path = "temp.pdf"
+                    st.session_state.pdf_path = "/mnt/data/temp.pdf"
                     with open(st.session_state.pdf_path, "wb") as f:
                         f.write(response.content)
                     st.session_state.pdf_loaded = False
@@ -79,11 +101,20 @@ if not st.session_state.pdf_loaded and "pdf_path" in st.session_state:
     with st.spinner("🔄 Processing document... Please wait."):
         loader = PDFPlumberLoader(st.session_state.pdf_path)
         docs = loader.load()
-        st.json(docs[0].metadata)
+
+        # Extract metadata
+        title, author, email_str, affiliation_str = extract_metadata(st.session_state.pdf_path)
+
+        # Display extracted metadata
+        st.subheader("📄 Extracted Document Metadata")
+        st.write(f"**Title:** {title}")
+        st.write(f"**Author:** {author}")
+        st.write(f"**Emails:** {email_str}")
+        st.write(f"**Affiliations:** {affiliation_str}")
 
         # Embedding Model
         model_name = "nomic-ai/modernbert-embed-base"
-        embedding_model = HuggingFaceEmbeddings(model_name=model_name, model_kwargs={"device": "cpu"}, encode_kwargs = {'normalize_embeddings': False})
+        embedding_model = HuggingFaceEmbeddings(model_name=model_name, model_kwargs={"device": "cpu"}, encode_kwargs={'normalize_embeddings': False})
 
         # Prevent unnecessary re-chunking
         if not st.session_state.chunked:
@@ -99,6 +130,7 @@ if not st.session_state.pdf_loaded and "pdf_path" in st.session_state:
 if not st.session_state.vector_created and st.session_state.processed_chunks:
     with st.spinner("🔄 Initializing Vector Store..."):
         st.session_state.vector_store = Chroma(
+            persist_directory=CHROMA_DB_DIR,  # <-- Ensures persistence
             collection_name="deepseek_collection",
             collection_metadata={"hnsw:space": "cosine"},
             embedding_function=embedding_model