Update app.py

app.py

@@ -12,7 +12,7 @@ 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")
 
-# ------------------ 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")
 
@@ -45,10 +45,10 @@ show_image = st.sidebar.checkbox("📖 Show Document Image", value=True)
 
 # ------------------ Load Structured Summary/FAQ ------------------
 with open("51940670-Manual-of-Surgical-Pathology-Third-Edition_1_structured_output.json", "r") as f:
-    structured_data = json.load(f)
+    structured_data = json.load(f)  # This is a list of dicts, not a dict
 
 # ------------------ Three-Column Layout ------------------
-left, center, right = st.columns([1, 2, 1])  # adjust as needed
+left, center, right = st.columns([1, 2, 1])
 
 # ------------------ Left Column: Document Image ------------------
 with left:
@@ -56,7 +56,7 @@ with left:
     if show_image and st.session_state.image_url:
         try:
             image = Image.open(requests.get(st.session_state.image_url, stream=True).raw)
-            st.image(image, caption="📑 Extracted Page", use_column_width=True)
+            st.image(image, caption="📑 Extracted Page", use_container_width=True)
             st.session_state.image_updated = False
         except Exception as e:
             st.warning("⚠️ Could not load image.")
@@ -79,20 +79,17 @@ with center:
 
             thread_id = st.session_state.thread_id
 
-            # Send user prompt
             client.beta.threads.messages.create(
                 thread_id=thread_id,
                 role="user",
                 content=prompt
             )
 
-            # Run assistant
             run = client.beta.threads.runs.create(
                 thread_id=thread_id,
                 assistant_id=ASSISTANT_ID
             )
 
-            # Poll until done
             with st.spinner("Assistant is thinking..."):
                 while True:
                     run_status = client.beta.threads.runs.retrieve(
@@ -103,7 +100,6 @@ with center:
                         break
                     time.sleep(1)
 
-            # Get assistant message
             messages = client.beta.threads.messages.list(thread_id=thread_id)
             assistant_message = None
             for message in reversed(messages.data):
@@ -114,7 +110,7 @@ with center:
             st.chat_message("assistant").write(assistant_message)
             st.session_state.messages.append({"role": "assistant", "content": assistant_message})
 
-            # Detect GitHub image in response
+            # Extract GitHub image URL
             image_match = re.search(
                 r'https://raw\.githubusercontent\.com/AndrewLORTech/surgical-pathology-manual/main/[\w\-/]*\.png',
                 assistant_message
@@ -131,17 +127,25 @@ with center:
 with right:
     st.subheader("📌 Summary")
 
+    # Parse page number from image URL if available
     if st.session_state.image_url:
         match = re.search(r'page_(\d+)', st.session_state.image_url)
-        page_number = int(match.group(1)) if match else None
+        page_number = int(match.group(1)) if match else 151
     else:
         page_number = 151  # default
 
-    summary_text = structured_data.get(str(page_number), {}).get("summary", "No summary available.")
+    # Get entry from structured data
+    page_entry = next((entry for entry in structured_data if entry.get("page") == page_number), None)
+    if page_entry:
+        summary_text = page_entry.get("summary", "No summary available.")
+        faq_list = page_entry.get("faqs", [])
+    else:
+        summary_text = "No summary available."
+        faq_list = []
+
     st.markdown(summary_text)
 
     st.subheader("❓ Auto-Generated FAQ")
-    faq_list = structured_data.get(str(page_number), {}).get("faqs", [])
     if faq_list:
         for faq in faq_list:
             st.markdown(f"**Q:** {faq.get('question', '')}\n\n**A:** {faq.get('answer', '')}")