2. 실습 ① — 통합 그래프 완성

import time
from langgraph.graph import StateGraph, START, END
from state import RAGState, init_state
import config

from nodes.classifier import classifier_node
from nodes.intents    import greeting_node, calc_node, scope_node
from nodes.retriever  import retriever_node
from nodes.generator  import generator_node
from nodes.verifier   import verifier_node
from nodes.rewriter   import rewrite_node
from nodes.fallback   import fallback_node

MAX_RETRY     = getattr(config, "MAX_RETRY", 2)
MAX_REWRITE   = getattr(config, "MAX_REWRITE", 2)
MAX_TOTAL_SEC = getattr(config, "MAX_TOTAL_SEC", 60)

# ── 경로 상수 (오타 방지) ──
OK, RESEARCH, REGEN, DONE, GIVEUP = \
    "ok", "research", "regenerate", "done", "giveup"

def bump_node(state) -> dict:
    n = state.get("retries", 0) + 1
    return {"retries": n, "log": [f"재생성 {n}회차 진입"]}

def _out_of_time(state) -> bool:
    return time.time() > state.get("deadline", float("inf"))

def route_by_intent(state) -> str:
    intent = state.get("intent", "document")
    return intent if intent in ("greeting", "calc", "scope") \
           else "document"

def route_after_retrieve(state) -> str:
    if _out_of_time(state):
        return GIVEUP
    if state.get("retrieval_ok"):
        return OK
    if state.get("rewrites", 0) >= MAX_REWRITE:
        return GIVEUP
    if state.get("fail_reason") in ("low_score", ""):
        return RESEARCH
    return GIVEUP           # search_error, no_result, empty_query

def route_after_verify(state) -> str:
    if _out_of_time(state):
        return GIVEUP
    if state.get("retries", 0) >= MAX_RETRY:
        return GIVEUP
    grade = state.get("grade", "retry")
    if grade == "pass":
        return DONE
    if grade == "giveup":
        return GIVEUP
    if grade == "research":
        return RESEARCH if state.get("rewrites", 0) < MAX_REWRITE \
               else GIVEUP
    return REGEN

def route_after_rewrite(state) -> str:
    if _out_of_time(state):
        return GIVEUP
    return "retry_search" if state.get("query") else GIVEUP

# ==============================================================
# 32차시. PPT 32차시_고도화ModularRAG. 코드 폴더는 src/32
# PPT: 32차시 경로 점검과 회귀 테스트
# 역할: 분류, 검색, 생성, 검증, 재생성, 재작성, fallback을 한 그래프로 모은 최종 연결입니다.
# 읽는 키: intent, retrieval_ok, fail_reason, rewrites, grade, retries, query
# 쓰는 키: retries, log, ask()는 최종 State
# 그래프:
#   START --> classify
#     |-- greeting|calc|scope --> END
#     +-- document --> retrieve
#           |-- ok --> generate --> verify
#           |         |-- done --> END
#           |         |-- regenerate --> bump --> generate
#           |         |-- research --> rewrite --> retrieve
#           |         +-- giveup --> fallback --> END
#           |-- research --> rewrite --> retrieve
#           +-- giveup --> 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에 적습니다.
# route_after_retrieve: 검색 성공은 ok. search_error, no_result, empty_query, 재작성 상한은 giveup.
# 그 외 품질 부족은 research로 rewrite에 갑니다.
# route_after_verify: pass는 done, giveup은 fallback, research는 재작성, retry 이고 횟수가 남으면 regenerate.
# route_after_rewrite: query가 비어 있지 않으면 retrieve로 돌아갑니다.
# 31차시와 달리 이 파일은 safe_node와 log_query를 붙이지 않습니다.
# ==============================================================

# ============================================================
# graph32.py
# 지금까지 만든 RAG 노드들을 하나의 LangGraph로 연결합니다.
# 질문 종류와 처리 결과에 따라 적절한 노드로 이동합니다.
# ============================================================

import os
import sys

from langgraph.graph import StateGraph, START, END

# ============================================================
# 이전 차시의 모듈을 사용할 수 있도록 경로 추가
# ============================================================

BASE_DIR = os.path.dirname(__file__)

folders = ["18", "19", "24", "25", "26", "29", "30", "31"]

for folder in folders:
    path = os.path.join(BASE_DIR, "..", folder)
    sys.path.insert(0, path)

# ============================================================
# 필요한 모듈 불러오기
# ============================================================

from graph_state2 import RAGState, make_initial_state

from classifier import classifier_node
from intents import greeting_node, calc_node, scope_node

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

import config

# ============================================================
# 최대 반복 횟수
# ============================================================

# 같은 자료로 답변을 다시 만드는 최대 횟수
MAX_RETRY = getattr(config, "MAX_RETRY", 2)

# 질문을 고쳐서 다시 검색하는 최대 횟수
MAX_REWRITE = getattr(config, "MAX_REWRITE", 2)

# ============================================================
# 재생성 횟수 증가 노드
# ============================================================

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

    retries = state.get("retries", 0)
    retries = retries + 1

    return {
        "retries": retries,
        "log": [f"재생성 {retries}회차"]
    }

# ============================================================
# 1. 질문 종류에 따른 이동 경로 결정
# ============================================================

def route_by_intent(state):
    """분류된 질문 종류에 따라 다음 경로를 결정합니다."""

    intent = state.get("intent", "document")

    if intent == "greeting":
        return "greeting"

    if intent == "calc":
        return "calc"

    if intent == "scope":
        return "scope"

    # 나머지 질문은 문서 검색으로 이동
    return "document"

# ============================================================
# 2. 검색 결과에 따른 이동 경로 결정
# ============================================================

