from modules.data_class import DataState
from modules.tools import data_node
from modules.nodes import chatbot_with_tools, human_node, maybe_exit_human_node, maybe_route_to_tools
from langgraph.graph import StateGraph, START, END

# Define the LangGraph chatbot
graph_builder = StateGraph(DataState)

# Add nodes
graph_builder.add_node("chatbot_healthassistant", chatbot_with_tools)
graph_builder.add_node("patient", human_node)
graph_builder.add_node("documenting", data_node)

# Define edges
graph_builder.add_conditional_edges("chatbot_healthassistant", maybe_route_to_tools)
graph_builder.add_conditional_edges("patient", maybe_exit_human_node)
graph_builder.add_edge("documenting", "chatbot_healthassistant")
graph_builder.add_edge(START, "chatbot_healthassistant")

# Compile the state graph into a runnable object
graph_with_order_tools = graph_builder.compile()

def invoke_our_graph(st_messages, callables):
    if not isinstance(callables, list):
        raise TypeError("callables must be a list")
    return graph_with_order_tools.invoke({"messages": st_messages}, config={"callbacks": callables})