Refactor RAGPipeline to extract page number from query and log requested page number

app.py

@@ -1,7 +1,10 @@
 # app.py
 
 import csv
+
 import datetime
+
+# from datetime import datetime
 import io
 import json
 import logging
@@ -377,6 +380,8 @@ def create_gr_interface() -> gr.Blocks:
 
             # Tab 2: PDF Chat Interface
             with gr.Tab("PDF Chat"):
+                pdf_processor = PDFProcessor()
+
                 with gr.Row():
                     # Left column: Chat and Input
                     with gr.Column(scale=7):
@@ -412,8 +417,6 @@ def create_gr_interface() -> gr.Blocks:
                         pdf_status = gr.Markdown()
                         current_collection = gr.State(value=None)
 
-        pdf_processor = PDFProcessor()
-
         # Event handlers for Study Analysis tab
         process_zotero_btn.click(
             process_zotero_library_items,
@@ -438,7 +441,6 @@ def create_gr_interface() -> gr.Blocks:
         # Event handlers for PDF Chat tab
 
         def handle_pdf_upload(files, name):
-            """Handle PDF upload and processing."""
             if not name:
                 return "Please provide a collection name", None
             if not files:
@@ -452,14 +454,20 @@ def create_gr_interface() -> gr.Blocks:
                 logger.error(f"Error in handle_pdf_upload: {str(e)}")
                 return f"Error: {str(e)}", None
 
+        upload_btn.click(
+            handle_pdf_upload,
+            inputs=[pdf_files, collection_name],
+            outputs=[pdf_status, current_collection],
+        )
+
         def add_message(history, message):
             """Add user message to chat history."""
             if not message.strip():
                 raise gr.Error("Please enter a message")
             history = history + [(message, None)]
-            return history, "", None  # Return empty preview
+            return history, "", None
 
-        def generate_chat_response(history, collection_id):
+        def generate_chat_response(history, collection_id, pdf_processor):
             """Generate response for the last message in history."""
             if not collection_id:
                 raise gr.Error("Please upload PDFs first")
@@ -468,13 +476,39 @@ def create_gr_interface() -> gr.Blocks:
 
             last_message = history[-1][0]
             try:
-                updated_history, preview_image = chat_response(
-                    last_message,
-                    history[:-1],
-                    collection_id,
-                    pdf_processor,
-                )
-                return updated_history, preview_image
+                # Get response and source info
+                rag = get_rag_pipeline(collection_id)
+                response, source_info = rag.query(last_message)
+
+                # Generate preview if source information is available
+                preview_image = None
+                if (
+                    source_info
+                    and source_info.get("source_file")
+                    and source_info.get("page_number") is not None
+                ):
+                    try:
+                        page_num = source_info["page_number"]
+                        logger.info(f"Attempting to render page {page_num}")
+                        preview_image = pdf_processor.render_page(
+                            source_info["source_file"], page_num
+                        )
+                        if preview_image:
+                            logger.info(
+                                f"Successfully generated preview for page {page_num}"
+                            )
+                        else:
+                            logger.warning(
+                                f"Failed to generate preview for page {page_num}"
+                            )
+                    except Exception as e:
+                        logger.error(f"Error generating PDF preview: {str(e)}")
+                        preview_image = None
+
+                # Update history with response
+                history[-1] = (last_message, response)
+                return history, preview_image
+
             except Exception as e:
                 logger.error(f"Error in generate_chat_response: {str(e)}")
                 history[-1] = (last_message, f"Error: {str(e)}")
@@ -493,7 +527,7 @@ def create_gr_interface() -> gr.Blocks:
             inputs=[chat_history, query_input],
             outputs=[chat_history, query_input, pdf_preview],
         ).success(
-            generate_chat_response,
+            lambda h, c: generate_chat_response(h, c, pdf_processor),
             inputs=[chat_history, current_collection],
             outputs=[chat_history, pdf_preview],
         )

rag/rag_pipeline.py

@@ -11,9 +11,12 @@ from llama_index.llms.openai import OpenAI
 from llama_index.vector_stores.chroma import ChromaVectorStore
 import chromadb
 from typing import Dict, Any, List, Tuple
+import re
+import logging
 
 
 logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 
 class RAGPipeline:
@@ -33,6 +36,23 @@ class RAGPipeline:
         self.load_documents()
         self.build_index()
 
+    def extract_page_number_from_query(self, query: str) -> int:
+        """Extract page number from query text."""
+        # Look for patterns like "page 3", "p3", "p. 3", etc.
+        patterns = [
+            r"page\s*(\d+)",
+            r"p\.\s*(\d+)",
+            r"p\s*(\d+)",
+            r"pg\.\s*(\d+)",
+            r"pg\s*(\d+)",
+        ]
+
+        for pattern in patterns:
+            match = re.search(pattern, query.lower())
+            if match:
+                return int(match.group(1))
+        return None
+
     def load_documents(self):
         if self.documents is None:
             with open(self.study_json, "r") as f:
@@ -55,7 +75,7 @@ class RAGPipeline:
                         "year": doc_data.get("date"),
                         "doi": doc_data.get("doi"),
                         "source_file": doc_data.get("source_file"),
-                        "page_number": page_num,  # Store single page number
+                        "page_number": int(page_num),  # Store as integer
                         "total_pages": len(pages),
                     }
 
@@ -103,14 +123,17 @@ class RAGPipeline:
                 "Given this information, please answer the question: {query_str}\n"
                 "Provide a detailed answer using the content from the context above. "
                 "If the question asks about specific page content, make sure to include that information. "
-                "Cite sources using square brackets for EVERY piece of information, e.g. [1], [2], etc. "
-                "Include page numbers in citations when available, e.g. [1, p.3]. "
+                "Cite sources using square brackets for EVERY piece of information, e.g. [1, p.3], [2, p.5], etc. "
                 "If you're unsure about something, say so rather than making assumptions."
             )
 
