Fimeg/letta-server
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

fix: Deprecate in memory function stores in agent.py (#2271)

Browse Source
This commit is contained in:
Matthew Zhou
2024-12-17 16:18:11 -08:00
committed by GitHub
parent 7fb8f16155
commit 9a0ffc84dd
7 changed files with 94 additions and 127 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
tests/helpers/endpoints_helper.py
View File

@@ -110,8 +110,7 @@ def check_first_response_is_valid_for_llm_endpoint(filename: str) -> ChatComplet
llm_config=agent_state.llm_config,
user_id=str(uuid.UUID(int=1)), # dummy user_id
messages=agent._messages,
functions=agent.functions,
functions_python=agent.functions_python,
functions=[t.json_schema for t in agent.agent_state.tools],
)
# Basic check

52
tests/integration_test_tool_execution_sandbox.py
View File

@@ -283,12 +283,12 @@ def test_local_sandbox_default(mock_e2b_api_key_none, add_integers_tool, test_us
# Mock and assert correct pathway was invoked
with patch.object(ToolExecutionSandbox, "run_local_dir_sandbox") as mock_run_local_dir_sandbox:
sandbox = ToolExecutionSandbox(add_integers_tool.name, args, user_id=test_user.id)
sandbox = ToolExecutionSandbox(add_integers_tool.name, args, user=test_user)
sandbox.run()
mock_run_local_dir_sandbox.assert_called_once()
# Run again to get actual response
sandbox = ToolExecutionSandbox(add_integers_tool.name, args, user_id=test_user.id)
sandbox = ToolExecutionSandbox(add_integers_tool.name, args, user=test_user)
result = sandbox.run()
assert result.func_return == args["x"] + args["y"]
@@ -297,7 +297,7 @@ def test_local_sandbox_default(mock_e2b_api_key_none, add_integers_tool, test_us
def test_local_sandbox_stateful_tool(mock_e2b_api_key_none, clear_core_memory_tool, test_user, agent_state):
args = {}
# Run again to get actual response
sandbox = ToolExecutionSandbox(clear_core_memory_tool.name, args, user_id=test_user.id)
sandbox = ToolExecutionSandbox(clear_core_memory_tool.name, args, user=test_user)
result = sandbox.run(agent_state=agent_state)
assert result.agent_state.memory.get_block("human").value == ""
assert result.agent_state.memory.get_block("persona").value == ""
@@ -306,7 +306,7 @@ def test_local_sandbox_stateful_tool(mock_e2b_api_key_none, clear_core_memory_to
@pytest.mark.local_sandbox
def test_local_sandbox_with_list_rv(mock_e2b_api_key_none, list_tool, test_user):
sandbox = ToolExecutionSandbox(list_tool.name, {}, user_id=test_user.id)
sandbox = ToolExecutionSandbox(list_tool.name, {}, user=test_user)
result = sandbox.run()
assert len(result.func_return) == 5
@@ -331,7 +331,7 @@ def test_local_sandbox_env(mock_e2b_api_key_none, get_env_tool, test_user):
args = {}
# Run the custom sandbox
sandbox = ToolExecutionSandbox(get_env_tool.name, args, user_id=test_user.id)
sandbox = ToolExecutionSandbox(get_env_tool.name, args, user=test_user)
result = sandbox.run()
assert long_random_string in result.func_return
@@ -349,7 +349,7 @@ def test_local_sandbox_e2e_composio_star_github(mock_e2b_api_key_none, check_com
actor=test_user,
)
result = ToolExecutionSandbox(composio_github_star_tool.name, {"owner": "letta-ai", "repo": "letta"}, user_id=test_user.id).run()
result = ToolExecutionSandbox(composio_github_star_tool.name, {"owner": "letta-ai", "repo": "letta"}, user=test_user).run()
assert result.func_return["details"] == "Action executed successfully"
@@ -359,7 +359,7 @@ def test_local_sandbox_external_codebase(mock_e2b_api_key_none, custom_test_sand
args = {"percentage": 10}
# Run again to get actual response
sandbox = ToolExecutionSandbox(external_codebase_tool.name, args, user_id=test_user.id)
sandbox = ToolExecutionSandbox(external_codebase_tool.name, args, user=test_user)
result = sandbox.run()
# Assert that the function return is correct
@@ -371,14 +371,14 @@ def test_local_sandbox_external_codebase(mock_e2b_api_key_none, custom_test_sand
def test_local_sandbox_with_venv_and_warnings_does_not_error(
mock_e2b_api_key_none, custom_test_sandbox_config, get_warning_tool, test_user
):
sandbox = ToolExecutionSandbox(get_warning_tool.name, {}, user_id=test_user.id)
sandbox = ToolExecutionSandbox(get_warning_tool.name, {}, user=test_user)
result = sandbox.run()
assert result.func_return == "Hello World"
@pytest.mark.e2b_sandbox
def test_local_sandbox_with_venv_errors(mock_e2b_api_key_none, custom_test_sandbox_config, always_err_tool, test_user):
sandbox = ToolExecutionSandbox(always_err_tool.name, {}, user_id=test_user.id)
sandbox = ToolExecutionSandbox(always_err_tool.name, {}, user=test_user)
# run the sandbox
result = sandbox.run()
@@ -397,12 +397,12 @@ def test_e2b_sandbox_default(check_e2b_key_is_set, add_integers_tool, test_user)
# Mock and assert correct pathway was invoked
with patch.object(ToolExecutionSandbox, "run_e2b_sandbox") as mock_run_local_dir_sandbox:
sandbox = ToolExecutionSandbox(add_integers_tool.name, args, user_id=test_user.id)
sandbox = ToolExecutionSandbox(add_integers_tool.name, args, user=test_user)
sandbox.run()
mock_run_local_dir_sandbox.assert_called_once()
# Run again to get actual response
sandbox = ToolExecutionSandbox(add_integers_tool.name, args, user_id=test_user.id)
sandbox = ToolExecutionSandbox(add_integers_tool.name, args, user=test_user)
result = sandbox.run()
assert int(result.func_return) == args["x"] + args["y"]
@@ -420,14 +420,14 @@ def test_e2b_sandbox_pip_installs(check_e2b_key_is_set, cowsay_tool, test_user):
SandboxEnvironmentVariableCreate(key=key, value=long_random_string), sandbox_config_id=config.id, actor=test_user
)
sandbox = ToolExecutionSandbox(cowsay_tool.name, {}, user_id=test_user.id)
sandbox = ToolExecutionSandbox(cowsay_tool.name, {}, user=test_user)
result = sandbox.run()
assert long_random_string in result.stdout[0]
@pytest.mark.e2b_sandbox
def test_e2b_sandbox_reuses_same_sandbox(check_e2b_key_is_set, list_tool, test_user):
sandbox = ToolExecutionSandbox(list_tool.name, {}, user_id=test_user.id)
sandbox = ToolExecutionSandbox(list_tool.name, {}, user=test_user)
# Run the function once
result = sandbox.run()
@@ -442,7 +442,7 @@ def test_e2b_sandbox_reuses_same_sandbox(check_e2b_key_is_set, list_tool, test_u
@pytest.mark.e2b_sandbox
def test_e2b_sandbox_stateful_tool(check_e2b_key_is_set, clear_core_memory_tool, test_user, agent_state):
sandbox = ToolExecutionSandbox(clear_core_memory_tool.name, {}, user_id=test_user.id)
sandbox = ToolExecutionSandbox(clear_core_memory_tool.name, {}, user=test_user)
# run the sandbox
result = sandbox.run(agent_state=agent_state)
@@ -458,7 +458,7 @@ def test_e2b_sandbox_inject_env_var_existing_sandbox(check_e2b_key_is_set, get_e
config = manager.create_or_update_sandbox_config(config_create, test_user)
# Run the custom sandbox once, assert nothing returns because missing env variable
sandbox = ToolExecutionSandbox(get_env_tool.name, {}, user_id=test_user.id, force_recreate=True)
sandbox = ToolExecutionSandbox(get_env_tool.name, {}, user=test_user, force_recreate=True)
result = sandbox.run()
# response should be None
assert result.func_return is None
@@ -471,7 +471,7 @@ def test_e2b_sandbox_inject_env_var_existing_sandbox(check_e2b_key_is_set, get_e
)
# Assert that the environment variable gets injected correctly, even when the sandbox is NOT refreshed
sandbox = ToolExecutionSandbox(get_env_tool.name, {}, user_id=test_user.id)
sandbox = ToolExecutionSandbox(get_env_tool.name, {}, user=test_user)
result = sandbox.run()
assert long_random_string in result.func_return
@@ -487,7 +487,7 @@ def test_e2b_sandbox_config_change_force_recreates_sandbox(check_e2b_key_is_set,
config = manager.create_or_update_sandbox_config(config_create, test_user)
# Run the custom sandbox once, assert a failure gets returned because missing environment variable
sandbox = ToolExecutionSandbox(list_tool.name, {}, user_id=test_user.id)
sandbox = ToolExecutionSandbox(list_tool.name, {}, user=test_user)
result = sandbox.run()
assert len(result.func_return) == 5
old_config_fingerprint = result.sandbox_config_fingerprint
@@ -497,7 +497,7 @@ def test_e2b_sandbox_config_change_force_recreates_sandbox(check_e2b_key_is_set,
config = manager.update_sandbox_config(config.id, config_update, test_user)
# Run again
result = ToolExecutionSandbox(list_tool.name, {}, user_id=test_user.id).run()
result = ToolExecutionSandbox(list_tool.name, {}, user=test_user).run()
new_config_fingerprint = result.sandbox_config_fingerprint
assert config.fingerprint() == new_config_fingerprint
@@ -507,7 +507,7 @@ def test_e2b_sandbox_config_change_force_recreates_sandbox(check_e2b_key_is_set,
@pytest.mark.e2b_sandbox
def test_e2b_sandbox_with_list_rv(check_e2b_key_is_set, list_tool, test_user):
sandbox = ToolExecutionSandbox(list_tool.name, {}, user_id=test_user.id)
sandbox = ToolExecutionSandbox(list_tool.name, {}, user=test_user)
result = sandbox.run()
assert len(result.func_return) == 5
@@ -524,7 +524,7 @@ def test_e2b_e2e_composio_star_github(check_e2b_key_is_set, check_composio_key_s
actor=test_user,
)
result = ToolExecutionSandbox(composio_github_star_tool.name, {"owner": "letta-ai", "repo": "letta"}, user_id=test_user.id).run()
result = ToolExecutionSandbox(composio_github_star_tool.name, {"owner": "letta-ai", "repo": "letta"}, user=test_user).run()
assert result.func_return["details"] == "Action executed successfully"
@@ -541,7 +541,7 @@ class TestCoreMemoryTools:
"""Test successful replacement of content in core memory - local sandbox."""
new_name = "Charles"
args = {"label": "human", "old_content": "Chad", "new_content": new_name}
sandbox = ToolExecutionSandbox(core_memory_tools["core_memory_replace"].name, args, user_id=test_user.id)
sandbox = ToolExecutionSandbox(core_memory_tools["core_memory_replace"].name, args, user=test_user)
result = sandbox.run(agent_state=agent_state)
assert new_name in result.agent_state.memory.get_block("human").value
@@ -552,7 +552,7 @@ class TestCoreMemoryTools:
"""Test successful appending of content to core memory - local sandbox."""
append_text = "\nLikes coffee"
args = {"label": "human", "content": append_text}
sandbox = ToolExecutionSandbox(core_memory_tools["core_memory_append"].name, args, user_id=test_user.id)
sandbox = ToolExecutionSandbox(core_memory_tools["core_memory_append"].name, args, user=test_user)
result = sandbox.run(agent_state=agent_state)
assert append_text in result.agent_state.memory.get_block("human").value
@@ -563,7 +563,7 @@ class TestCoreMemoryTools:
"""Test error handling when trying to replace non-existent content - local sandbox."""
nonexistent_name = "Alexander Wang"
args = {"label": "human", "old_content": nonexistent_name, "new_content": "Charles"}
sandbox = ToolExecutionSandbox(core_memory_tools["core_memory_replace"].name, args, user_id=test_user.id)
sandbox = ToolExecutionSandbox(core_memory_tools["core_memory_replace"].name, args, user=test_user)
result = sandbox.run(agent_state=agent_state)
assert len(result.stderr) != 0
@@ -575,7 +575,7 @@ class TestCoreMemoryTools:
"""Test successful replacement of content in core memory - e2b sandbox."""
new_name = "Charles"
args = {"label": "human", "old_content": "Chad", "new_content": new_name}
sandbox = ToolExecutionSandbox(core_memory_tools["core_memory_replace"].name, args, user_id=test_user.id)
sandbox = ToolExecutionSandbox(core_memory_tools["core_memory_replace"].name, args, user=test_user)
result = sandbox.run(agent_state=agent_state)
assert new_name in result.agent_state.memory.get_block("human").value
@@ -586,7 +586,7 @@ class TestCoreMemoryTools:
"""Test successful appending of content to core memory - e2b sandbox."""
append_text = "\nLikes coffee"
args = {"label": "human", "content": append_text}
sandbox = ToolExecutionSandbox(core_memory_tools["core_memory_append"].name, args, user_id=test_user.id)
sandbox = ToolExecutionSandbox(core_memory_tools["core_memory_append"].name, args, user=test_user)
result = sandbox.run(agent_state=agent_state)
assert append_text in result.agent_state.memory.get_block("human").value
@@ -597,7 +597,7 @@ class TestCoreMemoryTools:
"""Test error handling when trying to replace non-existent content - e2b sandbox."""
nonexistent_name = "Alexander Wang"
args = {"label": "human", "old_content": nonexistent_name, "new_content": "Charles"}
sandbox = ToolExecutionSandbox(core_memory_tools["core_memory_replace"].name, args, user_id=test_user.id)
sandbox = ToolExecutionSandbox(core_memory_tools["core_memory_replace"].name, args, user=test_user)
result = sandbox.run(agent_state=agent_state)
assert len(result.stderr) != 0

29
tests/test_server.py
View File

@@ -288,17 +288,18 @@ def org_id(server):
@pytest.fixture(scope="module")
def user_id(server, org_id):
# create user
def user(server, org_id):
user = server.user_manager.create_default_user()
print(f"Created user\n{user.id}")
yield user.id
# cleanup
yield user
server.user_manager.delete_user_by_id(user.id)
@pytest.fixture(scope="module")
def user_id(server, user):
# create user
yield user.id
@pytest.fixture(scope="module")
def base_tools(server, user_id):
actor = server.user_manager.get_user_or_default(user_id)
@@ -789,11 +790,11 @@ def ingest(message: str):
'''
def test_tool_run(server, mock_e2b_api_key_none, user_id, agent_id):
def test_tool_run(server, mock_e2b_api_key_none, user, agent_id):
"""Test that the server can run tools"""
result = server.run_tool_from_source(
user_id=user_id,
actor=user,
tool_source=EXAMPLE_TOOL_SOURCE,
tool_source_type="python",
tool_args=json.dumps({"message": "Hello, world!"}),
@@ -806,7 +807,7 @@ def test_tool_run(server, mock_e2b_api_key_none, user_id, agent_id):
assert not result.stderr
result = server.run_tool_from_source(
user_id=user_id,
actor=user,
tool_source=EXAMPLE_TOOL_SOURCE,
tool_source_type="python",
tool_args=json.dumps({"message": "Well well well"}),
@@ -819,7 +820,7 @@ def test_tool_run(server, mock_e2b_api_key_none, user_id, agent_id):
assert not result.stderr
result = server.run_tool_from_source(
user_id=user_id,
actor=user,
tool_source=EXAMPLE_TOOL_SOURCE,
tool_source_type="python",
tool_args=json.dumps({"bad_arg": "oh no"}),
@@ -835,7 +836,7 @@ def test_tool_run(server, mock_e2b_api_key_none, user_id, agent_id):
# Test that we can still pull the tool out by default (pulls that last tool in the source)
result = server.run_tool_from_source(
user_id=user_id,
actor=user,
tool_source=EXAMPLE_TOOL_SOURCE_WITH_DISTRACTOR,
tool_source_type="python",
tool_args=json.dumps({"message": "Well well well"}),
@@ -850,7 +851,7 @@ def test_tool_run(server, mock_e2b_api_key_none, user_id, agent_id):
# Test that we can pull the tool out by name
result = server.run_tool_from_source(
user_id=user_id,
actor=user,
tool_source=EXAMPLE_TOOL_SOURCE_WITH_DISTRACTOR,
tool_source_type="python",
tool_args=json.dumps({"message": "Well well well"}),
@@ -865,7 +866,7 @@ def test_tool_run(server, mock_e2b_api_key_none, user_id, agent_id):
# Test that we can pull a different tool out by name
result = server.run_tool_from_source(
user_id=user_id,
actor=user,
tool_source=EXAMPLE_TOOL_SOURCE_WITH_DISTRACTOR,
tool_source_type="python",
tool_args=json.dumps({}),