gabotorresruiz commented on code in PR #44575: URL: https://github.com/apache/superset/pull/44575#discussion_r4087613490
########## tests/unit_tests/mcp_service/test_tool_inventory.py: ########## @@ -0,0 +1,147 @@ +# 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. + +"""Per-tool size and schema fidelity budgets for the entire registered inventory.""" + +from copy import deepcopy + +import pytest +import tiktoken +from jsonschema import Draft202012Validator + +from superset.mcp_service.app import mcp +from superset.mcp_service.mcp_config import MCP_TOOL_SEARCH_CONFIG +from superset.mcp_service.server import _create_search_result_serializer, _strip_titles +from superset.utils import json + +# Compact JSON including tool metadata, measured independently as UTF-8 bytes +# and tiktoken 0.14.0 / cl100k_base tokens (BPE estimates, not Claude counts). +# Small tools round up to 100 bytes / 25 tokens; large chart tools retain their +# explicit delivery budgets. Keep both limits: bytes are not a token estimate. +TOOL_BUDGETS = { + "add_chart_to_existing_dashboard": (1_400, 325), + "apply_dashboard_filters": (2_800, 625), + "create_dataset": (1_700, 375), + "create_theme": (1_000, 225), + "create_virtual_dataset": (3_600, 800), + "delete_chart": (1_000, 225), + "delete_dashboard": (1_000, 225), + "duplicate_dashboard": (1_800, 375), + "execute_sql": (2_000, 450), + "find_users": (1_400, 325), + "generate_bug_report": (2_500, 575), + "generate_chart": (50_000, 20_000), + "generate_dashboard": (3_300, 725), + "generate_explore_link": (50_000, 20_000), + "get_annotation_layer_info": (900, 200), + "get_chart_data": (2_800, 625), + "get_chart_info": (3_500, 800), + "get_chart_preview": (3_300, 750), Review Comment: This block worries me a bit. The header says small tools round up to 100 bytes and 25 tokens, but rounding a measurement that already lands on the increment gives the measurement back, so six entries ship with zero headroom. Measured on this branch with the exact arithmetic `test_tool_inventory_size` uses: `get_chart_preview` is 3300 bytes against a 3300 byte budget, and `manage_dashboard_certification` (400), `apply_dashboard_filters` (625), `list_users` (650), `list_datasets` (1000) and `list_dashboards` (1000) each sit exactly on their token budget. Eleven tools have 10 bytes or less of slack. I checked what that costs the next contributor. On this branch I changed the `get_chart_preview` docstring from `Get chart preview by ID or UUID.` to `Get a chart preview by ID or UUID.` and got: ``` FAILED tests/unit_tests/mcp_service/test_tool_inventory.py::test_tool_inventory_size[get_chart_preview] AssertionError: ('get_chart_preview', 3301, 3300) ``` Two added bytes in a docstring, in a file that PR never opened. Green CI here will not warn us about it either. Could the ceiling always add a full increment before the table is regenerated, something like: ```python byte_budget = (measured_bytes // 100 + 1) * 100 token_budget = (measured_tokens // 25 + 1) * 25 ``` That keeps everything you are guarding (the chart tools stay pinned roughly 45 kB under master) without making every description edit a red build. ########## pyproject.toml: ########## @@ -327,6 +327,7 @@ development = [ "sqloxide", "statsd", "syntaqlite>=0.9.0,<0.10.0", + "tiktoken==0.14.0", # reproducible MCP inventory token budgets Review Comment: Not a blocker on its own, but I want to flag what this dependency costs at test time. `tiktoken` does not ship the vocabulary in the wheel, it fetches it on first use. I pointed `TIKTOKEN_CACHE_DIR` at an empty directory and ran `tiktoken.get_encoding("cl100k_base")` on this branch: it wrote a 1,681,126 byte blob pulled from `https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken`. With the network unreachable the same call raises, so `test_tool_inventory.py` and `test_chart_tool_inventory.py` error at collection time rather than fail with a useful message. That puts an external CDN in the path of `pytest ./tests/unit_tests` for every contributor and every CI run, and `pytest.importorskip("tiktoken")` would not help since the import is the part that succeeds. The byte budgets already catch the regression this PR is about. I confirmed `test_chart_tool_inventory_size` fails on `f8f293d2` for all three tools on the byte assertion alone, before the token one is reached. So the simplest option is to drop the token assertions and this dependency. If the token numbers are worth keeping, could we gate them on the encoding already being cached so an offline or network restricted run still passes? ########## docs/admin_docs/configuration/mcp-server.mdx: ########## @@ -775,7 +775,7 @@ MCP_TOOL_SEARCH_CONFIG = { | `max_results` | `5` | Maximum tools returned per search query | | `always_visible` | See above | Tools that always appear in `list_tools`, regardless of search | | `include_schemas` | `False` | When `False` (default, "summary mode"), search results omit `inputSchema` entirely and include a lightweight `parameters_hint` listing top-level parameter names. Set to `True` to include the full `inputSchema` in search results. Full schemas are always used when a tool is actually invoked via `call_tool`. | -| `compact_schemas` | `True` | Strip `$defs` / `$ref` and replace with `{"type": "object"}` in search results to reduce token cost. Only takes effect when `include_schemas=True` — ignored in summary mode. | +| `compact_schemas` | `True` | Legacy setting selecting the default description limit (300 when `True`, 0 when `False`) if `max_description_length` is omitted. Input schemas preserve `$defs`, `$ref`, nullable unions, and constraints in either mode; titles are omitted. Clients should resolve references within each tool's `inputSchema`. | Review Comment: Good catch updating this row. The same text lives in code and still describes the removed behaviour: `superset/mcp_service/mcp_config.py` has a `Schema Compaction:` block starting at line 466 saying search results `strip $defs sections and replace $ref pointers with {"type": "object"}`, a rollback bullet at line 476 offering `compact_schemas=False` to get `full $defs ... in search results`, and the inline comment on line 505 still reads `# Strip $defs/$ref (requires include_schemas=True)`. That module is what an operator reads when they copy `MCP_TOOL_SEARCH_CONFIG` into `superset_config.py`, so it is worth the same edit you made here. ########## superset/mcp_service/server.py: ########## @@ -535,22 +430,20 @@ def _create_search_result_serializer( ~80% vs compact mode while still conveying what parameters a tool accepts. - When ``include_schemas`` is True, the full ``compact_schemas``/ - ``max_description_length`` pipeline applies (existing behavior): - - * ``$defs`` sections and ``$ref`` pointers are collapsed when - ``compact_schemas`` is True (see :func:`_compact_schema`). - * Tool descriptions are truncated to ``max_description_length`` chars. + When ``include_schemas`` is True, input schemas retain their definitions, + references, and validation constraints. Inlining references duplicates shared + chart models and can make a single tool exceed client result limits. - Full schemas remain available when the tool is invoked via ``call_tool``. + Titles and output schemas are stripped by the base serializer. The legacy + ``compact_schemas`` setting only selects the default description limit; + ``max_description_length`` explicitly controls description truncation. """ include_schemas = config.get("include_schemas", False) if not include_schemas: max_desc = config.get("max_description_length", 300) return _build_summary_serializer(max_desc) - # include_schemas=True: apply full compact_schemas/max_description_length pipeline compact = config.get("compact_schemas", True) Review Comment: Just a small NIT: with the compaction gone, `compact` only picks the default for `max_desc`, and both return paths are now the same function. `compact_schemas=True` with `max_description_length=0` falls through to `_serializer`, which calls `_serialize_tools_without_output_schema` and then does nothing to the results. The whole block reads the same as: ```python compact = config.get("compact_schemas", True) max_desc = config.get("max_description_length", 300 if compact else 0) if not max_desc: return _serialize_tools_without_output_schema ``` Fewer moving parts around a setting we are now calling legacy. Non blocking either way. ########## docs/admin_docs/configuration/mcp-server.mdx: ########## @@ -775,7 +775,7 @@ MCP_TOOL_SEARCH_CONFIG = { | `max_results` | `5` | Maximum tools returned per search query | | `always_visible` | See above | Tools that always appear in `list_tools`, regardless of search | | `include_schemas` | `False` | When `False` (default, "summary mode"), search results omit `inputSchema` entirely and include a lightweight `parameters_hint` listing top-level parameter names. Set to `True` to include the full `inputSchema` in search results. Full schemas are always used when a tool is actually invoked via `call_tool`. | Review Comment: Just a small NIT, and it predates this PR, but you are already in this table. This row says `include_schemas` defaults to `False`, while `MCP_TOOL_SEARCH_CONFIG` in `superset/mcp_service/mcp_config.py:507` ships `"include_schemas": True`. Only the `config.get()` fallback inside `_create_search_result_serializer` is `False`, and nothing reaches it with the shipped config. So an operator reading this row concludes summary mode is the shipped default when full schemas actually are, which is exactly the mode your change affects. Cheap to correct in the same pass. -- 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]
