mirror of
https://github.com/apache/superset.git
synced 2026-07-25 16:12:39 +00:00
855 lines
33 KiB
Python
855 lines
33 KiB
Python
# 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.
|
|
|
|
"""
|
|
Unit tests for MCP service token utilities.
|
|
"""
|
|
|
|
from typing import Any, List
|
|
from unittest.mock import patch
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from superset.mcp_service.utils import token_utils
|
|
from superset.mcp_service.utils.token_utils import (
|
|
_replace_collections_with_summaries,
|
|
_summarize_large_dicts,
|
|
_truncate_lists,
|
|
_truncate_strings,
|
|
_truncate_strings_recursive,
|
|
CHARS_PER_TOKEN,
|
|
estimate_response_tokens,
|
|
estimate_token_count,
|
|
extract_query_params,
|
|
format_size_limit_error,
|
|
generate_size_reduction_suggestions,
|
|
get_response_size_bytes,
|
|
INFO_TOOLS,
|
|
truncate_oversized_response,
|
|
truncate_query_result,
|
|
)
|
|
|
|
|
|
class TestEstimateTokenCount:
|
|
"""Test estimate_token_count function."""
|
|
|
|
def test_estimate_string(self) -> None:
|
|
"""Should produce a positive non-zero estimate for a normal string.
|
|
|
|
We don't assert on a specific number because the result depends on
|
|
which tokenizer is loaded (tiktoken when available, char heuristic
|
|
otherwise).
|
|
"""
|
|
text = "Hello world"
|
|
result = estimate_token_count(text)
|
|
assert result > 0
|
|
|
|
def test_estimate_bytes(self) -> None:
|
|
"""Bytes input should be decoded and produce the same count as the
|
|
equivalent string."""
|
|
text = "Hello world"
|
|
assert estimate_token_count(text.encode("utf-8")) == estimate_token_count(text)
|
|
|
|
def test_empty_string(self) -> None:
|
|
"""Should return 0 for empty string and empty bytes."""
|
|
assert estimate_token_count("") == 0
|
|
assert estimate_token_count(b"") == 0
|
|
|
|
def test_json_like_content(self) -> None:
|
|
"""JSON content should produce a positive estimate."""
|
|
json_str = '{"name": "test", "value": 123, "items": [1, 2, 3]}'
|
|
assert estimate_token_count(json_str) > 0
|
|
|
|
def test_long_text_roughly_scales_with_length(self) -> None:
|
|
"""A doubled string should produce roughly double the token count
|
|
(within ±10%)."""
|
|
small = "the quick brown fox jumps over the lazy dog. " * 20
|
|
large = small * 2
|
|
small_n = estimate_token_count(small)
|
|
large_n = estimate_token_count(large)
|
|
# Within 10% of 2x — both tokenizers (tiktoken and the char
|
|
# fallback) preserve length monotonicity.
|
|
assert 1.8 * small_n <= large_n <= 2.2 * small_n
|
|
|
|
def test_fallback_uses_chars_per_token_when_tiktoken_unavailable(
|
|
self,
|
|
) -> None:
|
|
"""When the tiktoken encoding is None (not installed), the
|
|
function falls back to len/CHARS_PER_TOKEN math."""
|
|
text = "x" * 100
|
|
with patch.object(token_utils, "_ENCODING", None):
|
|
result = estimate_token_count(text)
|
|
assert result == int(100 / CHARS_PER_TOKEN)
|
|
|
|
def test_fallback_when_tiktoken_encode_raises(self) -> None:
|
|
"""A misbehaving encoding should fall back to the char heuristic
|
|
rather than raise — the size guard must never fail-open."""
|
|
|
|
class BoomEncoding:
|
|
def encode(self, text: str) -> list[int]:
|
|
raise ValueError("simulated tiktoken failure")
|
|
|
|
text = "abc" * 50
|
|
with patch.object(token_utils, "_ENCODING", BoomEncoding()):
|
|
result = estimate_token_count(text)
|
|
assert result == int(len(text) / CHARS_PER_TOKEN)
|
|
|
|
|
|
class TestEstimateResponseTokens:
|
|
"""Test estimate_response_tokens function."""
|
|
|
|
class MockResponse(BaseModel):
|
|
"""Mock Pydantic response model."""
|
|
|
|
name: str
|
|
value: int
|
|
items: List[Any]
|
|
|
|
def test_estimate_pydantic_model(self) -> None:
|
|
"""Should estimate tokens for Pydantic model."""
|
|
response = self.MockResponse(name="test", value=42, items=[1, 2, 3])
|
|
result = estimate_response_tokens(response)
|
|
assert result > 0
|
|
|
|
def test_estimate_dict(self) -> None:
|
|
"""Should estimate tokens for dict."""
|
|
response = {"name": "test", "value": 42}
|
|
result = estimate_response_tokens(response)
|
|
assert result > 0
|
|
|
|
def test_estimate_list(self) -> None:
|
|
"""Should estimate tokens for list."""
|
|
response = [{"name": "item1"}, {"name": "item2"}]
|
|
result = estimate_response_tokens(response)
|
|
assert result > 0
|
|
|
|
def test_estimate_string(self) -> None:
|
|
"""Should estimate tokens for string response."""
|
|
response = "Hello world"
|
|
result = estimate_response_tokens(response)
|
|
assert result > 0
|
|
|
|
def test_estimate_large_response(self) -> None:
|
|
"""Should estimate tokens for large response."""
|
|
response = {"items": [{"name": f"item{i}"} for i in range(1000)]}
|
|
result = estimate_response_tokens(response)
|
|
assert result > 1000 # Large response should have many tokens
|
|
|
|
|
|
class TestGetResponseSizeBytes:
|
|
"""Test get_response_size_bytes function."""
|
|
|
|
def test_size_dict(self) -> None:
|
|
"""Should return size in bytes for dict."""
|
|
response = {"name": "test"}
|
|
result = get_response_size_bytes(response)
|
|
assert result > 0
|
|
|
|
def test_size_string(self) -> None:
|
|
"""Should return size in bytes for string."""
|
|
response = "Hello world"
|
|
result = get_response_size_bytes(response)
|
|
assert result == len(response.encode("utf-8"))
|
|
|
|
def test_size_bytes(self) -> None:
|
|
"""Should return size for bytes."""
|
|
response = b"Hello world"
|
|
result = get_response_size_bytes(response)
|
|
assert result == len(response)
|
|
|
|
|
|
class TestExtractQueryParams:
|
|
"""Test extract_query_params function."""
|
|
|
|
def test_extract_pagination_params(self) -> None:
|
|
"""Should extract pagination parameters."""
|
|
params = {"page_size": 100, "limit": 50}
|
|
result = extract_query_params(params)
|
|
assert result["page_size"] == 100
|
|
assert result["limit"] == 50
|
|
|
|
def test_extract_column_selection(self) -> None:
|
|
"""Should extract column selection parameters."""
|
|
params = {"select_columns": ["name", "id"]}
|
|
result = extract_query_params(params)
|
|
assert result["select_columns"] == ["name", "id"]
|
|
|
|
def test_extract_from_nested_request(self) -> None:
|
|
"""Should extract from nested request object."""
|
|
params = {"request": {"page_size": 50, "filters": [{"col": "name"}]}}
|
|
result = extract_query_params(params)
|
|
assert result["page_size"] == 50
|
|
assert result["filters"] == [{"col": "name"}]
|
|
|
|
def test_empty_params(self) -> None:
|
|
"""Should return empty dict for empty params."""
|
|
assert extract_query_params(None) == {}
|
|
assert extract_query_params({}) == {}
|
|
|
|
def test_extract_filters(self) -> None:
|
|
"""Should extract filter parameters."""
|
|
params = {"filters": [{"col": "name", "opr": "eq", "value": "test"}]}
|
|
result = extract_query_params(params)
|
|
assert "filters" in result
|
|
|
|
|
|
class TestGenerateSizeReductionSuggestions:
|
|
"""Test generate_size_reduction_suggestions function."""
|
|
|
|
def test_suggest_reduce_page_size(self) -> None:
|
|
"""Should suggest reducing page_size when present."""
|
|
params = {"page_size": 100}
|
|
suggestions = generate_size_reduction_suggestions(
|
|
tool_name="list_charts",
|
|
params=params,
|
|
estimated_tokens=50000,
|
|
token_limit=25000,
|
|
)
|
|
assert any(
|
|
"page_size" in s.lower() or "limit" in s.lower() for s in suggestions
|
|
)
|
|
|
|
def test_suggest_add_limit_for_list_tools(self) -> None:
|
|
"""Should suggest adding limit for list tools."""
|
|
params: dict[str, Any] = {}
|
|
suggestions = generate_size_reduction_suggestions(
|
|
tool_name="list_charts",
|
|
params=params,
|
|
estimated_tokens=50000,
|
|
token_limit=25000,
|
|
)
|
|
assert any(
|
|
"limit" in s.lower() or "page_size" in s.lower() for s in suggestions
|
|
)
|
|
|
|
def test_suggest_select_columns(self) -> None:
|
|
"""Should suggest using select_columns."""
|
|
params: dict[str, Any] = {}
|
|
suggestions = generate_size_reduction_suggestions(
|
|
tool_name="list_charts",
|
|
params=params,
|
|
estimated_tokens=50000,
|
|
token_limit=25000,
|
|
)
|
|
assert any(
|
|
"select_columns" in s.lower() or "columns" in s.lower() for s in suggestions
|
|
)
|
|
|
|
def test_suggest_filters(self) -> None:
|
|
"""Should suggest adding filters."""
|
|
params: dict[str, Any] = {}
|
|
suggestions = generate_size_reduction_suggestions(
|
|
tool_name="list_charts",
|
|
params=params,
|
|
estimated_tokens=50000,
|
|
token_limit=25000,
|
|
)
|
|
assert any("filter" in s.lower() for s in suggestions)
|
|
assert not any("editor" in s.lower() for s in suggestions)
|
|
assert any("non-user attributes" in s for s in suggestions)
|
|
|
|
def test_tool_specific_suggestions_execute_sql(self) -> None:
|
|
"""Should provide SQL-specific suggestions for execute_sql."""
|
|
suggestions = generate_size_reduction_suggestions(
|
|
tool_name="execute_sql",
|
|
params={"sql": "SELECT * FROM table"},
|
|
estimated_tokens=50000,
|
|
token_limit=25000,
|
|
)
|
|
combined = " ".join(suggestions)
|
|
# Should suggest SQL LIMIT clause
|
|
assert "LIMIT" in combined
|
|
# Should suggest the tool's limit parameter
|
|
assert "'limit' parameter" in combined.lower() or "limit=" in combined.lower()
|
|
|
|
def test_execute_sql_with_limit_param_no_duplicate_suggestion(self) -> None:
|
|
"""When limit param is already set, should not suggest adding it again."""
|
|
suggestions = generate_size_reduction_suggestions(
|
|
tool_name="execute_sql",
|
|
params={"sql": "SELECT * FROM table", "limit": 500},
|
|
estimated_tokens=50000,
|
|
token_limit=25000,
|
|
)
|
|
combined = " ".join(suggestions)
|
|
# Should still suggest SQL LIMIT
|
|
assert "LIMIT" in combined
|
|
# Should suggest reducing the existing limit (from general suggestion)
|
|
assert "500" in combined or "limit" in combined.lower()
|
|
|
|
def test_tool_specific_suggestions_list_charts(self) -> None:
|
|
"""Should provide chart-specific suggestions for list_charts."""
|
|
suggestions = generate_size_reduction_suggestions(
|
|
tool_name="list_charts",
|
|
params={},
|
|
estimated_tokens=50000,
|
|
token_limit=25000,
|
|
)
|
|
# Should suggest excluding params or query_context
|
|
assert any(
|
|
"params" in s.lower() or "query_context" in s.lower() for s in suggestions
|
|
)
|
|
|
|
def test_suggests_search_parameter(self) -> None:
|
|
"""Should suggest using search parameter."""
|
|
suggestions = generate_size_reduction_suggestions(
|
|
tool_name="list_dashboards",
|
|
params={},
|
|
estimated_tokens=50000,
|
|
token_limit=25000,
|
|
)
|
|
assert any("search" in s.lower() for s in suggestions)
|
|
|
|
|
|
class TestFormatSizeLimitError:
|
|
"""Test format_size_limit_error function."""
|
|
|
|
def test_error_contains_token_counts(self) -> None:
|
|
"""Should include token counts in error message."""
|
|
error = format_size_limit_error(
|
|
tool_name="list_charts",
|
|
params={},
|
|
estimated_tokens=50000,
|
|
token_limit=25000,
|
|
)
|
|
assert "50,000" in error
|
|
assert "25,000" in error
|
|
|
|
def test_error_contains_tool_name(self) -> None:
|
|
"""Should include tool name in error message."""
|
|
error = format_size_limit_error(
|
|
tool_name="list_charts",
|
|
params={},
|
|
estimated_tokens=50000,
|
|
token_limit=25000,
|
|
)
|
|
assert "list_charts" in error
|
|
|
|
def test_error_contains_suggestions(self) -> None:
|
|
"""Should include suggestions in error message."""
|
|
error = format_size_limit_error(
|
|
tool_name="list_charts",
|
|
params={"page_size": 100},
|
|
estimated_tokens=50000,
|
|
token_limit=25000,
|
|
)
|
|
# Should have numbered suggestions
|
|
assert "1." in error
|
|
|
|
def test_error_contains_reduction_percentage(self) -> None:
|
|
"""Should include reduction percentage in error message."""
|
|
error = format_size_limit_error(
|
|
tool_name="list_charts",
|
|
params={},
|
|
estimated_tokens=50000,
|
|
token_limit=25000,
|
|
)
|
|
# 50% reduction needed
|
|
assert "50%" in error or "Reduction" in error
|
|
|
|
def test_error_limits_suggestions_to_five(self) -> None:
|
|
"""Should limit suggestions to 5."""
|
|
error = format_size_limit_error(
|
|
tool_name="list_charts",
|
|
params={},
|
|
estimated_tokens=100000,
|
|
token_limit=10000,
|
|
)
|
|
# Count numbered suggestions (1. through 5.)
|
|
suggestion_count = sum(1 for i in range(1, 10) if f"{i}." in error)
|
|
assert suggestion_count <= 5
|
|
|
|
def test_error_message_is_readable(self) -> None:
|
|
"""Should produce human-readable error message."""
|
|
error = format_size_limit_error(
|
|
tool_name="list_charts",
|
|
params={"page_size": 100},
|
|
estimated_tokens=75000,
|
|
token_limit=25000,
|
|
)
|
|
# Should be multi-line and contain key information
|
|
lines = error.split("\n")
|
|
assert len(lines) > 5
|
|
assert "Response too large" in error
|
|
assert "Please modify your query" in error
|
|
|
|
|
|
class TestCalculatedSuggestions:
|
|
"""Test that suggestions include calculated values."""
|
|
|
|
def test_suggested_limit_is_calculated(self) -> None:
|
|
"""Should calculate suggested limit based on reduction needed."""
|
|
params = {"page_size": 100}
|
|
suggestions = generate_size_reduction_suggestions(
|
|
tool_name="list_charts",
|
|
params=params,
|
|
estimated_tokens=50000, # 2x over limit
|
|
token_limit=25000,
|
|
)
|
|
# Find the page_size suggestion
|
|
page_size_suggestion = next(
|
|
(s for s in suggestions if "page_size" in s.lower()), None
|
|
)
|
|
assert page_size_suggestion is not None
|
|
# Should suggest reducing from 100 to approximately 50
|
|
assert "100" in page_size_suggestion
|
|
assert (
|
|
"50" in page_size_suggestion or "reduction" in page_size_suggestion.lower()
|
|
)
|
|
|
|
def test_reduction_percentage_in_suggestions(self) -> None:
|
|
"""Should include reduction percentage in suggestions."""
|
|
params = {"page_size": 100}
|
|
suggestions = generate_size_reduction_suggestions(
|
|
tool_name="list_charts",
|
|
params=params,
|
|
estimated_tokens=75000, # 3x over limit
|
|
token_limit=25000,
|
|
)
|
|
# Should mention ~66% reduction needed (int truncation of 66.6%)
|
|
combined = " ".join(suggestions)
|
|
assert "66%" in combined
|
|
|
|
|
|
class TestInfoToolsSet:
|
|
"""Test the INFO_TOOLS constant."""
|
|
|
|
def test_info_tools_contains_expected_tools(self) -> None:
|
|
"""Should contain all info tools."""
|
|
assert "get_chart_info" in INFO_TOOLS
|
|
assert "get_dataset_info" in INFO_TOOLS
|
|
assert "get_dashboard_info" in INFO_TOOLS
|
|
assert "get_instance_info" in INFO_TOOLS
|
|
|
|
def test_info_tools_does_not_contain_list_tools(self) -> None:
|
|
"""Should not contain list or write tools."""
|
|
assert "list_charts" not in INFO_TOOLS
|
|
assert "execute_sql" not in INFO_TOOLS
|
|
assert "generate_chart" not in INFO_TOOLS
|
|
|
|
|
|
class TestTruncateStrings:
|
|
"""Test _truncate_strings helper."""
|
|
|
|
def test_truncates_long_strings(self) -> None:
|
|
"""Should truncate strings exceeding max_chars."""
|
|
data: dict[str, Any] = {"description": "x" * 1000, "name": "short"}
|
|
notes: list[str] = []
|
|
changed = _truncate_strings(data, notes, max_chars=500)
|
|
assert changed is True
|
|
assert len(data["description"]) < 1000
|
|
assert "[truncated from 1000 chars]" in data["description"]
|
|
assert data["name"] == "short"
|
|
assert len(notes) == 1
|
|
|
|
def test_does_not_truncate_short_strings(self) -> None:
|
|
"""Should not truncate strings within limit."""
|
|
data: dict[str, Any] = {"name": "hello", "id": 123}
|
|
notes: list[str] = []
|
|
changed = _truncate_strings(data, notes, max_chars=500)
|
|
assert changed is False
|
|
assert data["name"] == "hello"
|
|
assert len(notes) == 0
|
|
|
|
|
|
class TestTruncateStringsRecursive:
|
|
"""Test _truncate_strings_recursive helper."""
|
|
|
|
def test_truncates_nested_strings_in_list_items(self) -> None:
|
|
"""Should truncate strings inside list items (e.g. charts[i].description)."""
|
|
data: dict[str, Any] = {
|
|
"id": 1,
|
|
"charts": [
|
|
{"id": 1, "description": "x" * 1000},
|
|
{"id": 2, "description": "short"},
|
|
],
|
|
}
|
|
notes: list[str] = []
|
|
changed = _truncate_strings_recursive(data, notes, max_chars=500)
|
|
assert changed is True
|
|
assert "[truncated" in data["charts"][0]["description"]
|
|
assert data["charts"][1]["description"] == "short"
|
|
assert len(notes) == 1
|
|
assert "charts[0].description" in notes[0]
|
|
|
|
def test_truncates_nested_strings_in_dicts(self) -> None:
|
|
"""Should truncate strings inside nested dicts."""
|
|
data: dict[str, Any] = {
|
|
"filter_state": {
|
|
"dataMask": {"some_filter": "y" * 2000},
|
|
},
|
|
}
|
|
notes: list[str] = []
|
|
changed = _truncate_strings_recursive(data, notes, max_chars=500)
|
|
assert changed is True
|
|
assert "[truncated" in data["filter_state"]["dataMask"]["some_filter"]
|
|
|
|
def test_respects_depth_limit(self) -> None:
|
|
"""Should stop recursing at depth 10."""
|
|
# Build a deeply nested structure (15 levels)
|
|
data: dict[str, Any] = {"level": "x" * 1000}
|
|
current = data
|
|
for _ in range(15):
|
|
current["nested"] = {"level": "x" * 1000}
|
|
current = current["nested"]
|
|
notes: list[str] = []
|
|
_truncate_strings_recursive(data, notes, max_chars=500)
|
|
# Should truncate levels 0-10 but stop before 15
|
|
assert len(notes) <= 11
|
|
|
|
def test_handles_empty_structures(self) -> None:
|
|
"""Should handle empty dicts and lists gracefully."""
|
|
data: dict[str, Any] = {"items": [], "meta": {}, "name": "ok"}
|
|
notes: list[str] = []
|
|
changed = _truncate_strings_recursive(data, notes, max_chars=500)
|
|
assert changed is False
|
|
|
|
def test_dashboard_with_many_charts_edge_case(self) -> None:
|
|
"""Simulate a dashboard with 30 charts each having long descriptions."""
|
|
data: dict[str, Any] = {
|
|
"id": 1,
|
|
"dashboard_title": "Big Dashboard",
|
|
"charts": [
|
|
{"id": i, "slice_name": f"Chart {i}", "description": "d" * 2000}
|
|
for i in range(30)
|
|
],
|
|
}
|
|
notes: list[str] = []
|
|
changed = _truncate_strings_recursive(data, notes, max_chars=500)
|
|
assert changed is True
|
|
# All 30 chart descriptions should be truncated
|
|
assert len(notes) == 30
|
|
for chart in data["charts"]:
|
|
assert len(chart["description"]) < 2000
|
|
assert "[truncated" in chart["description"]
|
|
|
|
|
|
class TestTruncateLists:
|
|
"""Test _truncate_lists helper."""
|
|
|
|
def test_truncates_long_lists(self) -> None:
|
|
"""Should truncate lists exceeding max_items without inline markers."""
|
|
data: dict[str, Any] = {
|
|
"columns": [{"name": f"col_{i}"} for i in range(50)],
|
|
"tags": [1, 2],
|
|
}
|
|
notes: list[str] = []
|
|
changed = _truncate_lists(data, notes, max_items=10)
|
|
assert changed is True
|
|
# Exactly 10 items — no marker appended (preserves type contract)
|
|
assert len(data["columns"]) == 10
|
|
assert all(isinstance(c, dict) and "name" in c for c in data["columns"])
|
|
assert data["tags"] == [1, 2] # Not truncated
|
|
assert len(notes) == 1
|
|
assert "50" in notes[0]
|
|
|
|
def test_does_not_truncate_short_lists(self) -> None:
|
|
"""Should not truncate lists within limit."""
|
|
data: dict[str, Any] = {"items": [1, 2, 3]}
|
|
notes: list[str] = []
|
|
changed = _truncate_lists(data, notes, max_items=10)
|
|
assert changed is False
|
|
|
|
|
|
class TestSummarizeLargeDicts:
|
|
"""Test _summarize_large_dicts helper."""
|
|
|
|
def test_summarizes_large_dicts(self) -> None:
|
|
"""Should replace large dicts with key summaries."""
|
|
big_dict = {f"key_{i}": f"value_{i}" for i in range(30)}
|
|
data: dict[str, Any] = {"form_data": big_dict, "id": 1}
|
|
notes: list[str] = []
|
|
changed = _summarize_large_dicts(data, notes, max_keys=20)
|
|
assert changed is True
|
|
assert data["form_data"]["_truncated"] is True
|
|
assert "30 keys" in data["form_data"]["_message"]
|
|
assert data["id"] == 1
|
|
|
|
def test_does_not_summarize_small_dicts(self) -> None:
|
|
"""Should not summarize dicts within limit."""
|
|
data: dict[str, Any] = {"params": {"a": 1, "b": 2}}
|
|
notes: list[str] = []
|
|
changed = _summarize_large_dicts(data, notes, max_keys=20)
|
|
assert changed is False
|
|
|
|
|
|
class TestReplaceCollectionsWithSummaries:
|
|
"""Test _replace_collections_with_summaries helper."""
|
|
|
|
def test_replaces_lists_and_dicts(self) -> None:
|
|
"""Should clear non-empty collections to reduce size."""
|
|
data: dict[str, Any] = {
|
|
"columns": [1, 2, 3],
|
|
"params": {"a": 1},
|
|
"name": "test",
|
|
"empty": [],
|
|
}
|
|
notes: list[str] = []
|
|
changed = _replace_collections_with_summaries(data, notes)
|
|
assert changed is True
|
|
# Lists become empty lists (preserves type)
|
|
assert data["columns"] == []
|
|
# Dicts become empty dicts (preserves type)
|
|
assert data["params"] == {}
|
|
# Scalars unchanged
|
|
assert data["name"] == "test"
|
|
# Empty collections unchanged
|
|
assert data["empty"] == []
|
|
assert len(notes) == 2
|
|
|
|
|
|
class TestTruncateOversizedResponse:
|
|
"""Test truncate_oversized_response function."""
|
|
|
|
def test_no_truncation_needed(self) -> None:
|
|
"""Should return original data when under limit."""
|
|
response = {"id": 1, "name": "test"}
|
|
result, was_truncated, notes = truncate_oversized_response(response, 10000)
|
|
assert was_truncated is False
|
|
assert notes == []
|
|
|
|
def test_truncates_large_string_fields(self) -> None:
|
|
"""Should truncate long strings to fit."""
|
|
response = {
|
|
"id": 1,
|
|
"description": "x" * 50000, # Very large description
|
|
}
|
|
result, was_truncated, notes = truncate_oversized_response(response, 500)
|
|
assert was_truncated is True
|
|
assert isinstance(result, dict)
|
|
assert "[truncated" in result["description"]
|
|
assert any("description" in n for n in notes)
|
|
|
|
def test_truncates_large_lists(self) -> None:
|
|
"""Should truncate lists when strings alone are not enough."""
|
|
response = {
|
|
"id": 1,
|
|
"columns": [{"name": f"col_{i}", "type": "VARCHAR"} for i in range(200)],
|
|
}
|
|
result, was_truncated, notes = truncate_oversized_response(response, 500)
|
|
assert was_truncated is True
|
|
assert isinstance(result, dict)
|
|
# Should have been truncated
|
|
assert len(result["columns"]) < 200
|
|
|
|
def test_handles_pydantic_model(self) -> None:
|
|
"""Should handle Pydantic model input."""
|
|
|
|
class FakeInfo(BaseModel):
|
|
id: int = 1
|
|
description: str = "x" * 5000
|
|
|
|
response = FakeInfo()
|
|
result, was_truncated, notes = truncate_oversized_response(response, 200)
|
|
assert was_truncated is True
|
|
assert isinstance(result, dict)
|
|
|
|
def test_returns_non_dict_unchanged(self) -> None:
|
|
"""Should return non-dict/model responses unchanged."""
|
|
result, was_truncated, notes = truncate_oversized_response("just a string", 100)
|
|
assert was_truncated is False
|
|
assert result == "just a string"
|
|
|
|
def test_progressive_truncation(self) -> None:
|
|
"""Should progressively apply truncation phases."""
|
|
# Build a response that's quite large
|
|
response = {
|
|
"id": 1,
|
|
"description": "x" * 2000,
|
|
"css": "y" * 2000,
|
|
"columns": [{"name": f"col_{i}"} for i in range(100)],
|
|
"form_data": {f"key_{i}": f"val_{i}" for i in range(50)},
|
|
}
|
|
result, was_truncated, notes = truncate_oversized_response(response, 300)
|
|
assert was_truncated is True
|
|
assert isinstance(result, dict)
|
|
assert result["id"] == 1 # Scalar fields preserved
|
|
assert len(notes) > 0
|
|
|
|
@staticmethod
|
|
def _build_large_dashboard_response() -> dict[str, Any]:
|
|
"""A dashboard with 463 charts and 48 native_filters, shared by the
|
|
default- and custom-max_list_items regression tests below."""
|
|
return {
|
|
"id": 1,
|
|
"dashboard_title": "x" * 2000, # forces Phase 2 to trigger
|
|
"charts": [{"id": i, "slice_name": f"chart_{i}"} for i in range(463)],
|
|
"native_filters": [{"id": i, "name": f"filter_{i}"} for i in range(48)],
|
|
}
|
|
|
|
def test_large_dashboard_respects_default_max_list_items(self) -> None:
|
|
"""Regression test for the Medialab large-dashboard report.
|
|
|
|
A dashboard with 463 charts and 48 native_filters should have
|
|
native_filters (48 items) left untouched under the new default cap
|
|
of 100, while charts (463 items) is truncated to 100 — a clear
|
|
improvement over the old flat 30-item cap, which truncated both.
|
|
"""
|
|
response: dict[str, Any] = self._build_large_dashboard_response()
|
|
result: Any
|
|
was_truncated: bool
|
|
notes: list[str]
|
|
result, was_truncated, notes = truncate_oversized_response(response, 3000)
|
|
assert was_truncated is True
|
|
assert isinstance(result, dict)
|
|
assert len(result["charts"]) == 100
|
|
assert len(result["native_filters"]) == 48
|
|
assert any("charts" in n and "463" in n for n in notes)
|
|
assert not any("native_filters" in n for n in notes)
|
|
|
|
def test_large_dashboard_respects_custom_max_list_items(self) -> None:
|
|
"""A custom max_list_items below both list sizes should truncate both fields."""
|
|
response: dict[str, Any] = self._build_large_dashboard_response()
|
|
result: Any
|
|
was_truncated: bool
|
|
notes: list[str]
|
|
result, was_truncated, notes = truncate_oversized_response(
|
|
response, 3000, max_list_items=30
|
|
)
|
|
assert was_truncated is True
|
|
assert isinstance(result, dict)
|
|
assert len(result["charts"]) == 30
|
|
assert len(result["native_filters"]) == 30
|
|
assert any("charts" in n and "30" in n for n in notes)
|
|
assert any("native_filters" in n and "30" in n for n in notes)
|
|
|
|
def test_custom_max_list_items_below_phase_four_survives_phase_four(self) -> None:
|
|
"""A max_list_items below Phase 4's hardcoded 10 should not be widened.
|
|
|
|
Phase 2 truncates ``charts`` to 5 first; the response is still over
|
|
budget because of the oversized ``form_data`` dict, so truncation
|
|
proceeds to Phase 4, whose ``_truncate_lists(..., max_items=10)``
|
|
call only shrinks lists larger than 10 — it must leave the
|
|
already-smaller 5-item list untouched rather than re-expanding it.
|
|
"""
|
|
response: dict[str, Any] = {
|
|
"id": 1,
|
|
"charts": [{"id": i, "slice_name": f"chart_{i}"} for i in range(300)],
|
|
"form_data": {f"key_{i}": f"val_{i}" for i in range(50)},
|
|
}
|
|
result: Any
|
|
was_truncated: bool
|
|
notes: list[str]
|
|
result, was_truncated, notes = truncate_oversized_response(
|
|
response, 200, max_list_items=5
|
|
)
|
|
assert was_truncated is True
|
|
assert isinstance(result, dict)
|
|
assert len(result["charts"]) == 5
|
|
assert any("form_data" in n for n in notes)
|
|
|
|
|
|
class TestTruncateQueryResult:
|
|
"""Tests for ``truncate_query_result`` (data-query row/scalar truncation)."""
|
|
|
|
def _rows_response(self, row_field: str, count: int = 200) -> dict[str, Any]:
|
|
row = {f"col_{i}": f"value_{i}" for i in range(10)}
|
|
return {
|
|
"status": "success",
|
|
row_field: [row] * count,
|
|
"row_count": count,
|
|
}
|
|
|
|
def test_no_truncation_needed(self) -> None:
|
|
response = self._rows_response("rows", count=3)
|
|
result, was_truncated, notes = truncate_query_result(response, 25000)
|
|
assert was_truncated is False
|
|
assert notes == []
|
|
assert result == response
|
|
|
|
def test_truncated_result_fits_under_limit(self) -> None:
|
|
"""The final payload (rows + note metadata) must itself fit.
|
|
|
|
Regression test: the note is built from the kept row count, but
|
|
that note text also consumes tokens. The bisection must reserve
|
|
room for it up front rather than measuring fit on bare rows and
|
|
appending the note afterward, which could push the final payload
|
|
back over the limit.
|
|
"""
|
|
response = self._rows_response("rows")
|
|
result, was_truncated, notes = truncate_query_result(response, 500)
|
|
assert was_truncated is True
|
|
assert isinstance(result, dict)
|
|
assert estimate_response_tokens(result) <= 500
|
|
assert result["row_count"] == len(result["rows"])
|
|
assert result["row_count"] < 200
|
|
|
|
def test_single_oversized_row_is_kept_anyway(self) -> None:
|
|
"""A single row that alone exceeds the limit is still returned.
|
|
|
|
``truncate_query_result`` always keeps >=1 row when the original
|
|
list is non-empty; it is the caller's (middleware) job to reject
|
|
a still-oversized result rather than ship it silently.
|
|
"""
|
|
response = {
|
|
"status": "success",
|
|
"rows": [{"col": "x" * 5000}] * 3,
|
|
"row_count": 3,
|
|
}
|
|
result, was_truncated, notes = truncate_query_result(response, 50)
|
|
assert was_truncated is True
|
|
assert isinstance(result, dict)
|
|
assert len(result["rows"]) == 1
|
|
assert notes
|
|
|
|
def test_truncates_csv_scalar_field_when_rows_empty(self) -> None:
|
|
"""CSV exports carry their payload in ``csv_data`` with ``data=[]``."""
|
|
response: dict[str, Any] = {
|
|
"chart_id": 1,
|
|
"data": [],
|
|
"csv_data": "col_0,col_1\n" + ("value,value\n" * 2000),
|
|
"format": "csv",
|
|
}
|
|
result, was_truncated, notes = truncate_query_result(
|
|
response, 500, tool_name="get_chart_data"
|
|
)
|
|
assert was_truncated is True
|
|
assert isinstance(result, dict)
|
|
assert len(result["csv_data"]) < len(response["csv_data"])
|
|
assert estimate_response_tokens(result) <= 500
|
|
assert any("CSV" in n for n in notes)
|
|
|
|
def test_does_not_truncate_excel_binary_field(self) -> None:
|
|
"""excel_data is base64 binary — truncating it would corrupt the file."""
|
|
response = {
|
|
"chart_id": 1,
|
|
"data": [],
|
|
"excel_data": "QUJDREVGRw==" * 5000,
|
|
"format": "excel",
|
|
}
|
|
result, was_truncated, notes = truncate_query_result(response, 500)
|
|
assert was_truncated is False
|
|
assert notes == []
|
|
assert isinstance(result, dict)
|
|
assert result["excel_data"] == response["excel_data"]
|
|
|
|
def test_get_chart_data_advice_mentions_limit_param(self) -> None:
|
|
response = self._rows_response("data")
|
|
_, _, notes = truncate_query_result(response, 500, tool_name="get_chart_data")
|
|
assert any("'limit' parameter" in n for n in notes)
|
|
assert not any("LIMIT clause" in n for n in notes)
|
|
|
|
def test_query_dataset_advice_mentions_row_limit_param(self) -> None:
|
|
response = self._rows_response("data")
|
|
_, _, notes = truncate_query_result(response, 500, tool_name="query_dataset")
|
|
assert any("'row_limit' parameter" in n for n in notes)
|
|
assert not any("LIMIT clause" in n for n in notes)
|
|
|
|
def test_execute_sql_advice_mentions_limit_clause(self) -> None:
|
|
response = self._rows_response("rows")
|
|
_, _, notes = truncate_query_result(response, 500, tool_name="execute_sql")
|
|
assert any("LIMIT clause" in n for n in notes)
|