Update app.py

app.py

@@ -177,24 +177,44 @@ def format_with_line_numbers(text):
     formatted_text = '\n'.join(f"{i+1}: {line}" for i, line in enumerate(lines))
     return formatted_text
 
-# 💾 Save responses to Cosmos DB
+
+def generate_unique_id():
+    return str(uuid.uuid4())
+
+def get_databases(client):
+    return [db['id'] for db in client.list_databases()]
+
+def get_containers(database):
+    return [container['id'] for container in database.list_containers()]
+
+def get_documents(container, limit=None):
+    query = "SELECT * FROM c ORDER BY c._ts DESC"
+    items = list(container.query_items(query=query, enable_cross_partition_query=True, max_item_count=limit))
+    return items
+
 def save_to_cosmos_db(query, response1, response2):
-    # 💾 Saving responses to Cosmos DB—because even the cosmos needs backups! 💿
-    cosmos_container = st.session_state.get("cosmos_container")
-    if cosmos_container:
-        record = {
-            "id": generate_unique_id(),
-            "query": query,
-            "response1": json.loads(response1),
-            "response2": json.loads(response2)
-        }
-        try:
-            cosmos_container.create_item(body=record)
-            st.success(f"Record saved successfully with ID: {record['id']}")
-        except exceptions.CosmosHttpResponseError as e:
-            st.error(f"Error saving record to Cosmos DB: {e}")
-    else:
-        st.error("Cosmos DB is not initialized.")
+    try:
+        cosmos_container = st.session_state.get("cosmos_container")
+        if cosmos_container:
+            record = {
+                "id": generate_unique_id(),
+                "query": query,
+                "response1": response1,
+                "response2": response2
+            }
+            try:
+                cosmos_container.create_item(body=record)
+                st.success(f"Record saved successfully with ID: {record['id']}")
+                # Refresh the documents display
+                st.session_state.documents = get_documents(cosmos_container)
+            except exceptions.CosmosHttpResponseError as e:
+                st.error(f"Error saving record to Cosmos DB: {e}")
+        else:
+            st.error("Cosmos DB is not initialized.")
+    except Exception as e:
+        st.error(f"An unexpected error occurred: {str(e)}")
+
+
 
 
 # Add dropdowns for model and database choices