🔥 LangGraph Code Assistant

📝 프로젝트 개요

LangGraph?

내가 LangGraph를 좋아하는 이유는 필요한 모듈을 온라인에서 가져와 내 마음대로 커스터마이징하기 편하기 때문이다. LangChain처럼 모든 기능이 일체형으로 묶여 있지 않고, Node 단위로 구성되어 있어 코드 수정과 확장이 훨씬 유연하다.

다른 사람들의 코드를 참고하며 많은 도움을 받은 만큼, 나도 누군가의 프로젝트에 작은 기여를 할 수 있으면 좋겠다는 생각이 들었다. 물론 코드 어시스턴트 예제는 이미 LangGraph 공식 문서에 나와 있지만, 한국어 자료는 부족하고 공식 문서는 내용이 다소 무겁게 느껴질 수 있다. 그래서 이 글에서는 처음 접하는 사람도 부담 없이 이해할 수 있도록 더 가볍고 실용적인 코드 쿡북을 만들어봤다. 특히 데이터 시각화 측면에 초점을 맞춰 코드를 짜봤다.

이 코드가 누군가의 프로젝트에 인용된다면 정말 뿌듯할 것 같다.
그럼 시작해보자! 🚀

---

🧩 LangGraph의 기본 요소

Agentic Code Generator를 만들기 전에, LangGraph의 핵심 개념부터 가볍게 훑어보자.

1. 상태(State): 작업 중 데이터를 저장하는 공간이야. 예를 들어 사용자 입력, 전처리 결과, 코드 실행 결과 같은 게 여기에 담긴다.
2. 노드(Node): 특정 작업을 처리하는 단위다. 데이터 로드, 전처리, 코드 생성, 시각화 등 단계별 작업이 여기서 이뤄진다.
3. 간선(Edge): 노드들끼리 연결해주는 다리 같은 존재다. 이걸로 작업 순서나 조건 분기를 설정할 수 있어.

이 구조 덕분에 작업 흐름을 시각적으로 파악하기도 쉽고, 코드 관리도 훨씬 편하다. 🙌

---

🛠️ Agentic Code Generator 플로우 설계

Agentic Code Generator는 아래 흐름으로 돌아간다.

1. 데이터 로드
2. 데이터 전처리
3. 사용자 요구사항 분석 및 코드 생성
4. 코드 실행 및 결과 확인
5. 에러 발생 시 코드 재생성 및 재실행

🔍 플로우 다이어그램

위의 그림은 Agentic Code Generator의 플로우 다이어그램이다. 최대한 간소화해서 만들어보았다.

flowchart LR
    A(START) --> B[데이터 로드]
    B --> C[데이터 전처리]
    C --> D[코드 생성]
    D --> E[시각화]
    E -->|에러 발생?| D
    E -->|에러 없음| F(END)

에러가 발생하면 자동으로 코드 생성 단계로 돌아가 반복 실행한다. 🔄

---

💻 코드 구현 과정

1️⃣ 상태(State) 정의

먼저 작업에 필요한 데이터를 저장할 공간을 만들어보자.

from typing import TypedDict

class State(TypedDict):
    input: str
    data: str
    processed_data: str
    code: str
    is_error: bool
    result: str

• input: 사용자의 요청 (input query)
• data: 로드된 데이터
• processed_data: 전처리된 데이터
• code: 생성된 코드
• is_error: 에러 발생 여부
• result: 실행 결과

각각의 쓰임새가 있는데, 아래에서 조금씩 설명해보겠다.

---

2️⃣ 노드(Node) 함수 정의

각 작업 단계별로 함수를 만들어보자.

📥 데이터 로드 노드

import pandas as pd

def load_data(state: State) -> State:
    data = pd.read_csv("data/example.csv")
    return {**state, "data": data}

데이터를 가져오는 간단한 코드다. 데이터 시각화를 위해서 CSV 파일을 가져온다. 참고로 해당 example.csv 파일은 Kaggle의 Titanic 생존자 예측 데이터를 사용했다.

🧹 데이터 전처리 노드

def preprocess_data(state: State) -> State:
    processed_data = state["data"].head(100)
    return {**state, "processed_data": processed_data}

데이터가 너무 많으면 처리 시간이 길어지니까 상위 100개만 추려보자.

🧠 코드 생성 노드

from langchain_openai import ChatOpenAI

def analyze_data(state: State) -> State:
    model = ChatOpenAI(model="gpt-4o-mini")
    prompt = f"{state['input']} 데이터: {state['processed_data']} 단, 오직 Code만 출력하세요."
    response = model.invoke(prompt)
    return {**state, "code": response.content}

