Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import os | |
| import time | |
| import re | |
| 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") | |
| # ------------------ 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) | |
| # ------------------ Session State Initialization ------------------ | |
| if "messages" not in st.session_state: | |
| st.session_state.messages = [] | |
| if "thread_id" not in st.session_state: | |
| st.session_state.thread_id = None | |
| if "image_url" not in st.session_state: | |
| st.session_state.image_url = None | |
| if "image_updated" not in st.session_state: | |
| st.session_state.image_updated = False | |
| # ------------------ Sidebar Controls ------------------ | |
| st.sidebar.header("🔧 Settings") | |
| if st.sidebar.button("🔄 Clear Chat"): | |
| st.session_state.messages = [] | |
| st.session_state.thread_id = None | |
| st.session_state.image_url = None | |
| st.session_state.image_updated = False | |
| st.rerun() | |
| show_image = st.sidebar.checkbox("📖 Show Document Image", value=True) | |
| # ------------------ Split Layout ------------------ | |
| col1, col2 = st.columns([1, 2]) | |
| # ------------------ 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 | |
| # ------------------ Chat Panel (Right) ------------------ | |
| with col2: | |
| # Pair user + assistant messages | |
| paired_messages = [] | |
| buffer = [] | |
| for msg in st.session_state.messages: | |
| buffer.append(msg) | |
| if msg["role"] == "assistant" and len(buffer) == 2: | |
| paired_messages.append(buffer.copy()) | |
| buffer.clear() | |
| # If last message is from user and no assistant reply yet | |
| if buffer: | |
| paired_messages.append(buffer.copy()) | |
| # Show pairs in reverse order (latest on top) | |
| for pair in reversed(paired_messages): | |
| for msg in pair: | |
| st.chat_message(msg["role"]).write(msg["content"]) | |
| # ------------------ Chat Input Handling ------------------ | |
| if prompt := st.chat_input("Type your question about the document..."): | |
| st.session_state.messages.append({"role": "user", "content": prompt}) | |
| try: | |
| # Create thread if not exists | |
| 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 user 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 to respond | |
| 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) | |
| # Retrieve 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 | |
| st.session_state.messages.append({"role": "assistant", "content": assistant_message}) | |
| # Extract image URL if present | |
| image_match = re.search( | |
| r'https://raw\.githubusercontent\.com/AndrewLORTech/surgical-pathology-manual/main/[\w\-/]*\.png', | |
| assistant_message | |
| ) | |
| if image_match: | |
| st.session_state.image_url = image_match.group(0) | |
| st.session_state.image_updated = True | |
| st.rerun() | |
| except Exception as e: | |
| st.error(f"❌ Error: {str(e)}") | |