app_hf.py

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'<span class="citation-pill" data-tooltip="Source {i}: {snippet}">{i}</span>'
            text = text.replace(marker, pill_html)
            
    return text

st.markdown("""
    <style>
        @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600&display=swap');
        
        html, body, p, h1, h2, h3, h4, h5, h6, li, ul, ol, div { 
            font-family: 'Inter', sans-serif; 
        }
        
        code, pre, code span, pre span {
            font-family: 'Courier New', Courier, monospace !important;
        }
        pre {
            background-color: #f1f5f9 !important;
            padding: 1rem !important;
            border-radius: 0.5rem !important;
            border: 1px solid #e2e8f0 !important;
            overflow-x: auto !important;
        }
        code {
            background-color: #f1f5f9 !important;
            color: #e11d48 !important;
            padding: 0.1rem 0.3rem !important;
            border-radius: 0.25rem !important;
        }
        pre code {
            background-color: transparent !important;
            color: #0f172a !important;
            padding: 0 !important;
        }

        .block-container { 
            padding-top: 1rem !important; padding-bottom: 0rem !important; 
            padding-left: 2rem !important; padding-right: 2rem !important; 
            max-width: 100% !important; 
        }
        
        header {visibility: hidden;}
        
        .stApp { overflow-y: hidden !important; }
        
        [data-testid="stVerticalBlockBorderWrapper"] {
            border-radius: 2rem !important; 
            border-color: #e5e7eb !important;
            box-shadow: 0 1px 2px 0 rgba(0, 0, 0, 0.05);
        }
        
        .welcome-text { text-align: center; margin-top: 10rem; margin-bottom: 2rem; }

        [data-testid="stChatMessageAvatarUser"] { display: none !important; }
        
        [data-testid="stChatMessageAvatarAssistant"] { 
            background-color: #374151 !important; 
            border-radius: 2px !important;
            width: 5px !important;
            height: 20px !important;
            min-width: 5px !important;
            margin-right: 12px !important;
            margin-top: 6px !important;
        }
        [data-testid="stChatMessageAvatarAssistant"] svg { 
            display: none !important;
        }

        ::-webkit-scrollbar { width: 6px; height: 6px; }
        ::-webkit-scrollbar-track { background: transparent !important; }
        ::-webkit-scrollbar-thumb { background: transparent !important; border-radius: 10px; }
        *:hover::-webkit-scrollbar-thumb { background: #d1d5db !important; }

        .citation-pill {
            display: inline-flex;
            align-items: center;
            justify-content: center;
            background-color: #f3f4f6;
            border: 1px solid #d1d5db;
            color: #4b5563;
            border-radius: 50%;
            width: 22px; height: 22px;
            font-size: 11px; font-weight: 600;
            cursor: help; position: relative; margin: 0 4px;
            transform: translateY(-2px); transition: all 0.2s ease;
        }
        .citation-pill:hover { background-color: #e5e7eb; color: #111827; border-color: #9ca3af; }
        .citation-pill::after {
            content: attr(data-tooltip); position: absolute; bottom: 145%; left: 50%; transform: translateX(-50%);
            width: max-content; max-width: 280px; background-color: #1f2937; color: #f9fafb; font-size: 12px; font-weight: 400;
            line-height: 1.4; padding: 8px 12px; border-radius: 8px; opacity: 0; visibility: hidden; transition: all 0.2s ease;
            z-index: 9999; box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1); pointer-events: none; white-space: normal; text-align: left;
        }
        .citation-pill:hover::after { opacity: 1; visibility: visible; }

        .terminal-loader {
            display: flex; align-items: center; gap: 8px;
            font-family: 'Inter', sans-serif; font-size: 15px; font-weight: 500; color: #4b5563; padding: 5px 0;
        }
        .terminal-cursor {
            width: 8px; height: 16px; background-color: #374151; animation: blink 1s step-end infinite;
        }
        .loading-text::after { content: "Let me learn it..."; animation: textSequence 12s infinite; }
        @keyframes blink { 0%, 100% { opacity: 1; } 50% { opacity: 0; } }
        @keyframes textSequence {
            0%, 15% { content: "Let me learn it..."; } 16%, 32% { content: "Scanning knowledge base..."; }
            33%, 49% { content: "Vectorizing context..."; } 50%, 66% { content: "Retrieving relevant data..."; }
            67%, 83% { content: "Synthesizing information..."; } 84%, 100% { content: "Formatting final response..."; }
        }
    </style>
""", unsafe_allow_html=True)

loading_html = """
<div class="terminal-loader">
    <div class="terminal-cursor"></div>
    <div class="loading-text"></div>
</div>
"""

st.markdown(
    """
    <div style='display: flex; align-items: center; gap: 12px; margin-bottom: 10px; padding-left: 5px;'>
        <div style='width: 36px; height: 36px; background-color: black; border-radius: 50%; display: flex; align-items: center; justify-content: center;'>
            <svg width="20" height="20" viewBox="0 0 24 24" fill="none" stroke="white" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><rect width="12" height="20" x="6" y="2" rx="2"></rect><rect width="20" height="12" x="2" y="6" rx="2"></rect></svg>
        </div>
        <span style='font-size: 22px; font-weight: 500; color: #1f2937;'>Omnibook</span>
    </div>
    """, 
    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("<h3 style='font-size: 16px; margin-bottom: 15px; padding-left: 8px;'>Sources</h3>", 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("<h3 style='font-size: 16px; padding-left: 8px; margin-top: 6px;'>Chat</h3>", 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("<div style='margin-bottom: 5px;'></div>", unsafe_allow_html=True)
        
        chat_area = st.container(height=480, border=False)
        
        with chat_area:
            if len(st.session_state.messages) == 0:
                st.markdown(
                    """
                    <div class="welcome-text">
                        <h1 style="font-size: 28px; font-weight: 600; color: #1f2937;">Let's learn it together</h1>
                    </div>
                    """,
                    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
                        })