import streamlit as st from azure.cosmos import CosmosClient, PartitionKey import os import pandas as pd # Streamlit page configuration for wide mode st.set_page_config(layout="wide") # Cosmos DB configuration ENDPOINT = "https://acae-afd.documents.azure.com:443/" SUBSCRIPTION_ID = "003fba60-5b3f-48f4-ab36-3ed11bc40816" DATABASE_NAME = os.environ.get("COSMOS_DATABASE_NAME") CONTAINER_NAME = os.environ.get("COSMOS_CONTAINER_NAME") Key = os.environ.get("Key") def insert_record(record): try: response = container.create_item(body=record) return True, response except Exception as e: return False, str(e) def call_stored_procedure(record): try: response = container.scripts.execute_stored_procedure( sproc="processPrompt", params=[record], partition_key=record['id'] ) return True, response except Exception as e: error_message = f"Error type: {type(e).__name__}\nError message: {str(e)}" if hasattr(e, 'sub_status'): error_message += f"\nSub-status: {e.sub_status}" if hasattr(e, 'response'): error_message += f"\nResponse: {e.response}" return False, error_message def fetch_all_records(): query = "SELECT * FROM c" items = list(container.query_items(query=query, enable_cross_partition_query=True)) return pd.DataFrame(items) def delete_records(ids): try: for id in ids: container.delete_item(item=id, partition_key=id) return True, f"Successfully deleted {len(ids)} records" except Exception as e: return False, f"Error deleting records: {str(e)}" def delete_all_records(): try: query = "SELECT c.id FROM c" items = list(container.query_items(query=query, enable_cross_partition_query=True)) for item in items: container.delete_item(item=item['id'], partition_key=item['id']) return True, f"Successfully deleted all records" except Exception as e: return False, f"Error deleting all records: {str(e)}" # Streamlit app st.title("🌟 Cosmos DB Record Management") # Login section if 'logged_in' not in st.session_state: st.session_state.logged_in = False if not st.session_state.logged_in: st.subheader("🔐 Login") input_key = Key if st.button("🚀 Login"): if input_key: st.session_state.primary_key = input_key st.session_state.logged_in = True st.rerun() else: st.error("Invalid key. Please check your environment variables.") else: # Initialize Cosmos DB client client = CosmosClient(ENDPOINT, credential=st.session_state.primary_key) database = client.get_database_client(DATABASE_NAME) container = database.get_container_client(CONTAINER_NAME) # Fetch all records records_df = fetch_all_records() # Sidebar for data editing st.sidebar.title("Edit Records") edited_df = st.sidebar.data_editor(records_df, num_rows="dynamic", key="sidebar_editor") # Main Content Area col1, col2 = st.columns([2, 1]) with col1: st.subheader("📝 Record Details") id = st.text_input("ID") name = st.text_input("Name") document = st.text_area("Document") evaluation_text = st.text_area("Evaluation Text") evaluation_score = st.number_input("Evaluation Score", min_value=0, max_value=100, step=1) col3, col4, col5 = st.columns(3) with col3: if st.button("💾 Insert Record"): record = { "id": id, "name": name, "document": document, "evaluationText": evaluation_text, "evaluationScore": evaluation_score } success, response = insert_record(record) if success: st.success("✅ Record inserted successfully!") st.json(response) else: st.error(f"❌ Failed to insert record: {response}") with col4: if st.button("🔧 Call Procedure"): record = { "id": id, "name": name, "document": document, "evaluationText": evaluation_text, "evaluationScore": evaluation_score } success, response = call_stored_procedure(record) if success: st.success("✅ Stored procedure executed successfully!") st.json(response) else: st.error(f"❌ Failed to execute stored procedure: {response}") with col5: if st.button("🗑️ Delete All"): success, message = delete_all_records() if success: st.success(message) else: st.error(message) st.rerun() with col2: st.subheader("🔍 Selected Record") selected_rows = edited_df[edited_df['select'] == True] if not selected_rows.empty: st.json(selected_rows.iloc[0].to_dict()) else: st.write("No record selected") st.subheader("📊 All Records") for i, record in edited_df.iterrows(): cols = st.columns([3, 2, 2, 1, 1, 1]) with cols[0]: st.markdown(f"**{record['name']}**") st.markdown(f"*{record['document']}*") with cols[1]: st.markdown(f"Evaluation: {record['evaluationText']}") with cols[2]: st.markdown(f"Score: {record['evaluationScore']}") with cols[3]: st.button("👍", key=f"like_{i}") with cols[4]: st.button("👎", key=f"dislike_{i}") with cols[5]: if st.button("🗑️", key=f"delete_{i}"): success, message = delete_records([record['id']]) if success: st.success(message) else: st.error(message) st.rerun() col6, col7 = st.columns(2) with col6: if st.button("🗑️ Delete Selected"): selected_ids = edited_df[edited_df['select'] == True]['id'].tolist() if selected_ids: success, message = delete_records(selected_ids) if success: st.success(message) else: st.error(message) st.rerun() else: st.warning("No records selected for deletion.") with col7: if st.download_button("📥 Download Data", edited_df.to_csv(index=False), "cosmos_db_data.csv", "text/csv"): st.success("Data downloaded successfully!") # Logout button if st.button("🚪 Logout"): st.session_state.logged_in = False st.rerun() # Display connection info st.sidebar.markdown("---") st.sidebar.subheader("🔗 Connection Information") st.sidebar.text(f"Endpoint: {ENDPOINT}") st.sidebar.text(f"Subscription ID: {SUBSCRIPTION_ID}") st.sidebar.text(f"Database: {DATABASE_NAME}") st.sidebar.text(f"Container: {CONTAINER_NAME}") # Update fields when a record is selected in the sidebar if st.session_state.logged_in and not edited_df[edited_df['select'] == True].empty: selected_record = edited_df[edited_df['select'] == True].iloc[0] st.session_state['id'] = selected_record['id'] st.session_state['name'] = selected_record['name'] st.session_state['document'] = selected_record['document'] st.session_state['evaluation_text'] = selected_record['evaluationText'] st.session_state['evaluation_score'] = selected_record['evaluationScore']