Update app.py

app.py

@@ -4,16 +4,13 @@ import nltk
 import fitz
 import random
 import base64
-import logging
 import pycountry
 from PIL import Image
 import streamlit as st
-from fastapi import FastAPI
 from langdetect import detect
 from config import load_config
 from dotenv import load_dotenv
 from nltk.corpus import stopwords
-from fastapi import FastAPI, Query
 from langchain_groq import ChatGroq
 from collections import defaultdict
 from log_utils import setup_logging
@@ -21,7 +18,6 @@ from nltk.tokenize import sent_tokenize
 from nltk.tokenize import word_tokenize
 from langchain.chains import RetrievalQA
 from upload_pdf import update_or_add_pdf 
-from fastapi.responses import JSONResponse
 from langchain.prompts import ChatPromptTemplate
 from langchain_community.vectorstores import Chroma
 from sklearn.metrics.pairwise import cosine_similarity
@@ -34,130 +30,7 @@ logger = setup_logging('app')
 nltk.download('punkt')
 nltk.download('punkt_tab')
 nltk.download('stopwords')
-
-app = FastAPI()
-
-@app.get("/pdf-details")
-async def get_pdf_details(
-    filename: str = Query(..., description="Filename of the PDF"),
-    page_number: int = Query(0, description="Page number (0-indexed)")
-):
-    logger.info(f"Processing PDF details request for file: {filename}, page: {page_number}")
-    try:
-        data_path = "/home/bacancy/Documents/Company/Smart PDF Search/data"
-        file_path = os.path.join(data_path, filename)
-        
-        # Open the PDF
-        logger.debug(f"Opening PDF file: {file_path}")
-        doc = fitz.open(file_path)
-        
-        # Extract full PDF text
-        full_text = ""
-        for page in doc:
-            full_text += page.get_text()
-            
-        # Get PDF metadata
-        pdf_metadata = doc.metadata or {}
-        
-        # Extract page text and render page image
-        page = doc.load_page(page_number)
-        page_text = page.get_text()
-        
-        # Render page as image
-        pix = page.get_pixmap()
-        page_image_base64 = base64.b64encode(pix.tobytes("png")).decode('utf-8')
-        
-        # Detect language
-        try:
-            lang_code = detect(page_text)
-            language = pycountry.languages.get(alpha_2=lang_code).name
-        except Exception as e:
-            logger.warning(f"Language detection failed: {str(e)}")
-            language = 'Unknown'
-        
-        # Prepare response
-        response = {
-            "file_path": file_path,
-            "filename": os.path.basename(file_path),
-            "total_pages": len(doc),
-            "current_page": page_number + 1,
-            "full_text": full_text,
-            "page_text": page_text,
-            "page_image": page_image_base64,
-            "file_size_bytes": os.path.getsize(file_path),
-            "file_size_kb": f"{os.path.getsize(file_path) / 1024:.2f} KB",
-            "language": language,
-            "metadata": {
-                "title": pdf_metadata.get('title', 'Unknown'),
-                "author": pdf_metadata.get('author', 'Unknown'),
-                "creator": pdf_metadata.get('creator', 'Unknown'),
-                "producer": pdf_metadata.get('producer', 'Unknown')
-            }
-        }
-        
-        logger.info(f"Successfully processed PDF details for {filename}")
-        return JSONResponse(content=response)
-    
-    except Exception as e:
-        logger.error(f"Error processing PDF details: {str(e)}", exc_info=True)
-        return JSONResponse(
-            content={"error": str(e)},
-            status_code=500
-        )
-
-@app.get("/romanized-text")
-async def get_romanized_text(
-    filename: str = Query(..., description="Filename of the PDF")
-):
-    logger.info(f"Processing romanized text request for file: {filename}")
-    try:
-        data_path = "/home/bacancy/Documents/Company/Smart PDF Search/data"
-        file_path = os.path.join(data_path, filename)
-        
-        # Open the PDF
-        logger.debug(f"Opening PDF file for romanization: {file_path}")
-        doc = fitz.open(file_path)
-        
-        # Extract full PDF text
-        full_text = ""
-        pages_text = []
     
-        for page in doc:
-            page_text = page.get_text()
-            full_text += page_text
-            # Add page info to pages_text list
-            pages_text.append({
-                    "page_number": page.number + 1,  # Adding 1 to make it 1-based instead of 0-based
-                    "text": page_text
-                })
-                
-        # Get PDF metadata
-        pdf_metadata = doc.metadata or {}
-        
-        response = {
-            "filename": os.path.basename(file_path),
-            "total_pages": len(doc),
-            "full_text": full_text,
-            "pages": pages_text, 
-            "file_size_kb": f"{os.path.getsize(file_path) / 1024:.2f} KB",
-            "metadata": {
-                "title": pdf_metadata.get('title', 'Unknown'),
-                "author": pdf_metadata.get('author', 'Unknown'),
-                "creator": pdf_metadata.get('creator', 'Unknown'),
-                "producer": pdf_metadata.get('producer', 'Unknown')
-            }
-        }
-        
-        logger.info(f"Successfully processed romanized text for {filename}")
-        return JSONResponse(content=response)
-        
-    except Exception as e:
-        logger.error(f"Error processing romanized text: {str(e)}", exc_info=True)
-        return JSONResponse(
-            content={"error": str(e)},
-            status_code=500
-        )
-
 # Load environment variables
 load_dotenv()
 
@@ -573,7 +446,7 @@ def display_source_documents_with_images(source_documents, query):
                     highlighted_snippet = highlight_query_words(snippet, query)
                     
                     st.markdown(f'<div class="source-content">{highlighted_snippet}</div>', unsafe_allow_html=True)
