From d7f81cbf55182175e2878bfb2f77c47c42e58dee Mon Sep 17 00:00:00 2001 From: Cristhian Zanforlin Lousa Date: Mon, 26 Jan 2026 17:59:37 -0300 Subject: [PATCH] perf: Limit enum options in tool schemas to reduce token usage (#11370) * fix current date tokens usage * Update src/lfx/src/lfx/io/schema.py * remove comment --------- Co-authored-by: Himavarsha <40851462+HimavarshaVS@users.noreply.github.com> --- src/lfx/src/lfx/io/schema.py | 25 +++-- .../unit/components/utilities/__init__.py | 0 .../components/utilities/test_current_date.py | 93 +++++++++++++++++++ 3 files changed, 112 insertions(+), 6 deletions(-) create mode 100644 src/lfx/tests/unit/components/utilities/__init__.py create mode 100644 src/lfx/tests/unit/components/utilities/test_current_date.py diff --git a/src/lfx/src/lfx/io/schema.py b/src/lfx/src/lfx/io/schema.py index f12bec75d1..bd42e130b1 100644 --- a/src/lfx/src/lfx/io/schema.py +++ b/src/lfx/src/lfx/io/schema.py @@ -15,6 +15,11 @@ from lfx.inputs.inputs import ( ) from lfx.schema.dotdict import dotdict +# Maximum number of options to include as enum in tool schemas. +# Dropdowns with more options will use string type with default value instead, +# avoiding token waste when sending tool schemas to LLMs. +MAX_OPTIONS_FOR_TOOL_ENUM = 50 + _convert_field_type_to_type: dict[FieldTypes, type] = { FieldTypes.TEXT: str, FieldTypes.INTEGER: int, @@ -217,10 +222,14 @@ def create_input_schema(inputs: list["InputTypes"]) -> type[BaseModel]: else: msg = f"Invalid field type: {field_type}" raise TypeError(msg) - if hasattr(input_model, "options") and isinstance(input_model.options, list) and input_model.options: + # Skip enum for large option lists to avoid token waste + if ( + hasattr(input_model, "options") + and isinstance(input_model.options, list) + and input_model.options + and len(input_model.options) <= MAX_OPTIONS_FOR_TOOL_ENUM + ): literal_string = f"Literal{input_model.options}" - # validate that the literal_string is a valid literal - field_type = eval(literal_string, {"Literal": Literal}) # noqa: S307 if hasattr(input_model, "is_list") and input_model.is_list: field_type = list[field_type] # type: ignore[valid-type] @@ -255,10 +264,14 @@ def create_input_schema_from_dict(inputs: list[dotdict], param_key: str | None = for input_model in inputs: # Create a Pydantic Field for each input field field_type = input_model.type - if hasattr(input_model, "options") and isinstance(input_model.options, list) and input_model.options: + # Skip enum for large option lists to avoid token waste + if ( + hasattr(input_model, "options") + and isinstance(input_model.options, list) + and input_model.options + and len(input_model.options) <= MAX_OPTIONS_FOR_TOOL_ENUM + ): literal_string = f"Literal{input_model.options}" - # validate that the literal_string is a valid literal - field_type = eval(literal_string, {"Literal": Literal}) # noqa: S307 if hasattr(input_model, "is_list") and input_model.is_list: field_type = list[field_type] # type: ignore[valid-type] diff --git a/src/lfx/tests/unit/components/utilities/__init__.py b/src/lfx/tests/unit/components/utilities/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/lfx/tests/unit/components/utilities/test_current_date.py b/src/lfx/tests/unit/components/utilities/test_current_date.py new file mode 100644 index 0000000000..a6c7751083 --- /dev/null +++ b/src/lfx/tests/unit/components/utilities/test_current_date.py @@ -0,0 +1,93 @@ +"""Tests for CurrentDateComponent tool schema optimization.""" + +import json + +from lfx.base.tools.component_tool import ComponentToolkit +from lfx.components.utilities.current_date import CurrentDateComponent +from lfx.io.schema import MAX_OPTIONS_FOR_TOOL_ENUM + + +class TestCurrentDateToolSchema: + """Tests to verify tool schema doesn't waste tokens on large option lists.""" + + def test_should_not_include_enum_when_options_exceed_limit(self): + """Verify schema uses string type instead of enum for large option lists.""" + # Arrange + component = CurrentDateComponent() + toolkit = ComponentToolkit(component) + + # Act + tools = toolkit.get_tools() + tool = tools[0] + schema = tool.args_schema.model_json_schema() + + # Assert + assert len(component.inputs[0].options) > MAX_OPTIONS_FOR_TOOL_ENUM + assert "enum" not in json.dumps(schema) + assert schema["properties"]["timezone"]["type"] == "string" + + def test_should_have_default_value_in_schema(self): + """Verify schema includes default value when enum is skipped.""" + # Arrange + component = CurrentDateComponent() + toolkit = ComponentToolkit(component) + + # Act + tools = toolkit.get_tools() + schema = tools[0].args_schema.model_json_schema() + + # Assert + assert schema["properties"]["timezone"]["default"] == "UTC" + + def test_should_reduce_schema_size_significantly(self): + """Verify schema size is reasonable (not wasting tokens).""" + # Arrange + component = CurrentDateComponent() + toolkit = ComponentToolkit(component) + max_acceptable_chars = 500 # Before fix was ~16000 + + # Act + tools = toolkit.get_tools() + schema_str = json.dumps(tools[0].args_schema.model_json_schema()) + + # Assert + assert len(schema_str) < max_acceptable_chars + + +class TestCurrentDateFunctionality: + """Tests to verify component still works correctly.""" + + def test_should_return_utc_time_by_default(self): + """Verify component returns UTC time when using default timezone.""" + # Arrange + component = CurrentDateComponent() + + # Act + result = component.get_current_date() + + # Assert + assert "UTC" in result.text + + def test_should_return_time_in_specified_timezone(self): + """Verify component respects timezone selection.""" + # Arrange + component = CurrentDateComponent() + component.timezone = "America/New_York" + + # Act + result = component.get_current_date() + + # Assert + assert "America/New_York" in result.text or "EST" in result.text or "EDT" in result.text + + def test_should_handle_invalid_timezone_gracefully(self): + """Verify component returns error for invalid timezone.""" + # Arrange + component = CurrentDateComponent() + component.timezone = "Invalid/Timezone" + + # Act + result = component.get_current_date() + + # Assert + assert "Error" in result.text