Update app.py

app.py

@@ -25,151 +25,105 @@ Key = os.environ.get("Key")  # 🔑 Don't forget your key!
 # 🏠 Your local app URL (Change this to your app's URL)
 LOCAL_APP_URL = "https://huggingface.co/spaces/awacke1/AzureCosmosDBUI"
 
-# 🐙 GitHub configuration
-def download_github_repo(url, local_path):
-    # 🚚 Let's download that GitHub repo!
-    if os.path.exists(local_path):
-        shutil.rmtree(local_path)
-    Repo.clone_from(url, local_path)
-
-def create_zip_file(source_dir, output_filename):
-    # 📦 Zipping up files like a pro!
-    shutil.make_archive(output_filename, 'zip', source_dir)
-
-def create_repo(g, repo_name):
-    # 🛠️ Creating a new GitHub repo. It's alive!
-    user = g.get_user()
-    return user.create_repo(repo_name)
-
-def push_to_github(local_path, repo, github_token):
-    # 🚀 Pushing code to GitHub. Blast off!
-    repo_url = f"https://{github_token}@github.com/{repo.full_name}.git"
-    local_repo = Repo(local_path)
-    if 'origin' in [remote.name for remote in local_repo.remotes]:
-        origin = local_repo.remote('origin')
-        origin.set_url(repo_url)
-    else:
-        origin = local_repo.create_remote('origin', repo_url)
-    
-    if not local_repo.heads:
-        local_repo.git.checkout('-b', 'main')
-        current_branch = 'main'
-    else:
-        current_branch = local_repo.active_branch.name
-    
-    local_repo.git.add(A=True)
-    
-    if local_repo.is_dirty():
-        local_repo.git.commit('-m', 'Initial commit')
-    
-    origin.push(refspec=f'{current_branch}:{current_branch}')
+# 📝 Initialize logging system
+if 'logs' not in st.session_state:
+    st.session_state.logs = []
 
-def get_base64_download_link(file_path, file_name):
-    # 🧙‍♂️ Generating a magical download link!
-    with open(file_path, "rb") as file:
-        contents = file.read()
-    base64_encoded = base64.b64encode(contents).decode()
-    return f'<a href="data:application/zip;base64,{base64_encoded}" download="{file_name}">⬇️ Download {file_name}</a>'
+def log_event(message):
+    """Adds a log entry to the log session state."""
+    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+    st.session_state.logs.append(f"[{timestamp}] {message}")
 
-# 🧭 Cosmos DB Functions
-def get_databases(client):
-    # 📚 Fetching list of databases!
-    return [db['id'] for db in client.list_databases()]
+# 🏠 Add logs to sidebar
+st.sidebar.title("📝 Log Viewer")
+st.sidebar.write("\n".join(st.session_state.logs))
 
-def get_containers(database):
-    # 📂 Listing containers within the selected database
-    return [container['id'] for container in database.list_containers()]
-
-def get_documents(container, limit=None):
-    # 📝 Fetching documents from Cosmos DB
-    query = "SELECT * FROM c ORDER BY c._ts DESC"
-    items = list(container.query_items(query=query, enable_cross_partition_query=True, max_item_count=limit))
-    return items
+# 🤖 OpenAI configuration
+# openai.api_key = os.environ.get("OPENAI_API_KEY")
+# MODEL = "gpt-3.5-turbo"  # Replace with your desired model
 
+# 🌟 Cosmos DB functions
 def insert_record(container, record):
-    # 📥 Inserting a record into the Cosmos DB
     try:
         container.create_item(body=record)
+        log_event("Record inserted successfully!")
         return True, "Record inserted successfully! 🎉"
     except exceptions.CosmosHttpResponseError as e:
+        log_event(f"HTTP error occurred: {str(e)} 🚨")
         return False, f"HTTP error occurred: {str(e)} 🚨"
     except Exception as e:
