Spaces:
Sleeping
Sleeping
File size: 4,956 Bytes
8c4492e b386f62 8c4492e b386f62 8c4492e c1043ca 8c4492e c1043ca 8c4492e c1043ca 8c4492e c1043ca 8c4492e c1043ca 8c4492e c1043ca 8c4492e c1043ca 8c4492e 3bbf4ab 57d0c38 3bbf4ab 57d0c38 3bbf4ab 8c4492e 57d0c38 8c4492e |
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 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 |
import streamlit as st
import os
import time
import re
import json
import requests
from PIL import Image
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 ------------------
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
ASSISTANT_ID = os.environ.get("ASSISTANT_ID")
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 = []
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)
# ------------------ Assistant Query Function ------------------
def query_assistant(prompt):
st.session_state.messages.insert(0, {"role": "user", "content": prompt}) # insert at top
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":
break
time.sleep(1)
messages = client.beta.threads.messages.list(thread_id=thread_id)
for message in reversed(messages.data):
if message.role == "assistant":
assistant_message = message.content[0].text.value
st.session_state.messages.insert(0, {"role": "assistant", "content": assistant_message})
# Extract GitHub image URL 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
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")
# 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.experimental_rerun()
# 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.")
|