Graphs

from graphai import Graph

# Create a graph with default settings
graph = Graph()

# Create a graph with custom max steps and initial state
graph = Graph(max_steps=20, initial_state={"history": []})

# Add a node to the graph
graph.add_node(my_node())

# Add multiple nodes
graph.add_node(node_a())
graph.add_node(node_b())
graph.add_node(node_c())

# Connect two nodes
graph.add_edge(source_node, destination_node)

# Can use node names instead of node objects
graph.add_edge("node_a", "node_b")

graph.add_edge(node_a, node_b)
graph.add_edge(node_b, node_c)

# Add a router with its sources and destinations
graph.add_router(
    sources=[node_a],  # Nodes that can lead to the router
    router=my_router(), # The router node itself
    destinations=[node_b, node_c]  # Possible destinations from the router
)

@router
async def my_router(input: dict):
    # Decision logic
    if some_condition:
        return {"choice": "node_b", "data": processed_data}
    else:
        return {"choice": "node_c", "data": processed_data}

import asyncio

async def run_graph():
    # Define initial input
    input_data = {"query": "Hello, world!"}
    
    # Execute the graph
    result = await graph.execute(input_data)
    
    return result

# Run the async function
result = asyncio.run(run_graph())

# Get the current state
state = graph.get_state()

# Set the state
graph.set_state({"history": [], "context": "some context"})

# Update the state
graph.update_state({"new_key": "new_value"})

# Reset the state
graph.reset_state()

@node
async def my_node(input: dict, state: dict):
    # Access the state
    history = state.get("history", [])
    
    # Update the state (changes won't persist outside this node)
    # For persistent changes, return them in the output
    return {"output": result, "history": history + [result]}

# Compile will raise exceptions if the graph is invalid
graph.compile()

# Visualize the graph
graph.visualize()