from langgraph.graph import StateGraph, START, END
from src.state.graph_state import TransportationState
from src.domain.node.input import create_input_node
from src.domain.node.call_llm import create_call_llm_node
from src.domain.node.save_cache import create_save_cache_node

def create_transportation_workflow() -> StateGraph:
    workflow = StateGraph(TransportationState)
    
    # Create nodes
    input_node = create_input_node()
    llm_node = create_call_llm_node()
    save_node = create_save_cache_node()
    
    # Add nodes
    workflow.add_node("input", input_node)
    workflow.add_node("llm", llm_node)
    workflow.add_node("save_cache", save_node)
    
    # Add edges
    workflow.add_edge(START, "input")
    workflow.add_edge("input", "llm")
    workflow.add_edge("llm", "save_cache")
    workflow.add_edge("save_cache", END)
    
    return workflow

app = create_transportation_workflow().compile()