LLMLLM API 개발 · 4중급

RAG 구현 실습 — ChromaDB로 문서 검색 시스템 만들기

LLMRAGChromaDB임베딩벡터DBPython

이론에서 코드로

LLM 기초 시리즈에서 RAG 개념을 배웠습니다. 이번에는 실제로 작동하는 RAG 시스템을 직접 구현합니다.

flowchart LR
    subgraph THIS["이번 편에서 구현"]
        direction TB
        I["문서 → 청킹 → 임베딩 → ChromaDB 저장"]
        Q["질문 → 임베딩 → 유사 문서 검색 → LLM 답변"]
    end

환경 설정

pip install openai chromadb python-dotenv
from openai import OpenAI
import chromadb
from chromadb.utils import embedding_functions

client = OpenAI()

# ChromaDB 로컬 클라이언트 (개발·테스트용)
chroma_client = chromadb.Client()

# OpenAI 임베딩 함수 설정
openai_ef = embedding_functions.OpenAIEmbeddingFunction(
    api_key=os.getenv("OPENAI_API_KEY"),
    model_name="text-embedding-3-small"  # 저렴하고 빠름
)

1단계: 문서 청킹

def chunk_text(text: str, chunk_size: int = 500, overlap: int = 50) -> list[str]:
    """
    텍스트를 chunk_size 크기로 분할, overlap만큼 앞뒤 중복
    """
    chunks = []
    start = 0
    
    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end]
        
        # 문장 경계에서 자르기 (마지막 마침표 찾기)
        if end < len(text):
            last_period = chunk.rfind(".")
            if last_period > chunk_size * 0.7:  # 70% 이상 위치면
                chunk = chunk[:last_period + 1]
        
        chunks.append(chunk.strip())
        start += len(chunk) - overlap
    
    return [c for c in chunks if c]  # 빈 청크 제거

# 테스트
sample_doc = """
파이썬은 1991년 귀도 반 로섬이 개발한 프로그래밍 언어입니다.
문법이 간결하고 읽기 쉬워 초보자부터 전문가까지 널리 사용됩니다.
데이터 분석, 머신러닝, 웹 개발 등 다양한 분야에서 활용됩니다.
NumPy, Pandas, TensorFlow 같은 강력한 라이브러리 생태계를 갖추고 있습니다.
"""

chunks = chunk_text(sample_doc, chunk_size=100, overlap=20)
for i, chunk in enumerate(chunks):
    print(f"청크 {i+1}: {chunk[:50]}...")

2단계: 문서 인덱싱 (ChromaDB 저장)

flowchart LR
    DOC["원본 문서"] --> CHUNK["청킹\n(500토큰 단위)"]
    CHUNK --> EMBED["OpenAI\n임베딩 API"]
    EMBED --> CHROMA["ChromaDB\n저장"]
def index_documents(docs: list[dict], collection_name: str = "my_docs"):
    """
    docs: [{"id": "doc1", "text": "...", "metadata": {...}}, ...]
    """
    collection = chroma_client.get_or_create_collection(
        name=collection_name,
        embedding_function=openai_ef
    )
    
    all_ids = []
    all_texts = []
    all_metadatas = []
    
    for doc in docs:
        chunks = chunk_text(doc["text"])
        
        for i, chunk in enumerate(chunks):
            all_ids.append(f"{doc['id']}_chunk_{i}")
            all_texts.append(chunk)
            all_metadatas.append({
                "source": doc["id"],
                "chunk_index": i,
                **doc.get("metadata", {})
            })
    
    # 배치 추가
    collection.add(
        ids=all_ids,
        documents=all_texts,
        metadatas=all_metadatas
    )
    
    print(f"✅ {len(all_texts)}개 청크 인덱싱 완료")
    return collection

# 실습: 회사 내부 문서 인덱싱
documents = [
    {
        "id": "policy_vacation",
        "text": """
        연차 휴가 정책: 입사 후 1년 미만은 월 1일 발생, 1년 이상은 15일.
        연차는 12월 31일 자동 소멸되며 이월되지 않습니다.
        연차 신청은 사용 3일 전 팀장 승인을 받아야 합니다.
        병가는 연간 10일이며 의사 소견서 제출 시 추가 5일 부여됩니다.
        """,
        "metadata": {"category": "HR", "updated": "2024-01"}
    },
    {
        "id": "policy_remote",
        "text": """
        재택근무 정책: 주 2회까지 허용, 금요일은 전 직원 출근.
        재택 근무 시 오전 9시까지 슬랙 체크인 필수.
        보안 상 이유로 카페 등 공공 WiFi 사용 금지.
        재택 장비는 회사 노트북만 사용 가능합니다.
        """,
        "metadata": {"category": "IT", "updated": "2024-03"}
    }
]

collection = index_documents(documents)

