fix/over-counts tool outputs and blocks legitimate conversations#3316
fix/over-counts tool outputs and blocks legitimate conversations#3316ezkemboi wants to merge 12 commits into
Conversation
![codescene-delta-analysis[bot]](https://avatars.githubusercontent.com/in/53921?s=60&v=4){kind=small}
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
|
@snopoke, I am doing more testing on this PR, later today and tomorrow, for it to be ready. |
📝 WalkthroughWalkthroughThis PR updates the message size validation middleware to focus token-budget enforcement on only the most recent Estimated code review effort🎯 2 (Simple) | ⏱️ ~12 minutes Possibly related issues
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@apps/pipelines/nodes/history_middleware.py`:
- Around line 148-153: The current early-exit uses "if not self._token_limit"
which treats 0 as "no limit" and skips validation; change that to only skip when
the limit is unknown (e.g., "if self._token_limit is None") so a token_limit of
0 is handled as a real (rejecting) limit, then ensure the validation logic
compares token usage against self._token_limit (using <= or < as appropriate)
and rejects when the limit is exceeded or equals zero; look for _token_limit and
_build_size_validation_middleware to update the guard, and ensure
count_tokens_approximately and HumanMessage handling remain unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 4cf33224-306d-4be3-be6b-a72b0267757e
📒 Files selected for processing (3)
apps/pipelines/nodes/history_middleware.pyapps/pipelines/nodes/llm_node.pyapps/pipelines/tests/test_llm_node.py
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
|
From the comment #3196 (comment), I have also removed counter on this end. |
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
|
@snopoke and @SmittieC, I’ve removed the frontend counter from this test. Also, on the current change, we only value the last message and not with history. The backend now uses _count_tokens, which relies on get_num_tokens_from_messages for accurate token counting, including proper handling of JSON-structured inputs. For models that don’t support get_num_tokens_from_messages, we fall back to count_tokens_approximately (langchain-based estimation). I have attached a screenshot of the test done. cc. @KipmurkorDev, please help retest the changes updated. |
|
|
||
| # Reserve space for the system message so trigger/keep thresholds reflect usable context | ||
| system_message_tokens = count_tokens_approximately([system_message]) | ||
| from apps.pipelines.nodes.history_middleware import _count_tokens # noqa: PLC0415 |
There was a problem hiding this comment.
can this import be moved to the top?
| if not max_token_limit: | ||
| return None | ||
| system_tokens = count_tokens_approximately([system_message]) | ||
| system_tokens = _count_tokens(model, [system_message]) |
There was a problem hiding this comment.
Rather make _count_tokens non-private if we're going to be importing it from another module
| system_message_tokens = count_tokens_approximately([system_message]) | ||
| from apps.pipelines.nodes.history_middleware import _count_tokens # noqa: PLC0415 | ||
|
|
||
| system_message_tokens = _count_tokens(model, [system_message]) if model is not None else 0 |
There was a problem hiding this comment.
Why assume 0 tokens if the model is missing? _count_tokens will fall back to approximating it if the model is missing, which seems like something we'd want?
| from apps.pipelines.exceptions import MessageTooLargeError | ||
|
|
||
|
|
||
| def _count_tokens(model, messages: list) -> int: |
There was a problem hiding this comment.
Suggestion:
You're passing the model from the node level down to here just to call .get_num_tokens_from_messages on it. What if you instead make a method that accepts an optional model param and return the token counter to use. Something like this:
def get_token_counter(model=None):
return model.get_num_tokens_from_messages if model else count_tokens_approximatelyThat way, the history middleware (and _build_size_validation_middleware) doesn't have to get a model from its callers, but just the token counter to use.
Regarding the try-catch behaviour: What scenarios might there be where model.get_num_tokens_from_messages raises and count_tokens_approximately doesn't? Maybe the caller should be responsible for doing any try-catching?
| f"Your message is too large for this model. " | ||
| f"It uses approximately {token_count} tokens, but only {self._token_limit} tokens are available." | ||
| f"It uses approximately {token_count} tokens, but only {self._token_limit} tokens are available " | ||
| f"after accounting for the system prompt." |
There was a problem hiding this comment.
after accounting for the system prompt
Nit: Just double checking, is there a specific reason why we need to include this?
| if not human_messages: | ||
| return None | ||
| token_count = count_tokens_approximately(state["messages"]) | ||
| token_count = _count_tokens(self._model, [human_messages[-1]]) |
There was a problem hiding this comment.
Do you know if the history middleware will remove all history when the human message is long enough that it effectively takes up all of the token space that is left over after subtracting the system message's tokens?
I have a feeling that a large human message might still fail at the API level, even though it passes here, unless the history middleware is smart enough to throw everything away except the last human and system message. Do you perhaps know if this is the case?
There was a problem hiding this comment.
Previously, the validator only checked the last human message against the token limit. The compressors (summarize/truncate) handled reducing history, but if they couldn't e.g. SUMMARIZE with fewer than 20 messages, or MAX_HISTORY_LENGTH which is count-based, a large message + existing history could slip through and fail silently at the API. That's exactly the edge case you caught.
The changes add a second check on the total conversation history (HumanMessages + AIMessages) that runs after the compressor. If the compressor couldn't reduce context enough, this raises a clear error before the API is called. ToolMessages are excluded because they sit between agent-loop iterations and aren't user-controlled, that was actually the root cause of #3303 where the previous total-context check was firing incorrectly.
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
|
@SmittieC, we are checking on the feedback and retesting based on the knowledge shared with @KipmurkorDev |
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
| f"It uses approximately {token_count} tokens, but only {self._token_limit} tokens are available." | ||
| ) | ||
|
|
||
| # Validate total history, excluding tool responses and the system prompt. |
There was a problem hiding this comment.
@SmittieC you can check on this approach and let me know what are your thoughts?
| and history size is managed by the compression middleware that runs before this one. | ||
| Two checks run on each call: the last HumanMessage is validated individually, and | ||
| the total conversation history (HumanMessages + AIMessages) is validated together. | ||
| ToolMessages and SystemMessages are excluded from the total — tool responses are |
There was a problem hiding this comment.
I don't quite follow why tool responses are excluded. They are also sent to the LLM along with the history + user's new message, so it should be included, no?
There was a problem hiding this comment.
I will include that as well in this PR.
Previously, I thought we were only validating the latest user message against the configured max_token_limit. Since some models support hundreds of thousands of tokens, whereas in our case we may configure a lower limit, such as 8192, I assumed the validation was only applied to the latest message and not to the entire payload.
From your explanation, I now understand that the validation should apply to the full request sent to the LLM, including the conversation history and tool responses, since all of that contributes to actual token usage.
| middleware.append(history_middleware) | ||
| if size_middleware := _build_size_validation_middleware(node, system_message, model): | ||
| middleware.append(size_middleware) | ||
| middleware = node.build_history_middleware(system_message=system_message, model=model) |
There was a problem hiding this comment.
Suggestion: Keep the history middleware as it was before and always include MessageSizeValidationMiddleware in the middleware, since there will always be a max token limit.
There was a problem hiding this comment.
Well noted. I will do so.
snopoke
left a comment
snopoke
left a comment
There was a problem hiding this comment.
@ezkemboi @SmittieC I'm wondering about the utility of this feature. The issue it is trying to resolve is quite rare and requires the user to upload a large document and paste a very large amount of text.
Given the complexity that validating this is turning out to be I'd like to propose we drop it entirely.
@snopoke, since @SmittieC was the creator of the same issue, his feedback carries more weight on the proposal above. |
I was starting to think along the same lines. The only benefit we get by checking and failing early is user feedback. If we still want the user feedback - which I do think is nice, because it tells users what they can do about it - we might as well let the API call fail, parse the failure and reraise a custom user friendly error. |
I think what you have mentioned above makes sense. |
|
+1 to parsing the error - Here is the error you get from OpenAI: I'm not sure what you get from the other providers. |
I will take a look at this PR this weekend so that we can get it out of line. |
for copilot model of Gpt is also the same as mentioned Reason: Request Failed: 400 {"error":{"message":"This model's maximum context length is 128000 tokens. However, your messages resulted in 132438 tokens (108594 in the messages, 23844 in the functions). Please reduce the length of the messages or functions.","code":"invalid_request_body"}} |
|
@snopoke noted. |
3 participants
Product Description
Users sending short messages were sometimes seeing errors like:
…even though their input was only a few words. This happened when a tool called during the conversation returned a large response (e.g. a Custom Action returning structured JSON). This fix resolves those false positives and re-enables the middleware that was temporarily disabled in #3302.
The frontend character counter (added in #3180) has also been removed. It was built on the same flawed approximation as the middleware and independently produced false positives — blocking the send button for messages the model would have accepted. The server-side error message now surfaces directly in the chat when a message is genuinely too large.
Technical Description
MessageSizeValidationMiddleware.before_modelpreviously rancount_tokens_approximately(state["messages"])— counting all messages in state on every agent loop iteration, includingToolMessagecontent from tool calls.Three problems stacked:
ToolMessagewas counted, blowing past the token limit even though the user's input was small.effective_limitalready deducts system tokens from the budget, butstate["messages"]can include theSystemMessage, so those tokens were subtracted twice.count_tokens_approximatelyuseschars / 4, which overshoots JSON and structured content by 1.3–2× under real BPE. Even when the count was "too large", the model may have accepted the request.Fix: filter
state["messages"]to only the lastHumanMessagebefore counting, and switch tomodel.get_num_tokens_from_messages(real BPE tokenizer via tiktoken for OpenAI models).Key changes:
history_middleware.py—before_modelcounts onlyhuman_messages[-1:]usingmodel.get_num_tokens_from_messages;token_limit=Noneskips validation,token_limit=0(system prompt exhausted full budget) correctly rejectsllm_node.py— middleware re-enabled (was commented out in fix(pipelines): temporarily disable MessageSizeValidationMiddleware #3302);model = node.get_chat_model()extracted once and passed to both the middleware andcreate_agentto avoid double instantiation; system token count also uses real tokenizertest_llm_node.py— regression tests for ToolMessage false positive, empty HumanMessage state, multi-turn history, zero budget, and None budgetutils.py,views/utils.py,experiment.py,chatbots/views.py,pipelines/views.py—compute_max_char_limit/get_max_char_limitand allmax_char_limittemplate context keys removedinput_bar.html—maxCharLimit,messageTooLong,:maxlength, error class bindings removedapps/pipelines/utils.py— deleted (emptied by counter removal)test_view_utils.py— deleted (all tests were for the removedget_max_char_limit)Migrations
No migrations.
Demo
No visible UI change on the happy path. Users who were previously blocked by false positives will now be able to send messages normally. When a message genuinely exceeds the model's token budget, a user-facing error appears inline in the chat.
Docs and Changelog