Update app.py

app.py

@@ -1,29 +1,41 @@
 import os
 import tempfile
+import uuid
 import streamlit as st
 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
-)
+from qdrant_client import models
+from langchain_community.vectorstores import Qdrant
+
+
+from utils import setup_openai_embeddings,setup_qdrant_client,delete_collection,is_document_embedded
+from embed import embed_documents_into_qdrant
+from preprocess import split_documents,update_metadata,load_documents_OCR
+from retrieve import retrieve_documents,retrieve_documents_from_collection
+from summarize import summarize_documents
+
 
 load_dotenv()
 
 def main():
     st.sidebar.title("PDF Management")
     uploaded_files = st.sidebar.file_uploader("Upload PDF files", type=["pdf"], accept_multiple_files=True)
-    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:
+
+    if 'uploaded_collection_name' not in st.session_state:
+        st.session_state['uploaded_collection_name'] = None
+
+    if uploaded_files:
+        if st.sidebar.button("Add Docs to Data Bank"):
+            files_info = save_uploaded_files(uploaded_files)
+            embed_documents_to_data_bank(files_info)
+        
+        if st.sidebar.button("Add Docs to Current Chat"):
             files_info = save_uploaded_files(uploaded_files)
-            embed_documents(files_info, model_name, use_ocr)
+            add_docs_to_current_chat(files_info)
 
     pages = {
         "Lex Document Summarization": page_summarization,
-        "Lex Q&A": page_qna
+        "Chat with Data Bank": page_qna,
+        "Chat with Uploaded Docs": page_chat_with_uploaded_docs
     }
 
     st.sidebar.title("Page Navigation")
@@ -35,10 +47,10 @@ def main():
 
     # Call the page function based on the user selection
     if page:
-        pages[page](uploaded_files, model_name, use_ocr)
+        pages[page](uploaded_files)
 
 def save_uploaded_files(uploaded_files):
-    """Save uploaded files to temporary directory and return their file paths along with original filenames."""
+    """Save uploaded files to a temporary directory and return their file paths along with original filenames."""
     files_info = []
     for uploaded_file in uploaded_files:
         with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmpfile:
@@ -46,7 +58,7 @@ def save_uploaded_files(uploaded_files):
             files_info.append((tmpfile.name, uploaded_file.name))
     return files_info
 
-def page_summarization(uploaded_files, model_name, use_ocr):
+def page_summarization(uploaded_files):
     """Page for document summarization."""
     st.title("Lex Document Summarization")
     if uploaded_files:
@@ -56,42 +68,84 @@ def page_summarization(uploaded_files, model_name, use_ocr):
             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
+                    if summary_button:  # Only summarize if the 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'))
+                            documents = load_documents_OCR(temp_path, os.getenv('UNSTRUCTURED_API'))
+                            summary = summarize_documents(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):
+def page_qna(uploaded_files):
     """Page for Q&A functionality."""
-    st.title("Lex Question and Answer")
-    user_query = st.text_area("Enter your question here:",height=300)
+    st.title("Chat with Data Bank")
+    user_query = st.text_area("Enter your question here:", height=300)
+    if st.button('Get Answer'):
+        if user_query:
+            answer = handle_query(user_query)
+            st.write(answer)
+        else:
+            st.error("Please enter a question to get an answer.")
+    
+def page_chat_with_uploaded_docs(uploaded_files):
+    """Page for chatting with uploaded documents."""
+    st.title("Chat with Uploaded Documents")
+    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)
+            answer = handle_uploaded_docs_query(user_query, st.session_state['uploaded_collection_name'])
             st.write(answer)
         else:
             st.error("Please enter a question to get an answer.")
+    
+    if st.session_state['uploaded_collection_name']:
+        if st.button('Delete Embedded Collection'):
+            collection_name = st.session_state['uploaded_collection_name']
+            delete_collection(collection_name, os.getenv('QDRANT_URL'), os.getenv('QDRANT_API_KEY'))
+            st.session_state['uploaded_collection_name'] = None
+            st.success(f"Deleted collection {collection_name}")
 
-def embed_documents(files_info, model_name, use_ocr):
-    """Function to embed documents."""
+def embed_documents_to_data_bank(files_info):
+    """Function to embed documents into the data bank."""
     for temp_path, original_name in files_info:
         if not is_document_embedded(original_name):
             try:
-                if use_ocr:
-                    documents = load_documents_OCR(temp_path, os.getenv('UNSTRUCTURED_API'))
-                else:
-                    documents = load_documents(temp_path)
+                documents = load_documents_OCR(temp_path, os.getenv('UNSTRUCTURED_API'))
                 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} into Data Bank")
+                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 add_docs_to_current_chat(files_info):
+    """Function to add documents to the current chat session."""
+    if not st.session_state['uploaded_collection_name']:
+        st.session_state['uploaded_collection_name'] = f"session-{uuid.uuid4()}"
+        client = setup_qdrant_client(os.getenv('QDRANT_URL'), os.getenv('QDRANT_API_KEY'))
+        client.create_collection(
+            collection_name=st.session_state['uploaded_collection_name'],
+            vectors_config=models.VectorParams(size=1536, distance=models.Distance.COSINE)
+        )
+    else:
+        client = setup_qdrant_client(os.getenv('QDRANT_URL'), os.getenv('QDRANT_API_KEY'))
+
+    embeddings_model = setup_openai_embeddings(os.getenv('OPENAI_API_KEY'))
+
+    for temp_path, original_name in files_info:
+        if not is_document_embedded(original_name):
+            try:
+                documents = load_documents_OCR(temp_path, os.getenv('UNSTRUCTURED_API'))
+                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'), collection_name=st.session_state['uploaded_collection_name'])
                     st.success(f"Embedded {original_name}")
                 else:
                     st.error(f"No documents found or extracted from {original_name}")
@@ -100,13 +154,29 @@ def embed_documents(files_info, model_name, use_ocr):
         else:
             st.info(f"{original_name} is already embedded.")
 
-def handle_query(query, model_name):
+def handle_query(query):
     """Retrieve answers based on the query."""
     try:
-        answer = retrieve_documents(query, os.getenv('OPENAI_API_KEY'), os.getenv('QDRANT_URL'), os.getenv('QDRANT_API_KEY'), model_name)
+        answer = retrieve_documents(query, os.getenv('OPENAI_API_KEY'), os.getenv('QDRANT_URL'), os.getenv('QDRANT_API_KEY'))
         return answer or "No relevant answer found."
     except Exception as e:
         return f"Error processing the query: {str(e)}"
 
+def handle_uploaded_docs_query(query, collection_name):
+    """Retrieve answers from the uploaded documents collection."""
+    try:
+        answer = retrieve_documents_from_collection(query, os.getenv('OPENAI_API_KEY'), os.getenv('QDRANT_URL'), os.getenv('QDRANT_API_KEY'), collection_name)
+        return answer or "No relevant answer found."
+    except Exception as e:
+        return f"Error processing the query: {str(e)}"
+
+def delete_collection(collection_name, qdrant_url, qdrant_api_key):
+    """Delete a Qdrant collection."""
+    client = setup_qdrant_client(qdrant_url, qdrant_api_key)
+    try:
+        client.delete_collection(collection_name=collection_name)
+    except Exception as e:
+        print("Failed to delete collection:", e)
+
 if __name__ == "__main__":
     main()