Update app.py

app.py

@@ -3,6 +3,7 @@ import os
 import requests
 import pdfplumber
 import chromadb
+import re
 from langchain.document_loaders import PDFPlumberLoader
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_experimental.text_splitter import SemanticChunker
@@ -10,7 +11,6 @@ from langchain_chroma import Chroma
 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 -----------------
@@ -21,7 +21,7 @@ st.title("Blah-1")
 os.environ["GROQ_API_KEY"] = st.secrets.get("GROQ_API_KEY", "")
 
 # ----------------- ChromaDB Persistent Directory -----------------
-CHROMA_DB_DIR = "/mnt/data/chroma_db" 
+CHROMA_DB_DIR = "/mnt/data/chroma_db"  # Ensure persistence
 os.makedirs(CHROMA_DB_DIR, exist_ok=True)
 
 # ----------------- Initialize Session State -----------------
@@ -36,9 +36,9 @@ if "processed_chunks" not in st.session_state:
 if "vector_store" not in st.session_state:
     st.session_state.vector_store = None
 
-# ----------------- Extract Metadata (Title, Author, Emails, Affiliations) -----------------
+# ----------------- Improved Metadata Extraction -----------------
 def extract_metadata(pdf_path):
-    """Extract metadata such as Title, Author, Emails, and Affiliations."""
+    """Extracts title, author, emails, and affiliations from PDF."""
     with pdfplumber.open(pdf_path) as pdf:
         metadata = pdf.metadata or {}
 
@@ -46,13 +46,14 @@ def extract_metadata(pdf_path):
         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"
+            title_match = re.search(r"(?i)title[:\-]?\s*(.*)", text or "")
+            title = title_match.group(1) if title_match else 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"
+            author_match = re.search(r"(?i)by\s+([A-Za-z\s,]+)", pdf.pages[0].extract_text() or "")
+            author = author_match.group(1).strip() if author_match 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 "")
@@ -77,25 +78,6 @@ if pdf_source == "Upload a PDF file":
         st.session_state.chunked = False
         st.session_state.vector_created = False
 
-elif pdf_source == "Enter a PDF URL":
-    pdf_url = st.text_input("Enter PDF URL:")
-    if pdf_url and not st.session_state.pdf_loaded:
-        with st.spinner("🔄 Downloading PDF..."):
-            try:
-                response = requests.get(pdf_url)
-                if response.status_code == 200:
-                    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
-                    st.session_state.chunked = False
-                    st.session_state.vector_created = False
-                    st.success("✅ PDF Downloaded Successfully!")
-                else:
-                    st.error("❌ Failed to download PDF. Check the URL.")
-            except Exception as e:
-                st.error(f"Error downloading PDF: {e}")
-
 # ----------------- Process PDF -----------------
 if not st.session_state.pdf_loaded and "pdf_path" in st.session_state:
     with st.spinner("🔄 Processing document... Please wait."):
@@ -117,10 +99,15 @@ if not st.session_state.pdf_loaded and "pdf_path" in st.session_state:
         model_name = "nomic-ai/modernbert-embed-base"
         embedding_model = HuggingFaceEmbeddings(model_name=model_name, model_kwargs={"device": "cpu"}, encode_kwargs={'normalize_embeddings': False})
 
+        # Convert metadata into a retrievable chunk
+        metadata_text = f"Title: {title}\nAuthor: {author}\nEmails: {email_str}\nAffiliations: {affiliation_str}"
+        metadata_doc = {"page_content": metadata_text, "metadata": {"source": "metadata"}}
+
         # Prevent unnecessary re-chunking
         if not st.session_state.chunked:
             text_splitter = SemanticChunker(embedding_model)
             document_chunks = text_splitter.split_documents(docs)
+            document_chunks.insert(0, metadata_doc)  # Insert metadata as a retrievable document
             st.session_state.processed_chunks = document_chunks
             st.session_state.chunked = True
 
@@ -140,6 +127,7 @@ if not st.session_state.vector_created and st.session_state.processed_chunks:
         st.session_state.vector_created = True
         st.success("✅ Vector store initialized successfully!")
 
+
 # ----------------- Query Input -----------------
 query = st.text_input("🔍 Ask a question about the document:")