-
Notifications
You must be signed in to change notification settings - Fork 1
/
database.py
42 lines (35 loc) · 1.24 KB
/
database.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
from pymongo import MongoClient
import os
password = os.environ.get('password')
uri = f"mongodb+srv://cidar:{password}@cluster0.yuhfxcj.mongodb.net/?retryWrites=true&w=majority"
# Create a new client and connect to the server
client = MongoClient(uri)
db = client['testDB']
collection = db['testCol']
def save_response(response_data):
"""Saves a response to the MongoDB collection."""
try:
# Insert the response data into the collection
print("right before insertion")
collection.insert_one(response_data)
print("right after insertion")
except Exception as e:
print(f"An error occurred: {e}")
exit(0)
def read_responses(filter_query=None):
"""Reads responses from the MongoDB collection based on an optional filter."""
try:
if filter_query is None:
filter_query = {} # An empty query will return all documents
ranks = {
'CIDAR': [0, 0, 0],
'ALPAGASUS': [0, 0, 0],
'CHAT': [0, 0, 0],
}
for doc in collection.find(filter_query):
for model in ranks:
ranks[model][doc[model] - 1] += 1
return ranks
except Exception as e:
print(f"An error occurred: {e}")
return []