모델에게 코드 생성을 요청하면, 사용자 입력과 전처리된 데이터를 기반으로 코드를 뚝딱 만들어준다.

만약 코드를 한번만 만든다면 위처럼 정의하면 되겠지만, 우리는 다이어그램에서 에러가 발생하면 다시 코드를 작성하기로 했으니까 코드를 조금 더 추가해야 한다.

def analyze_data(state: State) -> State:
    # 데이터 분석 로직 구현
    model = ChatOpenAI(model="gpt-4o-mini")
    if state["is_error"] is False:
        # 일반 코드 제작 프롬프트
        response = model.invoke(
            (
                f"{state['input']} "
                "단, 오직 Code만 출력하세요. "
                f"데이터: {state['processed_data']} "
            )
        )
    else:
        # Visualization 노드에서 코드 에러 발생시 다시 코드를 생성하는 프롬프트
        response = model.invoke(
            (
                f"{state['input']} "
                "단, 오직 Code만 출력하세요. "
                f"데이터: {state['processed_data']} "
                "====================== "
                f"[AI] {state['code']} "
                f"{state['result']} "
                "====================== "
                "실행 중 발생한 버그를 참고하여 다시 코드를 작성해주세요."
            )
        )
    code = response.content
    return State(code=code)

만약 처음으로 코드를 생성하는 상태라면 state["is_error"]가 False이므로 첫번째 if 문에서 "일반 코드 제작 프롬프트"로 LLM을 호출한다.

하지만 visualization 노드에서 코드를 실행했는데 error가 뜬 상황이라면 state["is_error"]가 True이므로 else 문에서 "다시 코드를 생성하는 프롬프트"로 LLM을 호출한다. 여기에는 에러 메시지와 다시 코드를 작성하라는 문구가 포함되어 있는 것을 볼 수 있다.

📊 시각화 노드

def visualization(state: State) -> State:
    import seaborn as sns
    import matplotlib.pyplot as plt
    from langchain_experimental.tools.python.tool import PythonAstREPLTool

    code = state["code"]
    repl_tool = PythonAstREPLTool(locals={"sns": sns, "plt": plt})
    
    try:
        exec_result = repl_tool.invoke(code)
        return {**state, "result": exec_result, "is_error": False}
    except Exception as e:
        return {**state, "result": str(e), "is_error": True}

AI가 생성한 코드를 실행해서 시각화를 출력한다. 실행 중 문제가 있으면 에러 메시지가 반환된다.

🚀 Tips

PythonAstREPLTool이 코드를 실행하는 Tool이다. PythonREPLTool과의 차이점은 보안이 조금 더 강하다는 점이다. PythonAstREPLTool은 코드 실행 전에 추상 구문 트리(Abstract Syntax Tree, AST) 분석을 통해 잠재적으로 위험한 코드나 금지된 연산을 사전에 차단할 수 있다. 이를 통해 외부 라이브러리 접근, 파일 시스템 조작, 무한 루프 실행 등 의도치 않은 동작을 방지하며 보다 안전한 코드 실행 환경을 제공한다. 따라서, 사용자 입력을 기반으로 코드를 실행해야 하는 시나리오에서 PythonAstREPLTool은 안정성과 보안을 모두 확보할 수 있는 적절한 선택지이다.

---

3️⃣ 플로우 구성하기

각 노드를 연결해 흐름을 완성하자.

from langgraph.graph import StateGraph, START, END

def build_flow():
    flow = StateGraph(State)
    
    flow.add_node("load_data", load_data)
    flow.add_node("preprocess_data", preprocess_data)
    flow.add_node("analyze_data", analyze_data)
    flow.add_node("visualization", visualization)

    flow.add_edge(START, "load_data")
    flow.add_edge("load_data", "preprocess_data")
    flow.add_edge("preprocess_data", "analyze_data")
    flow.add_edge("analyze_data", "visualization")

    flow.add_conditional_edges(
        "visualization",
        lambda state: "analyze_data" if state["is_error"] else "end",
        {"analyze_data": "analyze_data", "end": END},
    )

    return flow.compile()

에러가 발생하면 자동으로 코드 생성 단계로 돌아가 재실행할 수 있다. flow.add_conditional_edges에서 "visualization" 노드로부터 이어지는 conditional edges를 정한다. state["is_error"]가 True라면 "analyze_data" 노드로 이어지고, False라면 END 로 이어진다.

🚀 Tips

LangGraph의 State 객체는 직렬화를 거치기 때문에 pandas.DataFrame처럼 직렬화가 까다로운 객체는 직접 저장하기 어려울 수 있다. 특히 해당 Graph를 SubGraph로 추가할 때 직렬화와 관련된 Pydantic Error가 발생한다면 아래 솔루션을 시도해보자.
✅ 해결 방법:

