use set for better performance

This commit is contained in:
Pankaj Kaushal 2025-10-08 22:09:13 +05:30
parent 0ab4245462
commit 21035cfca0
3 changed files with 11 additions and 10 deletions

View File

@ -55,7 +55,7 @@ def decrypt_token_with_decoding(token: str, rsa_key, cipher_rsa):
@overload
def encrypt_secret_keys(
obj: Mapping[str, Any],
secret_variables: list[str] | None = None,
secret_variables: set[str] | None = None,
parent_key: str | None = None,
) -> Mapping[str, Any]: ...
@ -63,7 +63,7 @@ def encrypt_secret_keys(
@overload
def encrypt_secret_keys(
obj: list[Any],
secret_variables: list[str] | None = None,
secret_variables: set[str] | None = None,
parent_key: str | None = None,
) -> list[Any]: ...
@ -71,21 +71,22 @@ def encrypt_secret_keys(
@overload
def encrypt_secret_keys(
obj: Any,
secret_variables: list[str] | None = None,
secret_variables: set[str] | None = None,
parent_key: str | None = None,
) -> Any: ...
def encrypt_secret_keys(
obj: Any,
secret_variables: list[str] | None = None,
secret_variables: set[str] | None = None,
parent_key: str | None = None,
) -> Any:
"""
Recursively obfuscate the value if it belongs to a Secret Variable.
Preserves input type: dict -> dict, list -> list, scalar -> scalar.
"""
secret_variables = secret_variables or []
if secret_variables is None:
secret_variables = set()
if isinstance(obj, dict):
# recurse into dict

View File

@ -144,11 +144,11 @@ class AgentNode(Node):
env_vars = self.graph_runtime_state.variable_pool.variable_dictionary.get("env", {})
# get secret variables used
secret_variables = [
secret_variables = {
var.name
for var in env_vars.values() # iterate over the values directly
if isinstance(var, SecretVariable)
]
}
try:
message_stream = strategy.invoke(
@ -500,7 +500,7 @@ class AgentNode(Node):
node_type: NodeType,
node_id: str,
node_execution_id: str,
secret_variables: list[str] | None = None,
secret_variables: set[str] | None = None,
) -> Generator[NodeEventBase, None, None]:
"""
Convert ToolInvokeMessages into tuple[plain_text, files]

View File

@ -76,7 +76,7 @@ class CodeNode(Node):
code = self._node_data.code
# to store secret variables used in the code block.
secret_variables = []
secret_variables = set()
# Get variables
variables = {}
@ -84,7 +84,7 @@ class CodeNode(Node):
variable_name = variable_selector.variable
variable = self.graph_runtime_state.variable_pool.get(variable_selector.value_selector)
if isinstance(variable, SecretVariable):
secret_variables.append(variable_name)
secret_variables.add(variable_name)
if isinstance(variable, ArrayFileSegment):
variables[variable_name] = [v.to_dict() for v in variable.value] if variable.value else None