From 7930d01311a54e072688a009167a0f5738f07354 Mon Sep 17 00:00:00 2001 From: Carson Date: Tue, 6 Oct 2026 19:22:22 -0500 Subject: [PATCH] fix(chat): make get_tokens() robust to non-alternating turn histories get_tokens() computed per-turn token counts by pairing user and assistant turns positionally, which only works when turns strictly alternate. A trailing user turn with no completed assistant response (e.g. after a failed request or add_turn()), consecutive assistant turns, or a lone user turn raised ValueError. The computation is now driven off complete assistant turns: since each assistant turn's input token count is cumulative, implied user tokens fall out of consecutive differences without any alternation assumption. A history with no complete assistant turns now returns an empty list. Ports tidyverse/ellmer#1157 (fixes tidyverse/ellmer#1131). --- CHANGELOG.md | 1 + chatlas/_chat.py | 144 +++++++++++++++---------------------------- tests/test_tokens.py | 69 +++++++++++++++++++++ 3 files changed, 118 insertions(+), 96 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7501814e..fc295228 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,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. * `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) * `ChatBedrock()` (with the default `api="converse"`) no longer sends assistant turns with an empty `content` array, which Converse rejects. This happens when a response carries no content blocks, for example when a guardrail intervenes before the model produces any. A `"[empty string]"` placeholder is sent instead, matching how empty text content is already normalized. (#426) diff --git a/chatlas/_chat.py b/chatlas/_chat.py index 82d742bd..d9d549f5 100644 --- a/chatlas/_chat.py +++ b/chatlas/_chat.py @@ -484,17 +484,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: """ @@ -596,16 +585,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 = ( @@ -614,84 +609,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 diff --git a/tests/test_tokens.py b/tests/test_tokens.py index d3c19515..6034b2ee 100644 --- a/tests/test_tokens.py +++ b/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")