-                    # st.markdown(f"[View other results in this book](?page=pdf_details&filename={pdf_name}&page_number={page_number})", unsafe_allow_html=True)
+                    st.markdown(f"[View other results in this book](?page=pdf_details&filename={pdf_name}&page_number={page_number})", unsafe_allow_html=True)
                     
                     logger.debug(f"Successfully displayed content for {pdf_name}, page {page_number + 1}")
             
@@ -628,6 +501,110 @@ def is_query_relevant(question, source_documents, threshold=0.1):
         logger.error(f"Error checking query relevance: {str(e)}", exc_info=True)
         return False
 
+def get_pdf_details(filename, page_number):
+    """Get details of a specific PDF page."""
+    logger.info(f"Processing PDF details for file: {filename}, page: {page_number}")
+    try:
+        data_path = "/home/bacancy/Documents/Company/Smart PDF Search/data"
+        file_path = os.path.join(data_path, filename)
+        
+        # Open the PDF
+        logger.debug(f"Opening PDF file: {file_path}")
+        doc = fitz.open(file_path)
+        
+        # Extract full PDF text
+        full_text = ""
+        for page in doc:
+            full_text += page.get_text()
+            
+        # Get PDF metadata
+        pdf_metadata = doc.metadata or {}
+        
+        # Extract page text and render page image
+        page = doc.load_page(page_number)
+        page_text = page.get_text()
+        
+        # Render page as image
+        pix = page.get_pixmap()
+        img_bytes = pix.tobytes("png")
+        page_image_base64 = base64.b64encode(img_bytes).decode('utf-8')
+        
+        # Detect language
+        try:
+            lang_code = detect(page_text)
+            language = pycountry.languages.get(alpha_2=lang_code).name
+        except Exception as e:
+            logger.warning(f"Language detection failed: {str(e)}")
+            language = 'Unknown'
+        
+        # Prepare response
+        return {
+            "file_path": file_path,
+            "filename": os.path.basename(file_path),
+            "total_pages": len(doc),
+            "current_page": page_number + 1,
+            "full_text": full_text,
+            "page_text": page_text,
+            "page_image": page_image_base64,
+            "file_size_bytes": os.path.getsize(file_path),
+            "file_size_kb": f"{os.path.getsize(file_path) / 1024:.2f} KB",
+            "language": language,
+            "metadata": {
+                "title": pdf_metadata.get('title', 'Unknown'),
+                "author": pdf_metadata.get('author', 'Unknown'),
+                "creator": pdf_metadata.get('creator', 'Unknown'),
+                "producer": pdf_metadata.get('producer', 'Unknown')
+            }
+        }
+    
+    except Exception as e:
+        logger.error(f"Error processing PDF details: {str(e)}", exc_info=True)
+        raise
+
+def get_romanized_text(filename):
+    """Get romanized text from a PDF."""
+    logger.info(f"Processing romanized text for file: {filename}")
+    try:
+        data_path = "/home/bacancy/Documents/Company/Smart PDF Search/data"
+        file_path = os.path.join(data_path, filename)
+        
+        # Open the PDF
+        logger.debug(f"Opening PDF file for romanization: {file_path}")
+        doc = fitz.open(file_path)
+        
+        # Extract full PDF text
+        full_text = ""
+        pages_text = []
+    
+        for page in doc:
+            page_text = page.get_text()
+            full_text += page_text
+            pages_text.append({
+                "page_number": page.number + 1,
+                "text": page_text
+            })
+                
+        # Get PDF metadata
+        pdf_metadata = doc.metadata or {}
+        
+        return {
+            "filename": os.path.basename(file_path),
+            "total_pages": len(doc),
+            "full_text": full_text,
+            "pages": pages_text, 
+            "file_size_kb": f"{os.path.getsize(file_path) / 1024:.2f} KB",
+            "metadata": {
+                "title": pdf_metadata.get('title', 'Unknown'),
+                "author": pdf_metadata.get('author', 'Unknown'),
+                "creator": pdf_metadata.get('creator', 'Unknown'),
+                "producer": pdf_metadata.get('producer', 'Unknown')
+            }
+        }
+        
+    except Exception as e:
+        logger.error(f"Error processing romanized text: {str(e)}", exc_info=True)
+        raise
+    
 def main():
     logger.info("Starting Smart PDF Search application")
     
@@ -643,16 +620,24 @@ def main():
         logger.info(f"Displaying PDF details for {filename}, page {page_number}")
         
         if filename:
-            display_pdf_details(filename, page_number)
-        else:
-            logger.warning("No filename provided for PDF details")
-            st.error("No filename provided for PDF details")
+            try:
+                pdf_details = get_pdf_details(filename, page_number)
+                display_pdf_details(pdf_details, filename) 
+            except Exception as e:
+                logger.error(f"Error displaying PDF details: {str(e)}")
+                st.error(f"Error displaying PDF details: {str(e)}")
+                
     elif page == 'romanized_text':
         filename = query_params.get('filename', '')
         logger.info(f"Displaying romanized text for {filename}")
         
         if filename:
-            display_romanized_text_page(filename)
+            try:
+                romanized_data = get_romanized_text(filename)
+                display_romanized_text_page(romanized_data)
+            except Exception as e:
+                logger.error(f"Error displaying romanized text: {str(e)}")
+                st.error(f"Error displaying romanized text: {str(e)}")
         else:
             logger.warning("No filename provided for Romanized text")
             st.error("No filename provided for Romanized text")