5. 실습 ③ — 그래프 조립

MAX_RETRY = 2

def route_retrieve(state) -> str:
    return "ok" if state["documents"] else "empty"

def route_verify(state) -> str:
    if state["retries"] >= MAX_RETRY:    # ★ 안전장치 먼저
        return "giveup"
    if state["grade"] == "pass":
        return "done"
    return "retry"

g = StateGraph(MiniState)

# 노드 등록
for name, fn in [("retrieve", retrieve), ("generate", generate),
                 ("verify", verify), ("bump", bump),
                 ("fallback", fallback)]:
    g.add_node(name, fn)

# 연결
g.add_edge(START, "retrieve")
g.add_conditional_edges("retrieve", route_retrieve,
                        {"ok": "generate", "empty": "fallback"})
g.add_edge("generate", "verify")
g.add_conditional_edges("verify", route_verify,
                        {"done": END, "retry": "bump",
                         "giveup": "fallback"})
g.add_edge("bump", "generate")          # 루프
g.add_edge("fallback", END)

app = g.compile()
print(app.get_graph().draw_ascii())

# ============================================================
# 21차시 | 첫 번째 LangGraph 실습
# 교안: 21차시_첫번째LangGraph실습.pptx
# 역할: 더미 노드를 조건부 엣지로 연결하고, 질문 하나로 끝까지 실행한다.
# 입력: init("환불 규정은?")
# 출력: 문자 그래프, 노드 print, 최종 State
#
#   State  = MiniState. 이 파일의 State에는 log reducer 가 없다.
#            노드가 return 한 칸만 덮어쓴다.
#   Node   = retrieve, generate, verify, bump, fallback.
#   Edge   = START → retrieve, generate → verify, bump → generate, fallback → END.
#   조건부 엣지
#     after_retrieve: documents 가 있으면 "ok"→generate, 없으면 "empty"→fallback
#     after_verify: grade 가 pass 면 "done"→END,
#                   retries 가 MAX_RETRY 이상이면 "giveup"→fallback,
#                   아니면 "retry"→bump
#   compile = make_graph 가 앱을 반환. invoke = END까지 한 번에 실행.
#
# 흐름
#   START → retrieve
#              ├─ empty → fallback → END
#              └─ ok → generate → verify
#                                  ├─ done   → END
#                                  ├─ retry  → bump → generate (다시)
#                                  └─ giveup → fallback → END
# ============================================================

# 이 코드는 더미 데이터로 Modular RAG 그래프를 구성합니다.
# 검색, 답변 생성, 검증, 재시도, fallback 흐름을 확인합니다.

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

# ============================================================
# 1. State 정의
# ============================================================

class MiniState(TypedDict):
    # 사용자가 입력한 질문
    question: str

    # 검색된 문서 목록
    documents: list[str]

    # 생성된 답변
    answer: str

    # 검증 결과
    # pass: 통과, retry: 재시도, giveup: 포기
    grade: str

    # 재시도 횟수
    retries: int

# ============================================================
# 초기 State 생성
# ============================================================

# 노드가 아님. invoke 에 넣을 첫 노트를 만든다.
#   읽기: 인자 question
#   채우는 칸: question, documents, answer, grade, retries
def init(question: str) -> MiniState:
    """
    그래프 실행에 사용할 초기 상태를 만듭니다.
    """

    # 모든 항목을 기본값으로 설정합니다.
    return {
        "question": question,
        "documents": [],
        "answer": "",
        "grade": "",
        "retries": 0,
    }

# ============================================================
# 2. 노드 정의
# ============================================================

# 노드 retrieve
#   읽기: question
#   쓰기: documents
def retrieve(state: MiniState) -> dict:
    """
    질문에 대한 검색 결과를 만듭니다.

    질문에 '없는'이라는 단어가 있으면
    검색 실패 상황을 흉내 냅니다.
    """

    # 검색 결과를 저장할 리스트입니다.
    documents = []

    # 질문에 '없는'이 포함되어 있는지 확인합니다.
    if "없는" in state["question"]:
        # 검색 실패 상황입니다.
        documents = []
    else:
        # 검색 성공 상황을 흉내 냅니다.
        documents = ["조각A", "조각B", "조각C"]

    # 검색 결과 개수를 출력합니다.
    print(f"검색: {len(documents)}건")

    # 변경된 documents만 반환합니다.
    return {
        "documents": documents
    }

