Update app.py

app.py

@@ -1,103 +1,112 @@
 import os
-from dotenv import load_dotenv
 import tempfile
 import streamlit as st
-from langchain_community.document_loaders import PyPDFLoader
-
-from main import summarize_pdf_document
-from main import retrieve_documents
-from main import embed_document_data
-from main import is_document_embedded
-
-
-
+from dotenv import load_dotenv
+from main import (
+    load_and_split_documents, summarize_documents, embed_documents_into_qdrant, 
+    retrieve_documents, is_document_embedded, load_documents, split_documents, 
+    update_metadata, load_documents_OCR
+)
 
 load_dotenv()
 
 def main():
     st.sidebar.title("PDF Management")
     uploaded_files = st.sidebar.file_uploader("Upload PDF files", type=["pdf"], accept_multiple_files=True)
-    files_info = []  # Initialize files_info to an empty list before checking for uploaded files
+    model_name = st.sidebar.selectbox("Choose your model:", ["gpt-3.5-turbo", "gpt-4-turbo"])  # Model selection
+    use_ocr = st.sidebar.checkbox("Use OCR for document processing")
+    
+    if st.sidebar.button('Add Uploaded Documents in Q&A'):
+        if uploaded_files:
+            files_info = save_uploaded_files(uploaded_files)
+            embed_documents(files_info, model_name, use_ocr)
 
-    if uploaded_files:
-        files_info = save_uploaded_files(uploaded_files)
-        process_documents(files_info)
-        if st.button('Add Uploaded  Documents in Q nd A'):
-            embed_documents(files_info)
+    pages = {
+        "Lex Document Summarization": page_summarization,
+        "Lex Q&A": page_qna
+    }
 
-    # Call to display the Q&A section unconditionally
-    display_qna_section(files_info)
+    st.sidebar.title("Page Navigation")
+    page = st.sidebar.radio("Select a page", tuple(pages.keys()))
 
+    # Initialize session state for summarization results if not already set
+    if 'summaries' not in st.session_state:
+        st.session_state['summaries'] = {}
+
+    # Call the page function based on the user selection
+    if page:
+        pages[page](uploaded_files, model_name, use_ocr)
 
 def save_uploaded_files(uploaded_files):
     """Save uploaded files to temporary directory and return their file paths along with original filenames."""
     files_info = []
     for uploaded_file in uploaded_files:
-        # Create a temporary file
         with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmpfile:
-            # Write contents of the uploaded file to the temporary file
             tmpfile.write(uploaded_file.getvalue())
-            # Append both the temp file path and the original name
             files_info.append((tmpfile.name, uploaded_file.name))
     return files_info
 
-
-def process_documents(files_info):
-    """Handle document processing for summarization."""
-    st.header("Document Summaries")
-    summarize_button = st.button('Summarize Uploaded Documents')
-    
-    if summarize_button:
+def page_summarization(uploaded_files, model_name, use_ocr):
+    """Page for document summarization."""
+    st.title("Lex Document Summarization")
+    if uploaded_files:
+        files_info = save_uploaded_files(uploaded_files)
         for temp_path, original_name in files_info:
-            with st.container():  # Using container to better organize the display
-                st.write(f"Summary for {original_name}:")  # Display the original file name
-                try:
-                    summary = summarize_pdf_document(temp_path, os.getenv('OPENAI_API_KEY'))
-                    st.text_area("", value=summary, height=200, key=f"summary_{original_name}")
-                except Exception as e:
-                    st.error(f"Failed to summarize {original_name}: {str(e)}")
-                    
-    
-def embed_documents(files_info):
-    """Embed each document with correct metadata, replacing temp path with original filename."""
+            summary_button = st.button(f"Summarize {original_name}", key=original_name)
+            if summary_button or (original_name in st.session_state['summaries']):
+                with st.container():
+                    st.write(f"Summary for {original_name}:")
+                    if summary_button:  # Only summarize if button is pressed
+                        try:
+                            if use_ocr:
+                                documents = load_documents_OCR(temp_path, os.getenv('UNSTRUCTURED_API'))
+                            else:
+                                documents = load_and_split_documents(temp_path)
+                            summary = summarize_documents(model_name, documents, os.getenv('OPENAI_API_KEY'))
+                            st.session_state['summaries'][original_name] = summary  # Store summary in session state
+                        except Exception as e:
+                            st.error(f"Failed to summarize {original_name}: {str(e)}")
+                    st.text_area("", value=st.session_state['summaries'][original_name], height=200, key=f"summary_{original_name}")
+
+def page_qna(uploaded_files, model_name, use_ocr):
+    """Page for Q&A functionality."""
+    st.title("Lex Question and Answer")
+    user_query = st.text_area("Enter your question here:",height=300)
+    if st.button('Get Answer'):
+        if user_query:
+            answer = handle_query(user_query, model_name)
+            st.write(answer)
+        else:
+            st.error("Please enter a question to get an answer.")
+
+def embed_documents(files_info, model_name, use_ocr):
+    """Function to embed documents."""
     for temp_path, original_name in files_info:
-        if not is_document_embedded(original_name):  # Check if already embedded
+        if not is_document_embedded(original_name):
             try:
-                # Load the document
-                loader = PyPDFLoader(temp_path)
-                documents = loader.load()
-
-                # Update the metadata for each document
-                updated_documents = []
-                for doc in documents:
-                    doc.metadata['source'] = original_name  # Use original filename
-                    updated_documents.append(doc)
-
-                embed_document_data(updated_documents)  # Embed the documents
-                st.success(f"Embedded {original_name}")
+                if use_ocr:
+                    documents = load_documents_OCR(temp_path, os.getenv('UNSTRUCTURED_API'))
+                else:
+                    documents = load_documents(temp_path)
+                documents = update_metadata(documents, original_name)
+                documents = split_documents(documents)
+                if documents:
+                    embed_documents_into_qdrant(documents, os.getenv('OPENAI_API_KEY'), os.getenv('QDRANT_URL'), os.getenv('QDRANT_API_KEY'), 'Lex-v1')
+                    st.success(f"Embedded {original_name}")
+                else:
+                    st.error(f"No documents found or extracted from {original_name}")
             except Exception as e:
                 st.error(f"Failed to embed {original_name}: {str(e)}")
         else:
             st.info(f"{original_name} is already embedded.")
 
-
-def display_qna_section(files_info):
-    """Display Q&A section."""
-    st.header("Question and Answer")
-    with st.form("qa_form"):
-        user_query = st.text_input("Enter your question here:")
-        submit_button = st.form_submit_button('Get Answer')
-    
-    if submit_button and user_query:
-        answer = handle_query(user_query)
-        st.write(answer)
-    elif submit_button and not user_query:
-        st.error("Please enter a question to get an answer.")
-
-def handle_query(query):
+def handle_query(query, model_name):
     """Retrieve answers based on the query."""
-    answer = retrieve_documents(query)
-    return answer
+    try:
+        answer = retrieve_documents(query, os.getenv('OPENAI_API_KEY'), os.getenv('QDRANT_URL'), os.getenv('QDRANT_API_KEY'), model_name)
+        return answer or "No relevant answer found."
+    except Exception as e:
+        return f"Error processing the query: {str(e)}"
 
 if __name__ == "__main__":
-    main()
+    main()