From 7457dbc3955baa4d0772eac8b2a352e40619a4f0 Mon Sep 17 00:00:00 2001 From: "zhao.wang" Date: Thu, 20 Aug 2026 16:59:05 +0800 Subject: [PATCH] fix: avoid double-counting context token overhead --- .../atomic_agents/agents/atomic_agent.py | 16 ++- .../atomic_agents/utils/token_counter.py | 52 ++++++-- .../tests/utils/test_token_counter.py | 111 +++++++++++++++--- docs/api/agents.md | 18 ++- docs/api/utils.md | 33 ++++-- 5 files changed, 183 insertions(+), 47 deletions(-) diff --git a/atomic-agents/atomic_agents/agents/atomic_agent.py b/atomic-agents/atomic_agents/agents/atomic_agent.py index 10c77aa9..ab656e7b 100644 --- a/atomic-agents/atomic_agents/agents/atomic_agent.py +++ b/atomic-agents/atomic_agents/agents/atomic_agent.py @@ -496,14 +496,22 @@ def get_context_token_count(self) -> TokenCountResult: Returns: TokenCountResult: A named tuple containing: - - total: Total tokens in the context (including schema overhead) - - system_prompt: Tokens in the system prompt - - history: Tokens in the conversation history - - tools: Tokens in the tools/function definitions (TOOLS mode only) + - total: Total tokens in the complete message request; when tools + are provided without messages, this is schema-only overhead and + excludes request framing + - system_prompt: System message tokens, including request framing + when a system message is present + - history: Incremental tokens added by conversation history, + including request framing when no system message is present + - tools: Incremental tokens added by tool definitions (TOOLS mode only) - model: The model used for counting - max_tokens: Maximum context window (if known) - utilization: Percentage of context used (if max_tokens known) + The breakdown is additive: ``system_prompt + history + tools == total``. + When tools are provided without system or history messages, ``total`` + and ``tools`` report schema-only overhead and exclude request framing. + Example: ```python agent = AtomicAgent[InputSchema, OutputSchema](config) diff --git a/atomic-agents/atomic_agents/utils/token_counter.py b/atomic-agents/atomic_agents/utils/token_counter.py index 8215be85..f4571ce3 100644 --- a/atomic-agents/atomic_agents/utils/token_counter.py +++ b/atomic-agents/atomic_agents/utils/token_counter.py @@ -17,10 +17,15 @@ class TokenCountResult(NamedTuple): Result of a token count operation. Attributes: - total: Total number of tokens in the context (messages + tools). - system_prompt: Tokens in the system prompt (0 if no system prompt). - history: Tokens in the conversation history. - tools: Tokens in the tools/function definitions (0 if no tools). + total: Total number of tokens in the complete message request. When tools + are provided without messages, this is schema-only overhead and + excludes request framing. + system_prompt: Tokens in the system messages, including request framing + when a system message is present (0 otherwise). + history: Incremental tokens added by the history messages. Includes + request framing when there are no system messages. + tools: Incremental tokens added by tools to the complete message request. + When there are no messages, this is the schema-only token overhead. model: The model used for tokenization. max_tokens: Maximum context window for the model (None if unknown). utilization: Percentage of context window used (None if max_tokens unknown). @@ -170,6 +175,13 @@ def count_context( """ Count tokens with breakdown by system prompt, history, and tools. + The breakdown is additive: ``system_prompt + history + tools == total``. + Request framing is attributed to the system prompt when present, or to + history otherwise. Tool tokens are the increment over the full message + request without tools. When tools are provided without any messages, + ``total`` and ``tools`` report schema-only overhead and exclude request + framing. + Args: model: The model identifier. system_messages: System prompt messages (may be empty). @@ -183,18 +195,32 @@ def count_context( TokenCountError: If token counting fails. """ system_tokens = self.count_messages(model, system_messages) if system_messages else 0 - history_tokens = self.count_messages(model, history_messages) if history_messages else 0 + all_messages = system_messages + history_messages + + # Tokenizers include request-level framing, so independently counting the + # system prompt and history would count that overhead twice. Count the + # complete message list and attribute the incremental tokens to history. + if history_messages: + messages_tokens = self.count_messages(model, all_messages) + history_tokens = messages_tokens - system_tokens + else: + messages_tokens = system_tokens + history_tokens = 0 - # Count tool tokens separately if provided tools_tokens = 0 + total_tokens = messages_tokens if tools: - # To count just the tools overhead, we count empty messages with tools - # and subtract the base overhead - empty_with_tools = self.count_messages(model, [{"role": "user", "content": ""}], tools=tools) - empty_without_tools = self.count_messages(model, [{"role": "user", "content": ""}]) - tools_tokens = empty_with_tools - empty_without_tools - - total_tokens = system_tokens + history_tokens + tools_tokens + if all_messages: + total_tokens = self.count_messages(model, all_messages, tools=tools) + tools_tokens = total_tokens - messages_tokens + else: + # Preserve schema-only counting by subtracting the same placeholder + # request without tools. Keep the placeholder out of the total. + placeholder = [{"role": "user", "content": ""}] + with_tools_tokens = self.count_messages(model, placeholder, tools=tools) + without_tools_tokens = self.count_messages(model, placeholder) + tools_tokens = with_tools_tokens - without_tools_tokens + total_tokens = tools_tokens max_tokens = self.get_max_tokens(model) # Prevent division by zero diff --git a/atomic-agents/tests/utils/test_token_counter.py b/atomic-agents/tests/utils/test_token_counter.py index af53a5a0..0fc82638 100644 --- a/atomic-agents/tests/utils/test_token_counter.py +++ b/atomic-agents/tests/utils/test_token_counter.py @@ -1,5 +1,8 @@ +from unittest.mock import MagicMock, call, patch + import pytest -from unittest.mock import patch +from litellm import token_counter as litellm_token_counter + from atomic_agents.utils.token_counter import ( TokenCounter, TokenCountResult, @@ -203,7 +206,9 @@ def test_count_messages_raises_value_error_for_empty_model(self): @patch("litellm.get_model_info") @patch("litellm.token_counter") def test_count_context(self, mock_token_counter, mock_get_model_info): - mock_token_counter.side_effect = [30, 70] # system, then history + # Counting each group independently would report 30 + 70 = 100, but the + # complete request only has 95 tokens because framing overhead is shared. + mock_token_counter.side_effect = [30, 95] mock_get_model_info.return_value = {"max_input_tokens": 8192, "max_tokens": 4096} counter = TokenCounter() @@ -213,35 +218,109 @@ def test_count_context(self, mock_token_counter, mock_get_model_info): history_messages=[{"role": "user", "content": "Hello"}], ) - assert result.total == 100 + assert result.total == 95 assert result.system_prompt == 30 - assert result.history == 70 + assert result.history == 65 assert result.tools == 0 assert result.model == "gpt-4" assert result.max_tokens == 8192 - assert result.utilization == pytest.approx(100 / 8192) + assert result.utilization == pytest.approx(95 / 8192) @patch("litellm.get_model_info") @patch("litellm.token_counter") def test_count_context_with_tools(self, mock_token_counter, mock_get_model_info): - # system=30, history=70, empty_with_tools=60, empty_without_tools=10 -> tools=50 - mock_token_counter.side_effect = [30, 70, 60, 10] + # system=30, complete messages=95, complete request with tools=140 + mock_token_counter.side_effect = [30, 95, 140] mock_get_model_info.return_value = {"max_input_tokens": 8192, "max_tokens": 4096} counter = TokenCounter() + system_messages = [{"role": "system", "content": "You are helpful"}] + history_messages = [{"role": "user", "content": "Hello"}] tools = [{"type": "function", "function": {"name": "test_fn"}}] result = counter.count_context( model="gpt-4", - system_messages=[{"role": "system", "content": "You are helpful"}], - history_messages=[{"role": "user", "content": "Hello"}], + system_messages=system_messages, + history_messages=history_messages, tools=tools, ) assert result.system_prompt == 30 - assert result.history == 70 - assert result.tools == 50 - assert result.total == 150 # 30 + 70 + 50 + assert result.history == 65 + assert result.tools == 45 + assert result.total == 140 assert result.model == "gpt-4" + assert mock_token_counter.call_args_list == [ + call(model="gpt-4", messages=system_messages), + call(model="gpt-4", messages=system_messages + history_messages), + call(model="gpt-4", messages=system_messages + history_messages, tools=tools), + ] + + @patch("litellm.get_model_info") + def test_count_context_matches_complete_litellm_request(self, mock_get_model_info: MagicMock) -> None: + mock_get_model_info.return_value = {"max_input_tokens": 128000, "max_tokens": 16384} + + model = "gpt-4o-mini" + system_messages = [{"role": "system", "content": "You are a concise research assistant."}] + history_messages = [ + {"role": "user", "content": "Summarize retrieval-augmented generation."}, + {"role": "assistant", "content": "It grounds model responses in retrieved evidence."}, + ] + tools = [ + { + "type": "function", + "function": { + "name": "search_papers", + "description": "Search for relevant papers", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + }, + } + ] + complete_request_tokens = litellm_token_counter( + model=model, + messages=system_messages + history_messages, + tools=tools, + ) + + result = TokenCounter().count_context( + model=model, + system_messages=system_messages, + history_messages=history_messages, + tools=tools, + ) + + assert result.total == complete_request_tokens + assert result.total == result.system_prompt + result.history + result.tools + + @patch("litellm.get_model_info") + @patch("litellm.token_counter") + def test_count_context_with_tools_and_no_messages( + self, mock_token_counter: MagicMock, mock_get_model_info: MagicMock + ) -> None: + mock_token_counter.side_effect = [60, 10] + mock_get_model_info.return_value = {"max_input_tokens": 8192, "max_tokens": 4096} + + counter = TokenCounter() + tools = [{"type": "function", "function": {"name": "test_fn"}}] + result = counter.count_context( + model="gpt-4", + system_messages=[], + history_messages=[], + tools=tools, + ) + + placeholder = [{"role": "user", "content": ""}] + assert result.system_prompt == 0 + assert result.history == 0 + assert result.tools == 50 + assert result.total == 50 + assert mock_token_counter.call_args_list == [ + call(model="gpt-4", messages=placeholder, tools=tools), + call(model="gpt-4", messages=placeholder), + ] @patch("litellm.get_model_info") @patch("litellm.token_counter") @@ -265,7 +344,7 @@ def test_count_context_empty_system(self, mock_token_counter, mock_get_model_inf @patch("litellm.get_model_info") @patch("litellm.token_counter") def test_count_context_no_max_tokens(self, mock_token_counter, mock_get_model_info): - mock_token_counter.side_effect = [20, 30] + mock_token_counter.side_effect = [20, 45] mock_get_model_info.side_effect = Exception("Unknown model") counter = TokenCounter() @@ -275,7 +354,7 @@ def test_count_context_no_max_tokens(self, mock_token_counter, mock_get_model_in history_messages=[{"role": "user", "content": "Test"}], ) - assert result.total == 50 + assert result.total == 45 assert result.max_tokens is None assert result.utilization is None @@ -305,7 +384,7 @@ def test_count_messages_different_models(self, mock_token_counter): @patch("litellm.token_counter") def test_count_context_division_by_zero_prevention(self, mock_token_counter, mock_get_model_info): """Test that division by zero is prevented when max_tokens is 0.""" - mock_token_counter.side_effect = [20, 30] + mock_token_counter.side_effect = [20, 45] mock_get_model_info.return_value = {"max_input_tokens": 0, "max_tokens": 0} # Edge case counter = TokenCounter() @@ -315,7 +394,7 @@ def test_count_context_division_by_zero_prevention(self, mock_token_counter, moc history_messages=[{"role": "user", "content": "Test"}], ) - assert result.total == 50 + assert result.total == 45 assert result.max_tokens == 0 assert result.utilization is None # Should be None, not raise ZeroDivisionError diff --git a/docs/api/agents.md b/docs/api/agents.md index a794d755..d98d2ace 100644 --- a/docs/api/agents.md +++ b/docs/api/agents.md @@ -167,7 +167,7 @@ class BasicChatOutputSchema(BaseIOSchema): - `get_context_provider(provider_name: str)`: Get a registered context provider - `register_context_provider(provider_name: str, provider: BaseDynamicContextProvider)`: Register a new context provider - `unregister_context_provider(provider_name: str)`: Remove a context provider -- `get_context_token_count() -> TokenCountResult`: Get token count for current context (system prompt + history) +- `get_context_token_count() -> TokenCountResult`: Get the complete context token count with an additive breakdown ### Context Providers @@ -231,6 +231,7 @@ token_info = agent.get_context_token_count() print(f"Total tokens: {token_info.total}") print(f"System prompt (with schema): {token_info.system_prompt} tokens") print(f"History: {token_info.history} tokens") +print(f"Tools: {token_info.tools} tokens") print(f"Model: {token_info.model}") # Check context utilization if max tokens is known @@ -240,10 +241,19 @@ if token_info.utilization: print(f"Context utilization: {token_info.utilization:.1%}") ``` +The breakdown is additive: `system_prompt + history + tools == total`. Request +framing is attributed to `system_prompt` when a system message is present, or to +`history` otherwise. When tools are provided without system or history messages, +`total` and `tools` report schema-only overhead and exclude request framing. + The `TokenCountResult` contains: -- `total`: Total tokens in context (system + history + schema overhead) -- `system_prompt`: Tokens used by system prompt and output schema -- `history`: Tokens used by conversation history (including multimodal content) +- `total`: LiteLLM token count for the complete serialized context; with tools but + no messages, this is schema-only overhead and excludes request framing +- `system_prompt`: System message tokens, including request framing when a system + message is present; in JSON modes, this also includes the output schema +- `history`: Incremental tokens added by conversation history, including request + framing when no system message is present +- `tools`: Incremental tokens added by tool definitions in TOOLS mode - `model`: The model name used for counting - `max_tokens`: Maximum context window (if known) - `utilization`: Percentage of context used (if max_tokens known) diff --git a/docs/api/utils.md b/docs/api/utils.md index efb4135b..81cb9183 100644 --- a/docs/api/utils.md +++ b/docs/api/utils.md @@ -16,17 +16,25 @@ A named tuple containing token count information: .. py:attribute:: total :type: int - Total tokens in the context (system prompt + history + schema overhead). + LiteLLM token count for the complete serialized context. With tools but no + messages, this is schema-only overhead and excludes request framing. .. py:attribute:: system_prompt :type: int - Tokens used by the system prompt and output schema. + System message tokens, including request framing when a system message is + present. In JSON modes, this also includes the output schema. .. py:attribute:: history :type: int - Tokens used by conversation history (including multimodal content). + Incremental tokens added by conversation history, including request + framing when no system message is present. + + .. py:attribute:: tools + :type: int + + Incremental tokens added by tool definitions in TOOLS mode. .. py:attribute:: model :type: str @@ -76,13 +84,14 @@ The main utility class for counting tokens: :param model: The model name :return: Maximum tokens, or None if unknown - .. py:method:: count_context(model: str, system_messages: List[Dict], history_messages: List[Dict]) -> TokenCountResult + .. py:method:: count_context(model: str, system_messages: List[Dict], history_messages: List[Dict], tools: Optional[List[Dict]] = None) -> TokenCountResult - Count tokens for a complete context (system prompt + history). + Count tokens for a complete context with an additive breakdown. :param model: The model name :param system_messages: System prompt messages :param history_messages: Conversation history messages + :param tools: Optional tool definitions :return: TokenCountResult with detailed breakdown ``` @@ -126,16 +135,20 @@ token_info = agent.get_context_token_count() print(f"Total tokens: {token_info.total}") print(f"System prompt (with schema): {token_info.system_prompt} tokens") print(f"History: {token_info.history} tokens") +print(f"Tools: {token_info.tools} tokens") if token_info.utilization: print(f"Context utilization: {token_info.utilization:.1%}") ``` -The token count includes: -- System prompt content -- Output schema overhead (the JSON schema Instructor sends for structured output) -- Conversation history (including multimodal content like images, PDFs, audio) +The breakdown is additive: `system_prompt + history + tools == total`. Request +framing is attributed to `system_prompt` when a system message is present, or to +`history` otherwise. When tools are provided without system or history messages, +`total` and `tools` report schema-only overhead and exclude request framing. In +JSON modes, the output schema is included in `system_prompt`; in TOOLS mode, tool +definitions are reported through `tools`. -This gives you an accurate count that matches what would be sent to the API. +For requests with messages, the total matches LiteLLM's count for the complete +serialized request. ## Tool Message Formatting