maargdarshak/main.py

102 lines
3.2 KiB
Python
Raw Permalink Normal View History

2024-10-13 04:55:54 +00:00
import streamlit as st
import requests
2024-10-13 14:41:32 +00:00
import datetime
2024-10-13 15:47:29 +00:00
import json
import os
2024-10-13 14:41:32 +00:00
from streamlit_cookies_manager import CookieManager
2024-10-12 19:26:45 +00:00
2024-10-13 14:41:32 +00:00
# Initialize the Cookie Manager (no encryption)
cookies = CookieManager(prefix="ktosiek/streamlit-cookies-manager/")
if not cookies.ready():
st.stop()
2024-10-13 15:47:29 +00:00
# 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()
2024-10-13 14:41:32 +00:00
# Generate a user ID and store it in cookies if not already set
2024-10-13 15:47:29 +00:00
if cookies.get('user_id') is None: # Correct check for user_id in cookies
2024-10-13 14:41:32 +00:00
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}")
2024-10-13 15:47:29 +00:00
# Check if the user exists in the JSON file
if user_id not in user_data:
2024-10-13 14:41:32 +00:00
timestamp = str(datetime.datetime.now().isoformat()) # Unique timestamp for each user
2024-10-13 15:47:29 +00:00
user_data[user_id] = {
2024-10-13 14:41:32 +00:00
"timestamp": timestamp,
"messages": [] # Empty chat history
2024-10-13 15:47:29 +00:00
}
save_user_data(user_data)
st.write("New user detected! Created user entry in the JSON file.")
2024-10-13 14:41:32 +00:00
else:
2024-10-13 15:47:29 +00:00
# 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.")
2024-10-13 04:55:54 +00:00
2024-10-13 13:26:31 +00:00
# Function to send POST request to a local server
2024-10-13 04:55:54 +00:00
def send_post_request(prompt):
2024-10-13 13:26:31 +00:00
url = "http://localhost:19194/" # Update this to your actual server URL
2024-10-13 04:55:54 +00:00
payload = prompt.title()
response = requests.post(url, json=payload)
2024-10-13 14:25:06 +00:00
json = response.json()
if "response" in json:
return json["response"]
else:
return json
2024-10-13 04:55:54 +00:00
2024-10-13 14:41:32 +00:00
st.title('Ask Maargdarshak')
2024-10-13 13:26:31 +00:00
if 'messages' not in st.session_state:
st.session_state.messages = []
2024-10-13 14:41:32 +00:00
# Load chat history from session state
2024-10-13 13:26:31 +00:00
for message in st.session_state.messages:
st.chat_message(message['role']).markdown(message['content'])
2024-10-13 14:41:32 +00:00
# Prompt input
2024-10-13 13:26:31 +00:00
prompt = st.chat_input("Ask your question here")
2024-10-13 04:55:54 +00:00
2024-10-13 13:26:31 +00:00
if prompt:
2024-10-13 14:41:32 +00:00
# Display the user prompt in the chat UI
2024-10-13 13:26:31 +00:00
st.chat_message('user').markdown(prompt)
st.session_state.messages.append({'role': 'user', 'content': prompt})
2024-10-13 04:55:54 +00:00
2024-10-13 14:41:32 +00:00
# Get the bot's response via a POST request
2024-10-13 17:23:26 +00:00
with st.spinner("Fetching data..."):
try:
response = send_post_request(prompt)
except Exception as e:
response = f"An error occurred: {e}"
st.success("Data fetched successfully!")
2024-10-13 14:41:32 +00:00
# Display the bot's response in the chat UI
2024-10-13 13:26:31 +00:00
st.chat_message('bot').markdown(response)
st.session_state.messages.append({'role': 'bot', 'content': response})
2024-10-13 04:55:54 +00:00
2024-10-13 15:47:29 +00:00
# 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)
2024-10-13 14:41:32 +00:00
cookies['history'] = st.session_state.messages
cookies.save()