Update app.py

app.py

@@ -3,55 +3,51 @@ from openai import OpenAI
 import time
 import re
 
-# Streamlit setup
 st.set_page_config(page_title="documentaitest")
 st.title("documentaitest")
 st.caption("Chat with an Ai Assistant on your Documents")
 
-# Sidebar for API key input
+# Sidebar key entry
 with st.sidebar:
     OPENAI_API_KEY = st.text_input("Enter your C2 Group of Technologies Access Key", type="password")
 
-# Exit if no key provided
 if not OPENAI_API_KEY:
     st.error("Please enter your C2 Group of Technologies Access Key to continue.")
     st.stop()
 
-# Initialize OpenAI client
 client = OpenAI(api_key=OPENAI_API_KEY)
 ASSISTANT_ID = "asst_PJjxQftfz2IJUUMvnldK58lB"
 
-# Initialize session state
+# Session state setup
 if "messages" not in st.session_state:
     st.session_state.messages = []
 if "thread_id" not in st.session_state:
     st.session_state.thread_id = None
 
-# Button to clear chat
+# Clear button
 if st.button("Clear Chat", use_container_width=True):
     st.session_state.messages = []
     st.session_state.thread_id = None
     st.rerun()
 
-# Show message history
+# Display prior messages
 for message in st.session_state.messages:
     role, content = message["role"], message["content"]
     st.chat_message(role).write(content)
 
-# Handle new user message
+# Handle input
 if prompt := st.chat_input():
     st.session_state.messages.append({"role": "user", "content": prompt})
     st.chat_message("user").write(prompt)
 
     try:
-        # Create a new thread if none exists
         if st.session_state.thread_id is None:
             thread = client.beta.threads.create()
             st.session_state.thread_id = thread.id
 
         thread_id = st.session_state.thread_id
 
-        # Send user prompt to assistant
+        # Send prompt to assistant
         client.beta.threads.messages.create(
             thread_id=thread_id,
             role="user",
@@ -64,7 +60,7 @@ if prompt := st.chat_input():
             assistant_id=ASSISTANT_ID
         )
 
-        # Wait until completion
+        # Wait for assistant to finish
         while True:
             run_status = client.beta.threads.runs.retrieve(
                 thread_id=thread_id,
@@ -74,7 +70,7 @@ if prompt := st.chat_input():
                 break
             time.sleep(1)
 
-        # Retrieve assistant message
+        # Get assistant response
         messages = client.beta.threads.messages.list(thread_id=thread_id)
         assistant_message = None
         for message in reversed(messages.data):
@@ -82,18 +78,15 @@ if prompt := st.chat_input():
                 assistant_message = message.content[0].text.value
                 break
 
-        # Display assistant message
+        # Display message
         st.chat_message("assistant").write(assistant_message)
         st.session_state.messages.append({"role": "assistant", "content": assistant_message})
 
-        # Inject page image if assistant didn't include it
-        if "![Page Image](" not in assistant_message:
-            # Find the page number
-            page_match = re.search(r'Page (\d+)', assistant_message)
-            if page_match:
-                page_number = page_match.group(1).zfill(3)
-                image_url = f"https://raw.githubusercontent.com/AndrewLORTech/surgical-pathology-manual/main/51940670-Manual-of-Surgical-Pathology-Third-Edition_1_page_{page_number}.png"
-                st.image(image_url, caption=f"Page {int(page_number)}", use_container_width=True)
+        # 🔍 Extract and display GitHub image if present
+        image_match = re.search(r'https://raw\.githubusercontent\.com/AndrewLORTech/surgical-pathology-manual/main/[\w\-/]*\.png', assistant_message)
+        if image_match:
+            image_url = image_match.group(0)
+            st.image(image_url, caption="Page Image", use_container_width=True)
 
     except Exception as e:
         st.error(f"Error: {str(e)}")