Merge pull request #104 from AdiBak/main

Add dynamic page range slider for uploaded PDFs with validation

(Results in a run-time error for the offline mode; to be fixed separately)

app.py

@@ -222,6 +222,11 @@ with st.sidebar:
                 value='2024-05-01-preview',
             )
 
+        # Make slider with initial values
+        page_range_slider = st.slider('7: Specify a page range for the PDF file:',
+                  1, GlobalConfig.MAX_ALLOWED_PAGES, [1, GlobalConfig.MAX_ALLOWED_PAGES])
+        st.session_state['page_range_slider'] = page_range_slider
+
 
 def build_ui():
     """
@@ -255,6 +260,9 @@ def set_up_chat_ui():
     """
     Prepare the chat interface and related functionality.
     """
+    # Set start and end page
+    st.session_state['start_page'] = st.session_state['page_range_slider'][0]
+    st.session_state['end_page'] = st.session_state['page_range_slider'][1]
 
     with st.expander('Usage Instructions'):
         st.markdown(GlobalConfig.CHAT_USAGE_INSTRUCTIONS)
@@ -282,11 +290,38 @@ def set_up_chat_ui():
     ):
         prompt_text = prompt.text or ''
         if prompt['files']:
+            # Store uploaded pdf in session state
+            uploaded_pdf = prompt['files'][0]
+            st.session_state['pdf_file'] = uploaded_pdf  
             # Apparently, Streamlit stores uploaded files in memory and clears on browser close
             # https://docs.streamlit.io/knowledge-base/using-streamlit/where-file-uploader-store-when-deleted
-            st.session_state[ADDITIONAL_INFO] = filem.get_pdf_contents(prompt['files'][0])
-            print(f'{prompt["files"]=}')
 
+        # Check if pdf file is uploaded 
+        # (we can use the same file if the user doesn't upload a new one)
+        if 'pdf_file' in st.session_state:  
+            # Get validated page range 
+            st.session_state['start_page'], st.session_state['end_page'] = filem.validate_page_range(
+                                                                                    st.session_state['pdf_file'], 
+                                                                                    st.session_state['start_page'],
+                                                                                    st.session_state['end_page']
+                                                                                )
+            # Show sidebar text for page selection and file name
+            with st.sidebar:
+                if st.session_state['end_page'] is None:  # If the PDF has only one page
+                    st.text('Extracting page %d in %s' % (
+                        st.session_state['start_page'], st.session_state['pdf_file'].name
+                    ))
+                else:
+                    st.text('Extracting pages %d to %d in %s' % (
+                        st.session_state['start_page'], st.session_state['end_page'], st.session_state['pdf_file'].name
+                    ))
+
+            # Get pdf contents
+            st.session_state[ADDITIONAL_INFO] = filem.get_pdf_contents(
+                                                        st.session_state['pdf_file'], 
+                                                        (st.session_state['start_page'], 
+                                                        st.session_state['end_page'])
+                                                    )
         provider, llm_name = llm_helper.get_provider_model(
             llm_provider_to_use,
             use_ollama=RUN_IN_OFFLINE_MODE
@@ -593,4 +628,3 @@ def main():
 
 if __name__ == '__main__':
     main()
-

global_config.py

@@ -108,6 +108,7 @@ class GlobalConfig:
     DEFAULT_MODEL_INDEX = int(os.environ.get('DEFAULT_MODEL_INDEX', '4'))
     LLM_MODEL_TEMPERATURE = 0.2
     MAX_PAGE_COUNT = 50
+    MAX_ALLOWED_PAGES = 150
     LLM_MODEL_MAX_INPUT_LENGTH = 1000  # characters
 
     LOG_LEVEL = 'DEBUG'

helpers/file_manager.py

@@ -19,22 +19,55 @@ logger = logging.getLogger(__name__)
 
 def get_pdf_contents(
         pdf_file: st.runtime.uploaded_file_manager.UploadedFile,
-        max_pages: int = GlobalConfig.MAX_PAGE_COUNT
-) -> str:
+        page_range: tuple[int, int]) -> str:
     """
     Extract the text contents from a PDF file.
 
     :param pdf_file: The uploaded PDF file.
-    :param max_pages: The max no. of pages to extract contents from.
+    :param page_range: The range of pages to extract contents from.
     :return: The contents.
     """
 
     reader = PdfReader(pdf_file)
-    n_pages = min(max_pages, len(reader.pages))
+
+    start, end = page_range  # Set start and end per the range (user-specified values)
+
     text = ''
 
-    for page in range(n_pages):
-        page = reader.pages[page]
-        text += page.extract_text()
+    if end is None:
+        # If end is None (where PDF has only 1 page or start = end), extract start
+        end = start
 
+    # Get the text from the specified page range
+    for page_num in range(start - 1, end):
+        text += reader.pages[page_num].extract_text()
+
+    
     return text
+
+def validate_page_range(pdf_file: st.runtime.uploaded_file_manager.UploadedFile,
+                        start:int, end:int) -> tuple[int, int]:
+    """
+    Validate the page range.
+
+    :param pdf_file: The uploaded PDF file.
+    :param start: The start page 
+    :param end: The end page
+    :return: The validated page range tuple
+    """
+    n_pages = len(PdfReader(pdf_file).pages)
+
+    # Set start to max of 1 or specified start (whichever's higher)
+    start = max(1, start)
+
+    # Set end to min of pdf length or specified end (whichever's lower)
+    end = min(n_pages, end)
+
+    if start > end:  # If the start is higher than the end, make it 1
+        start = 1
+
+    if start == end:
+        # If start = end (including when PDF is 1 page long), set end to None
+        return start, None
+
+    return start, end