mirror of
https://github.com/langgenius/dify.git
synced 2025-12-19 17:27:16 -05:00
feat: _truncate_json_primitives support file (#29760)
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user