+        log_event(f"An unexpected error occurred: {str(e)} 😱")
         return False, f"An unexpected error occurred: {str(e)} 😱"
 
 def update_record(container, updated_record):
-    # 🛠️ Updating a record in the Cosmos DB
     try:
         container.upsert_item(body=updated_record)
+        log_event(f"Record with id {updated_record['id']} successfully updated.")
         return True, f"Record with id {updated_record['id']} successfully updated. 🛠️"
     except exceptions.CosmosHttpResponseError as e:
+        log_event(f"HTTP error occurred: {str(e)} 🚨")
         return False, f"HTTP error occurred: {str(e)} 🚨"
     except Exception as e:
+        log_event(f"An unexpected error occurred: {traceback.format_exc()} 😱")
         return False, f"An unexpected error occurred: {traceback.format_exc()} 😱"
 
 def delete_record(container, name, id):
-    # 🗑️ Deleting a record in Cosmos DB
     try:
         container.delete_item(item=id, partition_key=id)
+        log_event(f"Successfully deleted record with name: {name} and id: {id}")
         return True, f"Successfully deleted record with name: {name} and id: {id} 🗑️"
     except exceptions.CosmosResourceNotFoundError:
+        log_event(f"Record with id {id} not found. It may have been already deleted.")
         return False, f"Record with id {id} not found. It may have been already deleted. 🕵️‍♂️"
     except exceptions.CosmosHttpResponseError as e:
+        log_event(f"HTTP error occurred: {str(e)} 🚨")
         return False, f"HTTP error occurred: {str(e)} 🚨"
     except Exception as e:
+        log_event(f"An unexpected error occurred: {traceback.format_exc()} 😱")
         return False, f"An unexpected error occurred: {traceback.format_exc()} 😱"
 
-def archive_current_container(database_name, container_name, client):
-    # 📦 Archiving the container data from Cosmos DB
-    try:
-        base_dir = "./cosmos_archive_current_container"
-        if os.path.exists(base_dir):
-            shutil.rmtree(base_dir)
-        os.makedirs(base_dir)
-        
-        db_client = client.get_database_client(database_name)
-        container_client = db_client.get_container_client(container_name)
-        items = list(container_client.read_all_items())
-        
-        container_dir = os.path.join(base_dir, container_name)
-        os.makedirs(container_dir)
-        
-        for item in items:
-            item_id = item.get('id', f"unknown_{datetime.now().strftime('%Y%m%d%H%M%S')}")
-            with open(os.path.join(container_dir, f"{item_id}.json"), 'w') as f:
-                json.dump(item, f, indent=2)
-        
-        archive_name = f"{container_name}_archive_{datetime.now().strftime('%Y%m%d%H%M%S')}"
-        shutil.make_archive(archive_name, 'zip', base_dir)
-        
-        return get_base64_download_link(f"{archive_name}.zip", f"{archive_name}.zip")
-    except Exception as e:
-        return f"An error occurred while archiving data: {str(e)} 😢"
-
 # 🎲 Function to generate a unique UUID
 def generate_unique_id():
-    # 🔮 Generate a unique ID for new records
     return str(uuid.uuid4())
 
