Update app.py

app.py

@@ -1,329 +1,170 @@
 import openai
+import streamlit as st
 from openai import OpenAI
 import io
 import time
 import os
 from dotenv import load_dotenv
 
-# Load environment variables
+# Initialize the OpenAI client with your API key
+# Load environment variables from the .env file
 load_dotenv()
 
-# Securely fetch environment variables for Hugging Face Spaces
-def get_env_variable(name: str) -> str:
-    value = os.environ.get(name)
-    if value is None:
-        raise ValueError(f"Environment variable '{name}' is not set.")
-    return value
-
-# Initialize OpenAI client
-OPENAI_API_KEY = get_env_variable("OPENAI_API_KEY")
-VECTOR_STORE_ID = get_env_variable("VECTOR_STORE_ID")
-ASSISTANT_ID = get_env_variable("ASSISTANT_ID")
+# Get the OpenAI API key
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 
 client = OpenAI(api_key=OPENAI_API_KEY)
+vector_store_id = "vs_sHLBqA88ng66RDZYMgaBmDuX"  # Vector Store ID to use
 
-# Utility Functions
-def ensure_single_thread_id(session_state):
-    if "thread_id" not in session_state:
-        thread = client.beta.threads.create()
-        session_state.thread_id = thread.id
-    return session_state.thread_id
+# 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 = "asst_UeIHpkpVejAgrIz3RnisxSLm"  # 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:
-        return client.beta.threads.messages.create(
+        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):
-            time.sleep(1)
-            return safe_message_send(prompt, thread_id)
+            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):
-    safe_message_send(prompt, thread_id)
-    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
-    except Exception as e:
-        print(f"Streaming error: {str(e)}")
-
+    # 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
-        client.vector_stores.file_batches.upload_and_poll(
-            vector_store_id=VECTOR_STORE_ID,
+        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]
         )
-        print(f"File '{uploaded_file.name}' uploaded successfully.")
+
+
+        st.success(f"File '{uploaded_file.name}' processed and added to vector store. Status: {file_batch.status}")
     except Exception as e:
-        print(f"Upload failed: {str(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 = client.vector_stores.files.list(vector_store_id=VECTOR_STORE_ID)
-        for file in all_files.data:
-            print(f"File ID: {file.id}")
+        all_files = list(client.beta.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:
-        print(f"Failed to list files: {str(e)}")
-
+        st.error(f"Failed to list files: {str(e)}")
+        return {}
 
-def delete_file_from_vector_store(file_id):
+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,
+        client.beta.vector_stores.files.delete(
+            vector_store_id=vector_store_id,
             file_id=file_id
         )
-        print(f"File '{file_id}' deleted successfully.")
+        st.success(f"File with ID '{file_id}' deleted from vector store '{vector_store_id}'.")
     except Exception as e:
-        print(f"Delete failed: {str(e)}")
-
-# Simple CLI-based Interface for Testing
-def main():
-    import sys
-
-    session_state = {}
-
-    print("Welcome to the CLI Chatbot!")
-    while True:
-        print("\nOptions:")
-        print("1. Upload file")
-        print("2. List files")
-        print("3. Delete file")
-        print("4. Chat")
-        print("5. Exit")
-        choice = input("Choose an option: ")
+        st.error(f"Failed to delete file. File id is not Found.")
 
-        if choice == "1":
-            file_path = input("Enter path to PDF or DOCX file: ")
-            try:
-                with open(file_path, "rb") as f:
-                    class UploadedFile:
-                        def __init__(self, name, content):
-                            self.name = name
-                            self.content = content
+# 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)
 
-                        def getvalue(self):
-                            return self.content
+# Streamlit interface setup
+st.title("💬Chatbot")
+st.caption("🚀 A Streamlit Custom Chatbot")
 
-                    uploaded_file = UploadedFile(os.path.basename(file_path), f.read())
-                    upload_and_add_to_vector_store(uploaded_file)
-            except Exception as e:
-                print(f"File upload error: {str(e)}")
 
-        elif choice == "2":
-            list_all_files_in_vector_store()
-
-        elif choice == "3":
-            file_id = input("Enter File ID to delete: ")
-            delete_file_from_vector_store(file_id)
-
-        elif choice == "4":
-            if "messages" not in session_state:
-                session_state["messages"] = []
-            user_prompt = input("You: ")
-            session_state["messages"].append({"role": "user", "content": user_prompt})
-            thread_id = ensure_single_thread_id(session_state)
-            print("Assistant:", end=" ")
-            for part in stream_generator(user_prompt, thread_id):
-                print(part, end="", flush=True)
-            print()
-
-        elif choice == "5":
-            sys.exit()
+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:
-            print("Invalid choice, please try again.")
-
-
-if __name__ == "__main__":
-    main()
-
+            st.error("Please upload a file to process.")
 
-# import openai
-# import streamlit as st
-# from openai import OpenAI
-# import io
-# import time
-# import os
-# from dotenv import load_dotenv
-
-# # OpenAI keys
-# OPENAI_API_KEY = st.secrets["OPENAI_API_KEY"]
-# # VECTOR_STORE_ID = st.secrets["VECTOR_STORE_ID"]
-# # VECTOR_STORE_ID = os.getenv("VECTOR_STORE_ID")
-
-# client = OpenAI(api_key=OPENAI_API_KEY)
-# # vector_store_id =VECTOR_STORE_ID  # Vector Store ID to use
-# vector_store_id = "vs_sHLBqA88ng66RDZYMgaBmDuX"
-# st.write(vector_store_id)
-# # 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 = "asst_UeIHpkpVejAgrIz3RnisxSLm"  # 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.beta.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.beta.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 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})
+    # 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})
+        with st.chat_message("assistant"):
+            response = st.write_stream(stream_generator(prompt, thread_id))
+        st.session_state.messages.append({"role": "assistant", "content": response})