feat: Enhance Data Operations Component with JSON features (#9688)

* feat: Enhance Data Operations Component with JSON features

This PR enhances the DataOperationsComponent by introducing two new powerful operations:

JSON Path: Allows extracting values from nested JSON structures using path expressions, enabling more flexible data selection.

JSON Query: Integrates jq to support advanced JSON queries, including filters, projections, and transformations.

Additional improvements:

Input handling and validations for both new operations.

Enhanced error messages and logs for better debugging.

Extended UI field visibility logic in update_build_config to support dynamic path selection from JSON.

These updates significantly expand the component's capabilities for parsing, transforming, and querying structured JSON data.

Summary by CodeRabbit
New Features

Added Mapped JSON and jQuery actions to map and query JSON, including JSON path extraction and jq-based queries.
Automatically repairs and parses malformed JSON before processing.

New inputs for JSON operations: multiline JSON preview, dynamic key selector, and query field.
Documentation

Updated in-app help texts to clarify top-level key constraints and provide guidance for JSON mapping and querying.

* [autofix.ci] apply automated fixes

* feat: Refactor data_operations.py for clarity and updates

* [autofix.ci] apply automated fixes

* Update data_operations.py

* [autofix.ci] apply automated fixes

* Fix ruff check

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Eric Hare <[email protected]>

Author: 3 people

src/lfx/src/lfx/components/processing/data_operations.py

@@ -1,9 +1,13 @@
 import ast
+import json
 from typing import TYPE_CHECKING, Any
 
+import jq
+from json_repair import repair_json
+
 from lfx.custom import Component
 from lfx.inputs import DictInput, DropdownInput, MessageTextInput, SortableListInput
-from lfx.io import DataInput, Output
+from lfx.io import DataInput, MultilineInput, Output
 from lfx.log.logger import logger
 from lfx.schema import Data
 from lfx.schema.dotdict import dotdict
@@ -20,6 +24,8 @@
     "Append or Update": {"is_list": False, "log_msg": "setting Append or Update fields"},
     "Remove Keys": {"is_list": False, "log_msg": "setting remove keys fields"},
     "Rename Keys": {"is_list": False, "log_msg": "setting rename keys fields"},
+    "Path Selection": {"is_list": False, "log_msg": "setting mapped key extractor fields"},
+    "JQ Expression": {"is_list": False, "log_msg": "setting parse json fields"},
 }
 OPERATORS = {
     "equals": lambda a, b: str(a) == str(b),
@@ -33,7 +39,6 @@
 class DataOperationsComponent(Component):
     display_name = "Data Operations"
     description = "Perform various operations on a Data object."
-    documentation: str = "https://docs.langflow.org/components-processing#data-operations"
     icon = "file-json"
     name = "DataOperations"
     default_keys = ["operations", "data"]
@@ -59,6 +64,9 @@ class DataOperationsComponent(Component):
             "data filtering",
             "data selection",
             "data combination",
+            "Parse JSON",
+            "JSON Query",
+            "JQ Query",
         ],
     }
     actions_data = {
@@ -69,8 +77,47 @@ class DataOperationsComponent(Component):
         "Append or Update": ["append_update_data", "operations"],
         "Remove Keys": ["remove_keys_input", "operations"],
         "Rename Keys": ["rename_keys_input", "operations"],
+        "Path Selection": ["mapped_json_display", "selected_key", "operations"],
+        "JQ Expression": ["query", "operations"],
     }
 
