Update app.py

app.py

@@ -1,6 +1,8 @@
 import streamlit as st
 from azure.cosmos import CosmosClient
 import os
+import traceback
+import json
 
 # Cosmos DB configuration
 ENDPOINT = "https://acae-afd.documents.azure.com:443/"
@@ -10,26 +12,40 @@ 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):
+def insert_record(container, record):
     try:
         response = container.create_item(body=record)
         return True, response
     except Exception as e:
         return False, str(e)
 
-def call_stored_procedure(container, long_text):
+def call_stored_procedure(container, long_text, partition_key_type):
     try:
+        if partition_key_type == "first_word":
+            partition_key = long_text.split()[0]
+        elif partition_key_type == "first_two_words":
+            partition_key = " ".join(long_text.split()[:2])
+        elif partition_key_type == "first_line":
+            partition_key = long_text.split('\n')[0].strip()
+        elif partition_key_type == "json":
+            # Assume the text is a JSON string and extract a specific field
+            json_data = json.loads(long_text)
+            partition_key = json_data.get("partitionKey", "default")
+        else:
+            partition_key = "default"
+
         result = container.scripts.execute_stored_procedure(
-            sproc='processPrompt',  # Replace with your actual stored procedure name
+            sproc='processQTPrompts',  # Make sure this matches your stored procedure name
             params=[long_text],
-            partition_key=None  # You may need to adjust this based on your container's partition key
+            partition_key=partition_key
         )
         return True, result
     except Exception as e:
-        return False, str(e)
+        error_msg = f"Error: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
+        return False, error_msg
 
 # Streamlit app
-st.title("🌟 Cosmos DB Record Insertion and Stored Procedure")
+st.title("🌟 Cosmos DB Record Insertion and Stored Procedure Caller")
 
 # Login section
 if 'logged_in' not in st.session_state:
@@ -37,8 +53,7 @@ if 'logged_in' not in st.session_state:
 
 if not st.session_state.logged_in:
     st.subheader("🔐 Login")
-    #input_key = st.text_input("Enter your Cosmos DB Primary Key", type="password")
-    input_key = Key
+    input_key = st.text_input("Enter your Cosmos DB Primary Key", type="password", value=Key)
     if st.button("🚀 Login"):
         if input_key:
             st.session_state.primary_key = input_key
@@ -52,47 +67,52 @@ else:
     database = client.get_database_client(DATABASE_NAME)
     container = database.get_container_client(CONTAINER_NAME)
 
-    # Input fields for record insertion
-    st.subheader("📝 Enter Record Details")
-    id = st.text_input("ID")
-    name = st.text_input("Name")
-    age = st.number_input("Age", min_value=0, max_value=150)
-    city = st.text_input("City")
+    # Create tabs for different operations
+    tab1, tab2 = st.tabs(["Insert Record", "Call Stored Procedure"])
 
-    # Submit button for record insertion
-    if st.button("💾 Insert Record"):
-        record = {
-            "id": id,
-            "name": name,
-            "age": age,
-            "city": city
-        }
-        
-        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 tab1:
+        st.subheader("📝 Enter Record Details")
+        id = st.text_input("ID")
+        name = st.text_input("Name")
+        age = st.number_input("Age", min_value=0, max_value=150)
+        city = st.text_input("City")
+
+        if st.button("💾 Insert Record"):
+            record = {
+                "id": id,
+                "name": name,
+                "age": age,
+                "city": city
+            }
+            
+            success, response = insert_record(container, record)
+            if success:
+                st.success("✅ Record inserted successfully!")
+                st.json(response)
+            else:
+                st.error(f"❌ Failed to insert record: {response}")
 
-    # Separator
-    st.markdown("---")
+    with tab2:
+        st.subheader("📝 Enter Long Text for Stored Procedure")
+        long_text = st.text_area("Long Text", height=200)
 
-    # Input field for long text
-    st.subheader("📝 Enter Long Text for Stored Procedure")
-    long_text = st.text_area("Long Text", height=200)
+        partition_key_type = st.selectbox(
+            "Select Partition Key Type",
+            ["first_word", "first_two_words", "first_line", "json", "default"]
+        )
 
-    # Button to call stored procedure
-    if st.button("🚀 Call Stored Procedure"):
-        if long_text:
-            success, result = call_stored_procedure(container, long_text)
-            if success:
-                st.success("✅ Stored procedure executed successfully!")
-                st.json(result)
+        if st.button("🚀 Call Stored Procedure"):
+            if long_text:
+                with st.spinner('Calling stored procedure...'):
+                    success, result = call_stored_procedure(container, long_text, partition_key_type)
+                    if success:
+                        st.success("✅ Stored procedure executed successfully!")
+                        st.json(result)
+                    else:
+                        st.error("❌ Failed to execute stored procedure")
+                        st.text(result)  # This will show the detailed error message
             else:
-                st.error(f"❌ Failed to execute stored procedure: {result}")
-        else:
-            st.warning("Please enter some text before calling the stored procedure.")
+                st.warning("Please enter some text before calling the stored procedure.")
 
     # Logout button
     if st.button("🚪 Logout"):
@@ -104,4 +124,8 @@ else:
     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}")
+    st.sidebar.text(f"Container: {CONTAINER_NAME}")
+    
+    
+    
+