Update app.py

app.py

@@ -17,11 +17,24 @@ OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
 ASSISTANT_ID = os.environ.get("ASSISTANT_ID")
 
 if not OPENAI_API_KEY or not ASSISTANT_ID:
-    st.error("❌ Missing secrets. Please set both OPENAI_API_KEY and ASSISTANT_ID 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)
 
+# ------------------ Map document titles to exact GitHub folder names ------------------
+DOCUMENT_FOLDER_MAP = {
+    "Tender 5 of 2024 - Forrestdale": "Schlager Group Pty Ltd - Tender 5 of 2024 - Forrestdale ~ Pavilion and External Works Part 1 of 2",
+    "Tender 5 of 2024 - Forrestdale Part 2": "Schlager Group Pty Ltd - Tender 5 of 2024 - Forrestdale ~ Pavilion and External Works Part 2 of 2",
+    "Armadale Forrestdale Methodology": "Armadale Forrestdale Methodology",
+    "22.146.DS Electrical Specification IFC Rev_0": "22.146.DS Electrical Specification IFC Rev_0",
+    "2302-FORRESTDALE SPORTING PAVILION-TECHNICAL SPECIFICATION": "2302-FORRESTDALE SPORTING PAVILION-TECHNICAL SPECIFICATION",
+    "Forrrestdale Community Hub Program 23-04-2025": "Forrrestdale Community Hub Program 23-04-2025",
+    "Landscape Specification - Forrestdale Community Hub_02": "Landscape Specification - Forrestdale Community Hub_02",
+    "Mechanical_Spec_-_Forrestdale Sporting Club Room IFC": "Mechanical_Spec_-_Forrestdale Sporting Club Room IFC",
+    "V-23-021-HY-SPEC-R0": "V-23-021-HY-SPEC-R0"
+}
+
 # ------------------ Session State Initialization ------------------
 if "messages" not in st.session_state:
     st.session_state.messages = []
@@ -87,20 +100,19 @@ with col2:
                 thread = client.beta.threads.create()
                 st.session_state.thread_id = thread.id
 
-            # Send user prompt
+            # Send user message to assistant
             client.beta.threads.messages.create(
                 thread_id=st.session_state.thread_id,
                 role="user",
                 content=prompt
             )
 
-            # Run assistant
+            # Trigger assistant run
             run = client.beta.threads.runs.create(
                 thread_id=st.session_state.thread_id,
                 assistant_id=ASSISTANT_ID
             )
 
-            # Wait for completion
             with st.spinner("Assistant is thinking..."):
                 while True:
                     run_status = client.beta.threads.runs.retrieve(
@@ -111,7 +123,7 @@ with col2:
                         break
                     time.sleep(1)
 
-            # Get assistant reply
+            # Retrieve latest assistant message
             messages = client.beta.threads.messages.list(thread_id=st.session_state.thread_id)
             assistant_message = None
             for message in reversed(messages.data):
@@ -121,13 +133,23 @@ with col2:
 
             st.session_state.messages.append({"role": "assistant", "content": assistant_message})
 
-            # ✅ Extract GitHub raw image URL
-            image_match = re.search(
-                r'https://raw\.githubusercontent\.com/AndrewLORTech/[a-zA-Z0-9_\-/]+\.png',
+            # ✅ Extract image URL and correct document folder
+            match = re.search(
+                r'https://raw\.githubusercontent\.com/AndrewLORTech/[^/]+/([^/]+)/\1_page_(\d{4})\.png',
                 assistant_message
             )
-            if image_match:
-                st.session_state.image_url = image_match.group(0)
+            if match:
+                doc_display = match.group(1).replace("%20", " ")
+                page_num = match.group(2)
+
+                # Correct using folder map
+                correct_folder = DOCUMENT_FOLDER_MAP.get(doc_display, doc_display)
+                correct_url = (
+                    f"https://raw.githubusercontent.com/AndrewLORTech/c2ozschlaegerforrestdale/main/"
+                    f"{correct_folder}/{correct_folder}_page_{page_num}.png"
+                )
+
+                st.session_state.image_url = correct_url
                 st.session_state.image_updated = True
 
             st.rerun()