ByJJoon's Lab

앞서 간단한 예제를 통해 어떤 형식으로 PDF 파일을 읽어 질의하는지 확인했습니다. 이제 실제 벡터 DB를 구축하고 그걸 읽어 질의에 답변하는 코드를 작성해보겠습니다.

1. DB 구축 및 저장

PDF 파일을 읽어와 chroma_db라는 로컬 폴더를 만들고 그 안에 벡터 데이터를 영구 저장합니다.

# ingest.py
import os
from langchain_community.document_loaders import PyPDFLoader
from langchain_ollama import OllamaEmbeddings
from langchain_chroma import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 1. 환경 설정
PDF_PATH = "qrevo.pdf"
DB_DIR = "./chroma_db"  # 데이터가 영구 저장될 폴더명

# 2. 임베딩 모델 설정
embeddings = OllamaEmbeddings(model="bge-m3")

# 3. PDF 읽어오기
print("PDF 문서를 로드하는 중...")
loader = PyPDFLoader(PDF_PATH)
raw_documents = loader.load()

# 4. 텍스트 효율적으로 쪼개기
print("텍스트를 청크 단위로 분할하는 중...")
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
docs = text_splitter.split_documents(raw_documents)

# 5. Chroma DB를 하드디스크에 영구 저장 (persist_directory 지정)
print(f"임베딩 진행 후 '{DB_DIR}' 폴더에 벡터 데이터 저장 중...")
db = Chroma.from_documents(
    documents=docs,
    embedding=embeddings,
    persist_directory=DB_DIR  # 이 한 줄로 메모리가 아닌 실제 디스크에 저장됩니다.
)

print("벡터 데이터베이스 구축 완료! 이제 퀴즈/질의 코드를 실행하세요.")

2. 저장된 DB 로드 및 질의 스크립트

DB 구축후 해당 DB를 활용하여 질의하는 코드 입니다.

# query.py
import os
from langchain_ollama import OllamaLLM
from langchain_ollama import OllamaEmbeddings
from langchain_chroma import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

# 1. 환경 설정 및 엔진 로드
DB_DIR = "./chroma_db"  # 저장된 폴더 경로가 일치해야 합니다.

llm = OllamaLLM(model="qwen3:8b")
embeddings = OllamaEmbeddings(model="bge-m3")

# 2. 하드디스크에 저장된 Chroma DB 불러오기
if not os.path.exists(DB_DIR):
    print(f"에러: '{DB_DIR}' 폴더가 존재하지 않습니다. 먼저 ingest.py를 실행하세요.")
    exit()

print("영구 저장된 벡터 DB 로드 완료.")
db = Chroma(
    persist_directory=DB_DIR, 
    embedding_function=embeddings
)

# 3. 검색기(Retriever) 및 컨텍스트 유틸리티 설정
retriever = db.as_retriever(search_type="mmr", search_kwargs={"k": 5})

def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

# 4. 프롬프트 템플릿 설정
template = """당신은 주어진 문맥(Context)을 기반으로 질문에 답하는 친절한 AI 어시스턴트입니다.
질문에 답할 때 제공된 문맥을 최대한 활용하고, 문맥에 직접적인 단어가 없더라도 유추할 수 있다면 논리적으로 설명해 주세요.
만약 문맥을 아무리 찾아봐도 관련된 내용을 전혀 알 수 없다면, 억지로 지어내지 말고 "제공된 문서에서는 해당 내용을 찾을 수 없습니다"라고 정중하게 답하세요.

Context:
{context}

Question: {question}
Answer:"""
prompt = ChatPromptTemplate.from_template(template)

# 5. LCEL 파이프라인 구성
rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

# 6. 사용자 질의 응답 반복 루프
print("RAG AI가 준비되었습니다. 종료하려면 '종료'를 입력하세요.\n")
while True:
    query = input('질문을 입력하세요 : ')
    if query.strip() == "종료":
        print("프로그램을 종료합니다.")
        break

    if not query.strip():
        continue

    response = rag_chain.invoke(query)
    print("\n==== 진짜 RAG AI의 답변 ====")
    print(response)
    print("=" * 28 + "\n")

(venv) % python3 query.py
영구 저장된 벡터 DB 로드 완료.
RAG AI가 준비되었습니다. 종료하려면 '종료'를 입력하세요.

질문을 입력하세요 : 'Roborock Qrevo Curv 2 Flow' 제품에서 롤러가 하는 역할은?

==== 진짜 RAG AI의 답변 ====
제공된 문서에 따르면 **'Roborock Qrevo Curv 2 Flow'** 제품의 **롤러**는 **물걸레 청소 기능**을 수행하는 핵심 부품입니다. 구체적으로는 다음과 같은 역할을 합니다:

1. **물걸레 청소 기능**:  
   - 롤러 물걸레는 바닥을 닦는 데 사용되며, 물걸레를 롤러 축에 고정한 후 회전하면서 물을 퍼뜨려 바닥을 청소합니다.  
   - 사용 전에는 물걸레를 세척하고 자연 건조한 후, 1~3개월마다 교체해야 합니다.  

2. **청소 경로의 일부**:  
   - 로봇청소기의 청소 과정에서 롤러 물걸레는 바닥 표면을 닦는 주요 장치로, 먼지를 제거하고 깨끗한 상태로 유지하는 데 기여합니다.  

3. **청소 후 물 배출**:  
   - 롤러 물걸레가 사용된 후, 오수 통로(E3)를 통해 사용한 물과 쓰레기를 배출하는 과정에서 롤러의 역할이 연결됩니다.  

**결론**:  
롤러는 물걸레를 통해 바닥을 닦는 청소 기능을 수행하며, 물을 퍼뜨려 쓰레기를 제거하는 핵심 부품입니다. 문서에서는 롤러 자체의 구조나 추가 기능은 언급되지 않았으나, 물걸레 청소 과정에서 필수적인 역할을 합니다.
============================

질문을 입력하세요 : 종료
프로그램을 종료합니다.