LangGraph

LangGraph를 사용하면서 기억해 두면 좋을 만한 것들을 정리해 두었다.

1. Conditional Edges(조건 실행 엣지)

LangGraph를 구성하다보면 특정 Node의 결과에 따라서 A Node를 실행 할 수도, B Node를 실행 할 수도 혹은 모든 태스크를 끝내거나 실행했던 Node를 반복해서 실행할 필요가 있다. 아래의 코드와 같이 add_conditional_edges 함수를 사용해주면 된다. 

graph.add_conditional_edges("symbol_extraction", lambda state: "to_extra_mod" \
if state.error_file_cnt > AVAILABLE_ERROR_NUM and state.replay_cnt <LIMIT_REPLAY_CNT else "stop",\
{"to_extra_mod": "symbol_extraction_extra", "stop": END})

graph.add_conditional_edges("symbol_extraction_extra", lambda state: "to_extra_mod" \
if state.error_file_cnt > AVAILABLE_ERROR_NUM and state.replay_cnt < LIMIT_REPLAY_CNT \
else "stop", {"to_extra_mod": "symbol_extraction_extra", "stop": END})

 

2. State

LangGraph의  State는 Airflow의 xcom과 같이 node의 상태와 결과값을 다음 node에 넘겨 줄 수 있다. 즉 워크플로우(LangGrpah의 그래프, Airlfow의 DAG) 실행중 데이터를 공유하는 기능을 한다. 

Airflow에서는 node(task)에서 return을 한 후 xcom_pull(task_ids=id)를 통해서 해당 데이터에 접근할 수 있고 LangGraph에서는 node에서 "return state"를 해줌으로써 다른 node에서 해당 데이터에 접근 가능하다. "return state" 외에도 state.update()도 가능하다고 한다. return state를 안해주게 되면 state에 있는 값을 변경해도 다음 node에서 반영되지 않으니 꼭 return 해주자.

# Airflow에서의 xcom
ti = kwargs["ti"]
env, thread_ids = ti.xcom_pull(task_ids='insert_ddb')
embedding_save_path = ti.xcom_pull(task_ids='text_embedding')

# LangGraph에서의 state
def node_symbol_extraction(state):
    file_path_l = [[os.path.join(state.section_write_path, file_n), os.path.join(state.symbol_write_path, file_n.replace(".txt", ".json"))]\
                   for file_n in os.listdir(state.section_write_path)]
    for i in range(0, len(file_path_l), state.batch_size):
        print(f"{i/len(file_path_l)*100}% progressing...")
        cur_file_l = file_path_l[i:i+state.batch_size]
        symbol_temp_data = _3_symbol_extraction.make_data(cur_file_l, state.design_code)
        multi_thr.process_multi_thread_with_save(symbol_temp_data, _3_symbol_extraction.process, True)
    state.error_file_cnt = cnt_sympy_error(state.symbol_write_path) 
    print("Sympy 에러 개수", state.error_file_cnt)
    return state

 

 

3. 시각화 (Visualization)

pyviz를 설치 후에 아래와 같은 코드로 그래프를 도시할 수 있다. dot.layout(prog="dot") 부분에서 "dot"을 다른 스타일로 바꿀 수도 있다.

def graph_draw_pyviz(graph):
    import pygraphviz as pgv
    def visualize_graph(graph):
        dot = pgv.AGraph(strict=True, directed=True)
        for node in graph.nodes:
            dot.add_node(node)
        for edge in graph.edges:
            dot.add_edge(edge[0], edge[1])

        dot.layout(prog="dot")
        dot.draw("langgraph_visualization.png", format="png")
    visualize_graph(graph)
    
graph = graph.compile().get_graph()
graph_draw_pyviz(graph)