8. 오늘의 완성 코드

from indexer import get_store
TOP_K      = 3
MIN_SCORE  = 0.47
SEARCH_TYPE = "mmr"           # similarity / mmr
_store = get_store()

def get_retriever(k=TOP_K, search_type=SEARCH_TYPE, category=None):
    kwargs = {"k": k}
    if search_type == "mmr":
        kwargs.update({"fetch_k": k * 4, "lambda_mult": 0.5})
    if category:
        kwargs["filter"] = {"category": category}
    return _store.as_retriever(
        search_type=search_type, search_kwargs=kwargs)

def search(question, k=TOP_K, category=None):
    # 점수 필터가 필요하므로 store를 직접 쓰되, 형식은 통일
    if category:
        pairs = _store.similarity_search_with_relevance_scores(
            question, k=k, filter={"category": category})
    else:
        pairs = _store.similarity_search_with_relevance_scores(
            question, k=k)
    good = []
    for d, s in pairs:
        if s >= MIN_SCORE:
            good.append(d)
    return good

def build_context(docs):
    parts = []
    for i, d in enumerate(docs, 1):
        head = (f"[{i}] {d.metadata.get('filename','?')} "
                f"p.{d.metadata.get('page_no','?')}")
        parts.append(f"{head}\n{d.page_content}")
    return "\n\n---\n\n".join(parts)

if __name__ == "__main__":
    docs = search("환불과 교환의 차이는?")
    print(f"검색 결과 {len(docs)}개\n")
    print(build_context(docs))

# ============================================================
# 11차시 | PPT: 11차시_Retriever기초실습.pptx
# 역할: similarity와 MMR 검색, k별 점수, 점수 기준 필터, 자료 문자열 만들기를 비교
# 입력: 10차시 FAISS 인덱스, 환불/교환 관련 질문
# 출력: 검색기별 미리보기, k별 점수, 기준을 통과한 조각을 이어 붙인 자료
# 흐름:
#   get_store() -> FAISS
#     -> retriever.invoke 또는 similarity_search_with_relevance_scores
#     -> 점수 필터 -> build_context
# ============================================================
# retriever: 질문을 넣어 가까운 문서 조각을 받는 검색기
# FAISS: 벡터(숫자 목록)를 저장한 인덱스
# embedding: 질문 문장을 그 숫자 목록으로 바꾸는 과정 (검색 때 네트워크 호출)
# vector: embedding이 만든 숫자 목록
# chunk: 인덱스 안의 문서 조각. metadata는 쪽 번호, 파일명 같은 부가 정보
# citation: 뒤 차시에서 답변 문장 끝의 [1] 같은 근거 번호.
#   build_context가 붙이는 [1], [2]가 그 번호의 출발점입니다.

"""목표: FAISS 인덱스의 Retriever 검색 방식과 결과 후처리를 비교합니다.

similarity/MMR 검색, k별 점수 비교, 점수 임계값 필터링, context 포매팅을
확인하는 학습용 예제입니다.
"""

import os
import sys
import warnings

# indexer.py가 src/10에 있으므로 현재 파일 기준으로 모듈 검색 경로에 추가합니다.
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '10'))

# FAISS 인덱스를 로드하거나, 인덱스가 없을 때 생성하는 함수를 가져옵니다.
from indexer import get_store

# relevance score 범위 경고는 이 예제 출력에서 숨깁니다.
warnings.filterwarnings("ignore", message="Relevance scores must be between 0 and 1")

# ========== 설정 ==========
# 기본 검색에서 반환할 문서 조각 수입니다.
TOP_K      = 3
# relevance score가 이 값보다 낮은 검색 결과는 제외합니다.
MIN_SCORE  = 0.1
# 기본 검색 방식입니다. similarity 또는 mmr을 사용할 수 있습니다.
SEARCH_TYPE = "mmr"  # similarity / mmr

# 검색 함수들이 공통으로 사용하는 FAISS vector store를 준비합니다.
_store = get_store()

# ========== Part 1. 기본 함수 ==========
def show(name, retriever):
    """질문을 검색기에 전달하고 결과 수, 페이지, 본문 미리보기를 출력합니다."""
    # 이 예제에서 검색할 질문을 지정합니다.
    question = "환불과 교환은 어떻게 다른가요?"
    # Retriever를 실행해 관련 문서 조각을 가져옵니다.
    docs = retriever.invoke(question)
    # 검색기 이름과 반환된 문서 개수를 표시합니다.
    print(f"■ {name}  ({len(docs)}개)")
    # 각 문서의 줄바꿈을 정리하고 앞부분만 출력합니다.
    for d in docs:
        preview = d.page_content[:55].replace("\n", " ")         
        # 문서 metadata의 페이지 번호와 본문 미리보기를 표시합니다.
        print(f"   p.{d.metadata['page_no']}  {preview}...")     
    # 다음 검색 결과와 구분할 빈 줄을 출력합니다.
    print()

