Allow graph nodes to be reusable and/or extendable

[Finndersen]

Issue with current Graph implementation

The current graph framework implementation using type hints and nodes returning instances of other nodes is quite nice, however having the the node execution logic and conditional edge logic in the same place means that all nodes in a graph are effectively tightly coupled together, preventing extensibility and reuse.

I believe that the graph-based approach is most useful when nodes are re-usable building blocks that can be composed together into different configurations to customise behaviour of the system. For example, there is work underway to refactor theAgent class into a graph implementation. I would've thought that a major driver for this change would be the ability to extend & reuse the various useful nodes of the agent to customise behaviour as required (this would allow users to resolve issues like #675, #142, #677).

However with the current implementation I don't think this is possible. The new graph-based agent will still effectively be an immutable tightly coupled black box.

Potential solutions/improvements

Fully re-usable nodes

In order to properly allow graph nodes to be reusable building blocks, the edge routing logic would need to be de-coupled from the Node logic, which would be quite a drastic change to the framework. See example in my comment below.

Extensible/customisable nodes

At the very least, I think it would be quite valuable to allow overriding/substitution of nodes in a graph. One possible way to do this could be:

• Existing node is subclassed with a customised _run() implementation (attributes need to remain the same).
• Graph is initialised with subclassed node instead of original (along with other nodes).
• Whenever a node returns an instance of the original parent class node, the graph framework automatically up-casts it to the custom subclassed version before executing it

[YourTechBud]

I totally second this.

Technically, you can achieve reusable agents without adopting the node's and edges mechanism. My two-cents:

Here's an example of a reusable node:

• Let's say one of my nodes in the graph is query compresion, where I take a long chat history and compress it into a tiny context object. And this agent is smart enough to detect if the most recent question (user prompt) is a continuation of the previous discussion or if it starts a new discussion or not?
• I would want to reuse this node in several workflows, its so cool.

Why is it not reusable?

• Well, because this QueryCompressionNode will have to return the next node. That's literally the design.
• This is a problem cause it isn't possible to know the next node.

How to make it reusable?

Just accept a BaseNode as a constructor. So the node calling the query compression node can simply do:

class FirstNode(BaseNode):
    async run()
        # DO stuff
        return QueryCompressionNode(next_node=TheNodeWhichShouldComeInAfterQueryCompressor())

###########

# And in query compress class
class QueryCompressionNode(BaseNode):
    next_node: BaseNode


    async run(self) -> BaseNode:          # Can I return BaseNode here?
         return self.next_node

Problems I see with this approach:

• We are returning of type BaseNode. This wills crew up the mermaid graph stuff for people who care about it.
• Can't really use graph deps or graph state withing this query compression node anymore, cause every graph would have its own state.

What am I doing today?

Don't reuse nodes. Build agents which can be reused. Just create the node again. Annoying but it works.

Hope this helps @Finndersen

[Finndersen]

The main limitation of that approach is that it only supports a single next_node, which isn't very useful. And yeah you lose the ability to build the graph structure from the type hints.

I think a sensible approach is to seperate nodes and edges, where edges are functions which act as bridges between nodes, deciding the routing behaviour and translating between the inputs/outputs of different nodes. Using the graph implementation of Agent as an example (non-complete pseudocode):

# NODES
@dataclass
class UserPromptNode(BaseNode[...]):
    """Get the system & user prompt parts for initial message, or user prompt message if there is chat history"""
    user_prompt: str
    system_prompts: tuple[str, ...]
    ...

    def run(tx: GraphRunContext[...]) -> _messages.ModelRequest:
       
        ...


@dataclass
class ModelResponse:
    # Structure to store output of ModelRequestNode
    tool_calls: list[_messages.ToolCallPart]
    texts: list[str]


@dataclass
class ModelRequestNode(BaseNode[...]):
    """
    Make a request to the model using the last message in state.message_history (or a specified request).
    Returns the results of the model request.
    """
    request: _messages.ModelRequest
    model: models.Model   # IMO model should be configurable for this node

    def run(tx: GraphRunContext[...]) -> ModelResponse:
        ...


@dataclass
class ToolCallResult:
    tool_responses: list[_messages.ModelRequestPart]
    final_result: MarkFinalResult | None



@dataclass
class HandleToolCallsNode(BaseNode[...]):
    """
    Handles tool calls of a ModelResponse, including structured output via result_tool.
    """
    tool_calls: list[_messages.ModelRequestPart]

    def run(tx: GraphRunContext[...]) -> ToolCallResult:
        # Get results from calling tools
        # If result_tool is used, set final_result = True

