Callbacks

The Callback system in GraphAI provides a powerful mechanism for streaming data between nodes, particularly useful for handling streaming LLM outputs or other incremental data processing.

Callback Basics

At its core, the Callback is an asyncio-based system that:

Provides a queue for passing streaming data between components
Handles special tokens to mark node start/end events
Structures streaming content for easy consumption by downstream processes
Can be integrated with any async compatible streaming system

Creating a Callback

The Graph automatically creates a callback when needed, but you can also create and customize one:

from graphai import Callback

# Create a callback with default settings
callback = Callback()

# Create a callback with custom settings
callback = Callback(
    identifier="custom_id",  # Used for special tokens
    special_token_format="<{identifier}:{token}:{params}>",  # Format for special tokens
    token_format="{token}"  # Format for regular tokens
)

Callback In Nodes

To use callbacks in a node, mark it with stream=True:

from graphai import node

@node(stream=True)
async def streaming_node(input: dict, callback):
    """This node receives a callback parameter because stream=True."""
    # Process input
    for chunk in process_chunks(input["data"]):
        # Stream output chunks
        await callback.acall(chunk)
    
    # Return final result
    return {"result": "streaming complete"}

Important points:

The stream=True parameter tells GraphAI to inject a callback
The node must have a callback parameter
The callback can be used to stream output chunks

Streaming from LLMs

A common use case is streaming output from an LLM:

@node(stream=True)
async def llm_node(input: dict, callback):
    """Stream output from an LLM."""
    from openai import AsyncOpenAI
    
    client = AsyncOpenAI()
    
    # Start streaming response
    stream = await client.chat.completions.create(
        model="gpt-4",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": input["query"]}
        ],
        stream=True
    )
    
    response_text = ""
    
    # Stream chunks through the callback
    async for chunk in stream:
        if chunk.choices[0].delta.content:
            content = chunk.choices[0].delta.content
            response_text += content
            await callback.acall(content)
    
    # Return the complete response
    return {"response": response_text}

Callback Methods

The Callback provides several key methods:

Streaming Content

# Synchronous callback (use in sync contexts)
callback(token="chunk of text")

# Async callback (preferred)
await callback.acall(token="chunk of text")

Node Management

# Mark the start of a node
await callback.start_node(node_name="my_node")

# Mark the end of a node
await callback.end_node(node_name="my_node")

# Close the callback stream
await callback.close()

Consuming a Callback Stream

You can consume a callback’s stream using its async iterator:

async def consume_stream(callback):
    async for token in callback.aiter():
        # Process each token
        print(token, end="", flush=True)

This is especially useful for web applications that need to provide real-time updates.

Special Tokens

GraphAI uses special tokens to mark events in the stream:

<graphai:node_name:start>  # Marks the start of a node
<graphai:node_name>        # Identifies the node
<graphai:node_name:end>    # Marks the end of a node
<graphai:END>              # Marks the end of the stream

These tokens can be customized using the special_token_format parameter.

Example: Web Server with Streaming

Here’s how to use callbacks with a FastAPI server:

from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
from graphai import Graph, node
import asyncio

app = FastAPI()

@node(start=True, stream=True)
async def llm_stream(input: dict, callback):
    # Streaming LLM implementation...
    for chunk in get_llm_response(input["query"]):
        await callback.acall(chunk)
    return {"status": "complete"}

@node(end=True)
async def final_node(input: dict):
    return {"status": "success"}

# Create graph
graph = Graph()
graph.add_node(llm_stream())
graph.add_node(final_node())
graph.add_edge(llm_stream, final_node)

@app.post("/stream")
async def stream_endpoint(request: Request):
    data = await request.json()
    
    # Get a callback from the graph
    callback = graph.get_callback()
    
    # Start graph execution in background
    asyncio.create_task(graph.execute({"query": data["query"]}))
    
    # Return streaming response
    return StreamingResponse(
        callback.aiter(),
        media_type="text/event-stream"
    )

Callback Configuration

You can customize the callback’s behavior:

callback = Callback(
    # Custom identifier (default is "graphai")
    identifier="myapp",
    
    # Custom format for special tokens
    special_token_format="<[{identifier}|{token}|{params}]>",
    
    # Custom format for regular tokens
    token_format="TOKEN: {token}"
)

Advanced: Custom Processing of Special Tokens

You can implement custom processing of special tokens:

async def process_stream(callback):
    current_node = None
    buffer = ""
    
    async for token in callback.aiter():
        # Check if it's a special token
        if token.startswith("<graphai:") and token.endswith(">"):
            # Handle node start
            if ":start" in token:
                node_name = token.split(":")[1].split(":start")[0]
                current_node = node_name
                # Handle node start event
                
            # Handle node end
            elif ":end" in token:
                # Handle node end event
                current_node = None
                
            # Handle stream end
            elif token == "<graphai:END>":
                # Handle end of stream
                break
                
        # Regular token
        else:
            # Process regular token
            buffer += token
            # Maybe do something with buffer
            
    return buffer

Best Practices

Use async whenever possible: The callback system is built on asyncio
Close the callback when done: Always call await callback.close() when finished
Keep streaming chunks small: Don’t stream large objects; break them into manageable chunks
Handle special tokens correctly: When consuming streams, handle special tokens appropriately

Next Steps

Learn about Graphs for orchestrating node execution
Explore Nodes for processing logic
Check out State for maintaining context across nodes

Get Started

Client Reference

User Guide

Callback Basics

Creating a Callback

Callback In Nodes

Streaming from LLMs

Callback Methods

Streaming Content

Node Management

Consuming a Callback Stream

Special Tokens

Example: Web Server with Streaming

Callback Configuration

Advanced: Custom Processing of Special Tokens

Best Practices

Next Steps

​Callback Basics

​Creating a Callback

​Callback In Nodes

​Streaming from LLMs

​Callback Methods

​Streaming Content

​Node Management

​Consuming a Callback Stream

​Special Tokens

​Example: Web Server with Streaming

​Callback Configuration

​Advanced: Custom Processing of Special Tokens

​Best Practices

​Next Steps

Callback Basics

Creating a Callback

Callback In Nodes

Streaming from LLMs

Callback Methods

Streaming Content

Node Management

Consuming a Callback Stream

Special Tokens

Example: Web Server with Streaming

Callback Configuration

Advanced: Custom Processing of Special Tokens

Best Practices

Next Steps