app.py

from __future__ import annotations

import streamlit as st

from src.document_loader import load_pdf_text, load_text_files
from src.generator import generate_answer_with_openai
from src.retriever import TfidfRetriever
from src.text_splitter import split_text


st.set_page_config(page_title="AI Document Chatbot", page_icon="🤖", layout="wide")

st.title("AI Document Chatbot")
st.write(
    "Upload PDF or text documents, ask a question, and retrieve the most relevant document sections. "
    "If an OpenAI API key is configured, the app also generates a concise answer from the retrieved context."
)

with st.sidebar:
    st.header("Settings")
    chunk_size = st.slider("Chunk size", min_value=400, max_value=1600, value=900, step=100)
    overlap = st.slider("Chunk overlap", min_value=50, max_value=400, value=150, step=50)
    top_k = st.slider("Number of retrieved chunks", min_value=1, max_value=8, value=4)
    st.info("Set OPENAI_API_KEY in your environment to enable generated answers.")

uploaded_files = st.file_uploader(
    "Upload PDF or TXT files",
    type=["pdf", "txt"],
    accept_multiple_files=True,
)

if "retriever" not in st.session_state:
    st.session_state.retriever = None
    st.session_state.chunks = []

if uploaded_files:
    pdf_texts: list[str] = []
    txt_files: list[tuple[str, bytes]] = []

    for uploaded_file in uploaded_files:
        file_bytes = uploaded_file.read()
        if uploaded_file.name.lower().endswith(".pdf"):
            pdf_texts.append(load_pdf_text(file_bytes))
        else:
            txt_files.append((uploaded_file.name, file_bytes))

    combined_text = "\n\n".join(pdf_texts + [load_text_files(txt_files)])
    chunks = split_text(combined_text, chunk_size=chunk_size, overlap=overlap)

    if chunks:
        st.session_state.chunks = chunks
        st.session_state.retriever = TfidfRetriever(chunks)
        st.success(f"Loaded {len(uploaded_files)} file(s) and created {len(chunks)} searchable chunks.")
    else:
        st.error("No readable text was found. Try another PDF/TXT file.")

question = st.text_input("Ask a question about the uploaded documents")

if question and st.session_state.retriever:
    results = st.session_state.retriever.search(question, top_k=top_k)
    answer = generate_answer_with_openai(question, results)

    st.subheader("Answer")
    st.write(answer)

    st.subheader("Retrieved document sections")
    for result in results:
        with st.expander(f"Source {result.rank} — similarity score {result.score:.3f}"):
            st.write(result.text)

elif question and not st.session_state.retriever:
    st.warning("Please upload at least one document first.")

st.divider()
st.caption("Portfolio project: Python, Streamlit, scikit-learn, pypdf, optional OpenAI API.")