app.py

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']