실습 5 — 문서 챗봇 웹 화면

from pathlib import Path

APP_CODE = r'''
import re
import zlib
from pathlib import Path

import streamlit as st
from langchain_classic.chains import RetrievalQA
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.vectorstores import FAISS
from langchain_core.embeddings import Embeddings
from langchain_core.language_models.fake_chat_models import FakeListChatModel
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 앱 파일과 같은 폴더의 약관 PDF. 어디서 실행해도 같은 파일을 찾게 앱 파일의 자리를 기준으로 적는다
PDF = Path(__file__).parent / "travel_terms.pdf"
# 쪽마다 맨 위에 붙은 머리글. 조각마다 섞여 검색을 흐리므로 지운다
HEADER = "구름다리여행 여행 약관 (연습용 가상 문서)"


class BigramEmbeddings(Embeddings):
    """글자 두 개 묶음(바이그램)을 세어 숫자 목록을 만드는 연습용 임베딩"""

    def __init__(self, size=512):
        self.size = size

    def _vector(self, text):
        text = text.replace(" ", "").replace("\n", "")
        vec = [0.0] * self.size
        for i in range(len(text) - 1):
            # 글자 두 개를 crc32 로 0~511 사이의 칸 번호로 바꿔 그 칸을 1 올린다
            vec[zlib.crc32(text[i:i + 2].encode("utf-8")) % self.size] += 1.0
        norm = sum(v * v for v in vec) ** 0.5 or 1.0
        return [v / norm for v in vec]

    def embed_documents(self, texts):
        return [self._vector(t) for t in texts]

    def embed_query(self, text):
        return self._vector(text)


@st.cache_resource
def docbot_build_log():
    """색인을 만든 기록. 서버가 켜져 있는 동안 하나만 있다"""
    return []


# show_spinner: 처음 만드는 동안 화면에 보일 기다림 글
@st.cache_resource(show_spinner="약관으로 색인을 만드는 중입니다...")
def load_docbot():
    """6차시의 문서 검색: PDF 읽기 → 머리글 지우기 → 조각 내기 → FAISS 색인 → RetrievalQA"""
    docbot_build_log().append("색인")
    pages = PyPDFLoader(str(PDF)).load()
    for page in pages:
        page.page_content = page.page_content.replace(HEADER, "").strip()
    chunks = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=30).split_documents(pages)
    db = FAISS.from_documents(chunks, BigramEmbeddings())
    # 답을 쓰는 자리의 가짜 모델. 미리 적어 둔 답을 차례로 말한다
    llm = FakeListChatModel(responses=[
        "출발 3일 전부터 6일 전까지는 상품 금액의 30%를 취소 수수료로 뗍니다. 5일 전이면 30%입니다.",
        "반려동물은 '함께 여행' 메뉴의 동반 상품에서만 데려갈 수 있고, 7kg 이하의 개와 고양이만 됩니다.",
    ])
    qa = RetrievalQA.from_chain_type(
        llm=llm,
        retriever=db.as_retriever(search_kwargs={"k": 2}),   # 가까운 조각 2개를 찾는다
        return_source_documents=True,                         # 답과 함께 근거 조각을 돌려준다
    )
    return qa, len(chunks)


def source_label(doc):
    """근거 조각 하나를 '2쪽 · 제5조, 제6조' 꼴로 줄인다"""
    # re.findall: 글에서 '제(숫자)조' 꼴을 모두 찾는다. dict.fromkeys 로 겹친 것을 빼며 순서를 지킨다
    names = list(dict.fromkeys(re.findall(r"제\d+조", doc.page_content)))
    return f"{doc.metadata['page'] + 1}쪽 · {', '.join(names)}"


def show(message):
    """말풍선 하나를 그린다. 답에 근거가 있으면 접어 둔 칸에 넣는다"""
    with st.chat_message(message["role"]):
        st.markdown(message["content"])
        if message.get("sources"):
            # st.expander: 눌러야 펼쳐지는 접힌 칸
            with st.expander("근거 조항 보기"):
                for s in message["sources"]:
                    st.caption(s)


qa, n_chunks = load_docbot()

st.title("구름다리여행 약관 챗봇")
with st.sidebar:
    st.header("문서")
    st.caption(f"약관 조각: {n_chunks}개")
    st.caption(f"색인을 만든 횟수: {len(docbot_build_log())}")

if "doc_messages" not in st.session_state:
    st.session_state.doc_messages = []
for m in st.session_state.doc_messages:
    show(m)

if question := st.chat_input("약관에 대해 물어보세요"):
    user = {"role": "user", "content": question}
    show(user)
    # st.spinner: 기다리는 동안 돌아가는 표시와 글을 보인다
    with st.spinner("약관에서 찾는 중입니다..."):
        result = qa.invoke({"query": question})
    bot = {"role": "assistant", "content": result["result"],
           "sources": [source_label(d) for d in result["source_documents"]]}
    show(bot)
    st.session_state.doc_messages.extend([user, bot])
'''

Path("out_ch07_docbot.py").write_text(APP_CODE.lstrip(), encoding="utf-8")
print("저장: out_ch07_docbot.py,", len(APP_CODE.strip().splitlines()), "줄")

import streamlit as st
from streamlit.testing.v1 import AppTest

st.cache_resource.clear()

# 처음 켤 때 PDF 를 읽고 색인을 만들므로 기다리는 시간을 30초로 넉넉히 준다
at = AppTest.from_file("out_ch07_docbot.py", default_timeout=30).run()
print("제목:", at.title[0].value)
for c in at.sidebar.caption:
    print(c.value)
print("오류:", len(at.exception), "개")

at.chat_input[0].set_value("출발 5일 전에 취소하면 수수료가 얼마예요?").run()

bot = at.chat_message[-1]
print("답:", bot.markdown[0].value)
# 답 말풍선 안의 접힌 칸에서 근거를 꺼낸다
print("접힌 칸:", bot.expander[0].label)
for c in bot.expander[0].caption:
    print("  근거:", c.value)

at.chat_input[0].set_value("반려동물 데려가도 돼요?").run()

bot = at.chat_message[-1]
print("답:", bot.markdown[0].value)
print("근거:", [c.value for c in bot.expander[0].caption])
print("말풍선:", len(at.chat_message), "개 / 기록:", len(at.session_state["doc_messages"]), "개")
# 스크립트는 세 번 돌았지만 색인은 몇 번 만들었을까
print([c.value for c in at.sidebar.caption][-1])