From 565df927dcf8460a4f73c1f24adfb92474a4685e Mon Sep 17 00:00:00 2001 From: italojohnny Date: Sun, 9 Jun 2024 16:01:42 -0300 Subject: [PATCH] add exception to Credentials vars in session_id fields --- .../langflow/custom/custom_component/custom_component.py | 5 +++-- src/backend/base/langflow/interface/initialize/loading.py | 8 +++++--- src/backend/base/langflow/services/variable/service.py | 8 ++++++++ 3 files changed, 16 insertions(+), 5 deletions(-) diff --git a/src/backend/base/langflow/custom/custom_component/custom_component.py b/src/backend/base/langflow/custom/custom_component/custom_component.py index aeac9cae6..f6b90a368 100644 --- a/src/backend/base/langflow/custom/custom_component/custom_component.py +++ b/src/backend/base/langflow/custom/custom_component/custom_component.py @@ -378,13 +378,14 @@ class CustomComponent(Component): The variable for the current user with the specified name. """ - def get_variable(name: str): + def get_variable(name: str, field: str): if hasattr(self, "_user_id") and not self._user_id: raise ValueError(f"User id is not set for {self.__class__.__name__}") variable_service = get_variable_service() # Get service instance # Retrieve and decrypt the variable by name for the current user with session_scope() as session: - return variable_service.get_variable(user_id=self._user_id or "", name=name, session=session) + user_id = self._user_id or "" + return variable_service.get_variable(user_id=user_id, name=name, field=field, session=session) return get_variable diff --git a/src/backend/base/langflow/interface/initialize/loading.py b/src/backend/base/langflow/interface/initialize/loading.py index 6ff498c0c..7b336487e 100644 --- a/src/backend/base/langflow/interface/initialize/loading.py +++ b/src/backend/base/langflow/interface/initialize/loading.py @@ -71,7 +71,7 @@ def update_params_with_load_from_db_fields( try: key = None try: - key = custom_component.variables(params[field]) + key = custom_component.variables(params[field], field) except ValueError as e: # check if "User id is not set" is in the error message if "User id is not set" in str(e) and not fallback_to_env_vars: @@ -86,8 +86,10 @@ def update_params_with_load_from_db_fields( if key is None: logger.warning(f"Could not get value for {field}. Setting it to None.") - if field != "session_id": - params[field] = key + params[field] = key + + except TypeError as exc: + raise exc except Exception as exc: logger.error(f"Failed to get value for {field} from custom component. Setting it to None. Error: {exc}") diff --git a/src/backend/base/langflow/services/variable/service.py b/src/backend/base/langflow/services/variable/service.py index 84671e0f9..b2389e890 100644 --- a/src/backend/base/langflow/services/variable/service.py +++ b/src/backend/base/langflow/services/variable/service.py @@ -54,11 +54,19 @@ class VariableService(Service): self, user_id: Union[UUID, str], name: str, + field: str, session: Session = Depends(get_session), ) -> str: # we get the credential from the database # credential = session.query(Variable).filter(Variable.user_id == user_id, Variable.name == name).first() variable = session.exec(select(Variable).where(Variable.user_id == user_id, Variable.name == name)).first() + + if variable.type == "Credential" and field == "session_id": + raise TypeError( + f"variable {name} of type 'Credential' cannot be used in a Session ID field " + "because its purpose is to prevent the exposure of values." + ) + # we decrypt the value if not variable or not variable.value: raise ValueError(f"{name} variable not found.")