chore: switch to local json storage

This commit is contained in:
Sapan Shah 2024-10-13 21:17:29 +05:30
parent 1b842ff917
commit 45ab0d5480

63
main.py

@ -1,7 +1,8 @@
import streamlit as st
import requests
import datetime
from pymongo import MongoClient
import json
import os
from streamlit_cookies_manager import CookieManager
# Inject the scroll control HTML/JS
@ -14,35 +15,47 @@ cookies = CookieManager(prefix="ktosiek/streamlit-cookies-manager/")
if not cookies.ready():
st.stop()
# Connect to MongoDB
client = MongoClient("mongodb://localhost:27017/") # Adjust your MongoDB connection URI
db = client["chatbot_db"] # Database name
users_collection = db["users"] # Collection to store user data
# Define the file path for storing user data
user_data_file = "user_data.json"
# Helper function to load user data from JSON file
def load_user_data():
if os.path.exists(user_data_file):
with open(user_data_file, 'r') as file:
return json.load(file)
else:
save_user_data({})
return {}
# Helper function to save user data to JSON file
def save_user_data(data):
with open(user_data_file, 'w') as file:
json.dump(data, file, indent=4)
# Load existing user data
user_data = load_user_data()
# Generate a user ID and store it in cookies if not already set
if 'user_id' not in cookies:
if cookies.get('user_id') is None: # Correct check for user_id in cookies
cookies['user_id'] = str(datetime.datetime.now().timestamp()) # Use timestamp as a unique ID
cookies.save()
user_id = cookies.get('user_id')
# st.write(f"Current User ID: {user_id}")
# Check if the user exists in the database
user_data = users_collection.find_one({"user_id": user_id})
# If user doesn't exist, create a new entry with the current timestamp and empty chat history
if not user_data:
# Check if the user exists in the JSON file
if user_id not in user_data:
timestamp = str(datetime.datetime.now().isoformat()) # Unique timestamp for each user
users_collection.insert_one({
"user_id": user_id,
user_data[user_id] = {
"timestamp": timestamp,
"messages": [] # Empty chat history
})
st.write("New user detected! Created user entry in the database.")
}
save_user_data(user_data)
st.write("New user detected! Created user entry in the JSON file.")
else:
# If user exists, load the chat history from the database
st.session_state.messages = user_data["messages"]
# st.write("Returning user detected! Loaded chat history from the database.")
# If user exists, load the chat history from the JSON file
st.session_state.messages = user_data[user_id]["messages"]
# st.write("Returning user detected! Loaded chat history from the JSON file.")
# Function to send POST request to a local server
def send_post_request(prompt):
@ -121,16 +134,10 @@ if prompt:
st.chat_message('bot').markdown(response)
st.session_state.messages.append({'role': 'bot', 'content': response})
# Update the user's chat history in the database
users_collection.update_one(
{"user_id": user_id},
{"$push": {"messages": {'role': 'user', 'content': prompt}}},
)
users_collection.update_one(
{"user_id": user_id},
{"$push": {"messages": {'role': 'bot', 'content': response}}},
)
# Update the user's chat history in the JSON file
user_data[user_id]["messages"] = st.session_state.messages
save_user_data(user_data)
# Save updated history to cookies (optional, MongoDB handles persistence)
# Save updated history to cookies (optional)
cookies['history'] = st.session_state.messages
cookies.save()