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.")