ragflow/agent/component/variable_assigner.py

#
#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
from abc import ABC
import os
import numbers
from agent.component.base import ComponentBase, ComponentParamBase
from api.utils.api_utils import timeout

class VariableAssignerParam(ComponentParamBase):
    """
    Define the Variable Assigner component parameters.
    """
    def __init__(self):
        super().__init__()
        self.variables=[]

    def check(self):
        return True
    
    def get_input_form(self) -> dict[str, dict]:
        return {
            "items": {
                "type": "json",
                "name": "Items"
            }
        }

class VariableAssigner(ComponentBase,ABC):
    component_name = "VariableAssigner"
    _NO_PARAMETER_OPERATORS = {"clear", "remove_first", "remove_last"}

    @timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60)))
    def _invoke(self, **kwargs):
        if not isinstance(self._param.variables,list):
            return
        else:
            for item in self._param.variables:
                variable = item.get("variable")
                operator = item.get("operator")
                parameter = item.get("parameter")

                if any([not variable, not operator]):
                    raise ValueError("Variable is not complete.")
                if operator not in self._NO_PARAMETER_OPERATORS and parameter is None:
                    raise ValueError("Variable is not complete.")
                variable_value=self._canvas.get_variable_value(variable)
                new_variable=self._operate(variable_value,operator,parameter)
                self._canvas.set_variable_value(variable, new_variable)

    def _operate(self,variable,operator,parameter):
        if operator == "overwrite":
            return self._overwrite(parameter)
        elif operator == "clear":
            return self._clear(variable)
        elif operator == "set":
            return self._set(variable,parameter)
        elif operator == "append":
            return self._append(variable,parameter)
        elif operator == "extend":
            return self._extend(variable,parameter)
        elif operator == "remove_first":
            return self._remove_first(variable)
        elif operator == "remove_last":
            return self._remove_last(variable)
        elif operator == "+=":
            return self._add(variable,parameter)
        elif operator == "-=":
            return self._subtract(variable,parameter)
        elif operator == "*=":
            return self._multiply(variable,parameter)
        elif operator == "/=":
            return self._divide(variable,parameter)
        else:
            return
    
    def _overwrite(self,parameter):
        return self._canvas.get_variable_value(parameter)

    def _clear(self,variable):
        if isinstance(variable,list):
            return []
        elif isinstance(variable,str):
            return ""
        elif isinstance(variable,dict):
            return {}
        elif isinstance(variable,bool):
            return False
        elif isinstance(variable,int):
            return 0
        elif isinstance(variable,float):
            return 0.0
        else:
            return None

    def _set(self,variable,parameter):
        if variable is None:
            return self._canvas.get_value_with_variable(parameter)
        elif isinstance(variable,str):
            return self._canvas.get_value_with_variable(parameter)
        elif isinstance(variable,bool):
            return parameter
        elif isinstance(variable,int):
            return parameter
        elif isinstance(variable,float):
            return parameter
        else:
            return parameter

    def _append(self,variable,parameter):
        parameter=self._canvas.get_variable_value(parameter)
        if variable is None:
            variable=[]
        if not isinstance(variable,list):
            return "ERROR:VARIABLE_NOT_LIST"
        elif len(variable)!=0 and not isinstance(parameter,type(variable[0])):
            return "ERROR:PARAMETER_NOT_LIST_ELEMENT_TYPE"
        else:
            variable.append(parameter)
            return variable

    def _extend(self,variable,parameter):
        parameter=self._canvas.get_variable_value(parameter)
        if variable is None:
            variable=[]
        if not isinstance(variable,list):
            return "ERROR:VARIABLE_NOT_LIST"
        elif not isinstance(parameter,list):
            return "ERROR:PARAMETER_NOT_LIST"
        elif len(variable)!=0 and len(parameter)!=0 and not isinstance(parameter[0],type(variable[0])):
            return "ERROR:PARAMETER_NOT_LIST_ELEMENT_TYPE"
        else:
            return variable + parameter

    def _remove_first(self,variable):
        if not isinstance(variable,list):
            return "ERROR:VARIABLE_NOT_LIST"
        if len(variable)==0:
            return variable
        return variable[1:]

    def _remove_last(self,variable):
        if not isinstance(variable,list):
            return "ERROR:VARIABLE_NOT_LIST"
        if len(variable)==0:
            return variable
        return variable[:-1]

    def is_number(self, value):
        if isinstance(value, bool):
            return False
        return isinstance(value, numbers.Number)

    def _add(self,variable,parameter):
        if self.is_number(variable) and self.is_number(parameter):
            return variable + parameter
        else:
            return "ERROR:VARIABLE_NOT_NUMBER or PARAMETER_NOT_NUMBER"

    def _subtract(self,variable,parameter):
        if self.is_number(variable) and self.is_number(parameter):
            return variable - parameter
        else:
            return "ERROR:VARIABLE_NOT_NUMBER or PARAMETER_NOT_NUMBER"

    def _multiply(self,variable,parameter):
        if self.is_number(variable) and self.is_number(parameter):
            return variable * parameter
        else:
            return "ERROR:VARIABLE_NOT_NUMBER or PARAMETER_NOT_NUMBER"

    def _divide(self,variable,parameter):
        if self.is_number(variable) and self.is_number(parameter):
            if  parameter==0:
                return "ERROR:DIVIDE_BY_ZERO"
            else:
                return variable/parameter
        else:
            return  "ERROR:VARIABLE_NOT_NUMBER or PARAMETER_NOT_NUMBER"

    def thoughts(self) -> str:
        return "Assign variables from canvas."