# 노드 generate
#   읽기: documents, retries
#   쓰기: answer
def generate(state: MiniState) -> dict:
    """
    검색된 문서를 바탕으로 답변을 생성합니다.

    첫 번째 답변은 부실하게 만들고,
    재시도 후에는 좋은 답변을 만듭니다.
    """

    # 재시도 횟수를 확인합니다.
    retry_count = state["retries"]

    # 재시도 횟수에 따라 답변 품질을 결정합니다.
    if retry_count == 0:
        quality = "부실한"
    else:
        quality = "좋은"

    # 가짜 답변을 생성합니다.
    answer = f"{len(state['documents'])}건 근거로 만든 {quality} 답변"

    # 생성된 답변을 출력합니다.
    print(f"생성: {answer}")

    # 변경된 answer만 반환합니다.
    return {
        "answer": answer
    }

# 노드 verify
#   읽기: answer
#   쓰기: grade
def verify(state: MiniState) -> dict:
    """
    생성된 답변을 검증합니다.

    답변에 '좋은'이 포함되면 통과시키고,
    그렇지 않으면 재시도 대상으로 처리합니다.
    """

    # 답변에 '좋은'이 포함되어 있는지 확인합니다.
    if "좋은" in state["answer"]:
        # 검증을 통과시킵니다.
        print("검증: 통과")

        return {
            "grade": "pass"
        }

    # 검증에 실패하면 재시도합니다.
    print("검증: 실패")

    return {
        "grade": "retry"
    }

# 노드 bump
#   읽기: retries
#   쓰기: retries
def bump(state: MiniState) -> dict:
    """
    재시도 횟수를 1 증가시킵니다.
    """

    # 현재 재시도 횟수에 1을 더합니다.
    retry_count = state["retries"] + 1

    # 재시도 횟수를 출력합니다.
    print(f"재시도: {retry_count}회")

    # 변경된 retries만 반환합니다.
    return {
        "retries": retry_count
    }

# 노드 fallback
#   읽기: 없음
#   쓰기: answer, grade
def fallback(state: MiniState) -> dict:
    """
    검색에 실패했거나 재시도를 포기했을 때
    안내 답변을 반환합니다.
    """

    # fallback 노드가 실행되었음을 출력합니다.
    print("fallback")

    # 안내 답변과 포기 상태를 반환합니다.
    return {
        "answer": "자료를 찾지 못했습니다.",
        "grade": "giveup",
    }

# ============================================================
# 3. 분기 함수 정의
# ============================================================

# 허용할 최대 재시도 횟수입니다.
MAX_RETRY = 2

# 조건부 엣지 함수. 노드가 아니다.
#   읽기: documents
#   쓰기: 없음
#   "ok" → generate, "empty" → fallback
def after_retrieve(state: MiniState) -> str:
    """
    검색 결과가 있는지 확인하고
    다음 노드를 결정합니다.
    """

    # 검색 결과가 있으면 답변 생성으로 이동합니다.
    if state["documents"]:
        return "ok"

    # 검색 결과가 없으면 fallback으로 이동합니다.
    return "empty"

# 조건부 엣지 함수. 노드가 아니다.
#   읽기: grade, retries
#   쓰기: 없음
#   "done" → END, "giveup" → fallback, "retry" → bump
def after_verify(state: MiniState) -> str:
    """
    검증 결과와 재시도 횟수를 확인하고
    다음 노드를 결정합니다.
    """

    # 검증에 성공하면 그래프를 종료합니다.
    if state["grade"] == "pass":
        return "done"

    # 최대 재시도 횟수에 도달하면 fallback으로 이동합니다.
    if state["retries"] >= MAX_RETRY:
        return "giveup"

    # 아직 재시도할 수 있으면 다시 생성합니다.
    return "retry"

