🔨 refactor(process.py): refactor process_tweaks function to improve readability and maintainability

 feat(process.py): add input validation to process_tweaks function
The process_tweaks function has been refactored to improve readability and maintainability. The apply_tweaks function has been added to apply the tweaks to the node. The validate_input function has been added to validate the input parameters. The process_tweaks function now raises a ValueError if the input is not in the expected format.
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-06-19 11:59:34 -03:00
commit 83c28dcabe

View file

@ -11,7 +11,6 @@ from langflow.interface.run import (
from langflow.utils.logger import logger from langflow.utils.logger import logger
from langflow.graph import Graph from langflow.graph import Graph
from typing import Any, Dict, List, Optional, Tuple, Union from typing import Any, Dict, List, Optional, Tuple, Union
@ -179,23 +178,61 @@ def load_flow_from_json(
return graph return graph
def process_tweaks(graph_data: Dict, tweaks: Dict): def validate_input(
"""This function is used to tweak the graph data using the node id and the tweaks dict""" graph_data: Dict[str, Any], tweaks: Dict[str, Dict[str, Any]]
# the tweaks dict is a dict of dicts ) -> List[Dict[str, Any]]:
# the key is the node id and the value is a dict of the tweaks if not isinstance(graph_data, dict) or not isinstance(tweaks, dict):
# the dict of tweaks contains the name of a certain parameter and the value to be tweaked raise ValueError("graph_data and tweaks should be dictionaries")
nodes = graph_data.get("data", {}).get("nodes") or graph_data.get("nodes")
if not isinstance(nodes, list):
raise ValueError(
"graph_data should contain a list of nodes under 'data' key or directly under 'nodes' key"
)
return nodes
def apply_tweaks(node: Dict[str, Any], node_tweaks: Dict[str, Any]) -> None:
template_data = node.get("data", {}).get("node", {}).get("template")
if not isinstance(template_data, dict):
logger.warning(
f"Template data for node {node.get('id')} should be a dictionary"
)
return
for tweak_name, tweak_value in node_tweaks.items():
if tweak_name and tweak_value and tweak_name in template_data:
template_data[tweak_name]["value"] = tweak_value
def process_tweaks(
graph_data: Dict[str, Any], tweaks: Dict[str, Dict[str, Any]]
) -> Dict[str, Any]:
"""
This function is used to tweak the graph data using the node id and the tweaks dict.
:param graph_data: The dictionary containing the graph data. It must contain a 'data' key with
'nodes' as its child or directly contain 'nodes' key. Each node should have an 'id' and 'data'.
:param tweaks: A dictionary where the key is the node id and the value is a dictionary of the tweaks.
The inner dictionary contains the name of a certain parameter as the key and the value to be tweaked.
:return: The modified graph_data dictionary.
:raises ValueError: If the input is not in the expected format.
"""
nodes = validate_input(graph_data, tweaks)
# We need to process the graph data to add the tweaks
if "data" not in graph_data and "nodes" in graph_data:
nodes = graph_data["nodes"]
else:
nodes = graph_data["data"]["nodes"]
for node in nodes: for node in nodes:
node_id = node["id"] if isinstance(node, dict) and isinstance(node.get("id"), str):
if node_id in tweaks: node_id = node["id"]
node_tweaks = tweaks[node_id] if node_tweaks := tweaks.get(node_id):
template_data = node["data"]["node"]["template"] apply_tweaks(node, node_tweaks)
for tweak_name, tweake_value in node_tweaks.items(): else:
if tweak_name in template_data: logger.warning(
template_data[tweak_name]["value"] = tweake_value "Each node should be a dictionary with an 'id' key of type str"
)
return graph_data return graph_data