from agent.custom_state import AssistantState
from langgraph.graph import START, StateGraph, END
from agent.nodes import assisstant, tools, get_file
from langgraph.prebuilt import tools_condition, ToolNode

builder = StateGraph(AssistantState)

builder.add_node("get_file", get_file)
builder.add_node("assisstant", assisstant)
builder.add_node("tools", ToolNode(tools))

builder.add_edge(START, "get_file")
builder.add_edge("get_file", "assisstant")
builder.add_conditional_edges(
    "assisstant",
    tools_condition
)
builder.add_edge("tools", "assisstant")

graph = builder.compile()

png_bytes = graph.get_graph().draw_mermaid_png()

with open("graph.png", "wb") as f:
    f.write(png_bytes)