[AI/LLM] 직접 만들어보는 RAG — Chroma와 pgvector 실습편

지난 글에서 RAG, 임베딩, 벡터DB가 뭔지 개념을 잡았다. 이제 손으로 직접 만들어볼 차례다. 오늘은 아주 작은 RAG를 파이썬으로 처음부터 끝까지 만들어본다. 벡터DB는 두 가지, Chroma와 pgvector로 각각 구현해본다.

• Chroma: 설치가 간단하고 파이썬만으로 바로 돌아간다. 입문/프로토타이핑에 좋다.
• pgvector: 이미 쓰던 PostgreSQL에 얹는 방식이다. 실무에서 기존 DB랑 같이 쓰기 좋다.

같은 걸 두 방식으로 만들어보면 "벡터DB가 바뀌어도 RAG 흐름은 똑같다"는 걸 몸으로 알게 된다.

---

준비물

임베딩과 답변 생성에는 OpenAI API를 쓴다. (다른 임베딩 모델을 써도 흐름은 같다.) 먼저 공통 패키지부터 깔자.

pip install openai

그리고 API 키를 환경변수로 준비해둔다.

export OPENAI_API_KEY="sk-..."

실습에 쓸 예시 문서도 몇 개 준비하자. 여기서는 간단하게 리스트로 넣는다.

documents = [
    "우리 회사는 연차를 입사 첫해부터 15일 제공한다.",
    "퇴사 시 미사용 연차는 수당으로 정산해 지급한다.",
    "재택근무는 주 2회까지 신청 가능하다.",
    "점심 식대는 월 15만원까지 지원한다.",
]

---

공통: 임베딩 함수 만들기

RAG의 핵심은 "텍스트를 벡터로 바꾸는 것"이다. Chroma든 pgvector든 이 부분은 똑같이 쓴다.

from openai import OpenAI

client = OpenAI()

def embed(text: str) -> list[float]:
    res = client.embeddings.create(
        model="text-embedding-3-small",
        input=text,
    )
    return res.data[0].embedding

text-embedding-3-small 모델은 1536차원 벡터를 돌려준다. 즉 문장 하나가 숫자 1536개짜리 배열로 바뀐다. 이 숫자 뒤에 pgvector 테이블 만들 때 차원 수가 필요하니 기억해두자.

---

1부. Chroma로 만들기

설치

pip install chromadb

끝이다. 별도 서버 없이 파이썬 안에서 바로 돌아간다.

문서 저장하기 (준비 단계)

import chromadb

chroma_client = chromadb.Client()
collection = chroma_client.create_collection(name="company_docs")

# 문서를 임베딩해서 저장
for i, doc in enumerate(documents):
    collection.add(
        ids=[str(i)],
        embeddings=[embed(doc)],
        documents=[doc],
    )

collection.add()에 벡터(embeddings)와 원본 텍스트(documents)를 같이 넣는다. 나중에 검색하면 비슷한 벡터를 찾아서 원본 텍스트까지 돌려준다.

질문하고 검색하기

question = "퇴사하면 안 쓴 연차 어떻게 돼?"

results = collection.query(
    query_embeddings=[embed(question)],
    n_results=2,
)

related_docs = results["documents"][0]
print(related_docs)
# → ['퇴사 시 미사용 연차는 수당으로 정산해 지급한다.', '우리 회사는 연차를 입사 첫해부터 15일 제공한다.']

질문을 임베딩해서 query에 넘기면, 의미가 가장 비슷한 문서 2개를 찾아준다. 질문에는 "수당"이라는 단어가 없는데도 정산 관련 문서를 찾아온 게 포인트다. 키워드가 아니라 의미로 찾은 거다.

LLM으로 최종 답변 만들기

context = "\n".join(related_docs)

prompt = f"""아래 문서를 참고해서 질문에 답해줘.

문서:
{context}

질문: {question}
"""

answer = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[{"role": "user", "content": prompt}],
)

