app.py

"""Streamlit UI：聊天介面 + 來源引用 + 串流回應。

啟動：streamlit run src/app.py
"""

from __future__ import annotations

import streamlit as st

from src.config import config
from src.rag import build_chain

st.set_page_config(
    page_title="RAG 智能問答",
    layout="wide",
    initial_sidebar_state="expanded",
)


@st.cache_resource(show_spinner="載入 RAG 引擎中...")
def get_rag_engine():
    """快取 chain 與 retriever，避免每次互動都重建（重建會重讀整個向量資料庫）。"""
    return build_chain()


def render_sidebar() -> None:
    with st.sidebar:
        st.title("RAG AI")
        st.caption("基於教材附錄 A 實作的問答系統")

        st.subheader("目前設定")
        st.markdown(
            f"""
- **Embedding**：`{config.embedding_model}`
- **Chat 模型**：`{config.chat_model}`
- **Re-ranker**：`{'on' if config.use_reranker else 'off'}`
- **檢索 K**：`{config.top_k_retrieve}`
- **切塊大小**：`{config.chunk_size}` chars
"""
        )

        st.divider()
        st.subheader("使用步驟")
        st.markdown(
            """
1. 將文件放到 `data/` 目錄
2. 跑 `python -m src.ingest` 建索引
3. 在右側輸入問題

支援：PDF、Markdown、TXT、Word
"""
        )

        if st.button("清除對話紀錄", use_container_width=True):
            st.session_state.messages = []
            st.rerun()


def render_message(role: str, content: str, sources: list | None = None) -> None:
    with st.chat_message(role):
        st.markdown(content)
        if sources:
            with st.expander(f"查看 {len(sources)} 個引用來源"):
                for i, doc in enumerate(sources, 1):
                    src = doc.metadata.get("source", "未知")
                    page = doc.metadata.get("page", "-")
                    st.markdown(f"**來源 {i}**：`{src}` (頁 {page})")
                    st.caption(doc.page_content[:400] + ("..." if len(doc.page_content) > 400 else ""))


def main() -> None:
    render_sidebar()

    st.title("RAG 智能問答")
    st.caption("輸入問題，AI 會基於你已索引的文件回答")

    if not config.has_openai_key:
        st.error("找不到 `OPENAI_API_KEY`。請複製 `.env.example` 為 `.env` 並填入 API Key。")
        st.stop()

    chain, retriever = get_rag_engine()

    if "messages" not in st.session_state:
        st.session_state.messages = []

    for msg in st.session_state.messages:
        render_message(msg["role"], msg["content"], msg.get("sources"))

    if question := st.chat_input("請輸入問題..."):
        st.session_state.messages.append({"role": "user", "content": question})
        render_message("user", question)

        with st.chat_message("assistant"):
            sources = retriever.invoke(question)

            placeholder = st.empty()
            full_answer = ""
            for token in chain.stream(question):
                full_answer += token
                placeholder.markdown(full_answer + " ▌")
            placeholder.markdown(full_answer)

            with st.expander(f"查看 {len(sources)} 個引用來源"):
                for i, doc in enumerate(sources, 1):
                    src = doc.metadata.get("source", "未知")
                    page = doc.metadata.get("page", "-")
                    st.markdown(f"**來源 {i}**：`{src}` (頁 {page})")
                    preview = doc.page_content[:400] + ("..." if len(doc.page_content) > 400 else "")
                    st.caption(preview)

        st.session_state.messages.append(
            {"role": "assistant", "content": full_answer, "sources": sources}
        )


if __name__ == "__main__":
    main()