자바 개발자의 AI 입문기 (3편) - RAG 기초, 문서를 벡터로 저장하기

사용자 질문
    ↓
관련 문서 검색 (Retrieval)
    ↓
검색 결과 + 질문을 AI에게 전달
    ↓
AI가 답변 생성 (Generation)

죄송합니다. 귀사의 휴가 정책에 대한 정보가 없습니다.
일반적으로 연차 휴가는...

[검색된 문서: 휴가정책.pdf]
"연차는 입사 1년 후 15일 부여, 3년차부터 20일..."

→ AI 답변:
"귀사의 휴가 정책에 따르면, 연차는 입사 1년 후 15일이 부여되고,
3년차부터는 20일로 늘어납니다."

"자바 개발자" → [0.12, -0.34, 0.56, 0.78, ...]
"스프링 백엔드" → [0.15, -0.31, 0.52, 0.81, ...]
"맛있는 피자" → [-0.87, 0.22, -0.11, 0.03, ...]

코사인 유사도 = 두 벡터의 각도
- 1에 가까움 = 비슷함
- 0에 가까움 = 관련 없음
- -1에 가까움 = 반대

from langchain_openai import OpenAIEmbeddings
from dotenv import load_dotenv

load_dotenv()

# Embedding 모델 초기화
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

# 텍스트를 벡터로 변환
text = "자바 백엔드 개발자"
vector = embeddings.embed_query(text)

print(f"텍스트: {text}")
print(f"벡터 길이: {len(vector)}")
print(f"벡터 앞 5개: {vector[:5]}")

텍스트: 자바 백엔드 개발자
벡터 길이: 1536
벡터 앞 5개: [0.0234, -0.0156, 0.0412, ...]

pip install chromadb

# 샘플 문서 (실제로는 파일에서 로드)
documents = [
    "연차 휴가는 입사 1년 후 15일이 부여됩니다. 3년차부터는 20일로 늘어납니다.",
    "병가는 연간 10일까지 사용 가능합니다. 3일 이상 사용 시 진단서가 필요합니다.",
    "재택근무는 주 2회까지 가능합니다. 팀장 승인이 필요합니다.",
    "야근 수당은 오후 9시 이후 근무 시 지급됩니다. 시간당 기본급의 1.5배입니다.",
    "점심시간은 12시부터 1시까지입니다. 유연하게 조정 가능합니다.",
]

from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from dotenv import load_dotenv

load_dotenv()

# Embedding 모델
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

# 문서를 벡터로 변환하고 ChromaDB에 저장
vectorstore = Chroma.from_texts(
    texts=documents,
    embedding=embeddings,
    persist_directory="./chroma_db"  # 로컬 저장 경로
)

print(f"저장된 문서 수: {vectorstore._collection.count()}")

# 질문과 유사한 문서 검색
query = "휴가는 며칠이나 쓸 수 있나요?"
results = vectorstore.similarity_search(query, k=2)  # 상위 2개

print(f"질문: {query}\n")
print("검색된 문서:")
for i, doc in enumerate(results, 1):
    print(f"{i}. {doc.page_content}")

질문: 휴가는 며칠이나 쓸 수 있나요?

검색된 문서:
1. 연차 휴가는 입사 1년 후 15일이 부여됩니다. 3년차부터는 20일로 늘어납니다.
2. 병가는 연간 10일까지 사용 가능합니다. 3일 이상 사용 시 진단서가 필요합니다.

from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from dotenv import load_dotenv

load_dotenv()

# 샘플 문서
documents = [
    "연차 휴가는 입사 1년 후 15일이 부여됩니다. 3년차부터는 20일로 늘어납니다.",
    "병가는 연간 10일까지 사용 가능합니다. 3일 이상 사용 시 진단서가 필요합니다.",
    "재택근무는 주 2회까지 가능합니다. 팀장 승인이 필요합니다.",
    "야근 수당은 오후 9시 이후 근무 시 지급됩니다. 시간당 기본급의 1.5배입니다.",
    "점심시간은 12시부터 1시까지입니다. 유연하게 조정 가능합니다.",
]

# Embedding + Vector DB
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vectorstore = Chroma.from_texts(
    texts=documents,
    embedding=embeddings,
    persist_directory="./chroma_db"
)


def search_documents(query: str, k: int = 2) -> list[str]:
    """질문과 관련된 문서 검색"""
    results = vectorstore.similarity_search(query, k=k)
    return [doc.page_content for doc in results]


# 테스트
queries = [
    "휴가는 며칠이나 쓸 수 있나요?",
    "재택근무 하려면 어떻게 해야 하나요?",
    "야근하면 수당 나오나요?"
]

for query in queries:
    print(f"Q: {query}")
    results = search_documents(query)
    print(f"A: {results[0]}\n")