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