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