AabhavAmitabh/document-intelligence-chatbot

1# app.py
2# The web interface for your Document Intelligence Chatbot.
3# Run with: streamlit run app.py
4
5import streamlit as st
6import os
7import tempfile
8from pathlib import Path
9from ingestor import ingest_pdf
10from embedder import embed_and_store, load_collection
11from qa_chain import answer_question
12
13# ── Page configuration ──────────────────────────────────────────────────────
14st.set_page_config(
15    page_title="Document Intelligence Chatbot",
16    page_icon="📄",
17    layout="wide"
18)
19
20# ── Custom CSS for a cleaner look ────────────────────────────────────────────
21st.markdown("""
22<style>
23    .main-header {
24        font-size: 2rem;
25        font-weight: 600;
26        margin-bottom: 0.25rem;
27    }
28    .sub-header {
29        color: #666;
30        font-size: 1rem;
31        margin-bottom: 2rem;
32    }
33  .answer-box {
34        background-color: #1e3a5f;
35        border-left: 4px solid #2563eb;
36        padding: 1rem 1.25rem;
37        border-radius: 0 8px 8px 0;
38        margin: 1rem 0;
39        color: #ffffff;
40    }
41    .source-box {
42        background-color: #f9f9f9;
43        border: 1px solid #e0e0e0;
44        border-radius: 8px;
45        padding: 0.75rem 1rem;
46        font-size: 0.85rem;
47        color: #555;
48        margin-top: 0.5rem;
49    }
50    .status-ready {
51        color: #16a34a;
52        font-weight: 500;
53    }
54    .status-empty {
55        color: #dc2626;
56        font-weight: 500;
57    }
58</style>
59""", unsafe_allow_html=True)
60
61# ── Session state initialisation ─────────────────────────────────────────────
62# Streamlit reruns the entire script on every interaction.
63# st.session_state persists data between reruns — like a memory for the app.
64if "document_loaded" not in st.session_state:
65    st.session_state.document_loaded = False
66if "doc_name" not in st.session_state:
67    st.session_state.doc_name = ""
68if "chat_history" not in st.session_state:
69    st.session_state.chat_history = []
70
71# ── Header ───────────────────────────────────────────────────────────────────
72st.markdown('<div class="main-header">📄 Document Intelligence Chatbot</div>',
73            unsafe_allow_html=True)
74st.markdown('<div class="sub-header">Upload any PDF and ask questions about it in plain English.</div>',
75            unsafe_allow_html=True)
76
77# ── Layout: two columns ──────────────────────────────────────────────────────
78left_col, right_col = st.columns([1, 2])
79
80# ── LEFT COLUMN: Upload + status ─────────────────────────────────────────────
81with left_col:
82    st.subheader("Upload Document")
83
84    uploaded_file = st.file_uploader(
85        "Choose a PDF file",
86        type=["pdf"],
87        help="Upload any PDF — research papers, reports, manuals, articles."
88    )
89
90    if uploaded_file is not None:
91        # Show a process button once a file is selected
92        if st.button("Process Document", type="primary", use_container_width=True):
93            with st.spinner("Reading and indexing your document..."):
94                try:
95                    # Save uploaded file to a temp location so our
96                    # ingestor can read it from disk
97                    with tempfile.NamedTemporaryFile(
98                        delete=False, suffix=".pdf"
99                    ) as tmp_file:
100                        tmp_file.write(uploaded_file.getvalue())
101                        tmp_path = tmp_file.name
102
103                    # Get clean document name
104                    doc_name = Path(uploaded_file.name).stem
105
106                    # Run the full ingestion + embedding pipeline
107                    chunks = ingest_pdf(tmp_path)
108                    embed_and_store(chunks, doc_name=doc_name)
109
110                    # Clean up the temp file
111                    os.unlink(tmp_path)
112
113                    # Mark as ready in session state
114                    st.session_state.document_loaded = True
115                    st.session_state.doc_name = doc_name
116                    st.session_state.chat_history = []  # clear old chat
117
118                    st.success(f"Ready! Indexed {len(chunks)} passages.")
119
120                except Exception as e:
121                    st.error(f"Error processing document: {str(e)}")
122
123    # Status indicator
124    st.divider()
125    st.markdown("**Status**")
126    if st.session_state.document_loaded:
127        st.markdown(
128            f'<span class="status-ready">● Ready — {st.session_state.doc_name}</span>',
129            unsafe_allow_html=True
130        )
131    else:
132        st.markdown(
133            '<span class="status-empty">● No document loaded</span>',
134            unsafe_allow_html=True
135        )
136
137    # Sample questions to guide the user
138    st.divider()
139    st.markdown("**Try asking:**")
140    sample_questions = [
141        "What is this document about?",
142        "Summarise the key points",
143        "What are the main findings?",
144        "Who are the key people mentioned?",
145    ]
146    for q in sample_questions:
147        if st.button(q, use_container_width=True, disabled=not st.session_state.document_loaded):
148            st.session_state.pending_question = q
149
150# ── RIGHT COLUMN: Chat interface ─────────────────────────────────────────────
151with right_col:
152    st.subheader("Ask Questions")
153
154    # Display chat history
155    if st.session_state.chat_history:
156        for exchange in st.session_state.chat_history:
157            # User message
158            with st.chat_message("user"):
159                st.write(exchange["question"])
160            # Assistant answer
161            with st.chat_message("assistant"):
162                st.markdown(
163                    f'<div class="answer-box">{exchange["answer"]}</div>',
164                    unsafe_allow_html=True
165                )
166                # Show source passages in an expander
167                with st.expander(f"View {len(exchange['sources'])} source passages"):
168                    for i, src in enumerate(exchange["sources"]):
169                        st.markdown(
170                            f'<div class="source-box">'
171                            f'<strong>Passage {i+1}</strong> '
172                            f'(chunk {src["index"]} from <em>{src["source"]}</em>)<br><br>'
173                            f'{src["text"][:400]}{"..." if len(src["text"]) > 400 else ""}'
174                            f'</div>',
175                            unsafe_allow_html=True
176                        )
177    else:
178        if st.session_state.document_loaded:
179            st.info("Your document is ready. Ask anything below.")
180        else:
181            st.info("Upload and process a PDF on the left to get started.")
182
183    # Chat input box — always at the bottom
184    if st.session_state.document_loaded:
185        # Handle questions from sample buttons
186        default_q = st.session_state.get("pending_question", "")
187        if default_q:
188            st.session_state.pending_question = ""
189
190        user_question = st.chat_input("Ask a question about your document...")
191
192        # Process either typed question or button-selected question
193        question_to_answer = user_question or default_q
194
195        if question_to_answer:
196            with st.spinner("Searching document and generating answer..."):
197                try:
198                    result = answer_question(question_to_answer)
199                    # Add to chat history
200                    st.session_state.chat_history.append(result)
201                    st.rerun()  # refresh to show new message
202                except Exception as e:
203                    st.error(f"Error generating answer: {str(e)}")
204    else:
205        st.chat_input("Upload a document first...", disabled=True)