def route_after_retrieve(state):
    """검색 성공 여부에 따라 생성, 재검색, 포기를 결정합니다."""

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

    # 재검색 횟수를 모두 사용했으면 포기
    rewrites = state.get("rewrites", 0)

    if rewrites >= MAX_REWRITE:
        return "giveup"

    # 검색 시스템 자체의 문제라면 다시 검색하지 않음
    fail_reason = state.get("fail_reason", "")

    if fail_reason == "search_error":
        return "giveup"

    if fail_reason == "no_result":
        return "giveup"

    if fail_reason == "empty_query":
        return "giveup"

    # 검색 결과의 품질이 낮으면 질문을 고쳐서 재검색
    return "research"

# ============================================================
# 3. 답변 검증 결과에 따른 이동 경로 결정
# ============================================================

def route_after_verify(state):
    """검증 결과에 따라 종료, 재생성, 재검색을 결정합니다."""

    grade = state.get("grade", "retry")

    # 답변이 정상이면 종료
    if grade == "pass":
        return "done"

    # 더 이상 처리할 수 없으면 포기
    if grade == "giveup":
        return "giveup"

    # 검색 자료가 부족하면 질문을 고쳐서 재검색
    if grade == "research":

        rewrites = state.get("rewrites", 0)

        if rewrites >= MAX_REWRITE:
            return "giveup"

        return "research"

    # 답변에 문제가 있으면 같은 자료로 다시 생성
    retries = state.get("retries", 0)

    if retries < MAX_RETRY:
        return "regenerate"

    # 재생성 횟수를 모두 사용했으면 포기
    return "giveup"

# ============================================================
# 4. 질문 재작성 결과에 따른 이동 경로 결정
# ============================================================

def route_after_rewrite(state):
    """새 검색어가 만들어졌으면 다시 검색합니다."""

    query = state.get("query", "")

    if query:
        return "retry_search"

    return "giveup"

# ============================================================
# 전체 LangGraph 만들기
# ============================================================

def build_graph():
    """모든 노드를 등록하고 실행 순서를 연결합니다."""

    graph = StateGraph(RAGState)

    # --------------------------------------------------------
    # 1. 노드 등록
    # --------------------------------------------------------

    graph.add_node("classify", classifier_node)

    graph.add_node("greeting", greeting_node)
    graph.add_node("calc", calc_node)
    graph.add_node("scope", scope_node)

    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("rewrite", rewrite_node)
    graph.add_node("fallback", fallback_node)

    # --------------------------------------------------------
    # 2. 시작 → 질문 분류
    # --------------------------------------------------------

    graph.add_edge(
        START,
        "classify"
    )

    # --------------------------------------------------------
    # 3. 질문 종류에 따라 분기
    # --------------------------------------------------------

    graph.add_conditional_edges(
        "classify",
        route_by_intent,
        {
            "greeting": "greeting",
            "calc": "calc",
            "scope": "scope",
            "document": "retrieve"
        }
    )

    # 인사, 계산, 범위 밖 질문은 바로 종료
    graph.add_edge("greeting", END)
    graph.add_edge("calc", END)
    graph.add_edge("scope", END)

    # --------------------------------------------------------
    # 4. 검색 결과에 따라 분기
    # --------------------------------------------------------

    graph.add_conditional_edges(
        "retrieve",
        route_after_retrieve,
        {
            "ok": "generate",
            "research": "rewrite",
            "giveup": "fallback"
        }
    )

    # --------------------------------------------------------
    # 5. 답변 생성 → 검증
    # --------------------------------------------------------

    graph.add_edge(
        "generate",
        "verify"
    )

    # --------------------------------------------------------
    # 6. 검증 결과에 따라 분기
    # --------------------------------------------------------

    graph.add_conditional_edges(
        "verify",
        route_after_verify,
        {
            "done": END,
            "regenerate": "bump",
            "research": "rewrite",
            "giveup": "fallback"
        }
    )

    # --------------------------------------------------------
    # 7. 재생성 루프
    # --------------------------------------------------------
    # 같은 검색 자료를 사용하여 답변만 다시 만듭니다.

    graph.add_edge(
        "bump",
        "generate"
    )

    # --------------------------------------------------------
    # 8. 재검색 루프
    # --------------------------------------------------------
    # 질문을 고친 후 문서를 다시 검색합니다.

    graph.add_conditional_edges(
        "rewrite",
        route_after_rewrite,
        {
            "retry_search": "retrieve",
            "giveup": "fallback"
        }
    )

    # --------------------------------------------------------
    # 9. 실패 안내 후 종료
    # --------------------------------------------------------

    graph.add_edge(
        "fallback",
        END
    )

    # 완성된 그래프 반환
    return graph.compile()

# ============================================================
# 그래프 생성
# ============================================================

app = build_graph()

# ============================================================
# 질문 실행 함수
# ============================================================

def ask(question):
    """질문을 받아 초기 State를 만들고 그래프를 실행합니다."""

    # 질문을 이용하여 초기 State 생성
    state = make_initial_state(question)

    # 완성된 그래프 실행
    result = app.invoke(
        state,
        {"recursion_limit": 40}
    )

    return result

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

if __name__ == "__main__":

    # 테스트 질문
    question = "환불은 며칠 이내에 신청하고 대표이사 이름은?"

    # 그래프 실행
    result = ask(question)

    # 최종 결과 확인
    print("질문   :", question)
    print("답변   :", result.get("answer"))
    print("의도   :", result.get("intent"))
    print("판정   :", result.get("grade"))
    print("재생성 :", result.get("retries", 0))
    print("재검색 :", result.get("rewrites", 0))