Update app.py

app.py

@@ -169,81 +169,53 @@ def generate_unique_id():
     return str(uuid.uuid4())
     
 
-def save_to_cosmos_db(container, query, response1, response2):
-    max_retries = 5
-    base_wait_time = 0.1  # 100 ms
-
-    def generate_unique_id():
-        timestamp = datetime.utcnow().strftime('%Y%m%d%H%M%S%f')
-        unique_uuid = str(uuid.uuid4())
-        return f"{timestamp}-{unique_uuid}"
-
-    def generate_short_filename(query):
-        # Create a hash of the query
-        hash_object = hashlib.md5(query.encode())
-        query_hash = hash_object.hexdigest()[:8]
-        
-        # Use the first few words of the query (up to 50 characters)
-        short_query = ' '.join(query.split()[:5])[:50]
-        
-        # Remove any non-alphanumeric characters and replace spaces with underscores
-        safe_query = ''.join(c if c.isalnum() else '_' for c in short_query)
-        
-        # Combine timestamp, hash, and safe query
-        timestamp = datetime.utcnow().strftime('%Y%m%d%H%M%S')
-        return f"{timestamp}_{query_hash}_{safe_query}.md"
-
-    for attempt in range(max_retries):
-        try:
-            if container:
-                base_id = generate_unique_id()
-                record = {
-                    "id": base_id,
-                    "query": query,
-                    "response1": response1,
-                    "response2": response2,
-                    "timestamp": datetime.utcnow().isoformat()
-                }
-                
-                # Check if document with this ID exists
-                try:
-                    existing_doc = container.read_item(item=base_id, partition_key=base_id)
-                    # If it exists, append a random string to the ID
-                    record["id"] = f"{base_id}-{str(uuid.uuid4())[:8]}"
-                except exceptions.CosmosResourceNotFoundError:
-                    # If it doesn't exist, we can use the original ID
-                    pass
-                
-                container.upsert_item(body=record)
-                st.success(f"Record saved successfully with ID: {record['id']}")
-                st.session_state.documents = container.query_items(
-                    query="SELECT * FROM c ORDER BY c._ts DESC",
-                    enable_cross_partition_query=True
-                )
-
-                # Generate a short filename for local storage
-                filename = generate_short_filename(query)
-                with open(filename, 'w', encoding='utf-8') as file:
-                    file.write(f"Query: {query}\n\nResponse: {response1}")
-                st.success(f"File saved locally as: {filename}")
+def save_or_clone_to_cosmos_db(container, query=None, response=None, clone_id=None):
+    try:
+        if not container:
+            st.error("Cosmos DB container is not initialized.")
+            return
 
+        # Generate a new unique ID
+        new_id = str(uuid.uuid4())
+
+        # Create a new document
+        if clone_id:
+            # If clone_id is provided, we're cloning an existing document
+            try:
+                existing_doc = container.read_item(item=clone_id, partition_key=clone_id)
+                new_doc = existing_doc.copy()
+                new_doc['id'] = new_id
+                new_doc['cloned_from'] = clone_id
+                new_doc['cloned_at'] = datetime.utcnow().isoformat()
+            except exceptions.CosmosResourceNotFoundError:
+                st.error(f"Document with ID {clone_id} not found for cloning.")
                 return
-            else:
-                st.error("Cosmos DB container is not initialized.")
-                return
-        except exceptions.CosmosHttpResponseError as e:
-            if e.status_code == 409:  # Conflict error
-                wait_time = (2 ** attempt) * base_wait_time
-                st.warning(f"ID conflict occurred. Retrying in {wait_time:.2f} seconds... (Attempt {attempt + 1})")
-                time.sleep(wait_time)
-            else:
-                st.error(f"Error saving record to Cosmos DB: {e}")
-                return
-        except Exception as e:
-            st.error(f"An unexpected error occurred: {str(e)}")
-            return
+        else:
+            # If no clone_id, we're creating a new document
+            new_doc = {
+                'id': new_id,
+                'query': query,
+                'response': response,
+                'created_at': datetime.utcnow().isoformat()
+            }
+
+        # Insert the new document
+        container.create_item(body=new_doc)
+
+        st.success(f"{'Cloned' if clone_id else 'New'} document saved successfully with ID: {new_id}")
+
+        # Refresh the documents in the session state
+        st.session_state.documents = list(container.query_items(
+            query="SELECT * FROM c ORDER BY c._ts DESC",
+            enable_cross_partition_query=True
+        ))
+
+        return new_id
 
-    st.error("Failed to save record after maximum retries.")
+    except exceptions.CosmosHttpResponseError as e:
+        st.error(f"Error saving to Cosmos DB: {e}")
+    except Exception as e:
+        st.error(f"An unexpected error occurred: {str(e)}")
 
 # Add dropdowns for model and database choices
 def search_glossary(query):
@@ -320,11 +292,20 @@ def search_glossary(query):
 
     # When saving results, pass the container
     try:
-        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])
+        #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])
+
+        save_or_clone_to_cosmos_db(container=st.session_state.cosmos_container, query=query, response=result, clone_id=None)
+        save_or_clone_to_cosmos_db(container=st.session_state.cosmos_container, query=query, response=result2, clone_id=None)
+        save_or_clone_to_cosmos_db(container=st.session_state.cosmos_container, query=query, response=result3, clone_id=None)
+        save_or_clone_to_cosmos_db(container=st.session_state.cosmos_container, query=query, response=response2[0], clone_id=None)
+        save_or_clone_to_cosmos_db(container=st.session_state.cosmos_container, query=query, response=response2[1], clone_id=None)
+
+
+        
     except exceptions.CosmosHttpResponseError as e:
         return False, f"HTTP error occurred: {str(e)} 🚨"
     except Exception as e: