from IPython.display import Image, display
from langgraph.graph import StateGraph, START, END
from agent.agent_graph.Graph_Nodes import *
from agent.agent_graph.Graph_Routes import *
 
problem_graph = StateGraph(ProblemState)

# Add nodes
problem_graph.add_node("answer_question",answer_question)
problem_graph.add_node("update_context",update_context)
problem_graph.add_node("convertPriceToDollar",convertPriceToDollar)
problem_graph.add_node("step",step)


# Routes
problem_graph.add_conditional_edges(
    "update_context",
    is_question_clear,
    {
        True: "convertPriceToDollar",
        False: "step"
    }
)


# Edges
problem_graph.add_edge(START,"update_context")
problem_graph.add_edge("convertPriceToDollar","answer_question")

# Finalize
compiled_graph = problem_graph.compile()
#display(Image(compiled_graph.get_graph().draw_mermaid_png()))