def compare_k(question, ks=(1, 3, 5)):
    """k값을 바꾸며 검색 점수와 임계값 미달 개수를 비교합니다."""
    # 비교 대상 질문을 표시합니다.
    print(f"질문: {question}")
    # 질문과 결과 영역을 구분하는 선을 출력합니다.
    print("=" * 58)
    # 각 k값에 대해 별도의 유사도 검색을 수행합니다.
    for k in ks:
        pairs = _store.similarity_search_with_relevance_scores(question, k=k)
        # (문서, 점수) 쌍에서 점수만 모읍니다.
        scores = [s for _, s in pairs]
        # MIN_SCORE보다 낮은 결과의 개수를 계산합니다.
        low = sum(1 for s in scores if s < MIN_SCORE)
        # k별 점수와 기준 미달 문서 수를 보기 좋게 출력합니다.
        print(f"k={k:<3} 점수={[round(float(s), 3) for s in scores]}")
        print(f"     기준 미달({MIN_SCORE}) 조각: {low}개")
    # 다음 질문의 결과와 구분할 빈 줄을 출력합니다.
    print()

# ========== Part 2. 심화 함수 ==========
def get_retriever(k=TOP_K, search_type=SEARCH_TYPE, category=None):
    """k, 검색 방식, 선택적 category 필터를 적용한 Retriever를 만듭니다."""
    # 검색 결과 개수를 지정합니다.
    kwargs = {"k": k}
    # MMR은 더 넓은 후보군에서 관련성과 다양성을 함께 고려합니다.
    if search_type == "mmr":
        kwargs.update({"fetch_k": k * 4, "lambda_mult": 0.5})
    # category가 지정되면 해당 metadata 값을 가진 문서만 검색합니다.
    if category:
        kwargs["filter"] = {"category": category}
    # 구성한 옵션으로 LangChain Retriever 객체를 반환합니다.
    return _store.as_retriever(
        search_type=search_type, 
        search_kwargs=kwargs
    )

def search(question, k=TOP_K, category=None):
    """유사도 검색 후 relevance score 기준을 통과한 문서만 반환합니다."""
    # category가 있으면 해당 metadata로 먼저 후보 문서를 제한합니다.
    if category:
        pairs = _store.similarity_search_with_relevance_scores(
            question, k=k, filter={"category": category}
        )
    # 필터가 없으면 전체 인덱스에서 상위 k개를 검색합니다.
    else:
        pairs = _store.similarity_search_with_relevance_scores(question, k=k)
    
    # 검색 결과 중 relevance score가 MIN_SCORE 이상인 문서만 남깁니다.
    good = []
    for d, s in pairs:
        if s >= MIN_SCORE:
            good.append(d)
    # 필터링을 마친 문서 목록을 반환합니다.
    return good

def build_context(docs):
    """문서 목록을 출처와 페이지가 포함된 RAG context 문자열로 변환합니다."""
    # 포매팅한 문서 조각을 순서대로 저장합니다.
    parts = []
    # 문서마다 번호, 파일명, 페이지 번호를 제목으로 붙입니다.
    for i, d in enumerate(docs, 1):
        head = (f"[{i}] {d.metadata.get('filename','?')} "
                f"p.{d.metadata.get('page_no','?')}")
        # 제목과 본문을 하나의 context 조각으로 만듭니다.
        parts.append(f"{head}\n{d.page_content}")
    # 여러 조각 사이에 구분선을 넣어 하나의 문자열로 합칩니다.
    return "\n\n---\n\n".join(parts)

# ========== 테스트 ==========
if __name__ == "__main__":
    # 직접 실행한 경우에만 아래 검색 예제를 실행합니다.
    # Part 1에서는 기본 similarity와 MMR 결과를 비교합니다.
    print("=" * 58)
    print("Part 1. Retriever 기본 개념")
    print("=" * 58 + "\n")
    
    # similarity는 질문과 가까운 문서 조각 3개를 반환합니다.
    show("similarity k=3",
         _store.as_retriever(search_kwargs={"k": 3}))
    
        # MMR은 fetch_k 후보 중 관련성과 다양성을 고려해 3개를 선택합니다.
    show("MMR k=3 fetch_k=10",
         _store.as_retriever(search_type="mmr",
             search_kwargs={"k": 3, "fetch_k": 10}))
    
    # 같은 질문에 대해 k=1, 3, 5일 때 관련도 점수를 비교합니다.
    print("=" * 58)
    print("k값에 따른 유사도 점수")
    print("=" * 58 + "\n")
    
    compare_k("환불은 며칠 이내인가요?")
    compare_k("환불과 교환은 어떻게 다른가요?")
    
    # Part 2에서는 앞에서 정의한 점수 필터와 context 포매팅 함수를 사용합니다.
    print("=" * 58)
    print("Part 2. 통합 검색 함수")
    print("=" * 58 + "\n")
    
    # 질문과 관련된 문서를 검색하고 MIN_SCORE 기준을 적용합니다.
    docs = search("환불과 교환의 차이는?")
    # 검색 후 남은 문서 조각의 개수를 표시합니다.
    print(f"검색 결과: {len(docs)}개\n")
    # 결과가 있으면 LLM 입력용 context를 출력하고, 없으면 안내 문구를 표시합니다.
    if docs:
        print(build_context(docs))
    else:
        print("기준 이상의 관련 문서가 없습니다.")