Graph를 그림으로 보기

import matplotlib.pyplot as plt
from PIL import Image
import io

# Graph를 이미지로 변환
image_data = graph.get_graph().draw_mermaid_png()

# 이미지 표시
image = Image.open(io.BytesIO(image_data))
plt.imshow(image)
plt.show()

# ============================================================
# 17차시 | LangGraph 입문
# 교안: 17차시_LangGraph입문.pptx
# 역할: 조건부 그래프를 실행하지 않고, 연결 그림만 화면에 띄운다.
# 입력: 그래프 정의(State 키 query, results, answer). invoke 는 하지 않는다.
# 출력: matplotlib 창의 그래프 그림. 이 파일은 png 파일을 저장하지 않는다.
#
#   State  = 공유 노트. 이 파일의 노드는 그림용으로만 등록된다.
#   Node   = search, answer 함수.
#   Edge   = START → search, answer → END.
#   조건부 엣지 = search 다음 good 이면 answer, bad 이면 search.
#   compile = 그림을 그리려면 먼저 앱으로 확정해야 한다.
#   get_graph().draw_mermaid_png() = 확정된 연결을 그림 바이트로 만든다.
#
# 흐름 (그림에 나오는 연결)
#   START
#     |
#     v
#   search ---- bad ----+
#     |                  |
#     | good             v
#     v               search
#   answer
#     |
#     v
#    END
# ============================================================

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

import io
import matplotlib.pyplot as plt
from PIL import Image

class SimpleRAGState(TypedDict):
    query: str      # 사용자의 질문
    results: list   # 검색 결과 목록
    answer: str     # 최종 답변

# 노드 search
#   읽기: query
#   쓰기: results
def search_node(state: SimpleRAGState):
    query = state["query"]
    # 현재 State에서 질문을 읽습니다.
    print("[검색 노드] 질문:", state["query"])

    # 실제 검색 대신 고정된 문서 목록을 준비합니다.
    if "가격" in query:
        results = [
            "문서1: 제품 가격은 10,000원입니다.",
            "문서2: 배송비는 무료입니다.",
            "문서3: 현재 할인 행사가 진행 중입니다."
        ]
    else:
        results = []

    print("[검색] 결과:", results)

    # LangGraph가 results 항목을 갱신하도록 반환합니다.
    return {"results": results}

# 노드 answer
#   읽기: results
#   쓰기: answer
def answer_node(state: SimpleRAGState):
    # 앞의 검색 노드가 반환한 결과를 읽습니다.
    print("[답변 노드] 검색 결과:", state["results"])

    # 실제 LLM 대신 고정된 답변을 준비합니다.
    answer = "검색 결과를 바탕으로 만든 답변입니다."

    # LangGraph가 answer 항목을 갱신하도록 반환합니다.
    return {"answer": answer}

# 조건부 엣지 함수. 노드가 아니다.
#   읽기: results
#   쓰기: 없음
#   "good" → answer, "bad" → search
def check_search_quality(state: SimpleRAGState):
    # 검색 결과가 하나 이상이면 답변 경로를 선택합니다.
    if len(state["results"]) > 0:
        print("[판단] 검색 성공 → 답변")
        return "good"

    # 검색 결과가 없으면 재검색 경로를 선택합니다.
    print("[판단] 검색 실패 → 다시 검색")
    return "bad"

graph = StateGraph(SimpleRAGState)

graph.add_node("search", search_node)
graph.add_node("answer", answer_node)

graph.add_edge(START, "search")

graph.add_conditional_edges(
    "search",
    check_search_quality,
    {
        "good": "answer",  # 성공하면 답변 노드로 이동
        "bad": "search"    # 실패하면 검색 노드로 돌아감
    }
)

graph.add_edge("answer", END)

# compile 한 뒤에야 연결 그림을 뽑을 수 있다. 이 파일은 invoke 하지 않는다.
app = graph.compile()

print("Graph 이미지를 생성합니다...")

image_data = app.get_graph().draw_mermaid_png()

# PNG 바이트 데이터를 메모리상의 파일처럼 다룹니다.
image_file = io.BytesIO(image_data)

# 메모리상의 PNG 데이터를 PIL 이미지 객체로 엽니다.
image = Image.open(image_file)

# 8. 그래프 이미지를 화면에 표시합니다.
plt.figure(figsize=(8, 6))  # 그림의 가로·세로 크기를 지정합니다.
plt.imshow(image)          # PIL 이미지를 표시합니다.
plt.axis("off")            # 좌표축과 눈금을 숨깁니다.
plt.title("LangGraph - Conditional RAG Graph")
plt.show()                 # 그림 창을 화면에 띄웁니다.