rag update

app.py

@@ -1,9 +1,43 @@
 import streamlit as st
-from openai import OpenAI
+import os
+from typing_extensions import override
+from openai import AssistantEventHandler, OpenAI
+
+def list_files_in_directory(directory):
+    file_paths = []
+    for root, dirs, files in os.walk(directory):
+        for file in files:
+            file_path = os.path.join(root, file)
+            relative_path = os.path.relpath(file_path, directory)
+            file_paths.append(relative_path)
+    return file_paths
+
+class EventHandler(AssistantEventHandler):
+    @override
+    def on_text_created(self, text) -> None:
+        print(f"\nassistant > ", end="", flush=True)
+
+    @override
+    def on_tool_call_created(self, tool_call):
+        print(f"\nassistant > {tool_call.type}\n", flush=True)
+
+    @override
+    def on_message_done(self, message) -> None:
+        # print a citation to the file searched
+        message_content = message.content[0].text
+        annotations = message_content.annotations
+        citations = []
+        for index, annotation in enumerate(annotations):
+            message_content.value = message_content.value.replace(
+                annotation.text, f"[{index}]"
+            )
+            if file_citation := getattr(annotation, "file_citation", None):
+                cited_file = client.files.retrieve(file_citation.file_id)
+                citations.append(f"[{index}] {cited_file.filename}")
+
+        print(message_content.value)
+        print("\n".join(citations))
 
-# Set the assistant ID (replace with your own assistant ID)
-assistant_id = "asst_7Tpj3X81c2Z51nZ12GBWkT45"
-vectorstore_id = "vs_U7GH1VO9qLyIIJ21wE8Q7bVB"
 st.title("🤖 Chatbot")
 st.caption("🚀 A streamlit chatbot by Custom GPT AI Automation Academy")
 
@@ -13,17 +47,40 @@ api_key = st.text_input("Enter your OpenAI API key", type="password")
 # Initialize the OpenAI client with the provided API key
 if api_key:
     client = OpenAI(api_key=api_key)
+
+    assistant = client.beta.assistants.create(
+    name="Methylation Classification Assistant",
+    instructions="You are a helpful molecular pathologist, expert in CNS tumor classification using methylation array data.",
+    model="gpt-3.5-turbo-0125",
+    tools=[{"type": "file_search"}],
+    )
+
+    vector_store = client.beta.vector_stores.create(name="MAA_SOP")
+    
+    # Ready the files for upload to OpenAI
+    file_paths = list_files_in_directory('SOP')
+    file_streams = [open(path, "rb") for path in file_paths]
+    
+    # Use the upload and poll SDK helper to upload the files, add them to the vector store,
+    # and poll the status of the file batch for completion.
+    file_batch = client.beta.vector_stores.file_batches.upload_and_poll(
+    vector_store_id=vector_store.id, files=file_streams
+    )
+
+    assistant = client.beta.assistants.update(
+        assistant_id=assistant.id,
+        tool_resources={"file_search": {"vector_store_ids": [vector_store.id]}},
+    )
+
     thread = client.beta.threads.create()
 
     prompt = st.text_area("Enter your message")
     if prompt:
-        with client.beta.threads.runs.create_and_stream(
+        with client.beta.threads.runs.stream(
             thread_id=thread.id,
-            assistant_id=assistant_id,
-            model="gpt-4-turbo",
-            tools=[{"type": "file_search"}],
-            tool_resources={"file_search": {"vector_store_ids": vectorstore_id}},
-            instructions=prompt
+            assistant_id=assistant.id,
+            instructions=f"Please answer the users question: {prompt}",
+            event_handler=EventHandler(),
         ) as stream:
             with st.chat_message("assistant"):
                 response = st.write_stream(stream.text_deltas)