# ============================================================
# 4. 그래프 만들기
# ============================================================

# 노드가 아님. State, 노드, 엣지를 묶어 compile 된 앱을 반환한다.
def make_graph():
    """
    State, 노드, 엣지를 연결하여
    실행 가능한 그래프를 만듭니다.
    """

    # MiniState를 사용하는 그래프를 만듭니다.
    graph = StateGraph(MiniState)

    # 노드를 등록합니다.
    graph.add_node("retrieve", retrieve)
    graph.add_node("generate", generate)
    graph.add_node("verify", verify)
    graph.add_node("bump", bump)
    graph.add_node("fallback", fallback)

    # 그래프 시작점에서 retrieve로 이동합니다.
    graph.add_edge(START, "retrieve")

    # 검색 결과에 따라 경로를 나눕니다.
    graph.add_conditional_edges(
        "retrieve",
        after_retrieve,
        {
            "ok": "generate",
            "empty": "fallback",
        },
    )

    # 답변 생성 후 검증합니다.
    graph.add_edge("generate", "verify")

    # 검증 결과에 따라 경로를 나눕니다.
    graph.add_conditional_edges(
        "verify",
        after_verify,
        {
            "done": END,
            "retry": "bump",
            "giveup": "fallback",
        },
    )

    # 재시도 횟수를 증가시킨 뒤 다시 생성합니다.
    graph.add_edge("bump", "generate")

    # fallback 이후 그래프를 종료합니다.
    graph.add_edge("fallback", END)

    # compile: 연결을 실행 앱으로 확정해서 반환한다.
    return graph.compile()

# ============================================================
# 5. 그래프 실행
# ============================================================

if __name__ == "__main__":

    # 그래프를 만듭니다.
    app = make_graph()

    # 그래프 구조를 출력합니다.
    print("■ 그래프 구조")
    print(app.get_graph().draw_ascii())

    # 그래프를 실행합니다.
    print("\n■ 그래프 실행")
    # invoke: 초기 State를 넣고 END까지의 최종 State를 한 번에 받는다.
    result = app.invoke(init("환불 규정은?"))

    # 최종 결과를 출력합니다.
    print("\n최종 결과")
    print(result)

# ============================================================
# 21차시 | 첫 번째 LangGraph 실습
# 교안: 21차시_첫번째LangGraph실습.pptx
# 역할: 같은 더미 그래프를 stream 으로 노드마다 보고, 연결을 png 로 저장한다.
# 입력: init("환불 규정은?")
# 출력: 노드가 끝날 때마다의 부분 State, 그리고 graph_structure.png
#
#   State  = MiniState. 고친 칸만 덮어쓴다. log reducer 는 없다.
#   Node   = retrieve, generate, verify, bump, fallback.
#   Edge   = START → retrieve, generate → verify, bump → generate, fallback → END.
#   조건부 엣지
#     route_retrieve: documents 있으면 "ok"→generate, 없으면 "empty"→fallback
#     route_verify: pass → END, retries >= MAX_RETRY → fallback, 아니면 bump
#   compile = make_graph 의 반환 앱.
#   stream  = invoke 와 달리 END까지 모으지 않는다.
#            노드 하나가 끝날 때마다 {노드이름: 그 노드가 return 한 칸} 을 준다.
#
# 흐름
#   START → retrieve
#              ├─ empty → fallback → END
#              └─ ok → generate → verify
#                                  ├─ done   → END
#                                  ├─ retry  → bump → generate
#                                  └─ giveup → fallback → END
# ============================================================

# 이 코드는 더미 데이터로 Modular RAG 그래프를 구성합니다.
# 검색, 답변 생성, 검증, 재시도, fallback 흐름을 확인합니다.

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

# ============================================================
# 1. State 정의
# ============================================================

