From 8890d37a44fc974b760efc289eca548a703735aa Mon Sep 17 00:00:00 2001 From: Gabriel Almeida Date: Tue, 28 Mar 2023 20:41:18 -0300 Subject: [PATCH] feat: added new validation functions and tests --- .../langflow/interface/custom_types.py | 4 +- src/backend/langflow/utils/validate.py | 118 +++++++++++++++++- tests/test_validate_code.py | 69 +++++++++- 3 files changed, 183 insertions(+), 8 deletions(-) diff --git a/src/backend/langflow/interface/custom_types.py b/src/backend/langflow/interface/custom_types.py index a8d99c521..05d77fd1d 100644 --- a/src/backend/langflow/interface/custom_types.py +++ b/src/backend/langflow/interface/custom_types.py @@ -24,7 +24,9 @@ class Function(BaseModel): def get_function(self): """Get the function""" - return validate.eval_function(self.code) + function_name = validate.extract_function_name(self.code) + + return validate.create_function(self.code, function_name) class PythonFunction(Function): diff --git a/src/backend/langflow/utils/validate.py b/src/backend/langflow/utils/validate.py index 3dfde9a3b..d7b85ddc7 100644 --- a/src/backend/langflow/utils/validate.py +++ b/src/backend/langflow/utils/validate.py @@ -4,6 +4,15 @@ import types from typing import Dict +def add_type_ignores(): + if not hasattr(ast, "TypeIgnore"): + + class TypeIgnore(ast.AST): + _fields = () + + ast.TypeIgnore = TypeIgnore + + def validate_code(code): # Initialize the errors dictionary errors = {"imports": {"errors": []}, "function": {"errors": []}} @@ -16,12 +25,7 @@ def validate_code(code): return errors # Add a dummy type_ignores field to the AST - if not hasattr(ast, "TypeIgnore"): - - class TypeIgnore(ast.AST): - _fields = () - - ast.TypeIgnore = TypeIgnore + add_type_ignores() tree.type_ignores = [] # Evaluate the import statements @@ -61,3 +65,105 @@ def eval_function(function_string: str): if function_object is None: raise ValueError("Function string does not contain a function") return function_object + + +def execute_function(code, function_name, *args, **kwargs): + add_type_ignores() + + module = ast.parse(code) + exec_globals = globals().copy() + + for node in module.body: + if isinstance(node, ast.Import): + for alias in node.names: + try: + exec( + f"{alias.asname or alias.name} = importlib.import_module('{alias.name}')", + exec_globals, + locals(), + ) + exec_globals[alias.asname or alias.name] = importlib.import_module( + alias.name + ) + except ModuleNotFoundError as e: + raise ModuleNotFoundError( + f"Module {alias.name} not found. Please install it and try again." + ) from e + + function_code = next( + node + for node in module.body + if isinstance(node, ast.FunctionDef) and node.name == function_name + ) + function_code.parent = None + code_obj = compile( + ast.Module(body=[function_code], type_ignores=[]), "", "exec" + ) + try: + exec(code_obj, exec_globals, locals()) + except Exception as e: + # handle execution error here + pass + + # Add the function to the exec_globals dictionary + exec_globals[function_name] = locals()[function_name] + + return exec_globals[function_name](*args, **kwargs) + + +def create_function(code, function_name): + if not hasattr(ast, "TypeIgnore"): + + class TypeIgnore(ast.AST): + _fields = () + + ast.TypeIgnore = TypeIgnore + + module = ast.parse(code) + exec_globals = globals().copy() + + for node in module.body: + if isinstance(node, ast.Import): + for alias in node.names: + try: + exec_globals[alias.asname or alias.name] = importlib.import_module( + alias.name + ) + except ModuleNotFoundError as e: + raise ModuleNotFoundError( + f"Module {alias.name} not found. Please install it and try again." + ) from e + + function_code = next( + node + for node in module.body + if isinstance(node, ast.FunctionDef) and node.name == function_name + ) + function_code.parent = None + code_obj = compile( + ast.Module(body=[function_code], type_ignores=[]), "", "exec" + ) + try: + exec(code_obj, exec_globals, locals()) + except Exception as e: + pass + + exec_globals[function_name] = locals()[function_name] + + # Return a function that imports necessary modules and calls the target function + def wrapped_function(*args, **kwargs): + for module_name, module in exec_globals.items(): + if isinstance(module, type(importlib)): + globals()[module_name] = module + + return exec_globals[function_name](*args, **kwargs) + + return wrapped_function + + +def extract_function_name(code): + module = ast.parse(code) + for node in module.body: + if isinstance(node, ast.FunctionDef): + return node.name + raise ValueError("No function definition found in the code string") diff --git a/tests/test_validate_code.py b/tests/test_validate_code.py index 5fdfa1199..9cb47f6e1 100644 --- a/tests/test_validate_code.py +++ b/tests/test_validate_code.py @@ -1,4 +1,12 @@ -from langflow.utils.validate import validate_code +from langflow.utils.validate import ( + create_function, + extract_function_name, + validate_code, + execute_function, +) +import pytest +from requests.exceptions import MissingSchema +from unittest import mock def test_validate_code(): @@ -37,3 +45,62 @@ def square(x) "imports": {"errors": []}, "function": {"errors": ["expected ':' (, line 4)"]}, } + + +def test_execute_function_success(): + code = """ +import math + +def my_function(x): + return math.sin(x) + 1 + """ + result = execute_function(code, "my_function", 0.5) + assert result == 1.479425538604203 + + +def test_execute_function_missing_module(): + code = """ +import some_missing_module + +def my_function(x): + return some_missing_module.some_function(x) + """ + with pytest.raises(ModuleNotFoundError): + execute_function(code, "my_function", 0.5) + + +def test_execute_function_missing_function(): + code = """ +import math + +def my_function(x): + return math.some_missing_function(x) + """ + with pytest.raises(AttributeError): + execute_function(code, "my_function", 0.5) + + +def test_execute_function_missing_schema(): + code = """ +import requests + +def my_function(x): + return requests.get(x).text + """ + with mock.patch("requests.get", side_effect=MissingSchema): + with pytest.raises(MissingSchema): + execute_function(code, "my_function", "invalid_url") + + +def test_create_function(): + code = """ +import math + +def my_function(x): + return math.sin(x) + 1 +""" + + function_name = extract_function_name(code) + function = create_function(code, function_name) + result = function(0.5) + assert result == 1.479425538604203