app is now logging the question, results, vote and feedback in a mongodb collection.

app.py

@@ -3,8 +3,10 @@ import random
 from helpers import query_you_com, query_tavily, query_perplexity, query_brave
 from provider_info import search_providers
 from mongod_db import MongoDBHandler
+from bson.objectid import ObjectId
 # from swarms.utils.loguru_logger import logger
 import time
+import uuid
 
 # mongo = MongoDBHandler()
 
@@ -38,6 +40,7 @@ Join us in our mission to advance search technology through open collaboration a
 
 # Define the function to process the question
 def ProcessQuestion(question):
+    document_id = None
     # Randomly select two out of the four functions
     functions = [query_you_com, query_tavily, query_perplexity, query_brave]
     selected_functions = random.sample(functions, 2)
@@ -51,21 +54,42 @@ def ProcessQuestion(question):
     
     try: 
         # logger.info(f"Logging question: {question}")
-        mongo.add(
+        document_id = mongo.add(
             {
                 "question": question,
                 "answer_a": answer_a,
                 "answer_b": answer_b,
                 "selected_functions": [f.__name__ for f in selected_functions],
                 "query_time": time.time(),
+                "session_id": st.session_state.session_id,
             }
         )
         # logger.info("Successfully logged into mongodb")
     except Exception as e:
         # logger.error(f"Error logging into mongodb: {e}")
-        print("Error logging into mongodb: {e}")
+        print(f"Error logging into mongodb: {e}")
 
-    return answer_a, answer_b, selected_functions
+    return answer_a, answer_b, selected_functions, document_id
+
+def UpdateVote(session_id, vote):
+    mongo = MongoDBHandler()
+    try:
+        mongo.update(
+            {"session_id": session_id},
+            {"$set": {"vote": vote}}
+        )
+    except Exception as e:
+        print(f"Error updating vote in mongodb: {e}")
+
+def UpdateFeedback(session_id, feedback):
+    mongo = MongoDBHandler()
+    try:
+        mongo.update(
+            {"session_id": session_id},
+            {"$set": {"feedback": feedback}}
+        )
+    except Exception as e:
+        print(f"Error updating feedback in mongodb: {e}")
 
 
 # Initialize session state if not already done
@@ -77,7 +101,10 @@ default_values = {
     "source_a": "",
     "source_b": "",
     "winner": "",
-    "selected_button": ""
+    "selected_button": "",
+    "document_id": "",
+    "feedback": "",
+    "session_id": str(uuid.uuid4())
 }
 
 for key, value in default_values.items():
@@ -90,17 +117,19 @@ st.title("Search Engine Agent Comparison")
 def on_submit():
     question = st.session_state["question_input"]
     if question:
-        answer_a, answer_b, selected_functions = ProcessQuestion(question)
+        answer_a, answer_b, selected_functions, document_id = ProcessQuestion(question)
         st.session_state["question"] = question
         st.session_state["answer_a"] = answer_a
         st.session_state["answer_b"] = answer_b
         st.session_state["source_a"] = selected_functions[0].__name__
         st.session_state["source_b"] = selected_functions[1].__name__
         st.session_state["state"] = "arena_review"
+        st.session_state["document_id"] = document_id
 
 def handle_vote(vote):
     st.session_state["winner"] = vote
     st.session_state["state"] = "arena_results"
+    UpdateVote(st.session_state["session_id"], vote)
 
 def get_provider_info(provider_function_name):
     provider_name_map = {
@@ -172,8 +201,12 @@ def render_results_state():
         st.markdown(f"<div style='padding: 10px; border: 1px solid #ddd;'>{st.session_state['answer_b']}</div>", unsafe_allow_html=True)
     
     st.write("### Feedback")
-    st.text_area("Please provide feedback on why you chose the winner:", key="feedback")
+    feedback = st.text_area("Please provide feedback on why you chose the winner:")
+    if feedback:
+        UpdateFeedback(st.session_state["session_id"], feedback)
+    
     st.write("### About the search providers:")
+    
     col1, col2 = st.columns(2)
     with col1:
         st.write(f"**Website:** [{provider_info_a['website']}]({provider_info_a['website']})")

mongod_db.py

@@ -3,11 +3,10 @@ from pymongo import MongoClient
 from pymongo.errors import PyMongoError
 import os
 from dotenv import load_dotenv
-
+import unittest
 
 load_dotenv()
 
-
 def mongo_db_uri():
     return os.getenv("MONGODB_URI")
 
@@ -15,9 +14,9 @@ def mongo_db_uri():
 class MongoDBHandler:
     def __init__(
         self,
-        uri: str = mongo_db_uri(),
-        db_name: str = "search-arena",
-        collection_name: str = "search-arena-usage",
+        uri: str = os.getenv("MONGODB_URI") or mongo_db_uri(),
+        db_name: str = os.getenv("MONGODB_DB") or "search-arena",
+        collection_name: str = os.getenv("MONGODB_COLLECTION") or "search-arena-usage",
         *args,
         **kwargs
     ) -> None:
@@ -48,7 +47,23 @@ class MongoDBHandler:
         except PyMongoError as e:
             print(f"Error: {e}")
             raise
+        
+    def update(self, query: Dict[str, Any], update: Dict[str, Any]) -> None:
+        """
+        Updates documents in the MongoDB collection based on the provided query and update.
 
+        Args:
+            query (Dict[str, Any]): The query to filter documents.
+            update (Dict[str, Any]): The update to apply to the documents.
+
+        Raises:
+            PyMongoError: If an error occurs while updating the documents.
+        """
+        try:
+            self.collection.update_many(query, update)
+        except PyMongoError as e:
+            print(f"Error: {e}")
+            raise
     def query(self, query: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
         Queries documents from the MongoDB collection based on the provided query.
@@ -66,4 +81,4 @@ class MongoDBHandler:
             return list(self.collection.find(query))
         except PyMongoError as e:
             print(f"Error: {e}")
-            raise
+            raise

test_mongod_db.py

@@ -0,0 +1,32 @@
+import unittest
+from unittest.mock import MagicMock
+from mongod_db import MongoDBHandler
+
+class MongoDBHandlerTests(unittest.TestCase):
+    def setUp(self):
+        self.handler = MongoDBHandler(collection_name='search-arena-test')
+
+    def tearDown(self):
+        self.handler.collection.delete_many({})
+
+    def test_add_document(self):
+        document = {"name": "John Doe", "age": 30}
+        self.handler.add(document)
+        result = self.handler.collection.find_one(document)
+        self.assertIsNotNone(result)
+        self.assertEqual(result["name"], "John Doe")
+        self.assertEqual(result["age"], 30)
+
+    def test_query_documents(self):
+        document1 = {"name": "John Doe", "age": 30}
+        document2 = {"name": "Jane Smith", "age": 25}
+        self.handler.add(document1)
+        self.handler.add(document2)
+        query = {"age": {"$gt": 28}}
+        results = self.handler.query(query)
+        self.assertEqual(len(results), 1)
+        self.assertEqual(results[0]["name"], "John Doe")
+        self.assertEqual(results[0]["age"], 30)
+
+if __name__ == "__main__":
+    unittest.main()