class MiniState(TypedDict):
    # 사용자가 입력한 질문
    question: str

    # 검색된 문서 목록
    documents: list[str]

    # 생성된 답변
    answer: str

    # 검증 결과
    # pass: 통과, retry: 재시도, giveup: 포기
    grade: str

    # 재시도 횟수
    retries: int

# ============================================================
# 초기 State 생성
# ============================================================

# 노드가 아님. stream 에 넣을 첫 노트를 만든다.
#   읽기: 인자 question
#   채우는 칸: question, documents, answer, grade, retries
def init(question: str) -> MiniState:
    """
    그래프 실행에 사용할 초기 상태를 만듭니다.
    """

    # 모든 값은 기본값으로 시작합니다.
    return {
        "question": question,
        "documents": [],
        "answer": "",
        "grade": "",
        "retries": 0,
    }

# ============================================================
# 2. 노드 정의
# ============================================================

# 노드 retrieve
#   읽기: question
#   쓰기: documents
def retrieve(state: MiniState) -> dict:
    """
    질문에 대한 검색 결과를 만듭니다.

    질문에 '없는'이라는 단어가 있으면
    검색 결과가 없는 상황을 흉내 냅니다.
    """

    # 검색 결과를 저장할 리스트입니다.
    documents = []

    # 질문에 '없는'이 포함되어 있는지 확인합니다.
    if "없는" in state["question"]:
        # 검색 실패 상황입니다.
        documents = []
    else:
        # 검색 성공 상황을 흉내 냅니다.
        documents = ["조각A", "조각B", "조각C"]

    # 검색 결과의 개수를 출력합니다.
    print(f"→ 검색: {len(documents)}건")

    # 변경된 documents만 반환합니다.
    return {
        "documents": documents
    }

# 노드 generate
#   읽기: documents, retries
#   쓰기: answer
def generate(state: MiniState) -> dict:
    """
    검색 결과를 바탕으로 답변을 만듭니다.

    첫 번째 답변은 부실하게 만들고,
    재시도 후에는 좋은 답변을 만듭니다.
    """

    # 현재 재시도 횟수를 확인합니다.
    retry_count = state["retries"]

    # 답변의 품질을 저장할 변수입니다.
    quality = ""

    # 첫 번째 답변인지 확인합니다.
    if retry_count == 0:
        # 첫 번째 답변은 부실한 답변으로 처리합니다.
        quality = "부실한"
    else:
        # 재시도 후에는 좋은 답변으로 처리합니다.
        quality = "좋은"

    # 검색된 문서 개수를 확인합니다.
    document_count = len(state["documents"])

    # 가짜 답변을 만듭니다.
    answer = f"{document_count}건 근거로 만든 {quality} 답변"

    # 생성된 답변을 출력합니다.
    print(f"→ 생성: {answer}")

    # 변경된 answer만 반환합니다.
    return {
        "answer": answer
    }

# 노드 verify
#   읽기: answer
#   쓰기: grade
def verify(state: MiniState) -> dict:
    """
    생성된 답변을 검증합니다.

    답변에 '좋은'이라는 단어가 있으면 통과시키고,
    그렇지 않으면 재시도 대상으로 처리합니다.
    """

    # 현재 답변에 '좋은'이라는 단어가 있는지 확인합니다.
    if "좋은" in state["answer"]:
        # 검증을 통과시킵니다.
        print("→ 검증: 통과")

        return {
            "grade": "pass"
        }

    # 검증에 실패하면 재시도 상태를 반환합니다.
    print("→ 검증: 실패")

    return {
        "grade": "retry"
    }

# 노드 bump
#   읽기: retries
#   쓰기: retries
def bump(state: MiniState) -> dict:
    """
    재시도 횟수를 1 증가시킵니다.
    """

    # 현재 재시도 횟수에 1을 더합니다.
    retry_count = state["retries"] + 1

    # 재시도 횟수를 출력합니다.
    print(f"→ 재시도: {retry_count}회")

    # 변경된 retries만 반환합니다.
    return {
        "retries": retry_count
    }

