자바 개발자의 AI 입문기 (4편) - RAG 실전, Knowledge Base 검색 챗봇 만들기

graph LR
    A["사용자 질문"] --> B["관련 문서 검색"]
    B --> C["검색 결과 + 질문"]
    C --> D["LLM이 답변 생성"]
    D --> E["사용자에게 응답"]

from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from dotenv import load_dotenv

load_dotenv()

# 모델 설정
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.3)
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

# 3편에서 저장한 Vector DB 로드
vectorstore = Chroma(
    persist_directory="./chroma_db",
    embedding_function=embeddings
)

# RAG 프롬프트 템플릿
rag_prompt = ChatPromptTemplate.from_messages([
    ("system", """당신은 회사 정책 안내 도우미입니다.
    아래 참고 문서를 바탕으로 질문에 답변하세요.
    문서에 없는 내용은 "해당 정보를 찾을 수 없습니다"라고 답변하세요.
    
    참고 문서:
    {context}"""),
    ("user", "{question}")
])


def answer_with_rag(question: str) -> str:
    """RAG 기반 질의응답"""
    
    # 1. 관련 문서 검색
    docs = vectorstore.similarity_search(question, k=2)
    context = "\n".join([doc.page_content for doc in docs])
    
    # 2. 프롬프트 + LLM 체인
    chain = rag_prompt | llm | StrOutputParser()
    
    # 3. 답변 생성
    answer = chain.invoke({
        "context": context,
        "question": question
    })
    
    return answer


# 테스트
questions = [
    "연차는 언제부터 쓸 수 있나요?",
    "재택근무 하려면 뭐가 필요한가요?",
    "회사 주차장은 어디 있나요?"  # 없는 정보
]

for q in questions:
    print(f"Q: {q}")
    print(f"A: {answer_with_rag(q)}\n")

Q: 연차는 언제부터 쓸 수 있나요?
A: 연차 휴가는 입사 1년 후부터 사용할 수 있으며, 15일이 부여됩니다. 
   3년차부터는 20일로 늘어납니다.

Q: 재택근무 하려면 뭐가 필요한가요?
A: 재택근무를 하려면 팀장 승인이 필요합니다. 주 2회까지 가능합니다.

Q: 회사 주차장은 어디 있나요?
A: 해당 정보를 찾을 수 없습니다.

from langchain_community.document_loaders import TextLoader

# 텍스트 파일 로드
loader = TextLoader("./company_policy.txt", encoding="utf-8")
documents = loader.load()

print(f"로드된 문서 수: {len(documents)}")
print(f"첫 번째 문서 내용: {documents[0].page_content[:200]}...")

from langchain_community.document_loaders import PyPDFLoader

# PDF 파일 로드
loader = PyPDFLoader("./휴가정책.pdf")
documents = loader.load()

print(f"로드된 페이지 수: {len(documents)}")

pip install pypdf

from langchain.text_splitter import RecursiveCharacterTextSplitter

# 문서 분할기 설정
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,       # 청크 최대 길이
    chunk_overlap=50,     # 청크 간 중복 (문맥 유지용)
    separators=["\n\n", "\n", ".", " "]  # 분할 우선순위
)

# 문서 분할
chunks = text_splitter.split_documents(documents)

print(f"원본 문서 수: {len(documents)}")
print(f"분할 후 청크 수: {len(chunks)}")

from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_community.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from dotenv import load_dotenv
import os

load_dotenv()


