Update app.py

app.py

@@ -7,15 +7,14 @@ from openai import OpenAI
 # ------------------ App Configuration ------------------
 st.set_page_config(page_title="Document AI Assistant", layout="wide")
 st.title("📄 Document AI Assistant")
-st.caption("Chat with an AI Assistant on your medical/pathology documents")
+st.caption("Chat with an AI Assistant on your medical/pathology or construction documents")
 
-# ------------------ Load API Key and Assistant ID from Hugging Face Secrets ------------------
+# ------------------ Load API Key and Assistant ID ------------------
 OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
 ASSISTANT_ID = os.environ.get("ASSISTANT_ID")
 
-# ------------------ Error Handling for Missing Secrets ------------------
 if not OPENAI_API_KEY or not ASSISTANT_ID:
-    st.error("Missing secrets. Please ensure both OPENAI_API_KEY and ASSISTANT_ID are set in your Hugging Face Space secrets.")
+    st.error("❌ Missing secrets. Please set both OPENAI_API_KEY and ASSISTANT_ID in Hugging Face Space secrets.")
     st.stop()
 
 client = OpenAI(api_key=OPENAI_API_KEY)
@@ -41,21 +40,20 @@ if st.sidebar.button("🔄 Clear Chat"):
 
 show_image = st.sidebar.checkbox("📖 Show Document Image", value=True)
 
-# ------------------ Layout: Image (Left) and Chat (Right) ------------------
+# ------------------ Layout: Image + Chat ------------------
 col1, col2 = st.columns([1, 2])
 
-# ------------------ Left Panel: Document Image ------------------
+# ------------------ Left Panel: Image ------------------
 with col1:
     if show_image and st.session_state.image_url:
         st.image(st.session_state.image_url, caption="📑 Extracted Page", use_container_width=True)
         st.session_state.image_updated = False
 
-# ------------------ Right Panel: Chat Interface ------------------
+# ------------------ Right Panel: Chat ------------------
 with col2:
-    # -- Chat Input (Always at the Top)
     prompt = st.chat_input("Type your question about the document...")
 
-    # -- Pair messages (user + assistant) for display
+    # Display chat history
     paired_messages = []
     buffer = []
     for msg in st.session_state.messages:
@@ -66,37 +64,34 @@ with col2:
     if buffer:
         paired_messages.append(buffer.copy())
 
-    # -- Display chat history (latest first, older scroll down)
-    with st.container():
-        for pair in reversed(paired_messages):
-            for msg in pair:
-                with st.chat_message(msg["role"]):
-                    st.write(msg["content"])
+    for pair in reversed(paired_messages):
+        for msg in pair:
+            with st.chat_message(msg["role"]):
+                st.write(msg["content"])
 
-    # -- Process new prompt if entered
     if prompt:
         st.session_state.messages.append({"role": "user", "content": prompt})
 
         try:
-            # Initialize thread if not already done
+            # Create thread
             if st.session_state.thread_id is None:
                 thread = client.beta.threads.create()
                 st.session_state.thread_id = thread.id
 
-            # Send user message to assistant
+            # Send message to assistant
             client.beta.threads.messages.create(
                 thread_id=st.session_state.thread_id,
                 role="user",
                 content=prompt
             )
 
-            # Trigger assistant run
+            # Run assistant
             run = client.beta.threads.runs.create(
                 thread_id=st.session_state.thread_id,
                 assistant_id=ASSISTANT_ID
             )
 
-            # Wait for assistant to respond
+            # Wait for completion
             with st.spinner("Assistant is thinking..."):
                 while True:
                     run_status = client.beta.threads.runs.retrieve(
@@ -107,7 +102,7 @@ with col2:
                         break
                     time.sleep(1)
 
-            # Retrieve latest assistant message
+            # Get assistant reply
             messages = client.beta.threads.messages.list(thread_id=st.session_state.thread_id)
             assistant_message = None
             for message in reversed(messages.data):
@@ -117,16 +112,16 @@ with col2:
 
             st.session_state.messages.append({"role": "assistant", "content": assistant_message})
 
-            # -- Extract GitHub-hosted image URL if present
+            # ✅ Updated: Extract any GitHub-hosted .png image URL (flexible across repos)
             image_match = re.search(
-                r'https://raw\.githubusercontent\.com/AndrewLORTech/surgical-pathology-manual/main/[\w\-/]*\.png',
+                r'https://raw\.githubusercontent\.com/AndrewLORTech/[a-zA-Z0-9_\-/]+\.png',
                 assistant_message
             )
             if image_match:
                 st.session_state.image_url = image_match.group(0)
                 st.session_state.image_updated = True
 
-            st.rerun()  # Rerun to show updated image + message at top
+            st.rerun()
 
         except Exception as e:
             st.error(f"❌ Error: {str(e)}")