fix(chat): make get_tokens() robust to non-alternating turn histories

CHANGELOG.md

@@ -20,6 +20,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Bug fixes
 
+* `Chat.get_tokens()` no longer errors when the conversation ends with a user turn that has no completed assistant response yet (e.g. after a failed request or `add_turn()`), when assistant turns don't strictly alternate with user turns, or when the history holds a single user turn (which now returns an empty list). Token accounting is now driven off complete assistant turns instead of assuming strict user/assistant alternation.
 * `ChatBedrockAnthropic(reasoning=)` now accepts an effort level (`"low"`, `"medium"`, `"high"`, `"xhigh"` or `"max"`) to enable adaptive thinking, as `ChatAnthropic()` already did. It previously only took a token budget, which its default model (Claude Sonnet 5) rejects, so there was no way to turn on thinking with that model.
 * `ChatBedrock()` (with `api="messages"` or `api="responses"`) no longer sends a URL for `content_pdf_url()`/`content_document_url()` content, which `bedrock-mantle` rejects since it can't fetch URLs itself. The bytes are downloaded and sent instead, as they already were before URL passthrough was added. `ChatBedrockAnthropic()` gets the same fix, since Bedrock's InvokeModel API can't fetch `url` sources either. (#410)
 * Bedrock-backed providers now fail fast with a clear error for `content_image_url()` content (remote images were already rejected by Bedrock with an opaque API error, since it can't fetch URLs). Download the image and use `content_image_file()` instead. (#410)

chatlas/_chat.py

@@ -485,17 +485,6 @@ def _complete_assistant_turns(self) -> list[AssistantTurn]:
             t for t in self._turns if isinstance(t, AssistantTurn) and not t.is_partial
         ]
 