@dataclass
class FinalResultNode(BaseNode[...]):
    """
    For backwards compatibility, append a new ModelRequest to message history using the tool returns and retries.
    Also set logging run span attributes etc.
    """
    data: MarkFinalResult[NodeRunEndT]

    def run(tx: GraphRunContext[...]) -> MarkFinalResult:
        ...


# EDGES

def user_prompt_edge(ctx: GraphRunContext[...], user_prompt: _messages.ModelRequest) -> ModelRequestNode:
    return ModelRequestNode(model=ctx.deps.model, request=user_prompt)


def model_request_edge(ctx: GraphRunContext[...], model_response: ModelResponse) -> HandleToolCallsNode | ModelRequestNode | FinalResultNode:
    if model_response.tool_calls:
        return HandleToolCallsNode(tool_calls=model_response.tool_calls)
    elif model_response.texts:
        return _handle_text_response(ctx=ctx, texts=model_response.texts)   # Same as ModelRequestNode._handle_text_response() in linked PR
    else:
        raise exceptions.UnexpectedModelBehavior('Received empty model response')


def handle_tool_call_edge(ctx: GraphRunContext[...], tool_call_result: ToolCallResult) -> ModelRequestNode | FinalResultNode:
    if tool_call_result.final_result: 
        return FinalResultNode(data=tool_call_result.final_result, extra_parts=tool_call_result.tool_responses)
    else:
        return ModelRequestNode(messages=_messages.ModelRequest(parts=tool_call_result.tool_responses))

def final_result_edge(ctx: GraphRunContext[...], data: MarkFinalResult) -> End:
    return End(data)


## GRAPH

# The nodes and edges are provided together when building the graph

graph = Graph[...](
    nodes_and_edges=[
        (UserPromptNode, user_prompt_edge),
        (ModelRequestNode, model_request_edge),
        (HandleToolCallsNode, handle_tool_call_edge),
        (FinalResultNode, final_result_edge)
    ],
    ...
)

This means that Nodes are now completely re-usable and can be subclassed for extensibility, and outputs are more sensible primitive types that can be handled differently in different contexts. This allows people to use these building blocks to create their own custom agents.

Edge routing logic can also be customised, then various node and edge configurations combined together in different ways to build flexible graphs. And I believe it will still be possible to build graph diagrams and ensure type safety using the type hinting.

The next challenge would be to figure out how to deal with graph dependency and state typing, since they are also currently tightly coupled to the specific graph configuration

[asaf]

While I truly see your points in practice things are more complex than just 'throwing in nodes'.

This is why any flow engine (even the visualized one that aims for non tech personas) only encapsulates the logic in the node but you still can customize per graph at least:

1. Previous and next node.
2. Where the input of the node comes from.
3. To where it outputs (mostly)

I do see the point of the conditional edges, they are useful but they makes things more complex as you have another entity that encapsulates logic, I'll leave the decision whether to use them or not to the authors :)

My two cents: Generalize your logic, use it in different nodes per graph, in most cases that is easy enough to maintain and does not dictates the graph layout.

[Finndersen]

Yeah the edge functions should only serve to connect/route nodes based on the node output value.

Currently the useful logic of the graph Agent for example is contained within the nodes so can't really be re-used elsewhere. And with a class based approach, can break out pieces of functionality into methods that can be overridden as desired to customise individual pieces of behaviour of a node.

Curious to hear @samuelcolvin or @dmontagu thoughts on this

[izzyacademy]

Guys, it is possible for you to just define your static edge-unaware "Nodes" as classes that can be reused/extended/sub-classed in a "Router" pydantic Graph node? Ifyou want you may also use them in composition rather than via inheritance. I dont think you need edge-functions like what LangGraph is doing.

This approach is more elegant to me. I think we are probably stuck with the LangGraph mindset that is why we are seeing these complaints.

[samuelcolvin]

Sorry, I but don't think we need to change anything.

If you want to re-use code, just write a standalone function, and call it from nodes.

Obviously you could implement something with mixins, but I doubt that would be a good idea.

[Finndersen]

If that's the pattern you suggest then I think it should be followed for the graph Agent implementation, since its nodes contain a lot of logic that would be valuable to be reused or extended for custom agent implementations, but currently cannot be.

Like these components for example should be extracted into reusable functions so they can be used as desired in different contexts to build custom agents:

• dynamic system prompt construction
• making requests to model
• parsing model response into tool call actions
• handling/executing tool call actions and returning responses

Currently the graph Agent is still implemented as an immutable black box which I think is quite limiting