실습 1 — Retriever 기초 실습 실습 따라 하기

'''
번외편
이 파일은 AI 적용 사례를 위한 RAG(Retrieval-Augmented Generation) 프로젝트의 src/11/retriever_fast.py입니다.
속도를 높이기 위해 FAISS 인덱스를 로드하고, 질문에 대한 유사 문서를 검색하는 기능을 제공합니다.

'''
# ============================================================
# 11차시 | PPT: 11차시_Retriever기초실습.pptx
# 역할: 만들어 둔 FAISS 인덱스를 찾아 열고, 질문과 가까운 조각 3개를 출력
# 입력: src/11/faiss_index 또는 src/10/faiss_index
# 출력: 로드 시간, 검색 시간, 질문, 결과 개수, 순위·쪽·미리보기
# 흐름:
#   index.faiss가 있는 폴더 -> FAISS.load_local
#     -> similarity_search -> print
# ============================================================
# FAISS: 벡터(숫자 목록) 인덱스. embedding은 질문 문장을 그 숫자로 바꿉니다.
# vector: embedding 결과인 숫자 목록
# metadata: page_no, page, source 같은 부가 정보

# Path는 경로를 / 로 이어 붙일 수 있는 경로 객체입니다.
from pathlib import Path
from time import perf_counter

# .env를 읽는 도구, FAISS 저장소, OpenAI 임베딩을 가져옵니다.
from dotenv import load_dotenv
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings

# resolve()로 절대 경로를 만든 뒤 parent로 이 파일이 있는 폴더(src/11)를 얻습니다.
CURRENT_DIR = Path(__file__).resolve().parent
# 두 단계 위가 프로젝트 폴더입니다. .env 파일이 여기 있습니다.
PROJECT_DIR = CURRENT_DIR.parent.parent
# 11번 폴더를 먼저 보고, 없으면 10번 폴더의 인덱스를 씁니다.
INDEX_CANDIDATES = (
    CURRENT_DIR / "faiss_index",
    CURRENT_DIR.parent / "10" / "faiss_index",
)
# 후보를 차례로 보며 index.faiss가 있는 첫 폴더를 고릅니다. 하나도 없으면 None입니다.
INDEX_DIR = next(
    (path for path in INDEX_CANDIDATES if (path / "index.faiss").is_file()),
    None,
)
# 인덱스를 만들 때 쓴 임베딩 모델과 같아야 검색 결과가 맞습니다.
EMBED_MODEL = "text-embedding-3-small"
QUESTION = "환불과 교환은 어떻게 다른가요?"
RESULT_COUNT = 3

def load_store():
    # 인덱스가 없으면 무엇을 확인할지 알려 주는 오류를 냅니다.
    if INDEX_DIR is None:
        raise FileNotFoundError(
            "FAISS 인덱스를 찾을 수 없습니다. src/10 또는 src/11의 "
            "faiss_index 폴더를 확인하세요."
        )

    # .env와 인덱스 경로를 현재 작업 폴더가 아닌 프로젝트 파일 위치 기준으로 찾습니다.
    load_dotenv(PROJECT_DIR / ".env")
    # 질문을 벡터로 바꿀 임베딩 객체입니다. 이때부터 API 키가 필요합니다.
    embeddings = OpenAIEmbeddings(model=EMBED_MODEL)

    # 기존 인덱스만 로드해 매번 PDF 분할과 임베딩을 다시 하지 않습니다.
    # allow_dangerous_deserialization=True는 내가 만든 인덱스 파일이라 믿고 연다는 표시입니다.
    return FAISS.load_local(
        str(INDEX_DIR),
        embeddings,
        allow_dangerous_deserialization=True,
    )

def main():
    # perf_counter()는 초 단위 시계입니다. 시작 시각을 적어 두었다가 끝난 시각에서 뺍니다.
    load_started = perf_counter()
    store = load_store()
    # :.2f 는 소수점 둘째 자리까지 보여 줍니다.
    print(f"인덱스 로드: {perf_counter() - load_started:.2f}초 ({INDEX_DIR})")

    search_started = perf_counter()
    # 검색을 한 번만 수행해 질문 임베딩 API 요청도 한 번으로 제한합니다.
    documents = store.similarity_search(QUESTION, k=RESULT_COUNT)
    print(f"검색 및 질문 임베딩: {perf_counter() - search_started:.2f}초")
    print(f"질문: {QUESTION}\n검색 결과: {len(documents)}개")

    # enumerate(..., start=1)로 순위를 1부터 붙입니다.
    for rank, document in enumerate(documents, start=1):
        # page_no가 없으면 page를, 그것도 없으면 '-'를 씁니다.
        page = document.metadata.get("page_no", document.metadata.get("page", "-"))
        # split()과 join으로 줄바꿈과 여러 칸 공백을 한 칸으로 줄인 뒤 앞 120자만 자릅니다.
        preview = " ".join(document.page_content.split())[:120]
        print(f"{rank}. 페이지 {page}: {preview}")

# 직접 실행할 때만 main()을 부릅니다. retriever_test_fast.py는 import해서 load_store만 씁니다.
if __name__ == "__main__":
    main()