• CSV, JSON 등 문자열 형태로 변환하여 저장:

  state["data"] = df.to_json()  # 저장 시
  df = pd.read_json(state["data"])  # 불러올 때

  이렇게 하면 상태 전환 시 데이터 손실을 방지하면서도 안정적인 처리가 가능하다.

---

4️⃣ 실행 및 결과 확인

if __name__ == "__main__":
    flow = build_flow()
    initial_state = State(
        input="Titanic 데이터를 시각적으로 멋지게 보여줘!",
        data="",
        processed_data="",
        code="",
        is_error=False,
        result=""
    )

    for event in flow.stream(initial_state):
        for step, output in event.items():
            print(f"\n🚀 STEP: {step}\n{output}\n")

✅ 실행 결과 예시

🚀 STEP별 진행 로그

🚀 STEP: load_data
데이터 로드 완료!

🚀 STEP: preprocess_data
상위 100개 데이터 전처리 완료!

🚀 STEP: analyze_data
코드 생성 성공! 📝

🚀 STEP: visualization
시각화 성공 🎉 (에러 없음)

---

📄 실제 State 데이터 (Dump)

📥 STEP: load_data

데이터 로드 후 State에 저장된 내용:

{
  'data': 
      PassengerId  Survived  Pclass                                               Name     Sex  ...  Fare Cabin  Embarked
0              1         0       3                            Braund, Mr. Owen Harris    male  ...  7.2500   NaN      S
1              2         1       1  Cumings, Mrs. John Bradley (Florence Briggs Th...  female  ... 71.2833   C85      C
2              3         1       3                             Heikkinen, Miss. Laina  female  ...  7.9250   NaN      S
...          ...       ...     ...                                                ...     ...  ...     ...   ...    ...
890          891         0       3                                Dooley, Mr. Patrick    male  ...  7.7500   NaN      Q

[891 rows x 12 columns]
}

---

🧹 STEP: preprocess_data

상위 100개 데이터 전처리 후 State에 저장된 내용:

{
  'processed_data': 
      PassengerId  Survived  Pclass                                               Name  ...     Fare  Cabin Embarked
0              1         0       3                            Braund, Mr. Owen Harris  ...    7.2500   NaN       S
1              2         1       1  Cumings, Mrs. John Bradley (Florence Briggs Th...)  ...   71.2833    C85      C
2              3         1       3                             Heikkinen, Miss. Laina  ...    7.9250   NaN       S
...          ...       ...     ...                                                ...  ...      ...    ...     ...
99           100         0       2                                  Kantor, Mr. Sinai  ...   26.0000    NaN       S

[100 rows x 12 columns]
}

---

📊 STEP: analyze_data

분석 코드 생성 후 State에 저장된 내용:

{
  'code': 
"""
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

# 데이터 프레임 생성
data = {
    'PassengerId': range(1, 101),
    'Survived': [0, 1, 1, 1, 0] * 20,
    'Pclass': [3, 1, 3, 1, 3] * 20,
    'Fare': [7.25, 71.2833, 7.925, 53.1, 8.05] * 20,
    'Embarked': ['S', 'C', 'S', 'S', 'S'] * 20
}
df = pd.DataFrame(data)

# 시각화
plt.figure(figsize=(12, 8))

# 1. 생존자 수
plt.subplot(2, 2, 1)
sns.countplot(data=df, x='Survived', palette='pastel')
plt.title('Survival Count')

# 2. 클래스별 생존율
plt.subplot(2, 2, 2)
sns.barplot(data=df, x='Pclass', y='Survived', palette='pastel')
plt.title('Survival Rate by Class')

# 3. 요금 분포
plt.subplot(2, 2, 3)
sns.histplot(data=df, x='Fare', kde=True, bins=30, color='purple')
plt.title('Fare Distribution')

# 4. 탑승 항구별 생존자 수
plt.subplot(2, 2, 4)
sns.countplot(data=df, x='Embarked', hue='Survived', palette='pastel')
plt.title('Survival Count by Embarked')

plt.tight_layout()
plt.show()
"""
}

---

🖼️ STEP: visualization

시각화 결과 및 실행 상태:

{
  'result': '',
  'is_error': False
}

✅ 📈 시각화된 Figure:

Titanic 데이터 시각화 결과:
🧍‍♂️ 생존자 수 비교 | 🏷️ 클래스별 생존율
💰 요금 분포 | 🚢 탑승 항구별 생존자 수

---

🌊 플로우 시각화

from IPython.display import display, Image
from main import build_flow

flow = build_flow()
display(Image(flow.get_graph(xray=False).draw_mermaid_png()))

• 각 단계(load_data → preprocess_data → analyze_data → visualization)에서 처리 결과가 State에 저장된다.
• 시각화 단계에서 에러가 발생하면 analyze_data 단계로 되돌아가 코드가 자동으로 수정 및 재실행된다.
• 정상적으로 실행될 경우 위와 같은 시각화 결과가 출력된다.

전체 코드

from typing import TypedDict
import pandas as pd

from langchain_openai import ChatOpenAI
from langchain_experimental.tools.python.tool import PythonAstREPLTool
from langgraph.graph import StateGraph, START, END


# State 정의
class State(TypedDict):
    input: str
    data: str
    processed_data: str
    code: str
    is_error: bool
    result: str


# 각 Node에 해당하는 함수 정의
def load_data(state: State) -> State:
    # 데이터 로드 로직 구현
    data = pd.read_csv("./data/example.csv")
    return State(data=data)


def preprocess_data(state: State) -> State:
    # 데이터 전처리 로직 구현
    data = state["data"]
    processed_data = data.head(100)
    return State(processed_data=processed_data)


def analyze_data(state: State) -> State:
    # 데이터 분석 로직 구현
    model = ChatOpenAI(model="gpt-4o-mini")
    if state["is_error"] is False:
        response = model.invoke(
            (
                f"{state['input']} "
                "단, 오직 Code만 출력하세요. "
                f"데이터: {state['processed_data']} "
            )
        )
    else:
        response = model.invoke(
            (
                f"{state['input']} "
                "단, 오직 Code만 출력하세요. "
                f"데이터: {state['processed_data']} "
                "====================== "
                f"[AI] {state['code']} "
                f"{state['result']} "
                "====================== "
                "실행 중 발생한 버그를 참고하여 다시 코드를 작성해주세요."
            )
        )
    code = response.content
    return State(code=code)


def visualization(state: State) -> State:
    # 코드 실행
    import seaborn as sns
    import matplotlib.pyplot as plt

    code = state["code"]
    processed_code = None
    if code.startswith("```"):
        processed_code = "\n".join(code.split("\n")[1:-1])
    else:
        processed_code = code

    python_repl_tool = PythonAstREPLTool(locals={"sns": sns, "plt": plt})
    try:
        exec_result = str(python_repl_tool.invoke(processed_code))
        return State(result=exec_result, is_error=False)
    except Exception as e:
        return State(result=e, is_error=True)


def build_flow():
    # flow init
    flow = StateGraph(State)

    # node 정의
    flow.add_node("load_data", load_data)
    flow.add_node("preprocess_data", preprocess_data)
    flow.add_node("analyze_data", analyze_data)
    flow.add_node("visualization", visualization)

    # edge 연결
    flow.add_edge(START, "load_data")
    flow.add_edge("load_data", "preprocess_data")
    flow.add_edge("preprocess_data", "analyze_data")
    flow.add_edge("analyze_data", "visualization")

    # codintional edge 연결
    def is_error_check(state: State) -> str:
        return "analyze_data" if state["is_error"] else "end"

    flow.add_conditional_edges(
        "visualization",
        is_error_check,
        {"analyze_data": "analyze_data", "end": END},
    )

    # compile
    return flow.compile()


if __name__ == "__main__":
    flow = build_flow()
    initial_state = State(
        input="다음 데이터를 전문적으로 시각화하여 하나의 Figure로 표현해주세요.",
        is_error=False,
        code="",
        reulst="",
    )
    for event in flow.stream(initial_state):
        for node_name, value in event.items():
            print(f"\n==============\nSTEP: {node_name}\n==============\n")
            print(value)

---

🔔 마무리

이번 글에서는 LangGraph를 활용한 데이터 처리 파이프라인 구축 방법과 단계별 시각화 과정을 다뤘다.
각 단계별 에러 처리 흐름, 조건부 플로우 구성, 그리고 State 직렬화 팁 등을 포함해 실전 프로젝트에 적용할 수 있는 내용을 정리했다. 🚀

---

🔗 추가 자료 & 코드 확인하기

📂 전체 코드 GitHub에서 확인하기:
🔗 GitHub - Agentic Code Generator

📚 참고 자료:
🔗 LangGraph Code Assistant Tutorial

---

💬 읽어주셔서 감사합니다!
🙌 도움이 되셨다면 ⭐️ 스타를 눌러주시면 큰 힘이 됩니다.
✅ 추가로 다뤄보면 좋을 주제나 궁금한 내용이 있다면 댓글로 알려주세요. 관련 글을 작성해 보겠습니다. 😊