Update app.py

app.py

@@ -12,7 +12,6 @@ st.caption("Chat with an Ai Assistant on your Documents")
 with st.sidebar:
     OPENAI_API_KEY = st.text_input("Enter your C2 Group of Technologies Access Key", type="password")
 
-# Stop execution if no API key provided
 if not OPENAI_API_KEY:
     st.error("Please enter your C2 Group of Technologies Access Key to continue.")
     st.stop()
@@ -32,38 +31,37 @@ if st.button("Clear Chat", use_container_width=True):
     st.session_state.thread_id = None
     st.rerun()
 
-# Display message history
+# Display history
 for message in st.session_state.messages:
     role, content = message["role"], message["content"]
     st.chat_message(role).write(content)
 
-# Handle new user input
+# Handle new input
 if prompt := st.chat_input():
     st.session_state.messages.append({"role": "user", "content": prompt})
     st.chat_message("user").write(prompt)
 
     try:
-        # Create thread if not yet created
         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 message
+        # Send user input to OpenAI
         client.beta.threads.messages.create(
             thread_id=thread_id,
             role="user",
             content=prompt
         )
 
-        # Run assistant
+        # Start assistant run
         run = client.beta.threads.runs.create(
             thread_id=thread_id,
             assistant_id=ASSISTANT_ID
         )
 
-        # Poll until complete
+        # Wait until complete
         while True:
             run_status = client.beta.threads.runs.retrieve(
                 thread_id=thread_id,
@@ -73,7 +71,7 @@ if prompt := st.chat_input():
                 break
             time.sleep(1)
 
-        # Get assistant response
+        # Get the latest assistant message
         messages = client.beta.threads.messages.list(thread_id=thread_id)
         assistant_message = None
         for message in reversed(messages.data):
@@ -81,16 +79,18 @@ if prompt := st.chat_input():
                 assistant_message = message.content[0].text.value
                 break
 
-        # Display assistant message
+        # Display response text
         st.chat_message("assistant").write(assistant_message)
         st.session_state.messages.append({"role": "assistant", "content": assistant_message})
 
-        # 🔍 Detect and display GitHub-hosted image
-        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)
+        # Check if assistant already gave a full image URL
+        if "![Page Image](" not in assistant_message:
+            # Try to extract the referenced 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)
 
     except Exception as e:
         st.error(f"Error: {str(e)}")