-# 🔗 Helper function for hyperlink extraction
-def extract_hyperlinks(responses):
-    # 🔗 Extracting hyperlinks from response
-    hyperlinks = []
-    for response in responses:
-        parsed_response = json.loads(response)
-        links = [value for key, value in parsed_response.items() if isinstance(value, str) and value.startswith("http")]
-        hyperlinks.extend(links)
-    return hyperlinks
-
-# 📋 Helper function to format text with line numbers
-def format_with_line_numbers(text):
-    # 📋 Organizing lines with numbers
-    lines = text.splitlines()
-    formatted_text = '\n'.join(f"{i+1}: {line}" for i, line in enumerate(lines))
-    return formatted_text
+def get_databases(client):
+    return [db['id'] for db in client.list_databases()]
+
+def get_containers(database):
+    return [container['id'] for container in database.list_containers()]
+
+def get_documents(container, limit=None):
+    query = "SELECT * FROM c ORDER BY c._ts DESC"
+    items = list(container.query_items(query=query, enable_cross_partition_query=True, max_item_count=limit))
+    return items
+
+def save_to_cosmos_db(container, query, response1, response2):
+    try:
+        if container:
+            record = {
+                "id": generate_unique_id(),
+                "query": query,
+                "response1": response1,
+                "response2": response2
+            }
+            try:
+                container.create_item(body=record)
+                st.success(f"Record saved successfully with ID: {record['id']}")
+                log_event(f"Record saved to Cosmos DB with ID: {record['id']}")
+                # Refresh the documents display
+                st.session_state.documents = get_documents(container)
+            except exceptions.CosmosHttpResponseError as e:
+                st.error(f"Error saving record to Cosmos DB: {e}")
+                log_event(f"Error saving record to Cosmos DB: {e}")
+        else:
+            st.error("Cosmos DB container is not initialized.")
+            log_event("Error: Cosmos DB container is not initialized.")
+    except Exception as e:
+        st.error(f"An unexpected error occurred: {str(e)}")
+        log_event(f"An unexpected error occurred while saving to Cosmos DB: {str(e)}")
 
 # 🎤 Function to process text input
 def process_text(text_input):
-    # 🎤 Processing input for queries or searches
     if text_input:
         if 'messages' not in st.session_state:
             st.session_state.messages = []
@@ -182,22 +136,38 @@ def process_text(text_input):
         with st.chat_message("assistant"):
             search_glossary(text_input)
 
-# 🎈 Main function to drive the application
+# 🎈 Main function
 def main():
-    # 🎈 Let's kickstart the app with a fun introduction!
     st.title("🐙Git🌌Cosmos💫 - Azure Cosmos DB and Github Agent")
 
-    # 🚦 Initialize session state for managing components
     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:
         st.session_state.selected_database = None
     if 'selected_container' not in st.session_state:
         st.session_state.selected_container = None
+    if 'selected_document_id' not in st.session_state:
+        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
 
-    # 🔐 Login management
+    # ⚙️ Check query parameters for any action
+    try:
+        query_params = st.query_params
+        query = query_params.get('q') or query_params.get('query') or ''
+        if query:
+            process_text(query)
+            st.stop()
+    except Exception as e:
+        st.markdown(' ')
+    
+    # 🔐 Automatic Login
     if Key:
         st.session_state.primary_key = Key
         st.session_state.logged_in = True
@@ -206,20 +176,20 @@ def main():
         return
 
     if st.session_state.logged_in:
+        # 🌌 Initialize Cosmos DB client
         try:
-            # 🌌 Initialize Cosmos DB client
             if st.session_state.client is None:
                 st.session_state.client = CosmosClient(ENDPOINT, credential=st.session_state.primary_key)
             
-            # 🗄️ Sidebar for database and container selection
-            st.sidebar.title("🗄️ Cosmos DB Navigator")
-            
+            st.sidebar.title("🐙Git🌌Cosmos💫🗄️Navigator")
             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:
@@ -229,12 +199,14 @@ def main():
                 
                 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)
-
-                    # 📦 Export container data
+                    
+                    # 📦 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'):
@@ -242,57 +214,230 @@ def main():
                         else:
                             st.error(download_link)
                     
-                    # 📝 Display documents
+                    # Fetch documents
                     documents = get_documents(container)
-                    st.subheader(f"📊 Container: {st.session_state.selected_container}")
-                    if documents:
-                        df = pd.DataFrame(documents)
-                        st.dataframe(df)
+                    total_docs = len(documents)
+                    
+                    if total_docs > 5:
+                        documents_to_display = documents[:5]
+                        st.info("Showing top 5 most recent documents.")
                     else:
