LangGraph 상태 머신 설계

# pip install langgraph
from typing import TypedDict, Literal
from langgraph.graph import StateGraph, START, END

# 1. 상태 정의
class State(TypedDict):
    query: str
    category: str
    response: str

# 2. 노드 정의
def classify_query(state: State) -> State:
    """질문 분류"""
    query = state["query"].lower()
    if "결제" in query or "환불" in query:
        category = "billing"
    elif "기술" in query or "에러" in query:
        category = "tech"
    else:
        category = "general"
    return {"category": category}


def handle_billing(state: State) -> State:
    return {"response": f"결제 문의 처리: {state['query']}"}


def handle_tech(state: State) -> State:
    return {"response": f"기술 지원: {state['query']}"}


def handle_general(state: State) -> State:
    return {"response": f"일반 문의: {state['query']}"}


# 3. 그래프 구축
workflow = StateGraph(State)
workflow.add_node("classify", classify_query)
workflow.add_node("billing", handle_billing)
workflow.add_node("tech", handle_tech)
workflow.add_node("general", handle_general)


# 4. 조건 분기
def route_by_category(state: State) -> Literal["billing", "tech", "general"]:
    return state["category"]


workflow.add_edge(START, "classify")
workflow.add_conditional_edges(
    "classify",
    route_by_category,
    {
        "billing": "billing",
        "tech": "tech",
        "general": "general",
    },
)
workflow.add_edge("billing", END)
workflow.add_edge("tech", END)
workflow.add_edge("general", END)


# 5. 컴파일 + 실행
app = workflow.compile()
result = app.invoke({"query": "결제 환불 문의"})
print(result)
# {"query": "결제 환불 문의", "category": "billing", "response": "결제 문의 처리..."}

from langchain_anthropic import ChatAnthropic

llm = ChatAnthropic(model="claude-sonnet-4-6")


def classify_with_llm(state: State) -> State:
    """LLM으로 분류"""
    prompt = f"""다음 질문을 분류하세요. 답변은 단어 하나만:
- billing (결제, 환불, 가격)
- tech (오류, 기술 지원, 사용법)
- general (그 외)

질문: {state["query"]}
분류:"""
    category = llm.invoke(prompt).content.strip().lower()
    return {"category": category}


def handle_with_llm(template: str):
    def handler(state: State) -> State:
        prompt = template.format(query=state["query"])
        response = llm.invoke(prompt).content
        return {"response": response}
    return handler


# 워크플로우 재구축
workflow = StateGraph(State)
workflow.add_node("classify", classify_with_llm)
workflow.add_node("billing", handle_with_llm(
    "결제 문의 답변 (한국어, 친절): {query}"
))
workflow.add_node("tech", handle_with_llm(
    "기술 문제 진단 + 해결 단계: {query}"
))
workflow.add_node("general", handle_with_llm(
    "일반 질문 답변: {query}"
))
# 엣지는 위와 동일

class State(TypedDict):
    query: str
    category: str
    response: str
    confidence: float
    escalated: bool


def check_confidence(state: State) -> Literal["respond", "escalate"]:
    """확신도 낮으면 인간 상담사로 에스컬레이션"""
    if state.get("confidence", 0) < 0.7:
        return "escalate"
    return "respond"


def escalate_to_human(state: State) -> State:
    return {
        "response": "복잡한 문의입니다. 담당자가 곧 연락드립니다.",
        "escalated": True,
    }


workflow.add_conditional_edges(
    "classify",
    check_confidence,
    {
        "respond": "billing",
        "escalate": "human_handoff",
    },
)

# 검색 → 답변 → 평가 → 부족하면 다시 검색
class State(TypedDict):
    query: str
    docs: list
    answer: str
    iterations: int


def search_more(state: State) -> Literal["answer", "search_again"]:
    if state["iterations"] >= 3:
        return "answer"  # 3회 이상은 그만
    if len(state["docs"]) < 5:
        return "search_again"
    return "answer"


workflow.add_conditional_edges(
    "evaluate",
    search_more,
    {
        "search_again": "search",  # 다시 search 노드로
        "answer": "generate_answer",
    },
)

# Graph를 그림으로 출력
from IPython.display import Image, display

display(Image(app.get_graph().draw_mermaid_png()))
# 또는 Mermaid 문자열로
print(app.get_graph().draw_mermaid())