From 0c025977abf81571182fb006126131f23073bd02 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Wed, 27 Mar 2024 22:06:36 -0300 Subject: [PATCH] Refactor variable names and update variable service --- .../custom_component/custom_component.py | 14 +++--- .../langflow/interface/initialize/loading.py | 2 +- .../langflow/services/variable/service.py | 50 +++++++++++++++---- 3 files changed, 48 insertions(+), 18 deletions(-) diff --git a/src/backend/base/langflow/interface/custom/custom_component/custom_component.py b/src/backend/base/langflow/interface/custom/custom_component/custom_component.py index 4276d4616..c61b6ffb0 100644 --- a/src/backend/base/langflow/interface/custom/custom_component/custom_component.py +++ b/src/backend/base/langflow/interface/custom/custom_component/custom_component.py @@ -359,26 +359,26 @@ class CustomComponent(Component): return self.build_template_config() @property - def keys(self): + def variables(self): """ - Returns the credential for the current user with the specified name. + Returns the variable for the current user with the specified name. Raises: ValueError: If the user id is not set. Returns: - The credential for the current user with the specified name. + The variable for the current user with the specified name. """ - def get_credential(name: str): + def get_variable(name: 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 credential by name for the current user + # Retrieve and decrypt the variable by name for the current user with session_scope() as session: - return variable_service.get_credential(user_id=self._user_id or "", name=name, session=session) + return variable_service.get_variable(user_id=self._user_id or "", name=name, session=session) - return get_credential + return get_variable def list_key_names(self): """ diff --git a/src/backend/base/langflow/interface/initialize/loading.py b/src/backend/base/langflow/interface/initialize/loading.py index b1be1b6b6..bb64145ab 100644 --- a/src/backend/base/langflow/interface/initialize/loading.py +++ b/src/backend/base/langflow/interface/initialize/loading.py @@ -152,7 +152,7 @@ def update_params_with_load_from_db_fields(custom_component: "CustomComponent", for field in load_from_db_fields: if field in params: try: - key = custom_component.keys(params[field]) + key = custom_component.variables(params[field]) params[field] = key if key else params[field] except Exception as exc: logger.error(f"Failed to get value for {field} from custom component. Error: {exc}") diff --git a/src/backend/base/langflow/services/variable/service.py b/src/backend/base/langflow/services/variable/service.py index a555f3177..23d1cd806 100644 --- a/src/backend/base/langflow/services/variable/service.py +++ b/src/backend/base/langflow/services/variable/service.py @@ -19,18 +19,48 @@ class VariableService(Service): def __init__(self, settings_service: "SettingsService"): self.settings_service = settings_service - def get_credential(self, user_id: Union[UUID, str], name: str, session: Session = Depends(get_session)) -> str: + def get_variable(self, user_id: Union[UUID, str], name: 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() - credential = session.exec(select(Variable).where(Variable.user_id == user_id, Variable.name == name)).first() + variable = session.exec(select(Variable).where(Variable.user_id == user_id, Variable.name == name)).first() # we decrypt the value - if not credential or not credential.value: - raise ValueError(f"{name} credential not found.") - decrypted = auth_utils.decrypt_api_key(credential.value, settings_service=self.settings_service) + if not variable or not variable.value: + raise ValueError(f"{name} variable not found.") + decrypted = auth_utils.decrypt_api_key(variable.value, settings_service=self.settings_service) return decrypted - def list_credentials( - self, user_id: Union[UUID, str], session: Session = Depends(get_session) - ) -> list[Optional[str]]: - credentials = session.exec(select(Variable).where(Variable.user_id == user_id)).all() - return [credential.name for credential in credentials] + def list_variables(self, user_id: Union[UUID, str], session: Session = Depends(get_session)) -> list[Optional[str]]: + variables = session.exec(select(Variable).where(Variable.user_id == user_id)).all() + return [variable.name for variable in variables] + + def update_variable( + self, user_id: Union[UUID, str], name: str, value: str, session: Session = Depends(get_session) + ): + variable = session.exec(select(Variable).where(Variable.user_id == user_id, Variable.name == name)).first() + if not variable: + raise ValueError(f"{name} variable not found.") + encrypted = auth_utils.encrypt_api_key(value, settings_service=self.settings_service) + variable.value = encrypted + session.add(variable) + session.commit() + session.refresh(variable) + return variable + + def delete_variable(self, user_id: Union[UUID, str], name: str, session: Session = Depends(get_session)): + variable = session.exec(select(Variable).where(Variable.user_id == user_id, Variable.name == name)).first() + if not variable: + raise ValueError(f"{name} variable not found.") + session.delete(variable) + session.commit() + return variable + + def create_variable( + self, user_id: Union[UUID, str], name: str, value: str, session: Session = Depends(get_session) + ): + variable = Variable( + user_id=user_id, name=name, value=auth_utils.encrypt_api_key(value, settings_service=self.settings_service) + ) + session.add(variable) + session.commit() + session.refresh(variable) + return variable