Token count duplication in Invoke Agent API

When using the InvokeAgent API in Amazon Bedrock, you may indeed see two sets of input and output token counts being returned. This behavior is related to how the agent processes and maintains context throughout a conversation session.

The two sets of token counts you're observing likely represent different aspects of the agent's processing:

1. The first set of token counts probably represents the tokens processed for the current turn of the conversation. This includes the new user input and the agent's immediate response.

2. The second set, where you see increasing input token counts, likely represents the cumulative tokens processed for the entire conversation session. This includes all previous turns of the conversation, allowing the agent to maintain context.

The increasing input token count in the second set reflects the agent's ability to remember and consider the entire conversation history when generating responses. This is part of the session state management in Bedrock Agents.

Regarding your question about knowledge bases: While using a knowledge base with an agent can potentially increase token usage due to additional context being processed, the dual token count behavior you're observing is not specifically tied to knowledge base usage. It's more related to how the agent manages conversation context and session state.

It's important to note that for billing and usage tracking purposes, you should consider the sum of all token counts reported in each API response. This will give you the total number of tokens processed by the agent throughout the conversation.

If you need to manage token usage more precisely, you may want to consider strategies for limiting the conversation history or clearing the session state when appropriate for your use case.
Sources
Tokens processed in Amazon Bedrock | AWS re:Post
How Amazon Bedrock Agents works - Amazon Bedrock
InvokeAgent - Amazon Bedrock