4. 실습 ① — 그래프 완성

from langgraph.graph import StateGraph, START, END
from state import RAGState, init_state
from nodes.retriever import retriever_node
from nodes.generator import generator_node
from nodes.verifier  import verifier_node
from nodes.fallback  import fallback_node
import config

MAX_RETRY = getattr(config, "MAX_RETRY", 2)

# ── 재시도 카운터 노드 ──
def bump_node(state) -> dict:
    n = state.get("retries", 0) + 1
    return {"retries": n, "log": [f"재시도 {n}회차 진입"]}

# ── 라우팅 ──
def route_after_retrieve(state) -> str:
    return "ok" if state.get("retrieval_ok") else "empty"

def route_after_verify(state) -> str:
    if state.get("retries", 0) >= MAX_RETRY:
        return "giveup"
    grade = state.get("grade", "retry")
    if grade in ("pass", "giveup"):
        return "done" if grade == "pass" else "giveup"
    return "regenerate"

# ── 조립 ──
def build_graph():
    g = StateGraph(RAGState)

    g.add_node("retrieve", retriever_node)
    g.add_node("generate", generator_node)
    g.add_node("verify",   verifier_node)
    g.add_node("bump",     bump_node)
    g.add_node("fallback", fallback_node)

    g.add_edge(START, "retrieve")
    g.add_conditional_edges("retrieve", route_after_retrieve,
                            {"ok": "generate", "empty": "fallback"})
    g.add_edge("generate", "verify")
    g.add_conditional_edges("verify", route_after_verify,
                            {"done": END, "giveup": END,
                             "regenerate": "bump"})
    g.add_edge("bump", "generate")          # ★ 재생성 루프
    g.add_edge("fallback", END)

    # ─────────────────────────────────────────
    # 30차시에 rewrite 노드와 재검색 루프가 추가됩니다
    # ─────────────────────────────────────────

    return g.compile()

app = build_graph()

# ==============================================================
# 27차시. PPT 27차시_기본ModularRAG구현. 코드 폴더는 src/27
# PPT: 27차시 재생성 루프
# 역할: Retriever → Generator → Verifier. 실패하면 최대 MAX_RETRY까지 답을 다시 만들고, 넘으면 fallback입니다.
# 읽는 키: retrieval_ok, grade, retries, question
# 쓰는 키: retries, log, ask() 반환 answer, grade, reason, retries
# 그래프:
#   START --> retrieve
#     |-- ok ----> generate --> verify --+-- done --> END
#     |                                  |-- giveup -> fallback --> END
#     |                                  +-- retry --> bump --> generate
#     +-- empty -> fallback --> END
# --------------------------------------------------------------
# 파이썬만 아는 분을 위한 말
# State는 노드들이 함께 쓰는 딕셔너리입니다. 함수는 state를 읽고, 바꿀 키만 return 합니다.
# return에 적지 않은 키는 이전 값이 남습니다.
# log와 tried_queries는 리스트라서, 새로 return한 항목이 뒤에 이어 붙습니다.
# Retriever: 검색 담당. 질문과 비슷한 문서 조각만 고릅니다. 답 문장은 만들지 않습니다.
# Generator: 고른 문서를 근거로 답 문장을 만듭니다. 문서가 없으면 안내 문장만 돌려줍니다.
# Verifier: 답이 근거와 맞는지 검사하고 grade를 적습니다.
#   grade=pass     답을 확정하고 그래프를 끝냅니다.
#   grade=retry    같은 문서로 답 문장만 다시 씁니다. retries가 1 늘어납니다.
#   grade=research 문서가 부족합니다. 검색어를 바꾼 뒤 다시 찾습니다. rewrites가 늘어납니다.
#   grade=giveup   더 진행하지 않고 fallback으로 안내 문장만 돌려줍니다.
# fallback: 실패했을 때 답을 지어내지 않고 안내 문장을 반환하는 노드입니다.
# intent: 질문 종류. greeting 인사, calc 계산, scope 문서 밖, document 문서 검색.
# rewrite: 사용자가 친 question은 유지하고, 검색에 쓰는 query만 바꿉니다.
# safe_node: 노드를 try로 감쌉니다. 오류가 나도 그래프는 멈추지 않고 node_error에 적습니다.
# 24 폴더의 graph27.py는 재생성 갈래 이름이 regenerate 입니다.
# 이 파일은 같은 자리에 retry라는 이름을 쓰고, bump 다음에 generate로 돌아옵니다.
# grade가 pass면 done으로 END입니다. giveup이거나 retries>=MAX_RETRY면 fallback입니다.
# ask()는 answer, grade, reason, retries만 골라 반환합니다.
# ==============================================================

# 이 파일은 Retriever, Generator, Verifier를 하나의 그래프로 연결합니다.
# 답변 검증에 실패하면 최대 2회까지 다시 생성하고, 계속 실패하면 안내 메시지를 반환합니다.

import os
import sys

# ============================================================
# 필요한 폴더를 Python이 찾을 수 있도록 등록합니다.
# ============================================================

BASE_DIR = os.path.dirname(__file__)

