mirror of
https://github.com/apache/superset.git
synced 2026-07-26 16:42:32 +00:00
00c5e7957b
_bisect_row_limit already trims statements[*].data.rows in lockstep with the top-level rows field to keep the response under the token limit, but left each statement's row_count pointing at the pre-truncation count, so a truncated execute_sql response could claim more rows than statements[*].data actually contains. Track the statement dicts (not just their row lists) and update row_count alongside data.rows. Also log a full traceback (logger.exception) instead of a plain warning when data-query truncation hits an unexpected error, so a real bug isn't indistinguishable from routine size-limit enforcement. Adds regression tests for: total_rows surviving truncation distinct from row_count, the no-recognised-row-field fallback path, and the oversized-excel-export hard-failure path.
1040 lines
41 KiB
Python
1040 lines
41 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_csv_truncation_keeps_header_and_whole_records(self) -> None:
|
|
"""Truncated CSV must stay parseable: header intact, no partial rows.
|
|
|
|
Regression test: naive character-offset truncation can cut a CSV
|
|
payload mid-record (e.g. inside a quoted field), producing invalid
|
|
CSV. Truncation must cut on record boundaries instead.
|
|
"""
|
|
import csv
|
|
import io
|
|
|
|
header = "col_0,col_1,col_2"
|
|
# Include a quoted field containing a comma and an embedded newline,
|
|
# which would corrupt a naive character-offset cut.
|
|
rows = [
|
|
f'{i},"value, with comma","line1\nline2 for row {i}"' for i in range(500)
|
|
]
|
|
csv_data = header + "\n" + "\n".join(rows) + "\n"
|
|
response: dict[str, Any] = {
|
|
"chart_id": 1,
|
|
"data": [],
|
|
"csv_data": csv_data,
|
|
"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 estimate_response_tokens(result) <= 500
|
|
|
|
parsed_rows = list(csv.reader(io.StringIO(result["csv_data"])))
|
|
assert parsed_rows, "truncated CSV must remain parseable"
|
|
assert parsed_rows[0] == header.split(",")
|
|
# Every parsed record must be a complete, well-formed row (3 fields),
|
|
# proving no record was cut mid-field.
|
|
assert all(len(row) == 3 for row in parsed_rows)
|
|
assert len(parsed_rows) - 1 < len(rows)
|
|
assert any("data rows" 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)
|
|
|
|
def test_execute_sql_statements_rows_truncated_alongside_top_level(self) -> None:
|
|
"""Real ``ExecuteSqlResponse`` shapes duplicate rows under ``statements``.
|
|
|
|
Regression test: ``ExecuteSqlResponse.rows`` is a copy of the last
|
|
data-bearing statement's ``data.rows``, and the SAME rows are also
|
|
serialised under ``statements[0].data.rows``. If only the top-level
|
|
field is bisected, the nested copy stays at full size and the
|
|
response can remain oversized after "truncation".
|
|
"""
|
|
row = {f"col_{i}": f"value_{i}" for i in range(10)}
|
|
rows = [row] * 200
|
|
response: dict[str, Any] = {
|
|
"success": True,
|
|
"rows": rows,
|
|
"row_count": 200,
|
|
"statements": [
|
|
{
|
|
"original_sql": "SELECT * FROM t",
|
|
"executed_sql": "SELECT * FROM t",
|
|
"row_count": 200,
|
|
"data": {
|
|
"rows": rows,
|
|
"columns": [
|
|
{"name": f"col_{i}", "type": "TEXT"} for i in range(10)
|
|
],
|
|
},
|
|
}
|
|
],
|
|
}
|
|
result, was_truncated, notes = truncate_query_result(
|
|
response, 500, tool_name="execute_sql"
|
|
)
|
|
assert was_truncated is True
|
|
assert isinstance(result, dict)
|
|
# The overall payload — including the nested statements copy — must
|
|
# actually fit; this is what the middleware's post-truncation size
|
|
# check relies on to avoid falling back to a hard error.
|
|
assert estimate_response_tokens(result) <= 500
|
|
assert len(result["rows"]) < 200
|
|
assert len(result["statements"][0]["data"]["rows"]) == len(result["rows"])
|
|
assert notes
|
|
|
|
def test_execute_sql_statements_row_count_stays_in_sync(self) -> None:
|
|
"""Each statement's ``row_count`` must track its trimmed ``data.rows``.
|
|
|
|
Regression test: bisection already shrinks the nested
|
|
``statements[*].data.rows`` list in lockstep with the top-level
|
|
field (see the sibling ``..._rows_truncated_alongside_top_level``
|
|
test), but the per-statement ``row_count`` field is a separate
|
|
value copied at response-build time and must be updated too, or a
|
|
client would see a statement whose declared row count doesn't match
|
|
the row list actually returned.
|
|
"""
|
|
row = {f"col_{i}": f"value_{i}" for i in range(10)}
|
|
rows = [row] * 200
|
|
response: dict[str, Any] = {
|
|
"success": True,
|
|
"rows": rows,
|
|
"row_count": 200,
|
|
"statements": [
|
|
{
|
|
"original_sql": "SELECT * FROM t",
|
|
"executed_sql": "SELECT * FROM t",
|
|
"row_count": 200,
|
|
"data": {
|
|
"rows": rows,
|
|
"columns": [
|
|
{"name": f"col_{i}", "type": "TEXT"} for i in range(10)
|
|
],
|
|
},
|
|
}
|
|
],
|
|
}
|
|
result, was_truncated, _ = truncate_query_result(
|
|
response, 500, tool_name="execute_sql"
|
|
)
|
|
assert was_truncated is True
|
|
statement = result["statements"][0]
|
|
assert statement["row_count"] == len(statement["data"]["rows"])
|
|
assert statement["row_count"] < 200
|
|
|
|
def test_falls_back_to_generic_truncation_when_no_row_field(self) -> None:
|
|
"""Responses without a recognised row/data field truncate generically.
|
|
|
|
A DML-final ``execute_sql`` response (or any data-query response
|
|
with no ``rows``/``data`` list) has no row field to bisect on, so
|
|
``truncate_query_result`` delegates to the generic
|
|
``truncate_oversized_response`` field-truncation phases instead of
|
|
raising or returning the payload untouched.
|
|
"""
|
|
response: dict[str, Any] = {
|
|
"success": True,
|
|
"affected_rows": 5,
|
|
"multi_statement_warning": "x" * 5000,
|
|
}
|
|
result, was_truncated, notes = truncate_query_result(response, 200)
|
|
assert was_truncated is True
|
|
assert notes
|
|
assert len(result["multi_statement_warning"]) < 5000
|
|
assert result["affected_rows"] == 5
|
|
|
|
def test_execute_sql_multiple_statements_each_capped(self) -> None:
|
|
"""Every statement's rows are capped, not just the last (top-level) one."""
|
|
row = {f"col_{i}": f"value_{i}" for i in range(10)}
|
|
first_stmt_rows = [row] * 150
|
|
last_stmt_rows = [row] * 150
|
|
response: dict[str, Any] = {
|
|
"success": True,
|
|
"rows": last_stmt_rows,
|
|
"row_count": 150,
|
|
"multi_statement_warning": "2 data-bearing statements",
|
|
"statements": [
|
|
{
|
|
"original_sql": "SELECT * FROM a",
|
|
"executed_sql": "SELECT * FROM a",
|
|
"row_count": 150,
|
|
"data": {
|
|
"rows": first_stmt_rows,
|
|
"columns": [
|
|
{"name": f"col_{i}", "type": "TEXT"} for i in range(10)
|
|
],
|
|
},
|
|
},
|
|
{
|
|
"original_sql": "SELECT * FROM b",
|
|
"executed_sql": "SELECT * FROM b",
|
|
"row_count": 150,
|
|
"data": {
|
|
"rows": last_stmt_rows,
|
|
"columns": [
|
|
{"name": f"col_{i}", "type": "TEXT"} for i in range(10)
|
|
],
|
|
},
|
|
},
|
|
],
|
|
}
|
|
result, was_truncated, notes = truncate_query_result(
|
|
response, 2000, tool_name="execute_sql"
|
|
)
|
|
assert was_truncated is True
|
|
assert isinstance(result, dict)
|
|
assert estimate_response_tokens(result) <= 2000
|
|
assert len(result["statements"][0]["data"]["rows"]) < 150
|
|
assert len(result["statements"][1]["data"]["rows"]) < 150
|
|
assert notes
|