+    @staticmethod
+    def extract_all_paths(obj, path=""):
+        paths = []
+        if isinstance(obj, dict):
+            for k, v in obj.items():
+                new_path = f"{path}.{k}" if path else f".{k}"
+                paths.append(new_path)
+                paths.extend(DataOperationsComponent.extract_all_paths(v, new_path))
+        elif isinstance(obj, list) and obj:
+            new_path = f"{path}[0]"
+            paths.append(new_path)
+            paths.extend(DataOperationsComponent.extract_all_paths(obj[0], new_path))
+        return paths
+
+    @staticmethod
+    def remove_keys_recursive(obj, keys_to_remove):
+        if isinstance(obj, dict):
+            return {
+                k: DataOperationsComponent.remove_keys_recursive(v, keys_to_remove)
+                for k, v in obj.items()
+                if k not in keys_to_remove
+            }
+        if isinstance(obj, list):
+            return [DataOperationsComponent.remove_keys_recursive(item, keys_to_remove) for item in obj]
+        return obj
+
+    @staticmethod
+    def rename_keys_recursive(obj, rename_map):
+        if isinstance(obj, dict):
+            return {
+                rename_map.get(k, k): DataOperationsComponent.rename_keys_recursive(v, rename_map)
+                for k, v in obj.items()
+            }
+        if isinstance(obj, list):
+            return [DataOperationsComponent.rename_keys_recursive(item, rename_map) for item in obj]
+        return obj
+
     inputs = [
         DataInput(name="data", display_name="Data", info="Data object to filter.", required=True, is_list=True),
         SortableListInput(
@@ -86,6 +133,8 @@ class DataOperationsComponent(Component):
                 {"name": "Append or Update", "icon": "circle-plus"},
                 {"name": "Remove Keys", "icon": "eraser"},
                 {"name": "Rename Keys", "icon": "pencil-line"},
+                {"name": "Path Selection", "icon": "mouse-pointer"},
+                {"name": "JQ Expression", "icon": "terminal"},
             ],
             real_time_refresh=True,
             limit=1,
@@ -94,15 +143,18 @@ class DataOperationsComponent(Component):
         MessageTextInput(
             name="select_keys_input",
             display_name="Select Keys",
-            info="List of keys to select from the data.",
+            info="List of keys to select from the data. Only top-level keys can be selected.",
             show=False,
             is_list=True,
         ),
         # filter values inputs
         MessageTextInput(
             name="filter_key",
             display_name="Filter Key",
-            info="Key to filter by.",
+            info=(
+                "Name of the key containing the list to filter. "
+                "It must be a top-level key in the JSON and its value must be a list."
+            ),
             is_list=True,
             show=False,
         ),
@@ -126,7 +178,7 @@ class DataOperationsComponent(Component):
         DictInput(
             name="append_update_data",
             display_name="Append or Update",
-            info="Data to Append or Updatethe existing data with.",
+            info="Data to append or update the existing data with. Only top-level keys are checked.",
             show=False,
             value={"key": "value"},
             is_list=True,
@@ -148,6 +200,26 @@ class DataOperationsComponent(Component):
             is_list=True,
             value={"old_key": "new_key"},
         ),
+        MultilineInput(
+            name="mapped_json_display",
+            display_name="JSON to Map",
+            info="Paste or preview your JSON here to explore its structure and select a path for extraction.",
+            required=False,
+            refresh_button=True,
+            real_time_refresh=True,
+            placeholder="Add a JSON example.",
+            show=False,
+        ),
+        DropdownInput(
+            name="selected_key", display_name="Select Path", options=[], required=False, dynamic=True, show=False
+        ),
+        MessageTextInput(
+            name="query",
+            display_name="JQ Expression",
+            info="JSON Query to filter the data. Used by Parse JSON operation.",
+            placeholder="e.g., .properties.id",
+            show=False,
+        ),
     ]
     outputs = [
         Output(display_name="Data", name="data_output", method="as_data"),
@@ -156,10 +228,39 @@ class DataOperationsComponent(Component):
     # Helper methods for data operations
     def get_data_dict(self) -> dict:
         """Extract data dictionary from Data object."""
-        # TODO: rasie error if it s list of data objects
         data = self.data[0] if isinstance(self.data, list) and len(self.data) == 1 else self.data
         return data.model_dump()
 
+    def json_query(self) -> Data:
+        import json
+
+        import jq
+
+        if not self.query or not self.query.strip():
+            msg = "JSON Query is required and cannot be blank."
+            raise ValueError(msg)
+        raw_data = self.get_data_dict()
+        try:
+            input_str = json.dumps(raw_data)
+            repaired = repair_json(input_str)
+            data_json = json.loads(repaired)
+            jq_input = data_json["data"] if isinstance(data_json, dict) and "data" in data_json else data_json
+            results = jq.compile(self.query).input(jq_input).all()
+            if not results:
+                msg = "No result from JSON query."
+                raise ValueError(msg)
+            result = results[0] if len(results) == 1 else results
+            if result is None or result == "None":
+                msg = "JSON query returned null/None. Check if the path exists in your data."
+                raise ValueError(msg)
+            if isinstance(result, dict):
+                return Data(data=result)
+            return Data(data={"result": result})
+        except (ValueError, TypeError, KeyError, json.JSONDecodeError) as e:
+            logger.error(f"JSON Query failed: {e}")
+            msg = f"JSON Query error: {e}"
+            raise ValueError(msg) from e
+
     def get_normalized_data(self) -> dict:
         """Get normalized data dictionary, handling the 'data' key if present."""
         data_dict = self.get_data_dict()
@@ -204,34 +305,22 @@ def select_keys(self, *, evaluate: bool | None = None) -> Data:
         return Data(data=filtered)
 
     def remove_keys(self) -> Data:
-        """Remove specified keys from the data dictionary."""
+        """Remove specified keys from the data dictionary, recursively."""
         self.validate_single_data("Remove Keys")
         data_dict = self.get_normalized_data()
         remove_keys_input: list[str] = self.remove_keys_input
 
-        for key in remove_keys_input:
-            if key in data_dict:
-                data_dict.pop(key)
-            else:
-                logger.warning(f"Key '{key}' not found in data. Skipping removal.")
-
-        return Data(**data_dict)
+        filtered = DataOperationsComponent.remove_keys_recursive(data_dict, set(remove_keys_input))
+        return Data(data=filtered)
 
     def rename_keys(self) -> Data:
-        """Rename keys in the data dictionary."""
+        """Rename keys in the data dictionary, recursively."""
         self.validate_single_data("Rename Keys")
         data_dict = self.get_normalized_data()
         rename_keys_input: dict[str, str] = self.rename_keys_input
 
-        for old_key, new_key in rename_keys_input.items():
-            if old_key in data_dict:
-                data_dict[new_key] = data_dict[old_key]
-                data_dict.pop(old_key)
-            else:
-                msg = f"Key '{old_key}' not found in data. Skipping rename."
-                raise ValueError(msg)
-
-        return Data(**data_dict)
+        renamed = DataOperationsComponent.rename_keys_recursive(data_dict, rename_keys_input)
+        return Data(data=renamed)
 
     def recursive_eval(self, data: Any) -> Any:
         """Recursively evaluate string values in a dictionary or list.
@@ -299,13 +388,6 @@ def combine_data(self, *, evaluate: bool | None = None) -> Data:
 
         return Data(**combined_data)
 
-    def compare_values(self, item_value: Any, filter_value: str, operator: str) -> bool:
-        """Compare values based on the specified operator."""
-        comparison_func = OPERATORS.get(operator)
-        if comparison_func:
-            return comparison_func(item_value, filter_value)
-        return False
-
     def filter_data(self, input_data: list[dict[str, Any]], filter_key: str, filter_value: str, operator: str) -> list:
         """Filter list data based on key, value, and operator."""
         # Validate inputs
@@ -328,6 +410,12 @@ def filter_data(self, input_data: list[dict[str, Any]], filter_key: str, filter_
 
         return filtered_data
 
+    def compare_values(self, item_value: Any, filter_value: str, operator: str) -> bool:
+        comparison_func = OPERATORS.get(operator)
+        if comparison_func:
+            return comparison_func(item_value, filter_value)
+        return False
+
     def multi_filter_data(self) -> Data:
         """Apply multiple filters to the data."""
         self.validate_single_data("Filter Values")
@@ -366,57 +454,55 @@ def append_update(self) -> Data:
 
     # Configuration and execution methods
     def update_build_config(self, build_config: dotdict, field_value: Any, field_name: str | None = None) -> dotdict:
-        """Update build configuration based on selected action."""
-        if field_name != "operations":
-            return build_config
-
-        build_config["operations"]["value"] = field_value
-        selected_actions = [action["name"] for action in field_value]
-
-        # Handle single action case
-        if len(selected_actions) == 1 and selected_actions[0] in ACTION_CONFIG:
-            action = selected_actions[0]
-            config = ACTION_CONFIG[action]
-
-            build_config["data"]["is_list"] = config["is_list"]
-            logger.info(config["log_msg"])
-
-            return set_current_fields(
-                build_config=build_config,
-                action_fields=self.actions_data,
-                selected_action=action,
-                default_fields=self.default_keys,
-                func=set_field_display,
-            )
-
-        # Handle no operations case
-        if not selected_actions:
-            logger.info("setting default fields")
-            return set_current_fields(
-                build_config=build_config,
-                action_fields=self.actions_data,
-                selected_action=None,
-                default_fields=self.default_keys,
-                func=set_field_display,
-            )
+        if field_name == "operations":
+            build_config["operations"]["value"] = field_value
+            selected_actions = [action["name"] for action in field_value]
+            if len(selected_actions) == 1 and selected_actions[0] in ACTION_CONFIG:
+                action = selected_actions[0]
+                config = ACTION_CONFIG[action]
+                build_config["data"]["is_list"] = config["is_list"]
+                logger.info(config["log_msg"])
+                return set_current_fields(
+                    build_config, self.actions_data, action, ["operations", "data"], set_field_display
+                )
+
+        if field_name == "mapped_json_display":
+            try:
+                parsed_json = json.loads(field_value)
+                keys = DataOperationsComponent.extract_all_paths(parsed_json)
+                build_config["selected_key"]["options"] = keys
+                build_config["selected_key"]["show"] = True
+            except (json.JSONDecodeError, TypeError, ValueError) as e:
+                logger.error(f"Error parsing mapped JSON: {e}")
+                build_config["selected_key"]["show"] = False
 
         return build_config
 
+    def json_path(self) -> Data:
+        try:
+            if not self.data or not self.selected_key:
+                msg = "Missing input data or selected key."
+                raise ValueError(msg)
+            input_payload = self.data[0].data if isinstance(self.data, list) else self.data.data
+            compiled = jq.compile(self.selected_key)
+            result = compiled.input(input_payload).first()
+            if isinstance(result, dict):
+                return Data(data=result)
+            return Data(data={"result": result})
+        except (ValueError, TypeError, KeyError) as e:
+            self.status = f"Error: {e!s}"
+            self.log(self.status)
+            return Data(data={"error": str(e)})
+
     def as_data(self) -> Data:
-        """Execute the selected action on the data."""
         if not hasattr(self, "operations") or not self.operations:
             return Data(data={})
 
         selected_actions = [action["name"] for action in self.operations]
         logger.info(f"selected_actions: {selected_actions}")
-
-        # Only handle single action case for now
         if len(selected_actions) != 1:
             return Data(data={})
 
-        action = selected_actions[0]
-
-        # Explicitly type the action_map
         action_map: dict[str, Callable[[], Data]] = {
             "Select Keys": self.select_keys,
             "Literal Eval": self.evaluate_data,
@@ -425,14 +511,14 @@ def as_data(self) -> Data:
             "Append or Update": self.append_update,
             "Remove Keys": self.remove_keys,
             "Rename Keys": self.rename_keys,
+            "Path Selection": self.json_path,
+            "JQ Expression": self.json_query,
         }
-
-        handler: Callable[[], Data] | None = action_map.get(action)
+        handler: Callable[[], Data] | None = action_map.get(selected_actions[0])
         if handler:
             try:
                 return handler()
             except Exception as e:
-                logger.error(f"Error executing {action}: {e!s}")
+                logger.error(f"Error executing {selected_actions[0]}: {e!s}")
                 raise
-
         return Data(data={})