print(answer.choices[0].message.content)
# → "퇴사 시 사용하지 않은 연차는 수당으로 정산되어 지급됩니다."

이게 Chroma 버전 RAG의 전부다. 저장 → 검색 → LLM에 넘겨 답변, 딱 세 단계다.

---

2부. pgvector로 만들기

이번엔 같은 걸 PostgreSQL + pgvector로 해본다. 흐름은 똑같고, 저장/검색을 SQL로 한다는 점만 다르다.

설치와 확장 활성화

pgvector는 Postgres 확장이다. Docker로 빠르게 띄우는 게 편하다.

docker run -d --name pgvector-db \
  -e POSTGRES_PASSWORD=test \
  -p 5432:5432 \
  pgvector/pgvector:pg16

그리고 파이썬 드라이버를 깐다.

pip install psycopg

Postgres에 접속해서 확장을 켜고, 벡터를 저장할 테이블을 만든다. 아까 임베딩이 1536차원이었으니 vector(1536)으로 잡는다.

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE company_docs (
    id SERIAL PRIMARY KEY,
    content TEXT,
    embedding vector(1536)
);

문서 저장하기 (준비 단계)

import psycopg

conn = psycopg.connect(
    "host=localhost dbname=postgres user=postgres password=test"
)

with conn.cursor() as cur:
    for doc in documents:
        vector = embed(doc)
        cur.execute(
            "INSERT INTO company_docs (content, embedding) VALUES (%s, %s)",
            (doc, str(vector)),
        )
conn.commit()

pgvector는 벡터를 문자열 형태('[0.1, 0.2, ...]')로 받기 때문에 str(vector)로 넘긴다.

질문하고 검색하기

여기가 pgvector의 핵심이다. <=> 연산자가 코사인 거리를 계산해준다. 거리가 가까울수록 의미가 비슷한 거다.

question = "퇴사하면 안 쓴 연차 어떻게 돼?"
q_vector = str(embed(question))

with conn.cursor() as cur:
    cur.execute(
        """
        SELECT content
        FROM company_docs
        ORDER BY embedding <=> %s
        LIMIT 2
        """,
        (q_vector,),
    )
    related_docs = [row[0] for row in cur.fetchall()]

print(related_docs)

ORDER BY embedding <=> 질문벡터가 "질문이랑 가장 가까운 순서로 정렬"이고, LIMIT 2로 상위 2개만 가져온다. Chroma의 query(n_results=2)랑 정확히 같은 일을 SQL로 하는 거다.

LLM 답변은 동일

검색으로 related_docs를 얻은 다음은 Chroma 때랑 완전히 똑같다. 문서를 프롬프트에 넣고 LLM에 넘기면 끝이다. (위 1부의 "LLM으로 최종 답변 만들기" 코드 그대로 쓰면 된다.)

---

두 방식 비교

항목	Chroma	pgvector
설치	pip install 하나로 끝	Postgres + 확장 필요
검색 방식	파이썬 메서드 (query)	SQL (<=> 연산자)
적합한 상황	빠른 프로토타이핑, 입문	이미 Postgres 쓰는 실무
데이터 관리	별도 저장소	기존 DB랑 통합

---

정리

두 방식을 다 만들어보면 확실히 느껴지는 게 있다. 벡터DB가 바뀌어도 RAG의 뼈대는 똑같다는 것.

1. 문서를 임베딩해서 저장한다.
2. 질문을 임베딩해서 비슷한 문서를 찾는다.
3. 찾은 문서를 LLM한테 넘겨서 답하게 한다.

Chroma는 "일단 돌려보자" 할 때, pgvector는 "기존 Postgres에 얹자" 할 때 고르면 된다. 도구는 상황에 맞게 고르고, 핵심 흐름은 어디서나 똑같이 가져가면 된다.

다음엔 지금 코드에서 한 발 더 나가서, 긴 문서를 어떻게 쪼갤지(청킹 전략)랑 검색 정확도를 올리는 방법을 다뤄보자.