bito-code-review[bot] commented on PR #42244:
URL: https://github.com/apache/superset/pull/42244#issuecomment-5025467500
<!-- Bito Reply -->
The flagged issue is correct. In `_try_truncate_data_query_response`, the
token estimation is performed on the `truncated` payload alone, ignoring the
overhead of the `ToolResult` wrapper that is added back before returning. This
can result in a returned object that still exceeds the `token_limit`.
To resolve this, you should re-estimate the tokens on the final object after
rewrapping, or account for the wrapper overhead during the check. Here is the
corrected logic for the re-check:
```python
# ... (after truncation)
if not was_truncated:
return None
# Re-wrap to check total size including wrapper overhead
final_result = self._rewrap_as_tool_result(truncated, response) if
extracted is not None and isinstance(truncated, dict) else truncated
truncated_tokens = estimate_response_tokens(final_result)
if truncated_tokens > self.token_limit:
return None
# ... (proceed with logging and return final_result)
```
Would you like me to check the other comments on this PR and implement fixes
for them as well?
**superset/mcp_service/middleware.py**
```
# Re-wrap to check total size including wrapper overhead
final_result = self._rewrap_as_tool_result(truncated, response) if
extracted is not None and isinstance(truncated, dict) else truncated
truncated_tokens = estimate_response_tokens(final_result)
if truncated_tokens > self.token_limit:
return None
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
