SixthfactorBot / app.py
Deepakkori45's picture
Update app.py
64cd785 verified
import openai
import streamlit as st
from openai import OpenAI
import io
import time
import os
from dotenv import load_dotenv
# Initialize the OpenAI client with your API key
# Load environment variables from the .env file
load_dotenv()
# Get the OpenAI API key
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
client = OpenAI(api_key=OPENAI_API_KEY)
vector_store_id = os.getenv("VECTOR_STORE_ID") # Vector Store ID to use
# all_files = list(client.beta.vector_stores.files.list(vector_store_id))
# for file in all_files:
# # print(file)
# file_id = file.id
# st.write(file_id)
# Set the assistant ID
assistant_id = os.getenv("ASSISTANT_ID") # Replace with your own assistant ID
def ensure_single_thread_id():
if "thread_id" not in st.session_state:
thread = client.beta.threads.create()
st.session_state.thread_id = thread.id
return st.session_state.thread_id
def safe_message_send(prompt, thread_id):
try:
message = client.beta.threads.messages.create(
thread_id=thread_id,
role="user",
content=prompt
)
return message
except Exception as e:
if "active" in str(e):
print("Waiting for the current run to finish...")
time.sleep(1) # wait a bit before retrying
return safe_message_send(prompt, thread_id) # retry sending the message
else:
raise e
def stream_generator(prompt, thread_id):
# print(f'First time thread in the function {thread_id}')
message = safe_message_send(prompt, thread_id) # use the new safe send function
with st.spinner("Wait... Generating response..."):
try:
stream = client.beta.threads.runs.create(
thread_id=thread_id,
assistant_id=assistant_id,
stream=True
)
for event in stream:
if event.data.object == "thread.message.delta":
for content in event.data.delta.content:
if content.type == 'text':
yield content.text.value
time.sleep(0.01)
elif event.data.object == "thread.run.stop":
break # Break if the run stops
except Exception as e:
print(f"Error during streaming: {str(e)}")
def upload_and_add_to_vector_store(uploaded_file):
"""Upload a file to OpenAI and add it to the specified vector store."""
try:
# Convert the uploaded file to a BytesIO stream for uploading
file_stream = io.BytesIO(uploaded_file.getvalue())
file_stream.name = uploaded_file.name # Preserve the file name
# Upload the file to the vector store
file_batch = client.beta.vector_stores.file_batches.upload_and_poll(
vector_store_id=vector_store_id,
files=[file_stream]
)
st.success(f"File '{uploaded_file.name}' processed and added to vector store. Status: {file_batch.status}")
except Exception as e:
st.error(f"Failed to process file: {str(e)}")
def list_all_files_in_vector_store():
"""List all files in the specified vector store."""
try:
all_files = list(client.vector_stores.files.list(vector_store_id=vector_store_id))
# st.write(all_files)
for file in all_files:
file_id = file.id
st.write(file_id)
except Exception as e:
st.error(f"Failed to list files: {str(e)}")
return {}
def delete_file_from_vector_store(vector_store_id, file_id):
"""Delete a file from the specified vector store."""
try:
client.vector_stores.files.delete(
vector_store_id=vector_store_id,
file_id=file_id
)
st.success(f"File with ID '{file_id}' deleted from vector store '{vector_store_id}'.")
except Exception as e:
st.error(f"Failed to delete file. File id is not Found.")
# Interface to delete files from vector store
st.sidebar.subheader("Delete File from Vector Store")
file_id_to_delete = st.sidebar.text_input("Enter File ID to Delete", "")
if st.sidebar.button("Delete File"):
delete_file_from_vector_store(vector_store_id, file_id_to_delete)
# Streamlit interface setup
st.title("💬Chatbot")
st.caption("🚀 A Streamlit Custom Chatbot")
with st.sidebar:
st.write("Upload PDF File")
uploaded_file = st.file_uploader("Choose a file", type=['pdf', 'docx'], key='file_uploader')
if st.button('Upload File', key='process_file'):
if uploaded_file is not None:
upload_and_add_to_vector_store(uploaded_file)
st.success("File successfully uploaded and processed.")
else:
st.error("Please upload a file to process.")
# List all uploaded files
st.write("### Uploaded Files")
if 'uploaded_files' in st.session_state and st.session_state.uploaded_files:
for file_name, file_id in st.session_state.uploaded_files.items():
st.write(f"{file_name}: {file_id}")
# List all files in the vector store
st.write("## All Files in Vector Store")
all_files = list_all_files_in_vector_store()
# Initialize session state for chat
st.session_state.start_chat = True
if 'start_chat' not in st.session_state:
st.session_state.start_chat = False
# Main chat interface
if st.session_state.start_chat:
if "messages" not in st.session_state:
st.session_state.messages = []
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.markdown(message["content"])
prompt = st.chat_input("Enter your message")
if prompt:
thread_id = ensure_single_thread_id()
with st.chat_message("user"):
st.markdown(prompt)
st.session_state.messages.append({"role": "user", "content": prompt})
with st.chat_message("assistant"):
response = st.write_stream(stream_generator(prompt, thread_id))
st.session_state.messages.append({"role": "assistant", "content": response})