Update app.py

app.py

@@ -34,7 +34,54 @@ def display_cosmos_db_structure():
 # Button to Trigger Display of Cosmos DB Structure
 if st.button('Show Cosmos DB Structure'):
     display_cosmos_db_structure()
-        
+
+
+
+# Function to Add an Item
+def add_item_to_container(database_name, container_name, item):
+    container = cosmos_client.get_database_client(database_name).get_container_client(container_name)
+    container.create_item(item)
+
+# Function to Get All Item Keys from a Container
+def get_all_item_keys(database_name, container_name):
+    container = cosmos_client.get_database_client(database_name).get_container_client(container_name)
+    items = list(container.query_items(
+        query="SELECT c.id FROM c",
+        enable_cross_partition_query=True))
+    return [item['id'] for item in items]
+
+# Function to Retrieve and Display an Item
+def display_item(database_name, container_name, item_id):
+    container = cosmos_client.get_database_client(database_name).get_container_client(container_name)
+    item = container.read_item(item_id, partition_key=PartitionKey(item_id))
+    st.markdown(f"- **Item ID**: {item_id}, **Content**: {json.dumps(item)}")
+
+# Test Function to Add and Retrieve Image Prompts
+def test_image_prompts(database_name, container_name):
+    # Sample data
+    image_prompts = [{"id": f"image_{i}", "prompt": f"Image prompt {i}", "image_url": f"http://example.com/image_{i}.jpg"} for i in range(10)]
+    
+    # Add items
+    for item in image_prompts:
+        add_item_to_container(database_name, container_name, item)
+
+    # Retrieve and display items
+    item_keys = get_all_item_keys(database_name, container_name)
+    for key in item_keys:
+        display_item(database_name, container_name, key)
+
+# UI to Test Image Prompts Function
+if st.button('Test Image Prompts'):
+    test_database_name = st.text_input('Enter Test Database Name')
+    test_container_name = st.text_input('Enter Test Container Name')
+    if test_database_name and test_container_name:
+        test_image_prompts(test_database_name, test_container_name)
+    else:
+        st.error('Please enter the test database and container names.')
+
+
+
+    
 # Streamlit UI
 st.title('Azure Services Integration with Streamlit')