sys.path.insert(0, os.path.join(BASE_DIR, "..", "18"))
sys.path.insert(0, os.path.join(BASE_DIR, "..", "19"))
sys.path.insert(0, os.path.join(BASE_DIR, "..", "24"))
sys.path.insert(0, os.path.join(BASE_DIR, "..", "25"))
sys.path.insert(0, os.path.join(BASE_DIR, "..", "26"))

# ============================================================
# LangGraph와 각 모듈을 가져옵니다.
# ============================================================

from langgraph.graph import StateGraph, START, END

from graph_state import RAGState
from graph_state2 import make_initial_state

from retriever import retriever_node
from generator import generator_node
from verifier import verifier_node
from fallback import fallback_node

import config

# 최대 재시도 횟수를 설정합니다.
MAX_RETRY = getattr(config, "MAX_RETRY", 2)

# ============================================================
# 재시도 횟수 증가 노드
# ============================================================

def bump_node(state):
    """답변 재생성 횟수를 1 증가시킵니다."""

    # 현재 재시도 횟수를 가져옵니다.
    current_retries = state.get("retries", 0)

    # 재시도 횟수를 1 증가시킵니다.
    next_retries = current_retries + 1

    # 변경할 State 값만 반환합니다.
    return {
        "retries": next_retries,
        "log": [
            f"재시도 {next_retries}회차 진입"
        ],
    }

# ============================================================
# 검색 결과에 따른 이동 결정
# ============================================================

def route_after_retrieve(state):
    """검색 성공 여부에 따라 다음 노드를 결정합니다."""

    # 검색에 성공하면 답변을 생성합니다.
    if state.get("retrieval_ok"):
        return "ok"

    # 검색에 실패하면 안내 메시지를 반환합니다.
    return "empty"

# ============================================================
# 검증 결과에 따른 이동 결정
# ============================================================

def route_after_verify(state):
    """Verifier의 판정 결과에 따라 다음 노드를 결정합니다."""

    # 검증 결과를 가져옵니다.
    grade = state.get("grade", "retry")

    # 답변이 검증되었으면 그래프를 종료합니다.
    if grade == "pass":
        return "done"

    # Verifier가 포기 판정을 내리면 fallback으로 이동합니다.
    if grade == "giveup":
        return "giveup"

    # 최대 재시도 횟수에 도달했으면 fallback으로 이동합니다.
    retries = state.get("retries", 0)

    if retries >= MAX_RETRY:
        return "giveup"

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

# ============================================================
# LangGraph 조립
# ============================================================

def build_graph():
    """RAG 처리에 필요한 노드와 연결을 구성합니다."""

    # State 구조를 사용하는 그래프를 만듭니다.
    graph = StateGraph(RAGState)

    # 각 처리 함수를 그래프의 노드로 등록합니다.
    graph.add_node("retrieve", retriever_node)
    graph.add_node("generate", generator_node)
    graph.add_node("verify", verifier_node)
    graph.add_node("bump", bump_node)
    graph.add_node("fallback", fallback_node)

    # 그래프가 시작되면 검색부터 실행합니다.
    graph.add_edge(START, "retrieve")

    # 검색 결과에 따라 생성 또는 fallback으로 이동합니다.
    graph.add_conditional_edges(
        "retrieve",
        route_after_retrieve,
        {
            "ok": "generate",
            "empty": "fallback",
        },
    )

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

    # 검증 결과에 따라 종료, 재생성 또는 fallback으로 이동합니다.
    graph.add_conditional_edges(
        "verify",
        route_after_verify,
        {
            "done": END,
            "giveup": "fallback",
            "retry": "bump",
        },
    )

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

    # fallback 처리가 끝나면 그래프를 종료합니다.
    graph.add_edge("fallback", END)

    # 구성한 그래프를 실행 가능한 형태로 컴파일합니다.
    return graph.compile()

# 그래프를 한 번 만들어 여러 질문에 재사용합니다.
app = build_graph()

# ============================================================
# 외부에서 질문을 처리하는 함수
# ============================================================

def ask(question):
    """질문을 그래프에 전달하고 최종 결과를 반환합니다."""

    # 질문으로 초기 State를 만듭니다.
    initial_state = make_initial_state(question)

    # 전체 그래프를 실행합니다.
    result = app.invoke(
        initial_state,
        {
            "recursion_limit": 25
        },
    )

    # 사용자에게 필요한 결과만 정리해서 반환합니다.
    return {
        "answer": result.get("answer", ""),
        "grade": result.get("grade", ""),
        "reason": result.get("reason", ""),
        "retries": result.get("retries", 0),
    }

# ============================================================
# 직접 실행 테스트
# ============================================================

if __name__ == "__main__":

    # 테스트할 질문 목록입니다.
    questions = [
        "환불은 며칠 이내에 신청해야 하나요?",
        "대표이사가 누구인가요?",
        "환불 방법과 수수료를 알려주세요"
    ]

    # 질문을 하나씩 실행합니다.
    for question in questions:

        # 그래프에 질문을 전달합니다.
        result = ask(question)

        # 실행 결과를 출력합니다.
        print()
        print("Q:", question)
        print("A:", result["answer"])
        print("판정:", result["grade"])
        print("재시도:", result["retries"])