+        # Extract page number from query if present
+        requested_page = self.extract_page_number_from_query(context)
+        logger.info(f"Requested page number: {requested_page}")
+
         query_engine = self.index.as_query_engine(
             text_qa_template=prompt_template,
-            similarity_top_k=5,  # Reduced for more focused results
+            similarity_top_k=5,
             response_mode="tree_summarize",
             llm=OpenAI(model="gpt-4o-mini"),
         )
@@ -120,15 +143,23 @@ class RAGPipeline:
         # Extract source information from the response nodes
         source_info = {}
         if hasattr(response, "source_nodes") and response.source_nodes:
-            # Get the most relevant source
             source_node = response.source_nodes[0]
             metadata = source_node.metadata
+
+            # Use requested page number if available, otherwise use the page from metadata
+            page_number = (
+                requested_page
+                if requested_page is not None
+                else metadata.get("page_number", 0)
+            )
+
             source_info = {
                 "source_file": metadata.get("source_file"),
-                "page_number": metadata.get("page_number"),
+                "page_number": page_number,
                 "title": metadata.get("title"),
                 "authors": metadata.get("authors"),
-                "content": source_node.text,  # Include the actual content
+                "content": source_node.text,
             }
+            logger.info(f"Source info page number: {page_number}")
 
         return response.response, source_info

utils/pdf_processor.py

@@ -8,11 +8,12 @@ import os
 import fitz
 import logging
 from typing import Dict, List, Optional
-from datetime import datetime
+import datetime
 from slugify import slugify
 import json
 from PIL import Image
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -23,6 +24,60 @@ class PDFProcessor:
         os.makedirs(upload_dir, exist_ok=True)
         self.current_page = 0
 
+    def render_page(self, file_path: str, page_num: int) -> Optional[Image.Image]:
+        """Render a specific page from a PDF as an image."""
+        try:
+            logger.info(f"Attempting to render page {page_num} from {file_path}")
+            doc = fitz.open(file_path)
+
+            # Ensure page number is valid
+            if page_num < 0 or page_num >= len(doc):
+                logger.error(
+                    f"Invalid page number {page_num} for document with {len(doc)} pages"
+                )
+                return None
+
+            page = doc[page_num]
+            # Increase resolution for better quality
+            pix = page.get_pixmap(matrix=fitz.Matrix(300 / 72, 300 / 72))
+            image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+            doc.close()
+            logger.info(f"Successfully rendered page {page_num}")
+            return image
+        except Exception as e:
+            logger.error(f"Error rendering page {page_num} from {file_path}: {str(e)}")
+            return None
+
+    def process_pdfs(self, file_paths: List[str], collection_name: str) -> str:
+        """Process multiple PDF files and store their content."""
+        processed_docs = []
+
+        for file_path in file_paths:
+            try:
+                doc_data = self.extract_text_from_pdf(file_path)
+                processed_docs.append(doc_data)
+                logger.info(f"Successfully processed {file_path}")
+            except Exception as e:
+                logger.error(f"Error processing {file_path}: {str(e)}")
+                continue
+
+        if not processed_docs:
+            raise ValueError("No documents were successfully processed")
+
+        # Save to JSON file
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        output_filename = f"{slugify(collection_name)}_{timestamp}_documents.json"
+        output_path = os.path.join("data", output_filename)
+
+        # Ensure the data directory exists
+        os.makedirs("data", exist_ok=True)
+
+        with open(output_path, "w", encoding="utf-8") as f:
+            json.dump(processed_docs, f, indent=2, ensure_ascii=False)
+
+        logger.info(f"Saved processed documents to {output_path}")
+        return output_path
+
     def extract_text_from_pdf(self, file_path: str) -> Dict:
         """Extract text and metadata from a PDF file."""
         try:
@@ -33,7 +88,7 @@ class PDFProcessor:
             pages = {}
             for page_num in range(len(doc)):
                 page_text = doc[page_num].get_text()
-                pages[page_num] = page_text
+                pages[str(page_num)] = page_text  # Convert page_num to string for JSON
                 text += page_text + "\n"
 
             # Extract metadata
@@ -62,41 +117,3 @@ class PDFProcessor:
         except Exception as e:
             logger.error(f"Error processing PDF {file_path}: {str(e)}")
             raise
-
-    def process_pdfs(self, file_paths: List[str], collection_name: str) -> str:
-        """Process multiple PDF files and store their content."""
-        processed_docs = []
-
-        for file_path in file_paths:
-            try:
-                doc_data = self.extract_text_from_pdf(file_path)
-                processed_docs.append(doc_data)
-            except Exception as e:
-                logger.error(f"Error processing {file_path}: {str(e)}")
-                continue
-
-        if not processed_docs:
-            raise ValueError("No documents were successfully processed")
-
-        # Save to JSON file
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        output_filename = f"{slugify(collection_name)}_{timestamp}_documents.json"
-        output_path = f"data/{output_filename}"
-
-        with open(output_path, "w", encoding="utf-8") as f:
-            json.dump(processed_docs, f, indent=2, ensure_ascii=False)
-
-        return output_path
-
-    def render_page(self, file_path: str, page_num: int) -> Optional[Image.Image]:
-        """Render a specific page from a PDF as an image."""
-        try:
-            doc = fitz.open(file_path)
-            page = doc[page_num]
-            pix = page.get_pixmap(matrix=fitz.Matrix(300 / 72, 300 / 72))
-            image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-            doc.close()
-            return image
-        except Exception as e:
-            logger.error(f"Error rendering page {page_num} from {file_path}: {str(e)}")
-            return None