-                        st.info("No documents to display. 🧐")
+                        documents_to_display = documents
+                        st.info(f"Showing all {len(documents_to_display)} documents.")
+                    
+                    if documents_to_display:
+                        # 🎨 Add Viewer/Editor selection
+                        view_options = ['Show as Markdown', 'Show as Code Editor', 'Show as Edit and Save', 'Clone Document', 'New Record']
+                        selected_view = st.selectbox("Select Viewer/Editor", view_options, index=2)
+                        
+                        if selected_view == 'Show as Markdown':
+                            doc = documents[st.session_state.current_index]
+                            st.markdown(f"#### Document ID: {doc.get('id', '')}")
+                            content = json.dumps(doc, indent=2)
+                            st.markdown(f"```json\n{content}\n```")
+
+                            col_prev, col_next = st.columns([1, 1])
+                            with col_prev:
+                                if st.button("⬅️ Previous", key='prev_markdown'):
+                                    if st.session_state.current_index > 0:
+                                        st.session_state.current_index -= 1
+                                        st.rerun()
+                            with col_next:
+                                if st.button("➡️ Next", key='next_markdown'):
+                                    if st.session_state.current_index < total_docs - 1:
+                                        st.session_state.current_index += 1
+                                        st.rerun()
+
+                        elif selected_view == 'Show as Code Editor':
+                            doc = documents[st.session_state.current_index]
+                            st.markdown(f"#### Document ID: {doc.get('id', '')}")
+                            doc_str = st.text_area("Edit Document", value=json.dumps(doc, indent=2), height=300, key=f'code_editor_{st.session_state.current_index}')
+                            
+                            col_prev, col_next = st.columns([1, 1])
+                            with col_prev:
+                                if st.button("⬅️ Previous", key='prev_code'):
+                                    if st.session_state.current_index > 0:
+                                        st.session_state.current_index -= 1
+                                        st.rerun()
+                            with col_next:
+                                if st.button("➡️ Next", key='next_code'):
+                                    if st.session_state.current_index < total_docs - 1:
+                                        st.session_state.current_index += 1
+                                        st.rerun()
+                            
+                            if st.button("💾 Save Changes", key=f'save_button_{st.session_state.current_index}'):
+                                try:
+                                    updated_doc = json.loads(doc_str)
+                                    success, message = update_record(container, updated_doc)
+                                    if success:
+                                        st.success(f"Document {updated_doc['id']} saved successfully.")
+                                        st.session_state.selected_document_id = updated_doc['id']
+                                        st.rerun()
+                                    else:
+                                        st.error(message)
+                                except json.JSONDecodeError as e:
+                                    st.error(f"Invalid JSON: {str(e)} 🚫")
+                                    
+                        elif selected_view == 'Show as Edit and Save':
+                            st.markdown("#### Edit the document fields below:")
+                            num_cols = len(documents_to_display)
+                            cols = st.columns(num_cols)
+
+                            for idx, (col, doc) in enumerate(zip(cols, documents_to_display)):
+                                with col:
+                                    st.markdown(f"##### Document ID: {doc.get('id', '')}")
+                                    editable_id = st.text_input("ID", value=doc.get('id', ''), key=f'edit_id_{idx}')
+                                    editable_doc = doc.copy()
+                                    editable_doc.pop('id', None)
+                                    doc_str = st.text_area("Document Content (in JSON format)", value=json.dumps(editable_doc, indent=2), height=300, key=f'doc_str_{idx}')
+                                    
+                                    col_save, col_ai = st.columns(2)
+                                    with col_save:
+                                        if st.button("💾 Save Changes", key=f'save_button_{idx}'):
+                                            try:
+                                                updated_doc = json.loads(doc_str)
+                                                updated_doc['id'] = editable_id
+                                                success, message = update_record(container, updated_doc)
+                                                if success:
+                                                    st.success(f"Document {updated_doc['id']} saved successfully.")
+                                                    st.session_state.selected_document_id = updated_doc['id']
+                                                    st.rerun()
+                                                else:
+                                                    st.error(message)
+                                            except json.JSONDecodeError as e:
+                                                st.error(f"Invalid JSON: {str(e)} 🚫")
+                                    with col_ai:
+                                        if st.button("🤖 Run With AI", key=f'run_with_ai_button_{idx}'):
+                                            search_glossary(json.dumps(editable_doc, indent=2))
+
+                        elif selected_view == 'Clone Document':
+                            st.markdown("#### Clone a document:")
+                            for idx, doc in enumerate(documents_to_display):
+                                st.markdown(f"##### Document ID: {doc.get('id', '')}")
+                                if st.button("📄 Clone Document", key=f'clone_button_{idx}'):
+                                    cloned_doc = doc.copy()
+                                    cloned_doc['id'] = generate_unique_id()
+                                    st.session_state.cloned_doc = cloned_doc
+                                    st.session_state.cloned_doc_str = json.dumps(cloned_doc, indent=2)
+                                    st.session_state.clone_mode = True
+                                    st.rerun()
+
+                            if st.session_state.get('clone_mode', False):
+                                st.markdown("#### Edit Cloned Document:")
+                                cloned_doc_str = st.text_area("Cloned Document Content (in JSON format)", value=st.session_state.cloned_doc_str, height=300)
+                                if st.button("💾 Save Cloned Document"):
+                                    try:
+                                        new_doc = json.loads(cloned_doc_str)
+                                        success, message = insert_record(container, new_doc)
+                                        if success:
+                                            st.success(f"Cloned document saved with id: {new_doc['id']} 🎉")
+                                            st.session_state.selected_document_id = new_doc['id']
+                                            st.session_state.clone_mode = False
+                                            st.session_state.cloned_doc = None
+                                            st.session_state.cloned_doc_str = ''
+                                            st.rerun()
+                                        else:
+                                            st.error(message)
+                                    except json.JSONDecodeError as e:
+                                        st.error(f"Invalid JSON: {str(e)} 🚫")
+
+                        elif selected_view == 'New Record':
+                            st.markdown("#### Create a new document:")
+                            if st.button("🤖 Insert Auto-Generated Record"):
+                                success, message = insert_auto_generated_record(container)
+                                if success:
+                                    st.success(message)
+                                    st.rerun()
+                                else:
+                                    st.error(message)
+                            else:
+                                new_id = st.text_input("ID", value=generate_unique_id(), key='new_id')
+                                new_doc_str = st.text_area("Document Content (in JSON format)", value='{}', height=300)
+                                if st.button("➕ Create New Document"):
+                                    try:
+                                        new_doc = json.loads(new_doc_str)
+                                        new_doc['id'] = new_id
+                                        success, message = insert_record(container, new_doc)
+                                        if success:
+                                            st.success(f"New document created with id: {new_doc['id']} 🎉")
+                                            st.session_state.selected_document_id = new_doc['id']
+                                            st.rerun()
+                                        else:
+                                            st.error(message)
+                                    except json.JSONDecodeError as e:
+                                        st.error(f"Invalid JSON: {str(e)} 🚫")
+
+                    else:
+                        st.sidebar.info("No documents found in this container. 📭")
             
