Refactor operations count state setting (#631)

This commit is contained in:
Albert Villanova del Moral 2025-02-13 11:22:02 +01:00 committed by GitHub
parent 3316dd7b11
commit 41a388dac6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 2 additions and 4 deletions

View File

@ -1184,7 +1184,7 @@ def evaluate_ast(
The list of modules that can be imported by the code. By default, only a few safe modules are allowed.
If it contains "*", it will authorize any import. Use this at your own risk!
"""
if state["_operations_count"] >= MAX_OPERATIONS:
if state.setdefault("_operations_count", 0) >= MAX_OPERATIONS:
raise InterpreterError(
f"Reached the max number of operations of {MAX_OPERATIONS}. Maybe there is an infinite loop somewhere in the code, or you're just asking too many calculations."
)
@ -1363,7 +1363,6 @@ def evaluate_python_code(
custom_tools = custom_tools if custom_tools is not None else {}
result = None
state["_print_outputs"] = PrintContainer()
state["_operations_count"] = 0
def final_answer(value):
raise FinalAnswerException(value)

View File

@ -1149,7 +1149,6 @@ def test_evaluate_python_code_with_evaluate_delete(code, expected_error_message)
],
)
def test_evaluate_delete(code, state, expectation):
state["_operations_count"] = 0
delete_node = ast.parse(code).body[0]
if isinstance(expectation, Exception):
with pytest.raises(type(expectation)) as exception_info:
@ -1157,7 +1156,7 @@ def test_evaluate_delete(code, state, expectation):
assert str(expectation) in str(exception_info.value)
else:
evaluate_delete(delete_node, state, {}, {}, [])
del state["_operations_count"]
_ = state.pop("_operations_count", None)
assert state == expectation