Update app.py

app.py

@@ -17,20 +17,19 @@ cosmos_client = CosmosClient.from_connection_string(COSMOS_CONNECTION_STRING)
 # Initialize Azure Blob Storage Client
 blob_service = BlobServiceClient.from_connection_string(BLOB_STORAGE_CONNECTION_STRING)
 
-
 # Function to Retrieve and Display Cosmos DB Structure
 def display_cosmos_db_structure():
     st.subheader('Azure Cosmos DB Structure')
-    for db_properties in cosmos_client.list_databases():
+    db_properties = next(cosmos_client.list_databases(), None)
+    if db_properties:
         db_name = db_properties['id']
         st.markdown(f"#### Database: {db_name}")
         database_client = cosmos_client.get_database_client(db_name)
-        
-        for container_properties in database_client.list_containers():
+        container_properties = next(database_client.list_containers(), None)
+        if container_properties:
             container_name = container_properties['id']
             st.markdown(f"- **Container**: {container_name}")
             container_client = database_client.get_container_client(container_name)
-            
             items = list(container_client.read_all_items())
             for item in items:
                 item_desc = f"  - Item: `{item['id']}`"
@@ -40,51 +39,12 @@ def display_cosmos_db_structure():
                 else:
                     st.markdown(item_desc)
 
-                    
-# Function to Add or Update PNG Images
-def add_or_update_png_images(database_name, container_name):
-    container = cosmos_client.get_database_client(database_name).get_container_client(container_name)
-    existing_items = list(container.read_all_items())
-    existing_ids = {item['id'] for item in existing_items}
-
-    # Display and delete option for existing items
-    for item in existing_items:
-        if 'file_name' in item and item['file_name'].endswith('.png'):
-            st.markdown(f"Item: `{item['id']}`")
-            st.image(item['file_name'])
-            if st.button(f"🗑️ Delete {item['id']}", key=f"delete_{item['id']}"):
-                container.delete_item(item=item, partition_key=item['id'])
-
-    # Add or update PNG files from directory
-    png_files = glob.glob('*.png')
-    for file_name in png_files:
-        item_id = os.path.splitext(file_name)[0]
-        if item_id not in existing_ids:
-            item_data = {"id": item_id, "file_name": file_name}
-            container.create_item(body=item_data)
-            st.write(f"Added Item: {item_id}")
-        else:
-            st.write(f"Item already exists: {item_id}")
-
-# UI to manage PNG images
-st.title('Manage PNG Images in Cosmos DB')
-selected_db_name = st.text_input('Enter Database Name')
-selected_container_name = st.text_input('Enter Container Name')
-
-if st.button('Manage Images'):
-    if selected_db_name and selected_container_name:
-        add_or_update_png_images(selected_db_name, selected_container_name)
-    else:
-        st.error('Please enter both database and container names.')
-        
-                    
 # Button to Trigger Display of Cosmos DB Structure
 if st.button('Show Cosmos DB Structure'):
     display_cosmos_db_structure()
-        
-# Test Function to Insert PNG Images
-def test_insert_png_images():
-    # Get the first database and container
+
+# Function to Add or Update PNG Images
+def add_or_update_png_images():
     db_properties = next(cosmos_client.list_databases(), None)
     if db_properties:
         db_name = db_properties['id']
@@ -92,95 +52,26 @@ def test_insert_png_images():
         container_properties = next(database_client.list_containers(), None)
         if container_properties:
             container_name = container_properties['id']
+            container_client = database_client.get_container_client(container_name)
+            existing_items = list(container_client.read_all_items())
+            existing_ids = {item['id'] for item in existing_items}
 
-            # Insert PNG files
+            # Add or update PNG files from directory
             png_files = glob.glob('*.png')
             for file_name in png_files:
-                item_id = os.path.splitext(file_name)[0]  # Use file name without extension as ID
+                item_id = os.path.splitext(file_name)[0]
                 item_data = {"id": item_id, "file_name": file_name}
-                add_or_update_item(database_client, container_name, item_id, item_data)
-                st.write(f"Inserted: {file_name}")
-
-            # Displaying Images
-            st.subheader("Displaying Images in Container")
-            items = list(database_client.get_container_client(container_name).read_all_items())
-            for item in items:
-                if 'file_name' in item and item['file_name'].endswith('.png'):
-                    st.image(item['file_name'], caption=item['file_name'])
-        else:
-            st.error("No container found in the database.")
-    else:
-        st.error("No database found.")
-
-# UI to Test Image Insertion Function
-if st.button('Test Insert PNG Images'):
-    test_insert_png_images()
-
-
-
-# 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')
-
-# Azure Cosmos DB - CRUD Operations
-st.subheader('Azure Cosmos DB - CRUD Operations')
-cosmos_db = st.text_input('Database Name')
-cosmos_container = st.text_input('Container Name')
-item_id = st.text_input("Item ID (for Read, Update, Delete)")
-item_data = st.text_area("Item Data (JSON format, for Create and Update)")
-
-if st.button('Create Item in Cosmos DB'):
-    container = cosmos_client.get_database_client(cosmos_db).get_container_client(cosmos_container)
-    container.create_item(item_data)
+                if item_id not in existing_ids:
+                    container_client.create_item(body=item_data)
+                    st.write(f"Added Item: {item_id}")
+                else:
+                    if st.button(f"🗑️ Delete {item_id}", key=f"delete_{item_id}"):
+                        container_client.delete_item(item=item_data, partition_key=item_id)
+                    st.write(f"Item already exists: {item_id}")
 
-if st.button('Read Item from Cosmos DB'):
-    container = cosmos_client.get_database_client(cosmos_db).get_container_client(cosmos_container)
-    item = container.read_item(item_id, partition_key=PartitionKey(item_id))
-    st.json(item)
+# UI to Add or Update PNG Images
+if st.button('Manage PNG Images'):
+    add_or_update_png_images()
 
 # Azure Blob Storage - Upload/Download
 st.subheader('Azure Blob Storage - Upload/Download')