LangChain + LangGraph + Langfuse로 만드는 가장 작은 RAG

python -m venv .venv
source .venv/bin/activate
 
pip install \
  langchain \
  langchain-openai \
  langgraph \
  langfuse \
  python-dotenv

# .env.example
OPENAI_API_KEY=sk-...
LANGFUSE_PUBLIC_KEY=pk-lf-...
LANGFUSE_SECRET_KEY=sk-lf-...
LANGFUSE_HOST=https://cloud.langfuse.com

from typing import TypedDict
 
class RAGState(TypedDict):
    question: str            # 원래 사용자 질문
    rewritten_query: str     # 검색용으로 다듬은 쿼리
    retrieved_docs: list[str]  # 검색된 문서 조각
    final_answer: str        # 최종 답변

# app.py
import os
import sys
from typing import TypedDict
 
from dotenv import load_dotenv
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langfuse.callback import CallbackHandler
from langgraph.graph import END, StateGraph
 
load_dotenv()
 
# 1) State 정의
class RAGState(TypedDict):
    question: str
    rewritten_query: str
    retrieved_docs: list[str]
    final_answer: str
 
# 2) 더미 코퍼스 (실무에선 벡터 DB)
CORPUS = [
    "회사 휴가는 연 15일이며 분기당 5일까지 사용 가능하다.",
    "결제 모듈 에러 코드 4023은 한도 초과를 의미한다.",
    "주간 회의는 매주 월요일 10시에 진행한다.",
    "사내 위키 접근은 SSO 로그인 후 가능하다.",
    "재택근무는 주 2회까지 자율적으로 사용할 수 있다.",
]
 
# 3) LLM
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
 
# 4) 노드 함수
def query_rewrite_node(state: RAGState) -> RAGState:
    """사용자 질문을 검색용 키워드 중심 쿼리로 다듬는다."""
    msgs = [
        SystemMessage(
            content=(
                "사용자 질문을 한국어 검색용 키워드 쿼리로 다듬어라. "
                "한 줄, 부연 설명 없이 쿼리만 출력."
            )
        ),
        HumanMessage(content=state["question"]),
    ]
    rewritten = llm.invoke(msgs).content.strip()
    return {**state, "rewritten_query": rewritten}
 
def retrieve_node(state: RAGState) -> RAGState:
    """더미 키워드 매칭으로 top-3 문서를 가져온다."""
    query = state["rewritten_query"]
    scored = [
        (sum(1 for w in query.split() if w in doc), doc)
        for doc in CORPUS
    ]
    scored.sort(key=lambda x: x[0], reverse=True)
    top_docs = [doc for score, doc in scored[:3] if score > 0]
    return {**state, "retrieved_docs": top_docs}
 
def answer_node(state: RAGState) -> RAGState:
    """검색된 문서를 컨텍스트로 LLM이 답변한다."""
    context = "\n".join(f"- {d}" for d in state["retrieved_docs"]) or "(관련 문서 없음)"
    msgs = [
        SystemMessage(
            content=(
                "너는 사내 문서를 기반으로 답하는 도우미다. "
                "주어진 문서에 답이 없으면 '문서에 없습니다'라고만 답해라."
            )
        ),
        HumanMessage(
            content=f"[문서]\n{context}\n\n[질문]\n{state['question']}"
        ),
    ]
    answer = llm.invoke(msgs).content.strip()
    return {**state, "final_answer": answer}
 
# 5) Graph 빌드
graph = StateGraph(RAGState)
graph.add_node("query_rewrite", query_rewrite_node)
graph.add_node("retrieve", retrieve_node)
graph.add_node("answer", answer_node)
graph.set_entry_point("query_rewrite")
graph.add_edge("query_rewrite", "retrieve")
graph.add_edge("retrieve", "answer")
graph.add_edge("answer", END)
app = graph.compile()
 
# 6) Langfuse Callback
langfuse_handler = CallbackHandler(
    public_key=os.environ["LANGFUSE_PUBLIC_KEY"],
    secret_key=os.environ["LANGFUSE_SECRET_KEY"],
    host=os.environ.get("LANGFUSE_HOST", "https://cloud.langfuse.com"),
)
 
# 7) 실행
if __name__ == "__main__":
    question = sys.argv[1] if len(sys.argv) > 1 else "휴가 며칠이야?"
    result = app.invoke(
        {
            "question": question,
            "rewritten_query": "",
            "retrieved_docs": [],
            "final_answer": "",
        },
        config={"callbacks": [langfuse_handler]},
    )
    print("---")
    print("Rewritten:", result["rewritten_query"])
    print("Docs:", result["retrieved_docs"])
    print("Answer:", result["final_answer"])

python app.py "휴가 며칠이야?"