-    def _complete_turn_pairs(self) -> list[Turn]:
-        """Return turns with partial assistant turns (and their user turns) removed."""
-        turns = self.get_turns(include_system_prompt=False)
-        partial_indices: set[int] = set()
-        for i, t in enumerate(turns):
-            if isinstance(t, AssistantTurn) and t.is_partial:
-                partial_indices.add(i)
-                if i > 0:
-                    partial_indices.add(i - 1)
-        return [t for i, t in enumerate(turns) if i not in partial_indices]
-
     @property
     def system_prompt(self) -> str | None:
         """
@@ -597,16 +586,22 @@ def get_tokens(self) -> list[TokensDict]:
         Raises
         ------
         ValueError
-            If the chat's turns (i.e., `.get_turns()`) are not in an expected
-            format. This may happen if the chat history is manually set (i.e.,
-            `.set_turns()`). In this case, you can inspect the "raw" token
-            values via the `.get_turns()` method (each turn has a `.tokens`
-            attribute).
-        """
-
-        turns = self._complete_turn_pairs()
-
-        if len(turns) == 0:
+            If an assistant turn doesn't contain token counts. This may happen
+            if the chat history is manually set (i.e., `.set_turns()`). In
+            this case, you can inspect the "raw" token values via the
+            `.get_turns()` method (each turn has a `.tokens` attribute).
+        """
+
+        # Only complete assistant turns carry token counts, and each one's
+        # input count is cumulative over everything sent so far, so the
+        # implied user tokens for a round trip is the difference between
+        # consecutive requests. Driving the computation off assistant turns
+        # (rather than assuming strict user/assistant alternation) means
+        # trailing user turns, consecutive assistant turns, and partial
+        # (interrupted) assistant turns don't trip it up.
+        assistants = self._complete_assistant_turns()
+
+        if len(assistants) == 0:
             return []
 
         err_info = (
@@ -615,84 +610,41 @@ def get_tokens(self) -> list[TokensDict]:
             "(each turn has a `.tokens` attribute)."
         )
 
-        # Sanity checks for the assumptions made to figure out user token counts
-        if len(turns) == 1:
-            raise ValueError(
-                "Expected at least two turns in the chat history. " + err_info
-            )
-
-        if len(turns) % 2 != 0:
-            raise ValueError(
-                "Expected an even number of turns in the chat history. " + err_info
-            )
-
-        if turns[0].role != "user":
-            raise ValueError(
-                "Expected the 1st non-system turn to have role='user'. " + err_info
-            )
-
-        if not isinstance(turns[1], AssistantTurn):
-            raise ValueError(
-                "Expected the 2nd turn non-system to have role='assistant'. " + err_info
-            )
-
-        if turns[1].tokens is None:
-            raise ValueError(
-                "Expected the 1st assistant turn to contain token counts. " + err_info
-            )
-
-        res: list[TokensDict] = [
-            # Implied token count for the 1st user input
-            {
-                "role": "user",
-                "tokens": turns[1].tokens[0],
-                # Number of tokens currently cached (reduces input token usage)
-                "tokens_cached": turns[1].tokens[2],
-                "tokens_total": turns[1].tokens[0],
-            },
-            # The token count for the 1st assistant response
-            {
-                "role": "assistant",
-                "tokens": turns[1].tokens[1],
-                "tokens_cached": 0,
-                "tokens_total": turns[1].tokens[1],
-            },
-        ]
-
-        for i in range(1, len(turns) - 1, 2):
-            ti = turns[i]
-            tj = turns[i + 2]
-            if not isinstance(ti, AssistantTurn) or not isinstance(tj, AssistantTurn):
-                raise ValueError(
-                    "Expected even turns to have role='assistant'." + err_info
-                )
-            if ti.tokens is None or tj.tokens is None:
+        res: list[TokensDict] = []
+        prev: AssistantTurn | None = None
+        for curr in assistants:
+            if curr.tokens is None:
                 raise ValueError(
-                    "Expected role='assistant' turns to contain token counts."
-                    + err_info
+                    "Expected assistant turns to contain token counts. " + err_info
                 )
-            res.extend(
-                [
-                    {
-                        "role": "user",
-                        # Implied new token count for the user input (input tokens - context - cached reads)
-                        # Cached reads are only subtracted for particular providers
-                        "tokens": tj.tokens[0] - sum(ti.tokens),
-                        # Number of tokens currently cached (reduces input token usage depending on provider's API)
-                        "tokens_cached": tj.tokens[2],
-                        # Total tokens = Total User Tokens for the Turn = Distinct new tokens + context sent
-                        "tokens_total": tj.tokens[0],
-                    },
-                    {
-                        "role": "assistant",
-                        # The token count for the assistant response
-                        "tokens": tj.tokens[1],
-                        # Total tokens = Total Assistant tokens used in the turn
-                        "tokens_cached": 0,
-                        "tokens_total": tj.tokens[1],
-                    },
-                ]
+            input_tokens, output_tokens, cached_tokens = curr.tokens
+            # Implied new token count for the user input (input tokens -
+            # context - cached reads). Cached reads are only subtracted for
+            # particular providers.
+            implied_input = input_tokens
+            if prev is not None and prev.tokens is not None:
+                implied_input -= sum(prev.tokens)
+            res.append(
+                {
+                    "role": "user",
+                    "tokens": implied_input,
+                    # Number of tokens currently cached (reduces input token usage depending on provider's API)
+                    "tokens_cached": cached_tokens,
+                    # Total tokens = Total User Tokens for the Turn = Distinct new tokens + context sent
+                    "tokens_total": input_tokens,
+                }
+            )
+            res.append(
+                {
+                    "role": "assistant",
+                    # The token count for the assistant response
+                    "tokens": output_tokens,
+                    "tokens_cached": 0,
+                    # Total tokens = Total Assistant tokens used in the turn
+                    "tokens_total": output_tokens,
+                }
             )
+            prev = curr
 
         return res
 

tests/test_tokens.py

@@ -11,6 +11,7 @@
     tokens_log,
     tokens_reset,
 )
+from chatlas.types import ContentToolRequest, ContentToolResult
 from pydantic import BaseModel
 
 from .conftest import make_vcr_config
@@ -75,6 +76,74 @@ def test_tokens_method():
     ]
 
 
+def test_get_tokens_tolerates_trailing_user_turn():
+    # A conversation can end with a user turn that has no completed
+    # assistant response yet (e.g. after a failed request or add_turn())
+    chat = ChatOpenAI()
+    chat.set_turns(
+        [
+            UserTurn("Hi"),
+            AssistantTurn("Hello", tokens=(2, 10, 0)),
+            UserTurn("Are you there?"),
+        ]
+    )
+    assert chat.get_tokens() == [
+        {"role": "user", "tokens": 2, "tokens_cached": 0, "tokens_total": 2},
+        {"role": "assistant", "tokens": 10, "tokens_cached": 0, "tokens_total": 10},
+    ]
+
+
+def test_get_tokens_with_only_a_user_turn():
+    chat = ChatOpenAI()
+    chat.set_turns([UserTurn("Hi")])
+    assert chat.get_tokens() == []
+
+
+def test_get_tokens_tolerates_consecutive_assistant_turns():
+    chat = ChatOpenAI()
+    chat.set_turns(
+        [
+            UserTurn("Hi"),
+            AssistantTurn("Hello", tokens=(2, 10, 0)),
+            AssistantTurn("Hello again", tokens=(20, 5, 0)),
+        ]
+    )
+    assert chat.get_tokens() == [
+        {"role": "user", "tokens": 2, "tokens_cached": 0, "tokens_total": 2},
+        {"role": "assistant", "tokens": 10, "tokens_cached": 0, "tokens_total": 10},
+        {"role": "user", "tokens": 8, "tokens_cached": 0, "tokens_total": 20},
+        {"role": "assistant", "tokens": 5, "tokens_cached": 0, "tokens_total": 5},
+    ]
+
+
+def test_get_tokens_with_tool_loop():
+    # Tool calls produce user turns holding tool results; token accounting
+    # should treat them like any other user turn
+    request = ContentToolRequest(id="x1", name="my_tool", arguments={})
+    chat = ChatOpenAI()
+    chat.set_turns(
+        [
+            UserTurn("Hi"),
+            AssistantTurn([request], tokens=(10, 5, 0)),
+            UserTurn([ContentToolResult(value=1, request=request)]),
+            AssistantTurn("The answer is 1", tokens=(30, 10, 0)),
+        ]
+    )
+    assert chat.get_tokens() == [
+        {"role": "user", "tokens": 10, "tokens_cached": 0, "tokens_total": 10},
+        {"role": "assistant", "tokens": 5, "tokens_cached": 0, "tokens_total": 5},
+        {"role": "user", "tokens": 15, "tokens_cached": 0, "tokens_total": 30},
+        {"role": "assistant", "tokens": 10, "tokens_cached": 0, "tokens_total": 10},
+    ]
+
+
+def test_get_tokens_raises_without_token_counts():
+    chat = ChatOpenAI()
+    chat.set_turns([UserTurn("Hi"), AssistantTurn("Hello")])
+    with pytest.raises(ValueError, match="token counts"):
+        chat.get_tokens()
+
+
 @pytest.mark.vcr
 def test_token_count_method():
     chat = ChatOpenAI(model="gpt-4o-mini")