Update app.py

app.py

@@ -192,10 +192,9 @@ def get_documents(container, limit=None):
     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):
+def save_to_cosmos_db(container, query, response1, response2):
     try:
-        cosmos_container = st.session_state.get("cosmos_container")
-        if cosmos_container:
+        if container:
             record = {
                 "id": generate_unique_id(),
                 "query": query,
@@ -203,20 +202,17 @@ def save_to_cosmos_db(query, response1, response2):
                 "response2": response2
             }
             try:
-                cosmos_container.create_item(body=record)
+                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)
+                st.session_state.documents = get_documents(container)
             except exceptions.CosmosHttpResponseError as e:
                 st.error(f"Error saving record to Cosmos DB: {e}")
         else:
-            st.error("Cosmos DB is not initialized.")
+            st.error("Cosmos DB container is not initialized.")
     except Exception as e:
         st.error(f"An unexpected error occurred: {str(e)}")
 
-
-
-
 # Add dropdowns for model and database choices
 def search_glossary(query):
     st.markdown(f"### 🔍 Search Glossary for: `{query}`")
@@ -253,7 +249,6 @@ def search_glossary(query):
         )
         st.markdown(result)
         st.code(result, language="python", line_numbers=True)
-        save_to_cosmos_db(query, result, result) # Save both responses to Cosmos DB
         
         # 🔍 ArXiv RAG researcher expert ~-<>-~ Paper Summary & Ask LLM - api_name: /ask_llm
         result2 = client.predict(
@@ -264,7 +259,6 @@ def search_glossary(query):
         )
         st.markdown(result2)
         st.code(result2, language="python", line_numbers=True)
-        save_to_cosmos_db(query, result2, result2) # Save both responses to Cosmos DB
 
         # 🔍 ArXiv RAG researcher expert ~-<>-~ Paper Summary & Ask LLM - api_name: /ask_llm
         result3 = client.predict(
@@ -275,7 +269,6 @@ def search_glossary(query):
         )
         st.markdown(result3)
         st.code(result3, language="python", line_numbers=True)
-        save_to_cosmos_db(query, result3, result3) # Save both responses to Cosmos DB
 
         
         # 🔍 ArXiv RAG researcher expert ~-<>-~ Paper Summary & Ask LLM - api_name: /update_with_rag_md
@@ -289,12 +282,16 @@ def search_glossary(query):
 
         st.markdown(response2[0])
         st.code(response2[0], language="python", line_numbers=True, wrap_lines=True)  
-        save_to_cosmos_db(query, response2[0], response2[0]) # Save both responses to Cosmos DB
         
         st.markdown(response2[1])
         st.code(response2[1], language="python", line_numbers=True, wrap_lines=True)
-        save_to_cosmos_db(query, response2[1], response2[1]) # Save both responses to Cosmos DB
-        
+
+        # When saving results, pass the container
+        save_to_cosmos_db(st.session_state.cosmos_container, query, result, result)
+        save_to_cosmos_db(st.session_state.cosmos_container, query, result2, result2)
+        save_to_cosmos_db(st.session_state.cosmos_container, query, result3, result3)
+        save_to_cosmos_db(st.session_state.cosmos_container, query, response2[0], response2[0])
+        save_to_cosmos_db(st.session_state.cosmos_container, query, response2[1], response2[1])
         
         # Aggregate hyperlinks and show with emojis
         hyperlinks = extract_hyperlinks([response1, response2])