Update app.py

app.py

@@ -22,13 +22,13 @@ if not OPENAI_API_KEY or not ASSISTANT_ID:
 
 client = OpenAI(api_key=OPENAI_API_KEY)
 
-# ------------------ Map document titles to exact GitHub folder names ------------------
+# ------------------ Map document titles to 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",
+    "22.146.DS Electrical Specification IFC Rev_0": "22.146.DS Electrical Specification IFC Rev_0",
     "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",
@@ -100,14 +100,14 @@ with col2:
                 thread = client.beta.threads.create()
                 st.session_state.thread_id = thread.id
 
-            # Send user message to assistant
+            # Send user message
             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
@@ -123,7 +123,7 @@ with col2:
                         break
                     time.sleep(1)
 
-            # Retrieve latest assistant message
+            # Retrieve assistant message
             messages = client.beta.threads.messages.list(thread_id=st.session_state.thread_id)
             assistant_message = None
             for message in reversed(messages.data):
@@ -133,23 +133,22 @@ with col2:
 
             st.session_state.messages.append({"role": "assistant", "content": assistant_message})
 
-            # ✅ Extract image URL and correct document folder
-            match = re.search(
-                r'https://raw\.githubusercontent\.com/AndrewLORTech/[^/]+/([^/]+)/\1_page_(\d{4})\.png',
-                assistant_message
-            )
+            # ✅ Extract folder + page number from "Document Reference" line
+            match = re.search(r'Document Reference:\s+(.+?),\s+Page\s+(\d+)', assistant_message)
             if match:
-                doc_display = match.group(1).replace("%20", " ")
-                page_num = match.group(2)
+                doc_name_raw = match.group(1).strip()
+                page_num = int(match.group(2))
+                page_str = f"{page_num:04d}"
+
+                # Map to actual GitHub folder name
+                folder_name = DOCUMENT_FOLDER_MAP.get(doc_name_raw, doc_name_raw)
 
-                # Correct using folder map
-                correct_folder = DOCUMENT_FOLDER_MAP.get(doc_display, doc_display)
-                correct_url = (
+                corrected_url = (
                     f"https://raw.githubusercontent.com/AndrewLORTech/c2ozschlaegerforrestdale/main/"
-                    f"{correct_folder}/{correct_folder}_page_{page_num}.png"
+                    f"{folder_name}/{folder_name}_page_{page_str}.png"
                 )
 
-                st.session_state.image_url = correct_url
+                st.session_state.image_url = corrected_url
                 st.session_state.image_updated = True
 
             st.rerun()