자바 개발자의 AI 입문기 (5편) LangGraph 상태 기반 워크플로우

질문 → 검색 → 답변

질문 분석 
  → 카테고리 분류 (인사? 기술? 일반?)
  → 해당 DB에서 검색
  → 답변 생성
  → 답변 검증 (이상하면 다시 생성)
  → 최종 응답

pip install langgraph

def analyze_question(state):
    """질문을 분석하는 노드"""
    question = state["question"]
    # 분석 로직...
    return {"category": "hr", "question": question}

# A 노드 실행 후 B 노드로 이동
graph.add_edge("A", "B")

# 조건부 이동
graph.add_conditional_edges("A", route_function, {
    "hr": "hr_search",
    "tech": "tech_search"
})

from typing import TypedDict

class State(TypedDict):
    question: str
    category: str
    search_results: list
    answer: str

from typing import TypedDict, Literal

class ChatState(TypedDict):
    question: str
    category: Literal["hr", "tech", "general"]
    context: str
    answer: str

from langchain_openai import ChatOpenAI
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)


def classify_question(state: ChatState) -> ChatState:
    """질문을 카테고리로 분류"""
    
    prompt = ChatPromptTemplate.from_messages([
        ("system", """질문을 분류하세요.
        - hr: 휴가, 급여, 복지, 인사 관련
        - tech: 기술, 개발, 도구 관련
        - general: 그 외
        
        반드시 hr, tech, general 중 하나만 출력하세요."""),
        ("user", "{question}")
    ])
    
    chain = prompt | llm | StrOutputParser()
    category = chain.invoke({"question": state["question"]}).strip().lower()
    
    # 유효한 카테고리인지 확인
    if category not in ["hr", "tech", "general"]:
        category = "general"
    
    return {"category": category}


def search_hr(state: ChatState) -> ChatState:
    """인사 관련 검색 (시뮬레이션)"""
    # 실제로는 Vector DB 검색
    hr_docs = {
        "휴가": "연차 15일, 3년차부터 20일",
        "급여": "매월 25일 지급",
        "복지": "점심 식대 지원, 자기계발비 연 100만원"
    }
    
    question = state["question"]
    context = next((v for k, v in hr_docs.items() if k in question), 
                   "관련 인사 정책을 찾을 수 없습니다.")
    
    return {"context": context}


def search_tech(state: ChatState) -> ChatState:
    """기술 관련 검색 (시뮬레이션)"""
    tech_docs = {
        "배포": "GitHub Actions + ArgoCD 사용",
        "모니터링": "Datadog 대시보드 참고",
        "코드리뷰": "PR 올리면 2명 이상 승인 필요"
    }
    
    question = state["question"]
    context = next((v for k, v in tech_docs.items() if k in question),
                   "관련 기술 문서를 찾을 수 없습니다.")
    
    return {"context": context}


def search_general(state: ChatState) -> ChatState:
    """일반 검색"""
    return {"context": "일반 문의는 총무팀에 연락해주세요."}


def generate_answer(state: ChatState) -> ChatState:
    """최종 답변 생성"""
    
    prompt = ChatPromptTemplate.from_messages([
        ("system", """참고 정보를 바탕으로 친절하게 답변하세요.
        
        참고 정보:
        {context}"""),
        ("user", "{question}")
    ])
    
    chain = prompt | llm | StrOutputParser()
    answer = chain.invoke({
        "context": state["context"],
        "question": state["question"]
    })
    
    return {"answer": answer}

from langgraph.graph import StateGraph, END

def route_by_category(state: ChatState) -> str:
    """카테고리에 따라 다음 노드 결정"""
    category = state.get("category", "general")
    return category


# 그래프 생성
workflow = StateGraph(ChatState)

# 노드 추가
workflow.add_node("classify", classify_question)
workflow.add_node("hr", search_hr)
workflow.add_node("tech", search_tech)
workflow.add_node("general", search_general)
workflow.add_node("answer", generate_answer)

# 시작점 설정
workflow.set_entry_point("classify")

# 조건부 엣지: 분류 결과에 따라 검색 노드 선택
workflow.add_conditional_edges(
    "classify",
    route_by_category,
    {
        "hr": "hr",
        "tech": "tech",
        "general": "general"
    }
)

# 검색 노드들 → 답변 노드로
workflow.add_edge("hr", "answer")
workflow.add_edge("tech", "answer")
workflow.add_edge("general", "answer")

# 답변 노드 → 종료
workflow.add_edge("answer", END)

# 컴파일
app = workflow.compile()

# 테스트
questions = [
    "연차 며칠 받아요?",
    "배포는 어떻게 하나요?",
    "회사 주소가 어디예요?"
]

for q in questions:
    result = app.invoke({"question": q})
    print(f"Q: {q}")
    print(f"카테고리: {result['category']}")
    print(f"A: {result['answer']}\n")

Q: 연차 며칠 받아요?
카테고리: hr
A: 연차 휴가는 입사 1년 후 15일이 부여되고, 3년차부터는 20일로 늘어납니다.

Q: 배포는 어떻게 하나요?
카테고리: tech
A: 배포는 GitHub Actions와 ArgoCD를 사용합니다. 자세한 내용은 기술 문서를 참고해주세요.

Q: 회사 주소가 어디예요?
카테고리: general
A: 일반 문의는 총무팀에 연락해주세요.

from IPython.display import Image, display

# Mermaid 다이어그램 생성
print(app.get_graph().draw_mermaid())

graph TD
    __start__ --> classify
    classify --> hr
    classify --> tech
    classify --> general
    hr --> answer
    tech --> answer
    general --> answer
    answer --> __end__

def validate_answer(state: ChatState) -> str:
    """답변이 유효한지 검증"""
    answer = state.get("answer", "")
    
    # 답변이 너무 짧거나 "모르겠다"로 끝나면 재시도
    if len(answer) < 20 or "모르겠" in answer:
        return "retry"
    return "done"


# 재시도 노드
def regenerate_answer(state: ChatState) -> ChatState:
    """더 구체적인 프롬프트로 재생성"""
    # ... 재생성 로직
    return {"answer": "재생성된 답변..."}


# 조건부 엣지 추가
workflow.add_conditional_edges(
    "answer",
    validate_answer,
    {
        "retry": "regenerate",
        "done": END
    }
)