feat: _truncate_json_primitives support file (#29760)

This commit is contained in:
wangxiaolei
2025-12-17 07:10:43 +08:00
committed by GitHub
parent c2f2be6b08
commit dda7eb03c9
2 changed files with 56 additions and 1 deletions

View File

@@ -410,9 +410,12 @@ class VariableTruncator(BaseTruncator):
@overload
def _truncate_json_primitives(self, val: None, target_size: int) -> _PartResult[None]: ...
@overload
def _truncate_json_primitives(self, val: File, target_size: int) -> _PartResult[File]: ...
def _truncate_json_primitives(
self,
val: UpdatedVariable | str | list[object] | dict[str, object] | bool | int | float | None,
val: UpdatedVariable | File | str | list[object] | dict[str, object] | bool | int | float | None,
target_size: int,
) -> _PartResult[Any]:
"""Truncate a value within an object to fit within budget."""
@@ -425,6 +428,9 @@ class VariableTruncator(BaseTruncator):
return self._truncate_array(val, target_size)
elif isinstance(val, dict):
return self._truncate_object(val, target_size)
elif isinstance(val, File):
# File objects should not be truncated, return as-is
return _PartResult(val, self.calculate_json_size(val), False)
elif val is None or isinstance(val, (bool, int, float)):
return _PartResult(val, self.calculate_json_size(val), False)
else:

View File

@@ -518,6 +518,55 @@ class TestEdgeCases:
assert isinstance(result.result, StringSegment)
class TestTruncateJsonPrimitives:
"""Test _truncate_json_primitives method with different data types."""
@pytest.fixture
def truncator(self):
return VariableTruncator()
def test_truncate_json_primitives_file_type(self, truncator, file):
"""Test that File objects are handled correctly in _truncate_json_primitives."""
# Test File object is returned as-is without truncation
result = truncator._truncate_json_primitives(file, 1000)
assert result.value == file
assert result.truncated is False
# Size should be calculated correctly
expected_size = VariableTruncator.calculate_json_size(file)
assert result.value_size == expected_size
def test_truncate_json_primitives_file_type_small_budget(self, truncator, file):
"""Test that File objects are returned as-is even with small budget."""
# Even with a small size budget, File objects should not be truncated
result = truncator._truncate_json_primitives(file, 10)
assert result.value == file
assert result.truncated is False
def test_truncate_json_primitives_file_type_in_array(self, truncator, file):
"""Test File objects in arrays are handled correctly."""
array_with_files = [file, file]
result = truncator._truncate_json_primitives(array_with_files, 1000)
assert isinstance(result.value, list)
assert len(result.value) == 2
assert result.value[0] == file
assert result.value[1] == file
assert result.truncated is False
def test_truncate_json_primitives_file_type_in_object(self, truncator, file):
"""Test File objects in objects are handled correctly."""
obj_with_files = {"file1": file, "file2": file}
result = truncator._truncate_json_primitives(obj_with_files, 1000)
assert isinstance(result.value, dict)
assert len(result.value) == 2
assert result.value["file1"] == file
assert result.value["file2"] == file
assert result.truncated is False
class TestIntegrationScenarios:
"""Test realistic integration scenarios."""