from dotenv import load_dotenv
from embedding import openai_embedding_function as embedding_function
from langgraph.graph import StateGraph, END
from langchain.chains.query_constructor.base import AttributeInfo

load_dotenv()

from node.retrieve_node import (
    plan_retrieval_node, # 사용자의 질문을 분석해 어떤 방식으로 검색할지 계획
    execute_retrieval_node, # 계획에 따라 Chroma 벡터스토어에서 관련 문서를 임베딩 기반으로 검색
    grade_and_filter_node, # grade_and_filter_node
    rewrite_query_node, # rewrite_query_node
)
from node.decide_generate import decide_to_generate_or_rewrite, generate_node
from node.graph_state import GraphState
# --- 0. 환경 설정 (기존과 동일) ---

persist_directory= "chroma_db"

# 1. 메타데이터 필드에 대한 정보를 정의합니다.
metadata_field_info = [
    AttributeInfo(
        name="party",
        description="문서를 발표한 정당의 이름. (예: '더불어민주당', '국민의힘')",
        type="string",
    ),
    AttributeInfo(
        name="score",
        description="문서의 정치적 편향 점수. -1은 진보, 1은 보수.",
        type="float",
    ),
    AttributeInfo(
        name="date",
        description="문서가 발표된 날짜. 'YYYY-MM-DD' 형식의 문자열.",
        type="string",
    ),
]


# --- 3. 그래프 구성 및 연결 ---
graph_builder = StateGraph(GraphState) # 각 노드를 순차적으로 연결함

graph_builder.add_node("plan_retrieval", plan_retrieval_node)
graph_builder.add_node("execute_retrieval", execute_retrieval_node)
graph_builder.add_node("grade_and_filter", grade_and_filter_node)
graph_builder.add_node("rewrite_query", rewrite_query_node)
graph_builder.add_node("generate", generate_node)

graph_builder.set_entry_point("plan_retrieval")
graph_builder.add_edge("plan_retrieval", "execute_retrieval")
graph_builder.add_edge("execute_retrieval", "grade_and_filter")
graph_builder.add_conditional_edges(
    "grade_and_filter",
    decide_to_generate_or_rewrite,
    {"rewrite": "rewrite_query", "generate": "generate", "end": END}
)
graph_builder.add_edge("rewrite_query", "execute_retrieval") # 루프 생성!




app = graph_builder.compile()




# verbose=True로 설정하면, LLM이 생성한 쿼리와 필터를 출력해줍니다.
# query='대통령 평가' 
# filter=Expr(op='and', args=[Expr(op='eq', args=['score', -1.0]), Expr(op='gte', args=['date', '2024-01-01'])])


# --- 4. 메인 실행 루프 ---
def main():
    """메인 애플리케이션 루프"""
    print("\n--- LangGraph RAG 애플리케이션 시작 ---")
    print("질문을 입력하세요. 종료하려면 'exit' 또는 'quit'을 입력하세요.")

    while True:
        question = input("\n질문: ")
        if question.lower() in ["exit", "quit"]:
            print("애플리케이션을 종료합니다.")
            break
        
        # 그래프 실행
        inputs = {"question": question}
        final_state = app.invoke(inputs)
        
        # 결과 출력
        if not final_state.get("generation"):
            print("\n[AI 답변]\n죄송하지만, 제공된 정보만으로는 답변하기 어렵습니다.")
        else:
            print("\n[AI 답변]")
            print(final_state["generation"])
        
        # 근거 문서 출력
        if final_state.get("documents"):
            print("\n--- 검색된 근거 문서 ---")
            for i, doc in enumerate(final_state["documents"]):
                print(f"\n--- 근거 문서 {i+1} ---")
                print(f"내용: {doc.page_content}")
                print(f"출처 정당: {doc.metadata.get('party', 'N/A')}")
                print(f"정치 편향: {doc.metadata.get('score', 'N/A')}")
                print(f"작성 날짜: {doc.metadata.get('date', 'N/A')}")
            print("-" * 30)


if __name__ == "__main__":
    main()

1. 파이프라인 구조 개요

