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

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"  # 저렴하고 빠름
)

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]}...")

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)

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]}...")

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'])}")

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]