-            # 🐙 GitHub operations
+            # 🎉 Main content area
+            st.subheader(f"📊 Container: {st.session_state.selected_container}")
+            if st.session_state.selected_container:
+                if documents_to_display:
+                    df = pd.DataFrame(documents_to_display)
+                    st.dataframe(df)
+                else:
+                    st.info("No documents to display. 🧐")
+
+            # 🐙 GitHub section
             st.subheader("🐙 GitHub Operations")
-            github_token = os.environ.get("GITHUB")  # GitHub token for authentication
+            github_token = os.environ.get("GITHUB")  # Read GitHub token from environment variable
             source_repo = st.text_input("Source GitHub Repository URL", value="https://github.com/AaronCWacker/AIExamples-8-24-Streamlit")
-            new_repo_name = st.text_input("New Repository Name", value=f"AIExample-Clone-{datetime.now().strftime('%Y%m%d_%H%M%S')}")
+            new_repo_name = st.text_input("New Repository Name (for cloning)", value=f"AIExample-Clone-{datetime.now().strftime('%Y%m%d_%H%M%S')}")
 
             col1, col2 = st.columns(2)
             with col1:
                 if st.button("📥 Clone Repository"):
                     if github_token and source_repo:
-                        local_path = f"./temp_repo_{datetime.now().strftime('%Y%m%d%H%M%S')}"
-                        download_github_repo(source_repo, local_path)
-                        zip_filename = f"{new_repo_name}.zip"
-                        create_zip_file(local_path, zip_filename[:-4])
-                        st.markdown(get_base64_download_link(zip_filename, zip_filename), unsafe_allow_html=True)
-                        st.success("Repository cloned successfully! 🎉")
-                        shutil.rmtree(local_path)
+                        try:
+                            local_path = f"./temp_repo_{datetime.now().strftime('%Y%m%d%H%M%S')}"
+                            download_github_repo(source_repo, local_path)
+                            zip_filename = f"{new_repo_name}.zip"
+                            create_zip_file(local_path, zip_filename[:-4])
+                            st.markdown(get_base64_download_link(zip_filename, zip_filename), unsafe_allow_html=True)
+                            st.success("Repository cloned successfully! 🎉")
+                        except Exception as e:
+                            st.error(f"An error occurred: {str(e)} 😢")
+                        finally:
+                            if os.path.exists(local_path):
+                                shutil.rmtree(local_path)
+                            if os.path.exists(zip_filename):
+                                os.remove(zip_filename)
                     else:
