aminghadersohi commented on code in PR #37200: URL: https://github.com/apache/superset/pull/37200#discussion_r2775724706
########## superset/mcp_service/utils/token_utils.py: ########## @@ -0,0 +1,404 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +Token counting and response size utilities for MCP service. + +This module provides utilities to estimate token counts and generate smart +suggestions when responses exceed configured limits. This prevents large +responses from overwhelming LLM clients like Claude Desktop. +""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, List + +logger = logging.getLogger(__name__) + +# Approximate characters per token for estimation +# Claude tokenizer averages ~4 chars per token for English text +# JSON tends to be more verbose, so we use a slightly lower ratio +CHARS_PER_TOKEN = 3.5 + + +def estimate_token_count(text: str | bytes) -> int: + """ + Estimate the token count for a given text. + + Uses a character-based heuristic since we don't have direct access to + the actual tokenizer. This is conservative to avoid underestimating. + + Args: + text: The text to estimate tokens for (string or bytes) + + Returns: + Estimated number of tokens + """ + if isinstance(text, bytes): + text = text.decode("utf-8", errors="replace") + + # Simple heuristic: ~3.5 characters per token for JSON/code + return int(len(text) / CHARS_PER_TOKEN) + + +def estimate_response_tokens(response: Any) -> int: + """ + Estimate token count for an MCP tool response. + + Handles various response types including Pydantic models, dicts, and strings. + + Args: + response: The response object to estimate + + Returns: + Estimated number of tokens + """ + try: + from superset.utils import json + + # Convert response to JSON string for accurate estimation + if hasattr(response, "model_dump"): + # Pydantic model + response_str = json.dumps(response.model_dump()) + elif isinstance(response, (dict, list)): + response_str = json.dumps(response) + elif isinstance(response, (str, bytes)): + response_str = response if isinstance(response, str) else response.decode() + else: + response_str = str(response) + + return estimate_token_count(response_str) + except Exception as e: # noqa: BLE001 + logger.warning("Failed to estimate response tokens: %s", e) + # Return a high estimate to be safe (conservative fallback) + return 100000 + + +def get_response_size_bytes(response: Any) -> int: + """ + Get the size of a response in bytes. + + Args: + response: The response object + + Returns: + Size in bytes + """ + try: + from superset.utils import json + + if hasattr(response, "model_dump"): + response_str = json.dumps(response.model_dump()) + elif isinstance(response, (dict, list)): + response_str = json.dumps(response) + elif isinstance(response, bytes): + return len(response) + elif isinstance(response, str): + return len(response.encode("utf-8")) + else: + response_str = str(response) + + return len(response_str.encode("utf-8")) + except Exception as e: # noqa: BLE001 + logger.warning("Failed to get response size: %s", e) + # Return a conservative large value to avoid allowing oversized responses + # to bypass size checks (returning 0 would underestimate) + return 1_000_000 # 1MB fallback + + +def extract_query_params(params: Dict[str, Any] | None) -> Dict[str, Any]: + """ + Extract relevant query parameters from tool params for suggestions. + + Args: + params: The tool parameters dict + + Returns: + Extracted parameters relevant for size reduction suggestions + """ + if not params: + return {} + + # Handle nested request object (common pattern in MCP tools) + if "request" in params and isinstance(params["request"], dict): + params = params["request"] + + extracted = {} + + # Pagination params + if "page_size" in params: + extracted["page_size"] = params["page_size"] + if "limit" in params: + extracted["limit"] = params["limit"] + + # Column selection + if "select_columns" in params: + extracted["select_columns"] = params["select_columns"] + if "columns" in params: + extracted["columns"] = params["columns"] + + # Filters + if "filters" in params: + extracted["filters"] = params["filters"] + + # Search + if "search" in params: + extracted["search"] = params["search"] + + return extracted + + +def generate_size_reduction_suggestions( + tool_name: str, + params: Dict[str, Any] | None, + estimated_tokens: int, + token_limit: int, + response: Any = None, +) -> List[str]: + """ + Generate smart suggestions for reducing response size. + + Analyzes the tool and parameters to provide actionable recommendations. + + Args: + tool_name: Name of the MCP tool + params: The tool parameters + estimated_tokens: Estimated token count of the response + token_limit: Configured token limit + response: Optional response object for additional analysis + + Returns: + List of suggestion strings + """ + suggestions = [] + query_params = extract_query_params(params) + reduction_needed = estimated_tokens - token_limit + reduction_pct = int((reduction_needed / estimated_tokens) * 100) + + # Suggestion 1: Reduce page_size or limit + current_page_size = query_params.get("page_size") or query_params.get("limit") + if current_page_size: + # Calculate suggested new limit based on reduction needed + suggested_limit = max( + 1, int(current_page_size * (token_limit / estimated_tokens)) + ) + suggestions.append( + f"Reduce page_size/limit from {current_page_size} to {suggested_limit} " + f"(need ~{reduction_pct}% reduction)" + ) + else: + # No limit specified - suggest adding one + if "list_" in tool_name or tool_name.startswith("get_chart_data"): + suggestions.append( + f"Add a 'limit' or 'page_size' parameter (suggested: 10-25 items) " + f"to reduce response size by ~{reduction_pct}%" + ) + + # Suggestion 2: Use select_columns to reduce fields + current_columns = query_params.get("select_columns") or query_params.get("columns") + if current_columns and len(current_columns) > 5: + suggestions.append( + f"Reduce select_columns from {len(current_columns)} columns to only " + f"essential fields (currently: {current_columns[:3]}...)" + ) Review Comment: The logic is correct: `current_columns` only enters this block when `len(current_columns) > 5` (line 226), so slicing `[:3]` always returns exactly 3 items. The suffix logic `len(current_columns) > 3` is always true in this branch. Python slicing handles shorter lists gracefully without IndexError. -- 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]
