import streamlit as st
import streamlit_shadcn_ui as ui
import os
import tempfile
import html
from src.loader import load_pdf, split_documents
from src.embeddings import get_embedding_model
from src.vectorstore import create_vectorstore
from src.rag import answer_with_memory
st.set_page_config(page_title="Omnibook", layout="wide", initial_sidebar_state="collapsed")
@st.cache_resource
def load_embedding_model():
return get_embedding_model()
embedding_model = load_embedding_model()
if "messages" not in st.session_state:
st.session_state.messages = []
if "vectorstore" not in st.session_state:
st.session_state.vectorstore = None
def process_citations(text, sources):
if not sources:
return text
for i, doc in enumerate(sources, 1):
marker = f"[{i}]"
if marker in text:
snippet = html.escape(doc.page_content[:250].replace('\n', ' ')) + "..."
pill_html = f'{i}'
text = text.replace(marker, pill_html)
return text
st.markdown("""
""", unsafe_allow_html=True)
loading_html = """
"""
st.markdown(
"""
""",
unsafe_allow_html=True
)
col_left, col_right = st.columns([1, 2.8], gap="small")
with col_left:
with st.container(border=True):
st.markdown("Sources
", unsafe_allow_html=True)
uploaded_file = st.file_uploader("Upload PDF", type=["pdf"], label_visibility="collapsed")
process_clicked = ui.button("Process Document", key="btn_process")
if process_clicked:
if uploaded_file is None:
st.error("Please select a PDF file first.")
else:
with st.spinner("Indexing document..."):
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
tmp.write(uploaded_file.getvalue())
tmp_path = tmp.name
docs = load_pdf(tmp_path)
chunks = split_documents(docs)
st.session_state.vectorstore = create_vectorstore(chunks, embedding_model)
st.session_state.messages = []
os.unlink(tmp_path)
st.success(f"Indexed: {len(docs)} pages, {len(chunks)} chunks.")
with col_right:
with st.container(border=True):
head_c1, head_c2 = st.columns([8, 1.5])
with head_c1:
st.markdown("Chat
", unsafe_allow_html=True)
with head_c2:
clear_clicked = ui.button("Clear Chat", key="btn_clear")
if clear_clicked:
st.session_state.messages = []
st.rerun()
st.markdown("", unsafe_allow_html=True)
chat_area = st.container(height=480, border=False)
with chat_area:
if len(st.session_state.messages) == 0:
st.markdown(
"""
Let's learn it together
""",
unsafe_allow_html=True
)
else:
for msg in st.session_state.messages:
with st.chat_message(msg["role"]):
if msg["role"] == "assistant" and "sources" in msg and msg["sources"]:
formatted_text = process_citations(msg["content"], msg["sources"])
st.markdown(formatted_text, unsafe_allow_html=True)
else:
st.markdown(msg["content"])
prompt = st.chat_input("Ask a question about your document...")
if prompt:
if st.session_state.vectorstore is None:
guide = "Please upload a PDF from the left panel and click **Process Document** first."
st.session_state.messages.append({"role": "assistant", "content": guide})
with chat_area:
with st.chat_message("assistant"):
st.markdown(guide)
else:
st.session_state.messages.append({"role": "user", "content": prompt})
with chat_area:
with st.chat_message("user"):
st.markdown(prompt)
with chat_area:
with st.chat_message("assistant"):
loading_placeholder = st.empty()
loading_placeholder.markdown(loading_html, unsafe_allow_html=True)
history = []
msgs = st.session_state.messages[:-1]
for i in range(0, len(msgs) - 1, 2):
if msgs[i]["role"] == "user" and msgs[i+1]["role"] == "assistant":
history.append({
"question": msgs[i]["content"],
"answer": msgs[i+1]["content"],
})
answer, sources = answer_with_memory(
st.session_state.vectorstore, prompt, history
)
loading_placeholder.empty()
formatted_answer = process_citations(answer, sources)
st.markdown(formatted_answer, unsafe_allow_html=True)
st.session_state.messages.append({
"role": "assistant",
"content": answer,
"sources": sources
})