-                        st.error("Please provide GitHub token and source repo URL.")
-
+                        st.error("Please ensure GitHub token is set in environment variables and source repository URL is provided. 🔑❓")
+        
             with col2:
                 if st.button("📤 Push to New Repository"):
                     if github_token and source_repo:
-                        g = Github(github_token)
-                        new_repo = create_repo(g, new_repo_name)
-                        local_path = f"./temp_repo_{datetime.now().strftime('%Y%m%d%H%M%S')}"
-                        download_github_repo(source_repo, local_path)
-                        push_to_github(local_path, new_repo, github_token)
-                        st.success(f"Repository pushed successfully to {new_repo.html_url} 🚀")
-                        shutil.rmtree(local_path)
+                        try:
+                            g = Github(github_token)
+                            new_repo = create_repo(g, new_repo_name)
+                            local_path = f"./temp_repo_{datetime.now().strftime('%Y%m%d%H%M%S')}"
+                            download_github_repo(source_repo, local_path)
+                            push_to_github(local_path, new_repo, github_token)
+                            st.success(f"Repository pushed successfully to {new_repo.html_url} 🚀")
+                        except Exception as e:
+                            st.error(f"An error occurred: {str(e)} 😢")
+                        finally:
+                            if os.path.exists(local_path):
+                                shutil.rmtree(local_path)
                     else:
-                        st.error("Please provide GitHub token and source repo URL.")
+                        st.error("Please ensure GitHub token is set in environment variables and source repository URL is provided. 🔑❓")
 
         except exceptions.CosmosHttpResponseError as e:
             st.error(f"Failed to connect to Cosmos DB. HTTP error: {str(e)} 🚨")
         except Exception as e:
             st.error(f"An unexpected error occurred: {str(e)} 😱")
 
-    # 🚪 Logout option
+    # 🚪 Logout button
     if st.session_state.logged_in and st.sidebar.button("🚪 Logout"):
         st.session_state.logged_in = False
+        st.session_state.selected_records.clear()
         st.session_state.client = None
+        st.session_state.selected_database = None
+        st.session_state.selected_container = None
+        st.session_state.selected_document_id = None
+        st.session_state.current_index = 0
         st.rerun()
 
 if __name__ == "__main__":