본문 바로가기
Data & AI Intelligence/▶AI & Agent

Milvus + LangChain 기반 RAG 구성 예제

by 류딩이2025. 12. 10.

Milvus + LangChain 기반 RAG 구성 예제 정리

본 글에서는 Milvus 벡터DB를 활용하여 간단한 RAG(Retrieval-Augmented Generation) 파이프라인을 구성하는 전체 과정을 정리합니다.
OpenAI 임베딩 모델을 사용하여 문서를 벡터화하고, Milvus에 저장한 뒤 Retriever를 통해 유사도 검색을 수행하고 최종적으로 LLM이 응답을 생성하도록 구성합니다.


1. 전체 실행 흐름

  1. 임베딩 모델 초기화 OpenAI 임베딩(text-embedding-3-small) 모델을 불러옵니다.
  2. Milvus 벡터DB 연결 임베딩 함수를 포함한 Milvus 스토어를 초기화하고 컬렉션을 생성합니다.
  3. 문서를 임베딩 후 DB에 저장 Document 객체에 텍스트를 담아 add_documents()로 벡터를 저장합니다.
  4. Retriever 구성 저장된 문서를 기반으로 유사도 검색을 수행합니다(k=2).
  5. RAG 체인 구성 및 질의 실행 검색된 문서를 Prompt에 삽입하여 최종 답변을 생성합니다.

2. 전체 코드 (주석 포함)

from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_milvus import Milvus
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

# 1. OpenAI API KEY
API_KEY = "my_api_key"

# 2. 임베딩 모델 초기화
embeddings = OpenAIEmbeddings(
    model="text-embedding-3-small",
    api_key=API_KEY
)

# 3. Milvus 스토어 초기화
milvus_store = Milvus(
    embedding_function=embeddings,
    connection_args={"uri": "http://localhost:19530"},  # Milvus 서버 주소
    collection_name="rag_practice_collection",
    auto_id=True,     # 자동 ID 생성
    drop_old=True     # 기존 컬렉션 삭제 후 재생성
)

# 4. LLM 초기화
llm = ChatOpenAI(
    model_name="gpt-4o-mini",
    temperature=0.3,
    api_key=API_KEY
)

# 5. 벡터 DB에 적재할 텍스트 준비
print("데이터 적재중")

texts = [
    "이 반에는 일정륜, 이정륜, 삼정륜이 있다."
]

documents = [Document(page_content=text) for text in texts]

# Milvus에 문서 삽입 (임베딩 → 벡터 저장)
milvus_store.add_documents(documents)

print("데이터 적재 완료")

# 6. Retriever 생성 (유사도 기반 검색)
retriever = milvus_store.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 2}
)

# 7. RAG용 프롬프트 템플릿
template = """
너는 친절한 AI 어시스턴트입니다.
아래의 [Context]를 바탕으로 [Question]에 대해 한국어로 답변해주세요.

[Context]:
{context}

[Question]:
{question}

답변:
"""
prompt = ChatPromptTemplate.from_template(template)

# 8. RAG 체인 구성
rag_chain = (
    {"context": retriever, "question": RunnablePassthrough()} 
    | prompt
    | llm
    | StrOutputParser()
)

# 9. 질의 실행
query = "이 반에는 누가 있어?"
print(f"Query: {query}")
print("-" * 30)

response = rag_chain.invoke(query)
print(response)

3. 실행 결과 설명

코드는 “이 반에는 누가 있어?”라는 질문에 대해 Milvus에 저장된 문서 내용을 검색하여
“일정륜, 이정륜, 삼정륜이 있다”라는 정보를 기반으로 대답을 생성합니다.

RAG 구조의 핵심은 다음과 같습니다.

  • Retriever가 검색한 context만 LLM에 전달되므로 사실성이 높아짐
  • LLM 단독 추론이 아니라 DB 기반 검색을 포함하여 정확한 정보 제공 가능

4. 마무리

본 예제는 Milvus + LangChain + OpenAI를 통한 최소 형태의 RAG 파이프라인 구성 방식입니다.
실제 서비스에서는 데이터 청크 분할, Batch 임베딩, metadata 저장, hybrid search 적용 등이 추가적으로 필요합니다.