Update app.py

app.py

@@ -168,7 +168,6 @@ import uuid
 def generate_unique_id():
     return str(uuid.uuid4())
     
-
 def save_or_clone_to_cosmos_db(container, query=None, response=None, clone_id=None):
     try:
         if not container:
@@ -217,6 +216,7 @@ def save_or_clone_to_cosmos_db(container, query=None, response=None, clone_id=No
     except Exception as e:
         st.error(f"An unexpected error occurred: {str(e)}")
 
+
 # Add dropdowns for model and database choices
 def search_glossary(query):
     st.markdown(f"### 🔍 Search Glossary for: `{query}`")
@@ -449,10 +449,6 @@ def main():
     st.title("🐙Git🌌Cosmos💫 - Azure Cosmos DB and Github Agent")
 
     # Initialize session state
-    if 'logged_in' not in st.session_state:
-        st.session_state.logged_in = False
-    if 'selected_records' not in st.session_state:
-        st.session_state.selected_records = []
     if 'client' not in st.session_state:
         st.session_state.client = None
     if 'selected_database' not in st.session_state:
@@ -463,19 +459,68 @@ def main():
         st.session_state.selected_document_id = None
     if 'current_index' not in st.session_state:
         st.session_state.current_index = 0
-    if 'cloned_doc' not in st.session_state:
-        st.session_state.cloned_doc = None
+    if 'documents' not in st.session_state:
+        st.session_state.documents = []
 
-    # ⚙️ q= Run ArXiv search from query parameters
-    try:
-        query_params = st.query_params
-        query = query_params.get('q') or query_params.get('query') or ''
-        if query:
-            # 🕵️‍♂️ We have a query! Let's process it!
-            process_text(query)
-            st.stop()  # Stop further execution
-    except Exception as e:
-        st.markdown(' ')
+    # Cosmos DB Connection (assuming this is done earlier in your code)
+    if st.session_state.client is None:
+        st.session_state.client = CosmosClient(ENDPOINT, credential=st.session_state.primary_key)
+
+    # Database and Container Selection
+    databases = get_databases(st.session_state.client)
+    selected_db = st.sidebar.selectbox("🗃️ Select Database", databases)
+    
+    if selected_db != st.session_state.selected_database:
+        st.session_state.selected_database = selected_db
+        st.session_state.selected_container = None
+        st.session_state.selected_document_id = None
+        st.session_state.current_index = 0
+        st.rerun()
+    
+    if st.session_state.selected_database:
+        database = st.session_state.client.get_database_client(st.session_state.selected_database)
+        containers = get_containers(database)
+        selected_container = st.sidebar.selectbox("📁 Select Container", containers)
+        
+        if selected_container != st.session_state.selected_container:
+            st.session_state.selected_container = selected_container
+            st.session_state.selected_document_id = None
+            st.session_state.current_index = 0
+            st.rerun()
+        
+        if st.session_state.selected_container:
+            container = database.get_container_client(st.session_state.selected_container)
+            
+            # Add Export button
+            if st.button("📦 Export Container Data"):
+                download_link = archive_current_container(st.session_state.selected_database, st.session_state.selected_container, st.session_state.client)
+                if download_link.startswith('<a'):
+                    st.markdown(download_link, unsafe_allow_html=True)
+                else:
+                    st.error(download_link)
+            
+            # Fetch documents
+            st.session_state.documents = get_documents(container)
+            total_docs = len(st.session_state.documents)
+
+            # Add buttons for saving new document and cloning
+            if st.button("💾 Save New Document"):
+                user_query = st.session_state.get('user_query', '')
+                ai_response = st.session_state.get('ai_response', '')
+                save_or_clone_to_cosmos_db(container, query=user_query, response=ai_response)
+
+            if st.button("🔄 Clone Document"):
+                doc_to_clone = st.selectbox("Select document to clone", [doc['id'] for doc in st.session_state.documents])
+                if doc_to_clone:
+                    save_or_clone_to_cosmos_db(container, clone_id=doc_to_clone)
+
+            # Display documents
+            if st.session_state.documents:
+                for doc in st.session_state.documents:
+                    st.write(f"ID: {doc['id']}")
+                    st.write(f"Query: {doc.get('query', 'N/A')}")
+                    st.write(f"Response: {doc.get('response', 'N/A')}")
+                    st.write("---")
 
     # Automatic Login
     if Key:
@@ -847,6 +892,19 @@ def main():
         except Exception as e:
             st.error(f"An unexpected error occurred: {str(e)} 😱")
 
+
+    # ⚙️ q= Run ArXiv search from query parameters
+    try:
+        query_params = st.query_params
+        query = query_params.get('q') or query_params.get('query') or ''
+        if query:
+            # 🕵️‍♂️ We have a query! Let's process it!
+            process_text(query)
+            #st.stop()  # Stop further execution
+    except Exception as e:
+        st.markdown(' ')
+
+        
     # Logout button
     if st.session_state.logged_in and st.sidebar.button("🚪 Logout"):
         st.session_state.logged_in = False