cambio a atlas

app.py

@@ -5,7 +5,7 @@ import gradio as gr
 import pandas as pd
 from datetime import datetime, timedelta, timezone
 #import torch
-from config import groq_token, groq_model, QUESTION_PROMPT, init_google_sheets_client, groq_model, default_model_name, user_names, google_sheets_name
+from config import groq_token, groq_model, QUESTION_PROMPT, init_google_sheets_client, groq_model, default_model_name, user_names, google_sheets_name, AtlasClient
 
 #from config import  hugging_face_token, replicate_token
 #import replicate
@@ -74,64 +74,18 @@ model_history = []
 # Save all_answers to Google Drive
 FILE_ID = '1PwEiBxpHo0jRc6T1HixyC99UnP9iawbr'
 
+#init atlas client
+atlas_client = AtlasClient("chatbot_test")
 
-def save_answers(all_answers):
-    # Credenciales de la cuenta de servicio (reemplaza con tus credenciales)
-    SCOPES = ["https://spreadsheets.google.com/feeds", "https://www.googleapis.com/auth/drive"]
-    SERVICE_ACCOUNT_FILE = 'polar-land-440713-c4-bbc8d89804d8.json'
-    # Autentificación
-    credentials = service_account.Credentials.from_service_account_file(SERVICE_ACCOUNT_FILE, scopes=SCOPES)
-    service = build('drive', 'v3', credentials=credentials)
-
-    # Obtener el archivo existente
-    file = service.files().get(fileId=FILE_ID).execute()
-
-    # Download the content using get_media instead of export_media
-    request = service.files().get_media(fileId=FILE_ID)
-    fh = io.BytesIO()
-    downloader = MediaIoBaseDownload(fh, request)
-    done = False
-    while done is False:
-        status, done = downloader.next_chunk()
-        print("Download %d%%." % int(status.progress() * 100))
-
-    # Cargar el contenido JSON
-    content = fh.getvalue()
-    if content:
-        existing_data = json.loads(content)
-    else:
-        existing_data = {}
-
-    # Convert sets to lists before serialization if they exist
-    def convert_sets_to_lists(obj):
-        if isinstance(obj, set):
-            return list(obj)
-        if isinstance(obj, dict):
-            return {k: convert_sets_to_lists(v) for k, v in obj.items()}
-        if isinstance(obj, list):
-            return [convert_sets_to_lists(item) for item in obj]
-        return obj
-
-    existing_data = convert_sets_to_lists(existing_data)
-
-    # Agregar los nuevos datos al arreglo
-    if 'data' in existing_data:
-        existing_data['data'].append(all_answers)
-    else:
-        existing_data['data'] = [all_answers]
-    # Convertir los datos a formato JSON
-    new_content = json.dumps(existing_data)
-
-    # Create a temporary file to store the JSON data
-    with open('temp_data.json', 'w') as temp_file: 
-        temp_file.write(new_content)
-
-    media = MediaFileUpload('temp_data.json', mimetype='application/json')  
-    file = service.files().update(fileId=FILE_ID, 
-                                   media_body=media, 
-                                   fields='id').execute()
-    print('Archivo actualizado correctamente: %s' % file.get('id'))
 
+def save_answers(all_answers, user_name):
+    timestamp = datetime.now(timezone.utc) - timedelta(hours=3)  # Adjust to GMT-3
+    timestamp_str = timestamp.strftime("%Y-%m-%d %H:%M:%S")
+    data = {}
+    data["timestamp"] = timestamp_str
+    data["answers"] = all_answers
+    data['author'] = user_name
+    atlas_client.add(data, "dialogs")
 
 
     #Function to save comment and score
@@ -175,9 +129,9 @@ def save_comment_score(score, comment, story_name, user_name, system_prompt, mod
     user_sheet.append_row([timestamp_str, user_name, model_name, system_prompt, story_name, full_chat_history, score, comment])
 
 
-    # Save all answers to Google Drive as a JSON file
+    # Save all answers to AtlaS collection
     print(f"all answers...\n{all_answers}")
-    save_answers(all_answers)
+    save_answers(all_answers, user_name)
 
     #Append data and render the data table
     df = pd.DataFrame(data, columns=["Timestamp", "User Name", "Model Name", "System Prompt", "Story Name", "Chat History", "Score", "Comment"])

config.py

@@ -1,11 +1,28 @@
 import os
 import gspread
 from oauth2client.service_account import ServiceAccountCredentials
+from pymongo import MongoClient
 
 # Read the authentication token from the environment variable
 hugging_face_token = os.getenv("HUGGING_FACE_TOKEN")
 replicate_token = os.getenv("REPLICATE_TOKEN")
 groq_token = os.getenv("GROQ_TOKEN")
+atlas_token = os.getenv("ATLAS_TOKEN")
+
+#atlas configuration
+class AtlasClient:
+    def __init__(self, dbname):
+        self.mongodb_client = MongoClient(atlas_token)
+        self.database = self.mongodb_client[dbname]
+
+    # A quick way to test if we can connect to Atlas instance
+    def ping(self):
+        self.mongodb_client.admin.command("ping")
+
+
+    def add(self, item, collection_name):
+        collection = self.database[collection_name]
+        collection.insert_one(item)
 
 # Google Sheets configuration
 def init_google_sheets_client():