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
|
@overload
|
||||||
def _truncate_json_primitives(self, val: None, target_size: int) -> _PartResult[None]: ...
|
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(
|
def _truncate_json_primitives(
|
||||||
self,
|
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,
|
target_size: int,
|
||||||
) -> _PartResult[Any]:
|
) -> _PartResult[Any]:
|
||||||
"""Truncate a value within an object to fit within budget."""
|
"""Truncate a value within an object to fit within budget."""
|
||||||
@@ -425,6 +428,9 @@ class VariableTruncator(BaseTruncator):
|
|||||||
return self._truncate_array(val, target_size)
|
return self._truncate_array(val, target_size)
|
||||||
elif isinstance(val, dict):
|
elif isinstance(val, dict):
|
||||||
return self._truncate_object(val, target_size)
|
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)):
|
elif val is None or isinstance(val, (bool, int, float)):
|
||||||
return _PartResult(val, self.calculate_json_size(val), False)
|
return _PartResult(val, self.calculate_json_size(val), False)
|
||||||
else:
|
else:
|
||||||
|
|||||||
@@ -518,6 +518,55 @@ class TestEdgeCases:
|
|||||||
assert isinstance(result.result, StringSegment)
|
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:
|
class TestIntegrationScenarios:
|
||||||
"""Test realistic integration scenarios."""
|
"""Test realistic integration scenarios."""
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user