fix: avoid double-counting context token overhead

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)

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

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
 

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)

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