# 노드 fallback
#   읽기: 없음
#   쓰기: answer, grade
def fallback(state: MiniState) -> dict:
    """
    검색에 실패했거나 재시도 횟수를 초과했을 때
    안내 답변을 반환합니다.
    """

    # fallback 노드가 실행되었음을 출력합니다.
    print("→ fallback")

    # 안내 답변과 포기 상태를 반환합니다.
    return {
        "answer": "자료를 찾지 못했습니다.",
        "grade": "giveup",
    }

# ============================================================
# 3. 분기 함수 정의
# ============================================================

# 허용할 최대 재시도 횟수입니다.
MAX_RETRY = 2

# 조건부 엣지 함수. 노드가 아니다.
#   읽기: documents
#   쓰기: 없음
#   "ok" → generate, "empty" → fallback
def route_retrieve(state: MiniState) -> str:
    """
    검색 결과가 있는지 확인하고 다음 경로를 결정합니다.
    """

    # documents에 문서가 있으면 정상 경로로 이동합니다.
    if state["documents"]:
        return "ok"

    # documents가 비어 있으면 fallback 경로로 이동합니다.
    return "empty"

# 조건부 엣지 함수. 노드가 아니다.
#   읽기: grade, retries
#   쓰기: 없음
#   "done" → END, "giveup" → fallback, "retry" → bump
def route_verify(state: MiniState) -> str:
    """
    검증 결과와 재시도 횟수를 확인하고
    다음 실행 경로를 결정합니다.
    """

    # 답변 검증에 성공하면 그래프를 종료합니다.
    if state["grade"] == "pass":
        return "done"

    # 최대 재시도 횟수에 도달하면 포기합니다.
    if state["retries"] >= MAX_RETRY:
        return "giveup"

    # 아직 재시도할 수 있으면 bump로 이동합니다.
    return "retry"

# ============================================================
# 4. 그래프 만들기
# ============================================================

# 노드가 아님. 노드와 엣지를 compile 해서 앱을 반환한다.
def make_graph():
    """
    노드와 엣지를 연결하여 실행 가능한 그래프를 만듭니다.
    """

    # MiniState를 사용하는 그래프를 만듭니다.
    graph = StateGraph(MiniState)

    # 노드를 그래프에 등록합니다.
    graph.add_node("retrieve", retrieve)
    graph.add_node("generate", generate)
    graph.add_node("verify", verify)
    graph.add_node("bump", bump)
    graph.add_node("fallback", fallback)

    # 그래프 시작점에서 retrieve로 이동합니다.
    graph.add_edge(START, "retrieve")

    # 검색 결과에 따라 다음 노드를 결정합니다.
    graph.add_conditional_edges(
        "retrieve",
        route_retrieve,
        {
            "ok": "generate",
            "empty": "fallback",
        },
    )

    # 답변 생성 후 검증을 실행합니다.
    graph.add_edge("generate", "verify")

    # 검증 결과에 따라 다음 노드를 결정합니다.
    graph.add_conditional_edges(
        "verify",
        route_verify,
        {
            "done": END,
            "retry": "bump",
            "giveup": "fallback",
        },
    )

    # 재시도 횟수를 증가시킨 뒤 다시 답변을 생성합니다.
    graph.add_edge("bump", "generate")

    # fallback 이후 그래프를 종료합니다.
    graph.add_edge("fallback", END)

    # compile: 실행 가능한 앱으로 확정한다.
    return graph.compile()

# ============================================================
# 5. 그래프 실행 및 이미지 저장
# ============================================================

if __name__ == "__main__":

    # 그래프를 만듭니다.
    app = make_graph()

    print("■ 그래프 실행")

    # stream: 노드가 끝날 때마다 그 노드가 고친 칸만 딕셔너리로 받는다.
    for step in app.stream(init("환불 규정은?")):
        print(step)

    # 그래프 구조를 PNG 이미지로 변환합니다.
    image_data = app.get_graph().draw_mermaid_png()

    # 이미지 파일로 저장합니다.
    with open("graph_structure.png", "wb") as file:
        file.write(image_data)

    print("graph_structure.png 저장 완료")