3단계: 유사도 검색

def search_documents(query: str, collection_name: str = "my_docs", top_k: int = 3):
    collection = chroma_client.get_collection(
        name=collection_name,
        embedding_function=openai_ef
    )
    
    results = collection.query(
        query_texts=[query],
        n_results=top_k,
        include=["documents", "metadatas", "distances"]
    )
    
    docs = results["documents"][0]
    metadatas = results["metadatas"][0]
    distances = results["distances"][0]
    
    return [
        {
            "text": doc,
            "metadata": meta,
            "similarity": 1 - dist  # 거리 → 유사도 변환
        }
        for doc, meta, dist in zip(docs, metadatas, distances)
    ]

# 테스트
results = search_documents("연차 며칠이나 받아요?")
for r in results:
    print(f"[유사도: {r['similarity']:.2f}] {r['text'][:60]}...")

4단계: RAG 질의응답

def rag_query(question: str, collection_name: str = "my_docs") -> dict:
    # 1. 관련 문서 검색
    relevant_docs = search_documents(question, collection_name, top_k=3)
    
    if not relevant_docs:
        return {"answer": "관련 문서를 찾을 수 없습니다.", "sources": []}
    
    # 2. 컨텍스트 구성
    context = "\n\n".join([
        f"[문서 {i+1} - 출처: {doc['metadata']['source']}]\n{doc['text']}"
        for i, doc in enumerate(relevant_docs)
    ])
    
    # 3. LLM에 질문
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "system",
                "content": """당신은 회사 내부 문서 기반 Q&A 어시스턴트입니다.
아래 제공된 문서 내용만을 근거로 답변하세요.
문서에 없는 내용은 "해당 정보가 문서에 없습니다"라고 답하세요."""
            },
            {
                "role": "user",
                "content": f"[참고 문서]\n{context}\n\n[질문]\n{question}"
            }
        ],
        temperature=0
    )
    
    answer = response.choices[0].message.content
    sources = list(set(doc["metadata"]["source"] for doc in relevant_docs))
    
    return {
        "answer": answer,
        "sources": sources,
        "relevant_docs": relevant_docs
    }

# 테스트
result = rag_query("재택근무는 주에 몇 번 할 수 있나요?")
print(f"답변: {result['answer']}")
print(f"출처: {', '.join(result['sources'])}")

전체 파이프라인 확인

flowchart TB
    subgraph INDEX["인덱싱 (1회)"]
        D["문서"] --> C["청킹"] --> E["임베딩"] --> DB["ChromaDB"]
    end

    subgraph QUERY["검색 (매 질문)"]
        Q["질문"] --> QE["질문 임베딩"]
        QE --> SIM["유사도 검색\n(Top-3)"]
        DB --> SIM
        SIM --> PROMPT["프롬프트\n조합"]
        PROMPT --> LLM["GPT-4o-mini"]
        LLM --> ANS["답변 + 출처"]
    end

    INDEX --> QUERY

검색 품질 개선: 하이브리드 검색

벡터 검색만으로는 정확한 키워드 매칭이 약합니다.

from rank_bm25 import BM25Okapi  # pip install rank-bm25

def hybrid_search(query: str, all_chunks: list[str], top_k: int = 5):
    # BM25 키워드 검색
    tokenized = [doc.split() for doc in all_chunks]
    bm25 = BM25Okapi(tokenized)
    bm25_scores = bm25.get_scores(query.split())
    
    # 벡터 검색 결과와 점수 결합 (RRF: Reciprocal Rank Fusion)
    vector_results = search_documents(query, top_k=top_k * 2)
    
    combined_scores = {}
    for rank, doc in enumerate(vector_results):
        combined_scores[doc["text"]] = combined_scores.get(doc["text"], 0) + 1 / (rank + 60)
    
    return sorted(combined_scores.items(), key=lambda x: x[1], reverse=True)[:top_k]

실전 고려사항

항목개발 단계프로덕션
벡터 DBChromaDB (로컬)Pinecone, Weaviate
임베딩 모델text-embedding-3-small동일 또는 자체 모델
청크 크기500~800 토큰도메인에 맞게 튜닝
Top-K3~55~10 (리랭킹 후 3)
캐싱없음Redis로 동일 질문 캐시

정리

단계코드
청킹chunk_text(doc, size=500, overlap=50)
인덱싱collection.add(ids, documents, metadatas)
검색collection.query(query_texts, n_results)
RAG 조합검색 결과 → 프롬프트 컨텍스트 → LLM 답변

다음 편에서는 완성 프로젝트 — 지금까지 배운 모든 것을 합쳐서 실제 AI 챗봇 서비스를 만듭니다.

궁금한 점이 있으신가요?

협업·의뢰는 아래로, 가벼운 소통은 인스타그램 @bluefox._.hi도 환영이에요.