LLM & Generative AI/RAG in Practice

[LangChain RAG 구축 시리즈 Ep.17] 🔍 ChromaDB에서 문서 검색 및 Retriever 구성하기

ygtoken 2025. 4. 5. 21:33
728x90

이 글에서는 이전에 OpenAI 임베딩으로 저장한 벡터 데이터를 기반으로
사용자의 질문과 유사한 문서를 검색할 수 있는
Retriever(검색기) 를 구성해보겠습니다.


🎯 목표

  • ChromaDB에 저장된 벡터 데이터를 불러오기
  • 사용자의 자연어 질문을 임베딩하여 유사 문서 검색
  • Retriever 객체를 통해 top-k 검색 구현

🗂️ 디렉토리 구조 전제

rag-iceberg-project/
├── chroma_db/                         ← 벡터 저장소
└── src/search_documents.py           ← 검색기 구성 코드

🛠️ Step 1. ChromaDB에서 벡터 불러오기 + 검색기 구성

# src/search_documents.py

from langchain.vectorstores import Chroma                  # Chroma 벡터 DB
from langchain.embeddings import OpenAIEmbeddings          # OpenAI 임베딩 모델

def get_retriever(db_path: str, collection_name: str, k: int = 2):
    """
    저장된 ChromaDB 벡터 데이터를 기반으로 문서 검색기를 생성합니다.
    """

    # 1. OpenAI 임베딩 모델 생성
    embedding = OpenAIEmbeddings()

    # 2. ChromaDB 벡터 저장소 로드
    vectordb = Chroma(
        persist_directory=db_path,         # 저장된 벡터 위치
        embedding_function=embedding,      # 동일한 임베딩 모델 사용
        collection_name=collection_name    # 저장된 컬렉션 이름
    )

    # 3. Retriever 객체 반환 (top-k 검색 설정)
    retriever = vectordb.as_retriever(
        search_kwargs={"k": k}             # 유사도 상위 k개 문서 반환
    )

    # 4. 검색기 반환
    return retriever

🧪 Step 2. 사용자 질문으로 검색 테스트

# main.py (또는 테스트 스크립트에서 실행)

from src.search_documents import get_retriever  # 검색기 생성 함수 임포트

# ✅ 벡터 DB 정보 설정
db_path = "chroma_db"                      # 벡터 저장소 위치
collection_name = "iceberg_tables"         # 컬렉션 이름

# ✅ 검색기 생성
retriever = get_retriever(db_path, collection_name, k=2)

# ✅ 사용자 질문 정의
query = "상품 정보를 담고 있는 테이블이 뭐야?"

# ✅ 관련 문서 검색
results = retriever.get_relevant_documents(query)

# ✅ 결과 출력
print("📄 검색된 문서 목록:")
for i, doc in enumerate(results):
    print(f"[{i+1}] {doc.page_content.strip()}\n")

📤 출력 예시

📄 검색된 문서 목록:
[1] 📂 테이블명: products
🧾 컬럼 목록:
- product_id (long)
- product_name (string)
- category (string)
🧩 파티션 컬럼: category

✅ 사용자의 질문이 "products" 테이블과 관련 있다는 것을 벡터 유사도를 기반으로 정확히 찾아냅니다.


💡 참고: Retriever 내부 옵션들

 

옵션 키 설명
k top-k 유사 문서 개수
search_type 기본값: "similarity"
score_threshold 일정 유사도 이상일 경우만 반환 (선택)

 

예시:

retriever = vectordb.as_retriever(
    search_kwargs={
        "k": 3,
        "score_threshold": 0.3  # 점수 낮은 문서는 제외
    }
)

📎 요약 및 핵심 정리

  • 저장된 벡터 데이터를 Chroma에서 불러와 Retriever를 구성했습니다.
  • 사용자의 자연어 질문을 벡터화한 뒤, 유사도가 높은 문서를 top-k 방식으로 반환합니다.
  • 이 검색기는 LLM과 연결되는 RetrievalQA 체인의 핵심 요소입니다.

 

728x90