class KnowledgeBaseChatbot:
    """문서 기반 Q&A 챗봇"""
    
    def __init__(self, persist_dir: str = "./kb_chroma"):
        self.llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.3)
        self.embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
        self.persist_dir = persist_dir
        self.vectorstore = None
        
        # 프롬프트 템플릿
        self.prompt = ChatPromptTemplate.from_messages([
            ("system", """당신은 친절한 회사 정책 안내 도우미입니다.
            
            참고 문서:
            {context}
            
            규칙:
            1. 참고 문서를 바탕으로 답변하세요.
            2. 문서에 없는 내용은 "해당 정보를 찾을 수 없습니다"라고 답변하세요.
            3. 친절하고 자연스럽게 답변하세요.
            4. 한국어로 답변하세요."""),
            ("user", "{question}")
        ])
    
    def load_and_index(self, file_path: str) -> int:
        """파일을 로드하고 벡터 인덱싱"""
        
        # 1. 문서 로드
        loader = TextLoader(file_path, encoding="utf-8")
        documents = loader.load()
        
        # 2. 청킹
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=500,
            chunk_overlap=50
        )
        chunks = splitter.split_documents(documents)
        
        # 3. 벡터 저장
        self.vectorstore = Chroma.from_documents(
            documents=chunks,
            embedding=self.embeddings,
            persist_directory=self.persist_dir
        )
        
        return len(chunks)
    
    def load_existing_db(self):
        """기존 DB 로드"""
        if os.path.exists(self.persist_dir):
            self.vectorstore = Chroma(
                persist_directory=self.persist_dir,
                embedding_function=self.embeddings
            )
    
    def ask(self, question: str) -> dict:
        """질문에 답변"""
        if not self.vectorstore:
            return {"answer": "먼저 문서를 로드해주세요.", "sources": []}
        
        # 관련 문서 검색
        docs = self.vectorstore.similarity_search(question, k=3)
        context = "\n---\n".join([doc.page_content for doc in docs])
        
        # 답변 생성
        chain = self.prompt | self.llm | StrOutputParser()
        answer = chain.invoke({
            "context": context,
            "question": question
        })
        
        return {
            "answer": answer,
            "sources": [doc.page_content[:100] + "..." for doc in docs]
        }


# 사용 예시
if __name__ == "__main__":
    # 챗봇 초기화
    chatbot = KnowledgeBaseChatbot()
    
    # 샘플 문서 생성 (실제로는 파일에서 로드)
    sample_content = """회사 정책 안내

연차 휴가:
입사 1년 후 15일이 부여됩니다.
3년차부터는 20일로 늘어납니다.
연차는 1년 내 모두 소진해야 합니다.

병가:
연간 10일까지 사용 가능합니다.
3일 이상 사용 시 진단서가 필요합니다.

재택근무:
주 2회까지 가능합니다.
팀장 승인이 필요합니다.
코어타임(10시-4시)에는 연락 가능해야 합니다.

야근 수당:
오후 9시 이후 근무 시 지급됩니다.
시간당 기본급의 1.5배입니다.
사전 신청이 필요합니다.
"""
    
    # 샘플 파일 저장
    with open("company_policy.txt", "w", encoding="utf-8") as f:
        f.write(sample_content)
    
    # 문서 인덱싱
    chunk_count = chatbot.load_and_index("company_policy.txt")
    print(f"인덱싱 완료: {chunk_count}개 청크")
    
    # 대화
    questions = [
        "연차 몇 개 받아요?",
        "재택근무 하려면 뭐가 필요해요?",
        "야근 수당은 어떻게 계산되나요?",
        "회식비 지원되나요?"  # 없는 정보
    ]
    
    print("\n" + "="*50 + "\n")
    
    for q in questions:
        result = chatbot.ask(q)
        print(f"Q: {q}")
        print(f"A: {result['answer']}")
        print()

인덱싱 완료: 4개 청크

==================================================

Q: 연차 몇 개 받아요?
A: 입사 1년 후에 연차 15일이 부여됩니다. 3년차부터는 20일로 늘어나요.
   참고로 연차는 1년 내에 모두 소진해야 합니다.

Q: 재택근무 하려면 뭐가 필요해요?
A: 재택근무를 하려면 팀장 승인이 필요합니다. 주 2회까지 가능하고,
   코어타임(10시-4시)에는 연락 가능해야 해요.

Q: 야근 수당은 어떻게 계산되나요?
A: 오후 9시 이후 근무하시면 시간당 기본급의 1.5배가 지급됩니다.
   사전 신청이 필요하니 참고해주세요.

Q: 회식비 지원되나요?
A: 해당 정보를 찾을 수 없습니다.