Update app.py

app.py

@@ -85,16 +85,23 @@ def ensure_stored_procedure_exists(container):
         create_stored_procedure(container)
 
 def process_qt_prompts(container, prompt_text):
-    # Use a dummy partition key value
-    partition_key = "dummy_partition_key"
+    # Use the first prompt as the partition key
+    first_prompt = prompt_text.split('\n')[0].strip()
     return container.scripts.execute_stored_procedure(
         sproc='processQTPrompts',
         params=[prompt_text],
-        partition_key=partition_key
+        partition_key=first_prompt
     )
 
+def insert_record(container, record):
+    try:
+        response = container.create_item(body=record)
+        return True, response
+    except Exception as e:
+        return False, str(e)
+
 # Streamlit app
-st.title("🌟 QT Prompt Processor")
+st.title("🌟 Cosmos DB Record Insertion and QT Prompt Processor")
 
 # Login section
 if 'logged_in' not in st.session_state:
@@ -102,8 +109,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")
     if st.button("🚀 Login"):
         if input_key:
             st.session_state.primary_key = input_key
@@ -120,26 +126,50 @@ else:
     # Ensure the stored procedure exists
     ensure_stored_procedure_exists(container)
 
-    # Input field for QT prompts
-    st.subheader("📝 Enter QT Prompts")
-    default_text = "QT Crystal finders: Bioluminescent crystal caverns, quantum-powered explorers, prismatic hues, alien planet\nQT robot art: Cybernetic metropolis, sentient androids, rogue AI, neon-infused palette\nQT the Lava girl: Volcanic exoplanet, liquid metal rivers, heat-immune heroine, molten metallic palette"
-    qt_prompts = st.text_area("QT Prompts", value=default_text, height=300)
-
-    # Submit button
-    if st.button("🚀 Process QT Prompts"):
-        try:
-            results = process_qt_prompts(container, qt_prompts)
+    # Tab for different operations
+    tab1, tab2 = st.tabs(["Insert Record", "Process QT Prompts"])
+
+    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
+            }
             
-            # Display results in a dataframe
-            df_data = [{"Prompt": item['id'], "Occurrence Count": item['occurrenceCount']} for item in results]
-            st.dataframe(df_data)
-        except Exception as e:
-            st.error(f"An error occurred: {str(e)}")
+            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}")
+
+    with tab2:
+        st.subheader("📝 Enter QT Prompts")
+        default_text = "QT Crystal finders: Bioluminescent crystal caverns, quantum-powered explorers, prismatic hues, alien planet\nQT robot art: Cybernetic metropolis, sentient androids, rogue AI, neon-infused palette\nQT the Lava girl: Volcanic exoplanet, liquid metal rivers, heat-immune heroine, molten metallic palette"
+        qt_prompts = st.text_area("QT Prompts", value=default_text, height=300)
+
+        if st.button("🚀 Process QT Prompts"):
+            try:
+                results = process_qt_prompts(container, qt_prompts)
+                
+                # Display results in a dataframe
+                df_data = [{"Prompt": item['id'], "Occurrence Count": item['occurrenceCount']} for item in results]
+                st.dataframe(df_data)
+            except Exception as e:
+                st.error(f"An error occurred: {str(e)}")
 
     # Logout button
     if st.button("🚪 Logout"):
         st.session_state.logged_in = False
-        st.experimental_rerun()
+        st.rerun()
 
     # Display connection info
     st.sidebar.subheader("🔗 Connection Information")