• 입력: 사용자의 자연어 질문
• 과정: 질의 분석 → 검색 계획 → 벡터 검색 → 필터링/평가 → 필요시 쿼리 리라이트 → 반복 → 답변 생성
• 출력: AI가 생성한 답변 + 관련 근거 문서(정당, 점수, 날짜 등 메타데이터 포함)

---

2. 핵심 컴포넌트 및 흐름

(1) 그래프 노드 정의

각 노드는 node 폴더에 정의되어 있습니다.

• plan_retrieval_node: 사용자의 질문을 분석해 어떤 방식으로 검색할지 계획합니다.
• execute_retrieval_node: 계획에 따라 Chroma 벡터스토어에서 관련 문서를 임베딩 기반으로 검색합니다.
• grade_and_filter_node: 검색된 문서의 적합성을 평가하고, 메타데이터(정당, 점수, 날짜 등)로 필터링합니다.
• rewrite_query_node: 검색 결과가 부족하거나 부적절할 경우, LLM이 쿼리를 더 효과적으로 리라이트합니다.
• generate_node: 최종적으로 근거 문서와 함께 LLM이 답변을 생성합니다.

(2) 그래프 연결 및 조건 분기

• StateGraph를 사용해 각 노드를 순차적으로 연결합니다.
• grade_and_filter_node 이후,
  • 문서가 충분하면 → generate_node로 이동
  • 부족하면 → rewrite_query_node로 이동(질문을 다시 만듦)
  • 더 이상 반복이 필요 없으면 → 종료(END)
• rewrite_query_node에서 다시 execute_retrieval_node로 루프를 형성, 반복적으로 검색-리라이트를 수행할 수 있습니다.

(3) 메타데이터 활용

• 각 문서는 party(정당), score(정치적 편향), date(날짜) 등의 메타데이터를 가집니다.
• 필터링 단계에서 예를 들어 "2024년 이후 민주당 논평" 등 조건 검색이 가능합니다.

---

3. 실행 흐름 예시

1. 질문 입력

사용자가 "2024년 이후 민주당의 대통령 평가 논평은?"이라고 입력

2. plan_retrieval_node

→ "대통령 평가"라는 키워드와, "더불어민주당", "2024-01-01 이후"라는 필터 조건을 추출

3. execute_retrieval_node

→ Chroma에서 임베딩 검색 + 메타데이터 필터로 문서 추출

4. grade_and_filter_node

→ 문서가 충분하면 다음 단계, 부족하면 쿼리 리라이트로 분기

5. rewrite_query_node

→ LLM이 더 효과적인 검색 쿼리로 변환, 다시 검색

6. generate_node

→ 최종적으로 근거 문서와 함께 답변 생성

7. 출력

→ 답변 + 근거 문서(정당, 점수, 날짜 등) 출력

---

4. 코드에서의 주요 포인트

• StateGraph로 복잡한 검색-생성 플로우를 시각적으로 설계
• 조건 분기(add_conditional_edges)로, 검색 결과에 따라 반복/종료/생성 분기
• 메타데이터 기반 필터링으로, 단순 키워드 검색이 아닌 정교한 조건 검색 지원
• 루프 구조로, 검색 결과가 부족할 때 쿼리 리라이트 및 재검색 가능

---

5. 정리

이 파이프라인은 단순 검색이 아니라,

• 질문 분석 → 검색 → 평가/필터 → 필요시 쿼리 리라이트 → 반복 → 답변 생성
  의 구조로, 실제로 "근거 기반 AI 답변"을 구현합니다.
  LangChain의 StateGraph를 활용해 각 단계를 명확히 분리하고, 조건 분기와 반복을 자연스럽게 처리하는 것이 특징입니다.

---

본 후기는 [한글과컴퓨터x한국생산성본부x스나이퍼팩토리] 한컴 AI 아카데미 (B-log) 리뷰로 작성 되었습니다.

#한컴AI아카데미 #AI개발자 #AI개발자교육 #한글과컴퓨터 #한국생산성본부 #스나이퍼팩토리 #부트캠프 #AI전문가양성 #개발자교육 #개발자취업