📌 벡터 DB 저장하기: ChromaDB 활용

LLM 기반의 RAG(Retrieval-Augmented Generation) 시스템을 구축할 때, 대화 데이터를 효율적으로 저장하고 검색하기 위해 벡터 데이터베이스(Vector DB)를 사용한다. 이번 글에서는 ChromaDB를 활용하여 벡터 DB에 데이터를 저장하는 방법을 정리한다. 😭👩‍💻

📌 RAG 시스템에서 벡터 데이터베이스 선택하기

RAG(Retrieval-Augmented Generation) 시스템을 구축하려면, 검색 속도를 높이고 효율적인 문서 검색을 위해 벡터 데이터베이스(Vector DB)가 필수다.

📊 벡터 데이터베이스 선택 기준

✔ 데이터 규모 → 대량 데이터를 처리할 수 있는지
✔ 검색 성능 → 빠른 유사도 검색이 가능한지
✔ 확장성 → 클라우드 배포와 확장이 가능한지
✔ 호환성 → LangChain과 같은 프레임워크와 잘 통합되는지

🔎 주요 벡터 데이터베이스 비교

데이터베이스	특징	장점	단점
FAISS	Facebook AI 개발, 빠른 근사 최근접 이웃 검색	오픈소스, 빠른 검색 속도	클러스터링 기능 부족
ChromaDB	LangChain과 호환되는 벡터DB	간편한 사용, 오픈소스	대규모 데이터 처리 한계
Pinecone	클라우드 기반, AI 검색 최적화	확장성 뛰어남	유료 모델
Elasticsearch	키워드 + 벡터 검색 지원	복합 검색 가능	설정이 복잡

👉 내가 선택한 데이터베이스: ChromaDB

• LangChain과 호환성이 뛰어나고,
• 무료로 사용할 수 있으며,
• 구현이 간단해 초기 프로젝트에 적합하다.

⚠ 하지만 이후 프로젝트 확장 시 다른 DB로 교체할 수도 있음

1️⃣ 벡터 DB 설정하기

우선 ChromaDB를 사용할 준비를 해야 한다. 벡터화를 수행하기 위해 Ollama 임베딩 모델을 활용하며, 데이터 저장을 위한 디렉토리를 생성한다.

# Chroma 벡터 DB 설정
CHROMA_DB_DIR = "vectorstore"
embeddings = OllamaEmbeddings(model="llama3.1-instruct-8b:latest")
os.makedirs(CHROMA_DB_DIR, exist_ok=True)

✅ 설정 이유

• CHROMA_DB_DIR : 벡터 DB 데이터를 저장할 디렉토리 생성
• OllamaEmbeddings : 벡터화를 위한 임베딩 모델 지정
• os.makedirs() : 폴더가 없으면 자동 생성

---

2️⃣ 벡터 DB에 데이터 저장하기

메시지 데이터를 벡터화하여 ChromaDB에 저장하는 함수를 만들어 보자.

# 벡터 DB 저장 함수
def save_to_vector_db(messages):
    try:
        # ChromaDB 인스턴스 생성
        vector_db = Chroma(embedding_function=embeddings, persist_directory=CHROMA_DB_DIR)
        
        # 메시지를 하나의 텍스트로 결합하여 벡터화
        combined_text = " ".join(messages)
        vectors = embeddings.embed_documents([combined_text])

        # 고유 ID 생성 (문서 내용의 해시값 사용)
        doc_id = str(hash(combined_text))
        
        # Document 객체 생성
        document = Document(page_content=combined_text, metadata={"id": doc_id})
        
        # ChromaDB에 저장
        vector_db.add_documents(
            documents=[document],
            embeddings=vectors,
            ids=[doc_id]
        )
        
        logger.info("벡터 DB에 메시지 저장 성공")
    except Exception as e:
        logger.error(f"벡터 DB 저장 중 오류 발생: {e}")

✅ 핵심 개념

• 벡터화(Embedding): 텍스트 데이터를 수치 벡터로 변환
• 고유 ID: 해시값을 사용하여 중복 저장 방지
• 메타데이터: 문서 정보를 함께 저장하여 검색 시 활용

---

3️⃣ 저장된 데이터 확인하기

데이터가 정상적으로 저장되었는지 확인하는 방법이다.

save_to_vector_db(input_data.messages)  # 메시지 저장 실행

📌 저장된 데이터 확인 방법

logger.info(f"Document saved: {doc_id}, Message: {message}")
documents = vector_db.get()
logger.info(f"전체 조회: {documents}")

✅ 벡터 DB 조회 방법

• 저장된 문서의 ID와 내용을 확인 가능
• 전체 벡터 데이터를 로그로 출력하여 저장 상태 확인 가능

---

4️⃣ 확장성 고려: 메타데이터 추가

추후 검색과 관리를 용이하게 하기 위해, 문서 타입(텍스트, PDF, 웹 데이터 등)을 추가하여 저장하는 것이 좋다.

# 메타데이터 추가 (문서 타입 및 출처)
metadata = {
    "id": doc_id,
    "type": document_type,  # 'message', 'pdf', 'web'
    "source": "user_input"
}

✅ 메타데이터 추가 이유

• 문서 구분 가능 (텍스트, PDF, 웹 데이터 등)
• 출처 정보 추가 (어떤 유형의 데이터인지 기록)

---

5️⃣ OCR 데이터 저장 확인

OCR로 변환된 이미지 데이터를 벡터 DB에 저장할 경우, 아래와 같은 방식으로 처리된다.

✅ OCR 데이터 관리 방식

• 이미지에서 추출된 텍스트를 벡터화하여 저장
• 일반 텍스트와 동일한 방식으로 검색 가능

---

6️⃣ API에서 벡터 DB 활용하기

FastAPI를 사용하여, 벡터 DB를 대화형 API에 통합할 수 있다.

# 벡터 DB를 FastAPI 엔드포인트에 추가
async def chat(input: str = Form(...), file: Optional[UploadFile] = File(None),  document_type: str = "message", vector_db: Chroma = Depends(get_vector_db)):

📌 FastAPI 의존성 주입 (Dependency Injection)

# 벡터 DB 초기화 함수
def get_vector_db():
    global vector_db
    if vector_db is None:
        vector_db = Chroma(embedding_function=embeddings, persist_directory=CHROMA_DB_DIR)
        logger.info("벡터 DB 연결 완료")
    return vector_db

✅ 이점

• FastAPI에서 의존성 주입(Dependency Injection)을 통해 ChromaDB 객체를 관리 가능
• API 호출 시 자동으로 ChromaDB 연결 및 데이터 검색 가능

---

🚀 최종 구현 정리

✅ 현재까지 구현된 기능

• 벡터 DB 저장: 메시지를 벡터화하여 ChromaDB에 저장
• 메타데이터 추가: 문서 유형 및 출처 정보를 함께 저장
• API 연동: FastAPI를 활용하여 벡터 DB를 활용한 대화형 서비스 구축
• 로그 관리: 저장된 벡터 데이터 확인 가능

✅ 다음 목표

• 웹 페이지 및 PDF 내용을 벡터 DB에 저장하는 기능 추가
• 벡터 DB 최적화 및 검색 성능 향상
• PostgreSQL과 연동하여 AI 응답을 저장 및 관리

⚡ 다음 단계로 넘어가면서 최적화 및 검색 기능을 강화해보자....!!