Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: truncate "params" column on vertex_build table to prevent memory heap on database #4118

Merged
merged 5 commits into from
Oct 11, 2024
Merged
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,10 @@ def serialize_data(self, data: dict) -> dict:
def serialize_artifacts(self, data) -> dict:
return truncate_long_strings(data)

@field_serializer("params")
def serialize_params(self, data) -> dict:
return truncate_long_strings(data)


class VertexBuildTable(VertexBuildBase, table=True): # type: ignore[call-arg]
__tablename__ = "vertex_build"
Expand Down
7 changes: 6 additions & 1 deletion src/backend/base/langflow/utils/util_strings.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,12 @@ def truncate_long_strings(data, max_length=None):
if max_length is None:
max_length = constants.MAX_TEXT_LENGTH

if max_length < 0 or not isinstance(data, dict | list):
if max_length < 0:
return data

if not isinstance(data, dict | list):
if isinstance(data, str) and len(data) > max_length:
return data[:max_length] + "..."
return data

if isinstance(data, dict):
Expand Down
54 changes: 54 additions & 0 deletions src/backend/tests/unit/utils/test_truncate_long_strings.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
import pytest
from langflow.utils.util_strings import truncate_long_strings


@pytest.mark.parametrize(
"input_data, max_length, expected",
[
# Test case 1: String shorter than max_length
("short string", 20, "short string"),
# Test case 2: String exactly at max_length
("exact", 5, "exact"),
# Test case 3: String longer than max_length
("long string", 7, "long st..."),
# Test case 4: Empty string
("", 5, ""),
# Test case 5: Single character string
("a", 1, "a"),
# Test case 6: Unicode string
("こんにちは", 3, "こんに..."),
# Test case 7: Integer input
(12345, 3, 12345),
# Test case 8: Float input
(3.14159, 4, 3.14159),
# Test case 9: Boolean input
(True, 2, True),
# Test case 10: None input
(None, 5, None),
# Test case 11: Very long string
("a" * 1000, 10, "a" * 10 + "..."),
],
)
def test_truncate_long_strings_non_dict_list(input_data, max_length, expected):
result = truncate_long_strings(input_data, max_length)
assert result == expected


# Test for max_length of 0
def test_truncate_long_strings_zero_max_length():
assert truncate_long_strings("any string", 0) == "..."


# Test for negative max_length
def test_truncate_long_strings_negative_max_length():
assert truncate_long_strings("any string", -1) == "any string"


# Test for None max_length (should use default MAX_TEXT_LENGTH)
def test_truncate_long_strings_none_max_length():
from langflow.utils.constants import MAX_TEXT_LENGTH

long_string = "a" * (MAX_TEXT_LENGTH + 10)
result = truncate_long_strings(long_string, None)
assert len(result) == MAX_TEXT_LENGTH + 3 # +3 for "..."
assert result == "a" * MAX_TEXT_LENGTH + "..."
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ def test_truncate_long_strings_in_place_modification():
def test_truncate_long_strings_invalid_input():
input_string = "not a dict or list"
result = truncate_long_strings(input_string, 10)
assert result == input_string # The function should return the input unchanged
assert result == "not a dict..." # The function should return the input unchanged


# Updated test for negative max_length
Expand Down
Loading