Update app.py

app.py

@@ -4,16 +4,18 @@ import uuid
 import streamlit as st
 from dotenv import load_dotenv
 from qdrant_client import models
-
 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 retrieve import retrieve_documents_from_collection
 from summarize import summarize_documents
 
 # Load environment variables
 load_dotenv()
 
+
+
+
 def main():
     st.sidebar.title("PDF Management")
     uploaded_files = st.sidebar.file_uploader("Upload PDF files", type=["pdf"], accept_multiple_files=True)
@@ -26,6 +28,8 @@ def main():
         st.session_state['data_bank_answer'] = None
     if 'uploaded_docs_answer' not in st.session_state:
         st.session_state['uploaded_docs_answer'] = None
+    if 'voo_answer' not in st.session_state:
+        st.session_state['voo_answer'] = None
 
     if uploaded_files:
         if st.sidebar.button("Add Docs to Data Bank"):
@@ -38,8 +42,9 @@ def main():
 
     pages = {
         "Lex Document Summarization": page_summarization,
-        "Chat with RSCA": page_qna,
-        "Chat with Uploaded Docs": page_chat_with_uploaded_docs
+        "Chat with Data Bank": page_qna,
+        "Chat with Uploaded Docs": page_chat_with_uploaded_docs,
+        "Chat with VOO": page_chat_with_voo
     }
 
     st.sidebar.title("Page Navigation")
@@ -86,7 +91,7 @@ def page_summarization(uploaded_files):
 
 def page_qna(uploaded_files):
     """Page for Q&A functionality."""
-    st.title("Chat with RSCA")
+    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:
@@ -121,6 +126,21 @@ def page_chat_with_uploaded_docs(uploaded_files):
             st.session_state['uploaded_collection_name'] = None
             st.success(f"Deleted collection {collection_name}")
 
+def page_chat_with_voo(uploaded_files):
+    """Page for chatting with VOO documents."""
+    st.title("Chat with VOO")
+    user_query = st.text_area("Enter your question here:", height=300)
+    if st.button('Get Answer'):
+        if user_query:
+            answer = handle_voo_query(user_query)
+            st.session_state['voo_answer'] = answer  # Store the answer in session state
+            st.write(answer)
+        else:
+            st.error("Please enter a question to get an answer.")
+    # Display stored answer if it exists
+    if st.session_state['voo_answer']:
+        st.write(st.session_state['voo_answer'])
+
 def embed_documents_to_data_bank(files_info):
     """Function to embed documents into the data bank."""
     for temp_path, original_name in files_info:
@@ -172,7 +192,7 @@ def add_docs_to_current_chat(files_info):
 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'))
+        answer = retrieve_documents_from_collection(query, os.getenv('OPENAI_API_KEY'), os.getenv('QDRANT_URL'), os.getenv('QDRANT_API_KEY'),'Lex-v1')
         return answer or "No relevant answer found."
     except Exception as e:
         return f"Error processing the query: {str(e)}"
@@ -185,6 +205,14 @@ def handle_uploaded_docs_query(query, collection_name):
     except Exception as e:
         return f"Error processing the query: {str(e)}"
 
+def handle_voo_query(query):
+    """Retrieve answers from the VOO collection."""
+    try:
+        answer = retrieve_documents_from_collection(query, os.getenv('OPENAI_API_KEY'), os.getenv('QDRANT_URL'), os.getenv('QDRANT_API_KEY'), 'Lex-v2')
+        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)
@@ -194,4 +222,4 @@ def delete_collection(collection_name, qdrant_url, qdrant_api_key):
         print("Failed to delete collection:", e)
 
 if __name__ == "__main__":
-    main()
+    main()