LangGraph 기반 Agent DAG

from typing import TypedDict
from langgraph.graph import StateGraph, START, END

# 1. 전역 상태(State) 스키마 정의
class AgentState(TypedDict):
    input_text: str
    processed_data: str
    category: str

# 2. 개별 노드(Node) 함수 정의
def node_extract(state: AgentState):
    # 입력 텍스트를 처리하고 상태 업데이트 반환
    return {"processed_data": state["input_text"].strip().upper()}

def node_classify(state: AgentState):
    # 처리된 데이터를 기반으로 카테고리 분류 (LLM 모사)
    data = state["processed_data"]
    category = "A" if "URGENT" in data else "B"
    return {"category": category}

def node_handle_a(state: AgentState):
    return {"processed_data": state["processed_data"] + " [HANDLED AS A]"}

def node_handle_b(state: AgentState):
    return {"processed_data": state["processed_data"] + " [HANDLED AS B]"}

# 3. 라우팅 함수 (조건부 엣지용)
def route_by_category(state: AgentState) -> str:
    if state["category"] == "A":
        return "node_a"
    return "node_b"

# 4. StateGraph 선언 및 구축 (DAG 구성)
workflow = StateGraph(AgentState)

# 노드 등록
workflow.add_node("extractor", node_extract)
workflow.add_node("classifier", node_classify)
workflow.add_node("node_a", node_handle_a)
workflow.add_node("node_b", node_handle_b)

# 엣지 연결 (제어 흐름 정의)
workflow.add_edge(START, "extractor")
workflow.add_edge("extractor", "classifier")

# 조건부 엣지 등록 (classifier 노드 종료 후 route_by_category 함수 평가)
workflow.add_conditional_edges(
    "classifier",
    route_by_category,
    {
        "node_a": "node_a", # 라우팅 반환값이 "node_a"일 때 이동할 노드
        "node_b": "node_b"
    }
)

# 분기된 흐름을 END 노드로 수렴 (비순환 보장)
workflow.add_edge("node_a", END)
workflow.add_edge("node_b", END)

# 그래프 컴파일
app = workflow.compile()

# 실행
initial_state = {"input_text": " This is an urgent message "}
result = app.invoke(initial_state)
print(result)
# 출력: {'input_text': ' This is an urgent message ', 'processed_data': 'THIS IS AN URGENT MESSAGE [HANDLED AS A]', 'category': 'A'}

import operator
from typing import TypedDict, Annotated, Sequence
from pydantic import BaseModel, Field
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, END

# 1. Pydantic을 활용한 LLM 구조화 출력 스키마 정의 (라우팅용)
class OutputFormat(BaseModel):
    is_valid: bool = Field(..., description="문서가 요구사항을 충족하는지 여부")
    reason: str = Field(..., description="판단 사유")

# 2. 리듀서(Annotated + operator.add)가 적용된 프로덕션 State 스키마
# messages 리스트는 새로운 반환값이 들어올 때마다 기존 리스트에 append 됨
class GraphState(TypedDict):
    document: str
    messages: Annotated[Sequence[BaseMessage], operator.add]
    is_valid: bool

# 3. LLM 초기화 (구조화된 출력 강제)
llm = ChatOpenAI(model="gpt-4o", temperature=0)
evaluator_llm = llm.with_structured_output(OutputFormat)

# 4. 노드 정의
def document_analyzer_node(state: GraphState):
    """문서를 분석하고 유효성을 평가하는 노드"""
    doc = state["document"]
    prompt = f"다음 문서를 평가하여 완성도를 판단하세요. 문서: {doc}"
    
    # LLM 추론 및 구조화된 결과 파싱
    result: OutputFormat = evaluator_llm.invoke(prompt)
    
    ai_message = AIMessage(content=f"평가 완료. 사유: {result.reason}")
    
    # 상태 업데이트: 메시지 누적 및 유효성 플래그 세팅
    return {"messages": [ai_message], "is_valid": result.is_valid}

def success_handler_node(state: GraphState):
    """유효한 문서일 때 후속 처리를 담당하는 노드"""
    return {"messages": [AIMessage(content="[System] 문서가 승인되어 데이터베이스에 저장되었습니다.")]}

def rejection_handler_node(state: GraphState):
    """유효하지 않은 문서일 때 피드백을 생성하는 노드"""
    return {"messages": [AIMessage(content="[System] 문서가 반려되었습니다. 작성자에게 피드백을 전송합니다.")]}

# 5. 라우팅 로직
def route_based_on_validity(state: GraphState) -> str:
    """State의 is_valid 플래그를 검사하여 분기"""
    if state.get("is_valid"):
        return "success"
    return "rejection"

# 6. DAG 그래프 빌드
builder = StateGraph(GraphState)

builder.add_node("analyzer", document_analyzer_node)
builder.add_node("success", success_handler_node)
builder.add_node("rejection", rejection_handler_node)

builder.add_edge(START, "analyzer")

# 유효성에 따른 분기 (조건부 엣지)
builder.add_conditional_edges(
    "analyzer",
    route_based_on_validity,
    {
        "success": "success",
        "rejection": "rejection"
    }
)

builder.add_edge("success", END)
builder.add_edge("rejection", END)

# 런타임 컴파일
agent_app = builder.compile()

# 실행 예시
def run_agent(doc_text: str):
    inputs = {
        "document": doc_text,
        "messages": [HumanMessage(content="초기 문서 제출")]
    }
    
    # stream()을 사용하여 노드 실행 시점마다 실시간으로 상태 변경을 추적(Observability)
    for event in agent_app.stream(inputs):
        for node_name, state_update in event.items():
            print(f"--- [Node: {node_name}] Executed ---")
            print(state_update["messages"][-1].content)
            print("-" * 40)

# run_agent("이 문서는 완벽하게 작성된 최종 보고서입니다.")