Spaces:
Running
Running
File size: 4,554 Bytes
e760d91 1c296f6 e760d91 92001e7 e760d91 92001e7 e760d91 9d8abce 92001e7 e760d91 92001e7 e760d91 92001e7 e2e70fb 92001e7 e760d91 e2e70fb e760d91 92001e7 e760d91 e2e70fb e760d91 92001e7 e760d91 92001e7 e760d91 81240ab 92001e7 e2e70fb 92001e7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
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]) # 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
st.chat_message("assistant").write(assistant_message)
st.session_state.messages.append({"role": "assistant", "content": assistant_message})
# 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
)
if image_match:
st.session_state.image_url = image_match.group(0)
st.session_state.image_updated = True
st.rerun() # Trigger rerun to refresh image display
except Exception as e:
st.error(f"β Error: {str(e)}")
|