Token counting in AgentWorkflows

[martimfasantos]

Hi,
What is the best way of implementing a TokenCounter for an AgentWorkflow. My idea was to add a token counter to each agent and try to sum the counts in the end but that doesn't seem possible when I'm not the one calling the agent - the Workflow is.
It is also not possible to add TokenCountingHandler to the AgentWorkflow.
If anyone implemented this already, feel free to share your solution :)

[dosubot[bot]]

To implement a TokenCounter for an AgentWorkflow, you can utilize the existing utilities in the LlamaIndex codebase. The llama_index.core.callbacks.token_counting module provides a TokenCountingHandler that can be used to count tokens in LLM and Embedding events. This handler can be integrated into your workflow to track token usage across different agents.

Here's a general approach you can take:

1. Initialize a TokenCountingHandler: Create an instance of the TokenCountingHandler at the start of your workflow. This handler will be responsible for counting tokens throughout the execution of the workflow.

2. Attach the Handler to the Workflow: Ensure that the TokenCountingHandler is attached to the AgentWorkflow so that it can listen to events and count tokens as agents perform their tasks.

3. Aggregate Token Counts: At the end of the workflow, you can retrieve the total token counts from the TokenCountingHandler. This will give you the cumulative token usage across all agents in the workflow.

By using the TokenCountingHandler, you can effectively track and aggregate token usage without needing to manually sum counts from individual agents. This approach leverages the existing infrastructure in LlamaIndex for token counting, ensuring synchronization and accuracy.

For more details on the TokenCountingHandler, you can refer to the llama_index.core.callbacks.token_counting module ^[1].

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}

[martimfasantos]

But is it possible to add the token counter directly on the AgentWorkflow? It doesn't seem possible

[1-SarthakSarangi]

@dosu is it possible to add the token counter directly on the AgentWorkflow? It doesn't seem possible

[dosubot[bot]]

You're correct—it's not currently possible to attach a TokenCountingHandler (or similar token counter) directly to an AgentWorkflow. The design expects handlers like TokenCountingHandler to be integrated at a higher level, such as the global CallbackManager, rather than on the workflow itself. This is a known limitation and is reflected in both the code and community discussions—there's no documented or code-supported way to inject handlers directly into AgentWorkflow at this time.

If you need token counting for workflows, the recommended approach is to attach your TokenCountingHandler to the global CallbackManager (e.g., via Settings.callback_manager) so it observes all relevant events, or to implement a custom solution that parses token counts from LLM responses as they flow through the workflow. Some users have shared custom solutions for this in the community, since the built-in handler is aging and may not cover all use cases or new abstractions source.

Let me know if you want pointers to example custom solutions or more details on integrating with the CallbackManager.

To reply, just mention @dosu.

---

^{How did I do? Good | Irrelevant | Incorrect | Verbose | Hallucination | Report 🐛 | Other}