How to use on_agent_finish callback?

[pedrorjbr]

I suppose that the on_agent_finish callback should be called on END Node. but it seems not working.

Here is a sample of code:

from typing import Dict, Any, Optional, List
from uuid import UUID

from langchain_core.agents import AgentFinish, AgentAction
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.runnables import RunnableConfig
from langgraph.constants import END
from langgraph.graph import StateGraph

class MyCustomHandler(BaseCallbackHandler):

    def on_chain_start(
        self,
        serialized: Dict[str, Any],
        inputs: Dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Any:
        print(f"on_chain_start")

    def on_agent_action(
        self,
        action: AgentAction,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> Any:
        print("on_agent_action")

    def on_agent_finish(
        self,
        finish: AgentFinish,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> Any:
        print("on_agent_finish")

    async def on_chain_end(
        self,
        outputs: Dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> None:
        print("on_chain_end")


def node_a(state: dict, config: RunnableConfig):
    print("Node a")


# The second argument is optional
def node_b(state: dict, config: RunnableConfig):
    print("Node b")


builder = StateGraph(dict)
builder.add_node("node_a", node_a)
builder.add_node("node_b", node_b)
builder.set_entry_point("node_a")
builder.add_edge("node_a", "node_b")
builder.add_edge("node_b", END)
workflow = builder.compile()
config = {"callbacks": [MyCustomHandler()]}
workflow.invoke({}, config)

This is the output:

on_chain_start
on_chain_start
on_chain_end
on_chain_start
Node a
on_chain_start
on_chain_end
on_chain_end
on_chain_start
Node b
on_chain_start
on_chain_end
on_chain_end
on_chain_end