graph.py

"""
LangGraph 워크플로우
전체 분석 프로세스를 그래프로 정의합니다.
"""

from langgraph.graph import StateGraph, END
from state import AllerLensState
from nodes import (
    ocr_node,
    ingredient_extraction_node,
    rag_search_node,
    allergen_matching_node,
    risk_assessment_node,
    explanation_generation_node,
    error_handling_node
)


# ============================================
# 조건부 라우팅 함수
# ============================================

def should_continue_after_ocr(state: AllerLensState) -> str:
    """
    OCR 성공 여부에 따라 분기
    
    Args:
        state: 현재 상태
    
    Returns:
        "continue" 또는 "error"
    """
    # OCR 신뢰도 체크
    if state.get("ocr_confidence", 0) < 0.5:
        print("   ⚠️  OCR 신뢰도가 낮습니다.")
        return "error"
    
    # 추출된 텍스트 체크
    if not state.get("ocr_raw_text"):
        print("   ⚠️  텍스트가 추출되지 않았습니다.")
        return "error"
    
    # 에러 체크
    if state.get("errors"):
        return "error"
    
    return "continue"


def should_continue_after_extraction(state: AllerLensState) -> str:
    """
    성분 추출 성공 여부에 따라 분기
    
    Args:
        state: 현재 상태
    
    Returns:
        "continue" 또는 "error"
    """
    # 추출된 성분 체크
    if not state.get("extracted_ingredients"):
        print("   ⚠️  성분이 추출되지 않았습니다.")
        return "error"
    
    return "continue"


# ============================================
# 그래프 구성
# ============================================

def create_workflow() -> StateGraph:
    """
    AllerLens 워크플로우 그래프 생성
    
    Returns:
        컴파일된 StateGraph
    """
    
    # 그래프 생성
    workflow = StateGraph(AllerLensState)
    
    # ============================================
    # 노드 추가
    # ============================================
    
    workflow.add_node("ocr", ocr_node)
    workflow.add_node("extract_ingredients", ingredient_extraction_node)
    workflow.add_node("rag_search", rag_search_node)
    workflow.add_node("match_allergens", allergen_matching_node)
    workflow.add_node("assess_risk", risk_assessment_node)
    workflow.add_node("generate_explanation", explanation_generation_node)
    workflow.add_node("handle_error", error_handling_node)
    
    # ============================================
    # 엣지 정의
    # ============================================
    
    # 시작점
    workflow.set_entry_point("ocr")
    
    # OCR → 조건 분기
    workflow.add_conditional_edges(
        "ocr",
        should_continue_after_ocr,
        {
            "continue": "extract_ingredients",
            "error": "handle_error"
        }
    )
    
    # 성분 추출 → 조건 분기
    workflow.add_conditional_edges(
        "extract_ingredients",
        should_continue_after_extraction,
        {
            "continue": "rag_search",
            "error": "handle_error"
        }
    )
    
    # RAG 검색 → 알러지 매칭
    workflow.add_edge("rag_search", "match_allergens")
    
    # 알러지 매칭 → 위험도 분석
    workflow.add_edge("match_allergens", "assess_risk")
    
    # 위험도 분석 → 설명 생성
    workflow.add_edge("assess_risk", "generate_explanation")
    
    # 종료 엣지
    workflow.add_edge("generate_explanation", END)
    workflow.add_edge("handle_error", END)
    
    return workflow


# ============================================
# 컴파일된 앱
# ============================================

# 워크플로우 생성 및 컴파일
workflow = create_workflow().compile()


# ============================================
# 그래프 시각화 (옵션)
# ============================================

def visualize_graph():
    """
    그래프를 ASCII 아트로 출력
    """
    graph_ascii = """
    ┌─────────┐
    │  START  │
    └────┬────┘
         │
         ▼
    ┌──────────────┐
    │  OCR Node    │ ◄── CLOVA OCR
    └──┬───────┬───┘
       │       │
       │       └──────── (신뢰도 < 0.5 또는 에러)
       │                │
       │                ▼
       │           ┌──────────────┐
       │           │ Error Node   │──► END
       │           └──────────────┘
       │
       └─ (성공)
       │
       ▼
    ┌─────────────────────┐
    │ Extract Ingredients │ ◄── LLM: 성분 추출
    └──────────┬──────────┘
               │
               ▼
    ┌─────────────────────┐
    │   RAG Search Node   │ ◄── Pinecone 검색
    └──────────┬──────────┘
               │
               ▼
    ┌─────────────────────┐
    │  Match Allergens    │ ◄── RAG + 사용자 프로필
    └──────────┬──────────┘
               │
               ▼
    ┌─────────────────────┐
    │   Assess Risk       │ ◄── 위험도 계산
    └──────────┬──────────┘
               │
               ▼
    ┌─────────────────────┐
    │ Generate Explanation│ ◄── 설명 생성
    └──────────┬──────────┘
               │
               ▼
            ┌─────┐
            │ END │
            └─────┘
    """
    print(graph_ascii)


# ============================================
# 테스트 코드
# ============================================

if __name__ == "__main__":
    print("=== AllerLens 워크플로우 ===\n")
    visualize_graph()
    
    print("\n워크플로우 그래프가 생성되었습니다.")
    print("main.py를 실행하여 전체 파이프라인을 테스트하세요.")