Update app.py

app.py

@@ -2,9 +2,6 @@ import streamlit as st
 import os
 import time
 import re
-import json
-import requests
-from PIL import Image
 from openai import OpenAI
 
 # ------------------ App Configuration ------------------
@@ -12,25 +9,17 @@ 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 ------------------
+# ------------------ Load API Key and Assistant ID from Hugging Face Secrets ------------------
 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.stop()
 
 client = OpenAI(api_key=OPENAI_API_KEY)
 
-# ------------------ Load Structured JSON ------------------
-STRUCTURED_JSON_PATH = "51940670-Manual-of-Surgical-Pathology-Third-Edition_1_structured_output.json"
-try:
-    with open(STRUCTURED_JSON_PATH, "r") as f:
-        structured_data = json.load(f)
-except Exception as e:
-    st.error(f"❌ Failed to load structured summary file: {e}")
-    st.stop()
-
 # ------------------ Session State Initialization ------------------
 if "messages" not in st.session_state:
     st.session_state.messages = []
@@ -52,46 +41,69 @@ if st.sidebar.button("🔄 Clear Chat"):
 
 show_image = st.sidebar.checkbox("📖 Show Document Image", value=True)
 
-# ------------------ Assistant Query Function ------------------
-def query_assistant(prompt):
-    st.session_state.messages.insert(0, {"role": "user", "content": prompt})
-
-    try:
-        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
-
-        client.beta.threads.messages.create(
-            thread_id=thread_id,
-            role="user",
-            content=prompt
-        )
-
-        run = client.beta.threads.runs.create(
-            thread_id=thread_id,
-            assistant_id=ASSISTANT_ID
-        )
-
-        with st.spinner("Assistant is thinking..."):
-            while True:
-                run_status = client.beta.threads.runs.retrieve(
-                    thread_id=thread_id,
-                    run_id=run.id
-                )
-                if run_status.status == "completed":
+# ------------------ Split Layout ------------------
+col1, col2 = st.columns([1, 2])  # Adjust ratio as needed
+
+# ------------------ Image Panel (Left) ------------------
+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  # Reset flag after rendering
+
+# ------------------ Chat Panel (Right) ------------------
+with col2:
+    for message in st.session_state.messages:
+        role, content = message["role"], message["content"]
+        st.chat_message(role).write(content)
+
+    if prompt := st.chat_input("Type your question about the document..."):
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        st.chat_message("user").write(prompt)
+
+        try:
+            # Initialize thread if needed
+            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 to assistant
+            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
+            )
+
+            # Wait for assistant response
+            with st.spinner("Assistant is thinking..."):
+                while True:
+                    run_status = client.beta.threads.runs.retrieve(
+                        thread_id=thread_id,
+                        run_id=run.id
+                    )
+                    if run_status.status == "completed":
+                        break
+                    time.sleep(1)
+
+            # Get assistant response
+            messages = client.beta.threads.messages.list(thread_id=thread_id)
+            assistant_message = None
+            for message in reversed(messages.data):
+                if message.role == "assistant":
+                    assistant_message = message.content[0].text.value
                     break
-                time.sleep(1)
 
-        # ✅ Only fetch the latest assistant message to avoid infinite loop
-        messages = client.beta.threads.messages.list(thread_id=thread_id, limit=1)
-        latest_message = messages.data[0]
-        if latest_message.role == "assistant":
-            assistant_message = latest_message.content[0].text.value
-            st.session_state.messages.insert(0, {"role": "assistant", "content": assistant_message})
+            st.chat_message("assistant").write(assistant_message)
+            st.session_state.messages.append({"role": "assistant", "content": assistant_message})
 
-            # Extract GitHub image URL if available
+            # Extract GitHub image from response if available
             image_match = re.search(
                 r'https://raw\.githubusercontent\.com/AndrewLORTech/surgical-pathology-manual/main/[\w\-/]*\.png',
                 assistant_message
@@ -99,39 +111,7 @@ def query_assistant(prompt):
             if image_match:
                 st.session_state.image_url = image_match.group(0)
                 st.session_state.image_updated = True
-            return assistant_message
-
-    except Exception as e:
-        st.error(f"❌ Error: {str(e)}")
-        return None
-
-# ------------------ Layout ------------------
-left, center = st.columns([1, 2])
-
-# ------------------ Center Column: Chat UI with Static Input on Top ------------------
-with center:
-    st.subheader("💬 Document AI Assistant")
+                st.rerun()  # Trigger rerun to refresh image display
 
-    # Static Chat Input Bar
-    with st.container():
-        prompt = st.text_input("💡 Ask a question about the document:", key="chat_input")
-        if prompt:
-            query_assistant(prompt)
-            st.rerun()  # Rerun to refresh with new message
-
-    # Show messages: latest at top
-    for message in st.session_state.messages:
-        role = message["role"]
-        with st.chat_message(role):
-            st.markdown(message["content"])
-
-# ------------------ Left Column: Document Image ------------------
-with left:
-    st.subheader("📄 Document Image")
-    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_container_width=True)
-            st.session_state.image_updated = False
         except Exception as e:
-            st.warning("⚠️ Could not load image.")
+            st.error(f"❌ Error: {str(e)}")