Merge pull request 'chore: switch to local json storage' (#8) from chore/switch-to-local-json into main
Reviewed-on: #8
This commit is contained in:
commit
2c15c026c8
63
main.py
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()
|
||||
|
Loading…
Reference in New Issue
Block a user