Update app.py

app.py

@@ -36,12 +36,13 @@ if 'assistant_id' not in st.session_state:
     st.session_state.assistant_id = None
 if 'thread_id' not in st.session_state:
     st.session_state.thread_id = None
+if 'file_ids' not in st.session_state:
+    st.session_state.file_ids = []
 
 # ---------------------- Helper Functions ----------------------
 
 def get_vector_stores():
     try:
-        client = OpenAI()
         vector_stores = client.beta.vector_stores.list()
         return vector_stores
     except Exception as e:
@@ -103,14 +104,16 @@ def create_assistant():
     return assistant.id
 
 def chat_with_assistant(file_ids, user_message):
-    for file_id in file_ids:
-        print("File ID:", file_id)
-    # Create a thread and attach the file to the message
-    print("final file id:", file_id)
+    print("----- Starting chat_with_assistant -----")
+    print("Received file_ids:", file_ids)
+    print("Received user_message:", user_message)
+    
+    # Create attachments for each file_id
     attachments = [{"file_id": file_id, "tools": [{"type": "file_search"}]} for file_id in file_ids]
-    print("attachments:", attachments)
-
+    print("Attachments created:", attachments)
+    
     if st.session_state.thread_id is None:
+        print("No existing thread_id found. Creating a new thread.")
         thread = client.beta.threads.create(
             messages=[
                 {
@@ -121,36 +124,83 @@ def chat_with_assistant(file_ids, user_message):
             ]
         )
         st.session_state.thread_id = thread.id
+        print("New thread created with id:", st.session_state.thread_id)
     else:
-        thread = client.beta.threads.messages.create(
+        print(f"Existing thread_id found: {st.session_state.thread_id}. Adding message to the thread.")
+        # Add a message to the existing thread without updating thread_id
+        message = client.beta.threads.messages.create(
             thread_id=st.session_state.thread_id,
             role="user",
             content=user_message,
             attachments=attachments
         )
-
-    # The thread now has a vector store with that file in its tool resources.
-    print(thread.tool_resources.file_search)
-    print("assistant_id:", st.session_state.assistant_id)
-    print("thread_id:", thread.id)
-    run = client.beta.threads.runs.create_and_poll(
-        thread_id=thread.id, assistant_id=st.session_state.assistant_id
-    )
-    print("run:", run)
-    messages = list(client.beta.threads.messages.list(thread_id=thread.id, run_id=run.id))
-
-    message_content = messages[0].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))
-    return message_content.value, citations
+        print("Message added to thread with id:", message.id)
+        # Do NOT update st.session_state.thread_id here
+    
+    # Retrieve the thread object using the thread_id
+    try:
+        thread = client.beta.threads.retrieve(thread_id=st.session_state.thread_id)
+        print("Retrieved thread:", thread)
+    except Exception as e:
+        print(f"Error retrieving thread with id {st.session_state.thread_id}: {e}")
+        return "An error occurred while processing your request.", []
+    
+    # Debugging tool resources
+    try:
+        tool_resources = thread.tool_resources.file_search
+        print("Thread tool resources (file_search):", tool_resources)
+    except AttributeError:
+        print("No tool_resources.file_search found in thread.")
+    
+    print("Assistant ID:", st.session_state.assistant_id)
+    print("Thread ID:", thread.id)
+    
+    # Create and poll the run
+    try:
+        run = client.beta.threads.runs.create_and_poll(
+            thread_id=thread.id, assistant_id=st.session_state.assistant_id
+        )
+        print("Run created and polled:", run)
+    except Exception as e:
+        print("Error during run creation and polling:", e)
+        return "An error occurred while processing your request.", []
+    
+    # Retrieve messages
+    try:
+        messages = list(client.beta.threads.messages.list(thread_id=thread.id, run_id=run.id))
+        print("Retrieved messages:", messages)
+    except Exception as e:
+        print("Error retrieving messages:", e)
+        return "An error occurred while retrieving messages.", []
+    
+    # Process the first message content
+    if messages and messages[0].content:
+        message_content = messages[0].content[0].text
+        print("Raw message content:", message_content)
+        
+        annotations = message_content.annotations
+        print("Annotations found:", annotations)
+        
+        citations = []
+        for index, annotation in enumerate(annotations):
+            print(f"Processing annotation {index}: {annotation.text}")
+            message_content.value = message_content.value.replace(annotation.text, f"[{index}]")
+            if file_citation := getattr(annotation, "file_citation", None):
+                try:
+                    cited_file = client.files.retrieve(file_citation.file_id)
+                    citation_entry = f"[{index}] {cited_file.filename}"
+                    citations.append(citation_entry)
+                    print(f"Citation added: {citation_entry}")
+                except Exception as e:
+                    print(f"Error retrieving cited file for annotation {index}: {e}")
+        
+        print("Final message content after replacements:", message_content.value)
+        print("All citations:", citations)
+        print("----- Ending chat_with_assistant -----")
+        return message_content.value, citations
+    else:
+        print("No messages or content found in the retrieved messages.")
+        return "No response received from the assistant.", []
 
 
 
@@ -201,7 +251,7 @@ if page == "Home":
     if submit and user_input.strip() != "":
         # Add user message to chat history
         st.session_state.chat_history.append({"role": "user", "content": user_input})
-        
+        print("chat history:", st.session_state.chat_history)
         if not st.session_state.file_ids:
             st.error("Please process PDFs first.")
         else:
@@ -209,7 +259,10 @@ if page == "Home":
                 try:
                     response, citations = chat_with_assistant(st.session_state.file_ids, user_input)
                     # Add assistant response to chat history
+                    print("response:", response)
+                    print("citations:", citations)
                     st.session_state.chat_history.append({"role": "assistant", "content": response+"\n\n"+"\n".join(citations)})
+                    print("chat history:", st.session_state.chat_history)
                 except Exception as e:
                     st.error(f"Error generating response: {str(e)}")