chore: switch to local json storage #8

Merged
sapan merged 1 commits from chore/switch-to-local-json into main 2024-10-13 15:48:08 +00:00
Showing only changes of commit 45ab0d5480 - Show all commits

63
main.py

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