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

feat: Rewrite agents (#2232)

Browse Source
This commit is contained in:
Matthew Zhou
2024-12-13 14:43:19 -08:00
committed by GitHub
parent 65fd731917
commit 7908b8a15f
86 changed files with 2495 additions and 3980 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/conftest.py
View File

@@ -2,8 +2,6 @@ import logging
import pytest
from letta.settings import tool_settings
def pytest_configure(config):
logging.basicConfig(level=logging.DEBUG)
@@ -11,6 +9,8 @@ def pytest_configure(config):
@pytest.fixture
def mock_e2b_api_key_none():
from letta.settings import tool_settings
# Store the original value of e2b_api_key
original_api_key = tool_settings.e2b_api_key

14
tests/helpers/endpoints_helper.py
View File

@@ -61,7 +61,7 @@ def setup_agent(
filename: str,
memory_human_str: str = get_human_text(DEFAULT_HUMAN),
memory_persona_str: str = get_persona_text(DEFAULT_PERSONA),
tools: Optional[List[str]] = None,
tool_ids: Optional[List[str]] = None,
tool_rules: Optional[List[BaseToolRule]] = None,
agent_uuid: str = agent_uuid,
) -> AgentState:
@@ -77,7 +77,7 @@ def setup_agent(
memory = ChatMemory(human=memory_human_str, persona=memory_persona_str)
agent_state = client.create_agent(
name=agent_uuid, llm_config=llm_config, embedding_config=embedding_config, memory=memory, tools=tools, tool_rules=tool_rules
name=agent_uuid, llm_config=llm_config, embedding_config=embedding_config, memory=memory, tool_ids=tool_ids, tool_rules=tool_rules
)
return agent_state
@@ -103,7 +103,6 @@ def check_first_response_is_valid_for_llm_endpoint(filename: str) -> ChatComplet
cleanup(client=client, agent_uuid=agent_uuid)
agent_state = setup_agent(client, filename)
tools = [client.get_tool(client.get_tool_id(name=name)) for name in agent_state.tool_names]
full_agent_state = client.get_agent(agent_state.id)
agent = Agent(agent_state=full_agent_state, interface=None, user=client.user)
@@ -171,19 +170,18 @@ def check_agent_uses_external_tool(filename: str) -> LettaResponse:
client = create_client()
cleanup(client=client, agent_uuid=agent_uuid)
tool = client.load_composio_tool(action=Action.GITHUB_STAR_A_REPOSITORY_FOR_THE_AUTHENTICATED_USER)
tool_name = tool.name
# Set up persona for tool usage
persona = f"""
My name is Letta.
I am a personal assistant who answers a user's questions about a website `example.com`. When a user asks me a question about `example.com`, I will use a tool called {tool_name} which will search `example.com` and answer the relevant question.
I am a personal assistant who answers a user's questions about a website `example.com`. When a user asks me a question about `example.com`, I will use a tool called {tool.name} which will search `example.com` and answer the relevant question.
Don’t forget - inner monologue / inner thoughts should always be different than the contents of send_message! send_message is how you communicate with the user, whereas inner thoughts are your own personal inner thoughts.
"""
agent_state = setup_agent(client, filename, memory_persona_str=persona, tools=[tool_name])
agent_state = setup_agent(client, filename, memory_persona_str=persona, tool_ids=[tool.id])
response = client.user_message(agent_id=agent_state.id, message="What's on the example.com website?")
@@ -191,7 +189,7 @@ def check_agent_uses_external_tool(filename: str) -> LettaResponse:
assert_sanity_checks(response)
# Make sure the tool was called
assert_invoked_function_call(response.messages, tool_name)
assert_invoked_function_call(response.messages, tool.name)
# Make sure some inner monologue is present
assert_inner_monologue_is_present_and_valid(response.messages)
@@ -334,7 +332,7 @@ def check_agent_summarize_memory_simple(filename: str) -> LettaResponse:
client.user_message(agent_id=agent_state.id, message="Does the number 42 ring a bell?")
# Summarize
agent = client.server.load_agent(agent_id=agent_state.id)
agent = client.server.load_agent(agent_id=agent_state.id, actor=client.user)
agent.summarize_messages_inplace()
print(f"Summarization succeeded: messages[1] = \n\n{json_dumps(agent.messages[1])}\n")

55
tests/helpers/utils.py
View File

@@ -3,6 +3,7 @@ from typing import Union
from letta import LocalClient, RESTClient
from letta.functions.functions import parse_source_code
from letta.functions.schema_generator import generate_schema
from letta.schemas.agent import AgentState, CreateAgent, UpdateAgent
from letta.schemas.tool import Tool
@@ -24,3 +25,57 @@ def create_tool_from_func(func: callable):
source_code=parse_source_code(func),
json_schema=generate_schema(func, None),
)
def comprehensive_agent_checks(agent: AgentState, request: Union[CreateAgent, UpdateAgent]):
# Assert scalar fields
assert agent.system == request.system, f"System prompt mismatch: {agent.system} != {request.system}"
assert agent.description == request.description, f"Description mismatch: {agent.description} != {request.description}"
assert agent.metadata_ == request.metadata_, f"Metadata mismatch: {agent.metadata_} != {request.metadata_}"
# Assert agent type
if hasattr(request, "agent_type"):
assert agent.agent_type == request.agent_type, f"Agent type mismatch: {agent.agent_type} != {request.agent_type}"
# Assert LLM configuration
assert agent.llm_config == request.llm_config, f"LLM config mismatch: {agent.llm_config} != {request.llm_config}"
# Assert embedding configuration
assert (
agent.embedding_config == request.embedding_config
), f"Embedding config mismatch: {agent.embedding_config} != {request.embedding_config}"
# Assert memory blocks
if hasattr(request, "memory_blocks"):
assert len(agent.memory.blocks) == len(request.memory_blocks) + len(
request.block_ids
), f"Memory blocks count mismatch: {len(agent.memory.blocks)} != {len(request.memory_blocks) + len(request.block_ids)}"
memory_block_values = {block.value for block in agent.memory.blocks}
expected_block_values = {block.value for block in request.memory_blocks}
assert expected_block_values.issubset(
memory_block_values
), f"Memory blocks mismatch: {expected_block_values} not in {memory_block_values}"
# Assert tools
assert len(agent.tools) == len(request.tool_ids), f"Tools count mismatch: {len(agent.tools)} != {len(request.tool_ids)}"
assert {tool.id for tool in agent.tools} == set(
request.tool_ids
), f"Tools mismatch: {set(tool.id for tool in agent.tools)} != {set(request.tool_ids)}"
# Assert sources
assert len(agent.sources) == len(request.source_ids), f"Sources count mismatch: {len(agent.sources)} != {len(request.source_ids)}"
assert {source.id for source in agent.sources} == set(
request.source_ids
), f"Sources mismatch: {set(source.id for source in agent.sources)} != {set(request.source_ids)}"
# Assert tags
assert set(agent.tags) == set(request.tags), f"Tags mismatch: {set(agent.tags)} != {set(request.tags)}"
# Assert tool rules
if request.tool_rules:
assert len(agent.tool_rules) == len(
request.tool_rules
), f"Tool rules count mismatch: {len(agent.tool_rules)} != {len(request.tool_rules)}"
assert all(
any(rule.tool_name == req_rule.tool_name for rule in agent.tool_rules) for req_rule in request.tool_rules
), f"Tool rules mismatch: {agent.tool_rules} != {request.tool_rules}"

2
tests/test_agent_tool_graph.py → tests/integration_test_agent_tool_graph.py
View File

@@ -99,7 +99,7 @@ def test_single_path_agent_tool_call_graph(mock_e2b_api_key_none):
]
# Make agent state
agent_state = setup_agent(client, config_file, agent_uuid=agent_uuid, tools=[t.name for t in tools], tool_rules=tool_rules)
agent_state = setup_agent(client, config_file, agent_uuid=agent_uuid, tool_ids=[t.id for t in tools], tool_rules=tool_rules)
response = client.user_message(agent_id=agent_state.id, message="What is the fourth secret word?")
# Make checks

2
tests/test_o1_agent.py → tests/integration_test_o1_agent.py
View File

@@ -17,7 +17,7 @@ def test_o1_agent():
agent_state = client.create_agent(
agent_type=AgentType.o1_agent,
tools=[thinking_tool.name, final_tool.name],
tool_ids=[thinking_tool.id, final_tool.id],
llm_config=LLMConfig.default_config("gpt-4"),
embedding_config=EmbeddingConfig.default_config("text-embedding-ada-002"),
memory=ChatMemory(human=get_human_text(DEFAULT_HUMAN), persona=get_persona_text("o1_persona")),

17
tests/test_offline_memory_agent.py → tests/integration_test_offline_memory_agent.py
View File

@@ -32,8 +32,10 @@ def clear_agents(client):
for agent in client.list_agents():
client.delete_agent(agent.id)
def test_ripple_edit(client, mock_e2b_api_key_none):
trigger_rethink_memory_tool = client.create_or_update_tool(trigger_rethink_memory)
send_message = client.server.tool_manager.get_tool_by_name(tool_name="send_message", actor=client.user)
conversation_human_block = Block(name="human", label="human", value=get_human_text(DEFAULT_HUMAN), limit=2000)
conversation_persona_block = Block(name="persona", label="persona", value=get_persona_text(DEFAULT_PERSONA), limit=2000)
@@ -64,7 +66,7 @@ def test_ripple_edit(client, mock_e2b_api_key_none):
system=gpt_system.get_system_text("memgpt_convo_only"),
llm_config=LLMConfig.default_config("gpt-4"),
embedding_config=EmbeddingConfig.default_config("text-embedding-ada-002"),
tools=["send_message", trigger_rethink_memory_tool.name],
tool_ids=[send_message.id, trigger_rethink_memory_tool.id],
memory=conversation_memory,
include_base_tools=False,
)
@@ -81,7 +83,7 @@ def test_ripple_edit(client, mock_e2b_api_key_none):
memory=offline_memory,
llm_config=LLMConfig.default_config("gpt-4"),
embedding_config=EmbeddingConfig.default_config("text-embedding-ada-002"),
tools=[rethink_memory_tool.name, finish_rethinking_memory_tool.name],
tool_ids=[rethink_memory_tool.id, finish_rethinking_memory_tool.id],
tool_rules=[TerminalToolRule(tool_name=finish_rethinking_memory_tool.name)],
include_base_tools=False,
)
@@ -111,16 +113,16 @@ def test_chat_only_agent(client, mock_e2b_api_key_none):
)
conversation_memory = BasicBlockMemory(blocks=[conversation_persona_block, conversation_human_block])
client = create_client()
send_message = client.server.tool_manager.get_tool_by_name(tool_name="send_message", actor=client.user)
chat_only_agent = client.create_agent(
name="conversation_agent",
agent_type=AgentType.chat_only_agent,
llm_config=LLMConfig.default_config("gpt-4"),
embedding_config=EmbeddingConfig.default_config("text-embedding-ada-002"),
tools=["send_message"],
tool_ids=[send_message.id],
memory=conversation_memory,
include_base_tools=False,
metadata={"offline_memory_tools": [rethink_memory.name, finish_rethinking_memory.name]},
metadata={"offline_memory_tools": [rethink_memory.id, finish_rethinking_memory.id]},
)
assert chat_only_agent is not None
assert set(chat_only_agent.memory.list_block_labels()) == {"chat_agent_persona", "chat_agent_human"}
@@ -135,6 +137,7 @@ def test_chat_only_agent(client, mock_e2b_api_key_none):
# Clean up agent
client.delete_agent(chat_only_agent.id)
def test_initial_message_sequence(client, mock_e2b_api_key_none):
"""
Test that when we set the initial sequence to an empty list,
@@ -150,8 +153,6 @@ def test_initial_message_sequence(client, mock_e2b_api_key_none):
initial_message_sequence=[],
)
assert offline_memory_agent is not None
assert len(offline_memory_agent.message_ids) == 1 # There should just the system message
assert len(offline_memory_agent.message_ids) == 1 # There should just the system message
client.delete_agent(offline_memory_agent.id)

41
tests/test_autogen_integration.py
View File

@@ -1,41 +0,0 @@
# TODO: add back
# import os
# import subprocess
#
# import pytest
#
#
# @pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="Missing OpenAI API key")
# def test_agent_groupchat():
#
# # Define the path to the script you want to test
# script_path = "letta/autogen/examples/agent_groupchat.py"
#
# # Dynamically get the project's root directory (assuming this script is run from the root)
# # project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
# # print(project_root)
# # project_root = os.path.join(project_root, "Letta")
# # print(project_root)
# # sys.exit(1)
#
# project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
# project_root = os.path.join(project_root, "letta")
# print(f"Adding the following to PATH: {project_root}")
#
# # Prepare the environment, adding the project root to PYTHONPATH
# env = os.environ.copy()
# env["PYTHONPATH"] = f"{project_root}:{env.get('PYTHONPATH', '')}"
#
# # Run the script using subprocess.run
# # Capture the output (stdout) and the exit code
# # result = subprocess.run(["python", script_path], capture_output=True, text=True)
# result = subprocess.run(["poetry", "run", "python", script_path], capture_output=True, text=True)
#
# # Check the exit code (0 indicates success)
# assert result.returncode == 0, f"Script exited with code {result.returncode}: {result.stderr}"
#
# # Optionally, check the output for expected content
# # For example, if you expect a specific line in the output, uncomment and adapt the following line:
# # assert "expected output" in result.stdout, "Expected output not found in script's output"
#

23
tests/test_base_functions.py
View File

@@ -23,7 +23,7 @@ def agent_obj():
agent_state = client.create_agent()
global agent_obj
agent_obj = client.server.load_agent(agent_id=agent_state.id)
agent_obj = client.server.load_agent(agent_id=agent_state.id, actor=client.user)
yield agent_obj
client.delete_agent(agent_obj.agent_state.id)
@@ -35,49 +35,50 @@ def query_in_search_results(search_results, query):
return True
return False
def test_archival(agent_obj):
"""Test archival memory functions comprehensively."""
# Test 1: Basic insertion and retrieval
base_functions.archival_memory_insert(agent_obj, "The cat sleeps on the mat")
base_functions.archival_memory_insert(agent_obj, "The dog plays in the park")
base_functions.archival_memory_insert(agent_obj, "Python is a programming language")
# Test exact text search
results, _ = base_functions.archival_memory_search(agent_obj, "cat")
assert query_in_search_results(results, "cat")
# Test semantic search (should return animal-related content)
results, _ = base_functions.archival_memory_search(agent_obj, "animal pets")
assert query_in_search_results(results, "cat") or query_in_search_results(results, "dog")
# Test unrelated search (should not return animal content)
results, _ = base_functions.archival_memory_search(agent_obj, "programming computers")
assert query_in_search_results(results, "python")
# Test 2: Test pagination
# Insert more items to test pagination
for i in range(10):
base_functions.archival_memory_insert(agent_obj, f"Test passage number {i}")
# Get first page
page0_results, next_page = base_functions.archival_memory_search(agent_obj, "Test passage", page=0)
# Get second page
page1_results, _ = base_functions.archival_memory_search(agent_obj, "Test passage", page=1, start=next_page)
assert page0_results != page1_results
assert query_in_search_results(page0_results, "Test passage")
assert query_in_search_results(page1_results, "Test passage")
# Test 3: Test complex text patterns
base_functions.archival_memory_insert(agent_obj, "Important meeting on 2024-01-15 with John")
base_functions.archival_memory_insert(agent_obj, "Follow-up meeting scheduled for next week")
base_functions.archival_memory_insert(agent_obj, "Project deadline is approaching")
# Search for meeting-related content
results, _ = base_functions.archival_memory_search(agent_obj, "meeting schedule")
assert query_in_search_results(results, "meeting")
assert query_in_search_results(results, "2024-01-15") or query_in_search_results(results, "next week")
# Test 4: Test error handling
# Test invalid page number
try:
@@ -85,7 +86,7 @@ def test_archival(agent_obj):
assert False, "Should have raised ValueError"
except ValueError:
pass
def test_recall(agent_obj):
base_functions.conversation_search(agent_obj, "banana")

112
tests/test_client.py
View File

@@ -1,5 +1,4 @@
import asyncio
import json
import os
import threading
import time
@@ -42,8 +41,8 @@ def run_server():
@pytest.fixture(
# params=[{"server": False}, {"server": True}], # whether to use REST API server
params=[{"server": False}], # whether to use REST API server
params=[{"server": False}, {"server": True}], # whether to use REST API server
# params=[{"server": True}], # whether to use REST API server
scope="module",
)
def client(request):
@@ -69,6 +68,7 @@ def client(request):
@pytest.fixture(scope="module")
def agent(client: Union[LocalClient, RESTClient]):
agent_state = client.create_agent(name=f"test_client_{str(uuid.uuid4())}")
yield agent_state
# delete agent
@@ -86,6 +86,47 @@ def clear_tables():
session.commit()
def test_shared_blocks(mock_e2b_api_key_none, client: Union[LocalClient, RESTClient]):
# _reset_config()
# create a block
block = client.create_block(label="human", value="username: sarah")
# create agents with shared block
from letta.schemas.block import Block
from letta.schemas.memory import BasicBlockMemory
# persona1_block = client.create_block(label="persona", value="you are agent 1")
# persona2_block = client.create_block(label="persona", value="you are agent 2")
# create agents
agent_state1 = client.create_agent(
name="agent1", memory=BasicBlockMemory([Block(label="persona", value="you are agent 1")]), block_ids=[block.id]
)
agent_state2 = client.create_agent(
name="agent2", memory=BasicBlockMemory([Block(label="persona", value="you are agent 2")]), block_ids=[block.id]
)
## attach shared block to both agents
# client.link_agent_memory_block(agent_state1.id, block.id)
# client.link_agent_memory_block(agent_state2.id, block.id)
# update memory
client.user_message(agent_id=agent_state1.id, message="my name is actually charles")
# check agent 2 memory
assert "charles" in client.get_block(block.id).value.lower(), f"Shared block update failed {client.get_block(block.id).value}"
client.user_message(agent_id=agent_state2.id, message="whats my name?")
assert (
"charles" in client.get_core_memory(agent_state2.id).get_block("human").value.lower()
), f"Shared block update failed {client.get_core_memory(agent_state2.id).get_block('human').value}"
# assert "charles" in response.messages[1].text.lower(), f"Shared block update failed {response.messages[0].text}"
# cleanup
client.delete_agent(agent_state1.id)
client.delete_agent(agent_state2.id)
def test_sandbox_config_and_env_var_basic(client: Union[LocalClient, RESTClient]):
"""
Test sandbox config and environment variable functions for both LocalClient and RESTClient.
@@ -137,15 +178,15 @@ def test_sandbox_config_and_env_var_basic(client: Union[LocalClient, RESTClient]
client.delete_sandbox_config(sandbox_config_id=sandbox_config.id)
def test_add_and_manage_tags_for_agent(client: Union[LocalClient, RESTClient], agent: AgentState):
def test_add_and_manage_tags_for_agent(client: Union[LocalClient, RESTClient]):
"""
Comprehensive happy path test for adding, retrieving, and managing tags on an agent.
"""
tags_to_add = ["test_tag_1", "test_tag_2", "test_tag_3"]
# Step 0: create an agent with tags
tagged_agent = client.create_agent(tags=tags_to_add)
assert set(tagged_agent.tags) == set(tags_to_add), f"Expected tags {tags_to_add}, but got {tagged_agent.tags}"
# Step 0: create an agent with no tags
agent = client.create_agent()
assert len(agent.tags) == 0
# Step 1: Add multiple tags to the agent
client.update_agent(agent_id=agent.id, tags=tags_to_add)
@@ -175,6 +216,9 @@ def test_add_and_manage_tags_for_agent(client: Union[LocalClient, RESTClient], a
final_tags = client.get_agent(agent_id=agent.id).tags
assert len(final_tags) == 0, f"Expected no tags, but found {final_tags}"
# Remove agent
client.delete_agent(agent.id)
def test_update_agent_memory_label(client: Union[LocalClient, RESTClient], agent: AgentState):
"""Test that we can update the label of a block in an agent's memory"""
@@ -255,35 +299,33 @@ def test_add_remove_agent_memory_block(client: Union[LocalClient, RESTClient], a
# client.delete_agent(new_agent.id)
def test_update_agent_memory_limit(client: Union[LocalClient, RESTClient], agent: AgentState):
def test_update_agent_memory_limit(client: Union[LocalClient, RESTClient]):
"""Test that we can update the limit of a block in an agent's memory"""
agent = client.create_agent(name=create_random_username())
agent = client.create_agent()
try:
current_labels = agent.memory.list_block_labels()
example_label = current_labels[0]
example_new_limit = 1
current_block = agent.memory.get_block(label=example_label)
current_block_length = len(current_block.value)
current_labels = agent.memory.list_block_labels()
example_label = current_labels[0]
example_new_limit = 1
current_block = agent.memory.get_block(label=example_label)
current_block_length = len(current_block.value)
assert example_new_limit != agent.memory.get_block(label=example_label).limit
assert example_new_limit < current_block_length
assert example_new_limit != agent.memory.get_block(label=example_label).limit
assert example_new_limit < current_block_length
# We expect this to throw a value error
with pytest.raises(ValueError):
client.update_agent_memory_block(agent_id=agent.id, label=example_label, limit=example_new_limit)
# Now try the same thing with a higher limit
example_new_limit = current_block_length + 10000
assert example_new_limit > current_block_length
# We expect this to throw a value error
with pytest.raises(ValueError):
client.update_agent_memory_block(agent_id=agent.id, label=example_label, limit=example_new_limit)
updated_agent = client.get_agent(agent_id=agent.id)
assert example_new_limit == updated_agent.memory.get_block(label=example_label).limit
# Now try the same thing with a higher limit
example_new_limit = current_block_length + 10000
assert example_new_limit > current_block_length
client.update_agent_memory_block(agent_id=agent.id, label=example_label, limit=example_new_limit)
finally:
client.delete_agent(agent.id)
updated_agent = client.get_agent(agent_id=agent.id)
assert example_new_limit == updated_agent.memory.get_block(label=example_label).limit
client.delete_agent(agent.id)
def test_messages(client: Union[LocalClient, RESTClient], agent: AgentState):
@@ -316,7 +358,7 @@ def test_function_return_limit(client: Union[LocalClient, RESTClient]):
padding = len("[NOTE: function output was truncated since it exceeded the character limit (100000 > 1000)]") + 50
tool = client.create_or_update_tool(func=big_return, return_char_limit=1000)
agent = client.create_agent(name="agent1", tools=[tool.name])
agent = client.create_agent(tool_ids=[tool.id])
# get function response
response = client.send_message(agent_id=agent.id, message="call the big_return function", role="user")
print(response.messages)
@@ -330,10 +372,14 @@ def test_function_return_limit(client: Union[LocalClient, RESTClient]):
assert response_message, "FunctionReturn message not found in response"
res = response_message.function_return
assert "function output was truncated " in res
res_json = json.loads(res)
assert (
len(res_json["message"]) <= 1000 + padding
), f"Expected length to be less than or equal to 1000 + {padding}, but got {len(res_json['message'])}"
# TODO: Re-enable later
# res_json = json.loads(res)
# assert (
# len(res_json["message"]) <= 1000 + padding
# ), f"Expected length to be less than or equal to 1000 + {padding}, but got {len(res_json['message'])}"
client.delete_agent(agent_id=agent.id)
@pytest.mark.asyncio

37
tests/test_client_legacy.py
View File

@@ -583,43 +583,6 @@ def test_list_llm_models(client: RESTClient):
assert has_model_endpoint_type(models, "anthropic")
def test_shared_blocks(mock_e2b_api_key_none, client: Union[LocalClient, RESTClient], agent: AgentState):
# _reset_config()
# create a block
block = client.create_block(label="human", value="username: sarah")
# create agents with shared block
from letta.schemas.block import Block
from letta.schemas.memory import BasicBlockMemory
# persona1_block = client.create_block(label="persona", value="you are agent 1")
# persona2_block = client.create_block(label="persona", value="you are agent 2")
# create agnets
agent_state1 = client.create_agent(name="agent1", memory=BasicBlockMemory([Block(label="persona", value="you are agent 1"), block]))
agent_state2 = client.create_agent(name="agent2", memory=BasicBlockMemory([Block(label="persona", value="you are agent 2"), block]))
## attach shared block to both agents
# client.link_agent_memory_block(agent_state1.id, block.id)
# client.link_agent_memory_block(agent_state2.id, block.id)
# update memory
response = client.user_message(agent_id=agent_state1.id, message="my name is actually charles")
# check agent 2 memory
assert "charles" in client.get_block(block.id).value.lower(), f"Shared block update failed {client.get_block(block.id).value}"
response = client.user_message(agent_id=agent_state2.id, message="whats my name?")
assert (
"charles" in client.get_core_memory(agent_state2.id).get_block("human").value.lower()
), f"Shared block update failed {client.get_core_memory(agent_state2.id).get_block('human').value}"
# assert "charles" in response.messages[1].text.lower(), f"Shared block update failed {response.messages[0].text}"
# cleanup
client.delete_agent(agent_state1.id)
client.delete_agent(agent_state2.id)
@pytest.fixture
def cleanup_agents(client):
created_agents = []

142
tests/test_concurrent_connections.py
View File

@@ -1,142 +0,0 @@
# TODO: add back when messaging works
# import os
# import threading
# import time
# import uuid
#
# import pytest
# from dotenv import load_dotenv
#
# from letta import Admin, create_client
# from letta.config import LettaConfig
# from letta.credentials import LettaCredentials
# from letta.settings import settings
# from tests.utils import create_config
#
# test_agent_name = f"test_client_{str(uuid.uuid4())}"
## test_preset_name = "test_preset"
# test_agent_state = None
# client = None
#
# test_agent_state_post_message = None
# test_user_id = uuid.uuid4()
#
#
## admin credentials
# test_server_token = "test_server_token"
#
#
# def _reset_config():
#
# # Use os.getenv with a fallback to os.environ.get
# db_url = settings.letta_pg_uri
#
# if os.getenv("OPENAI_API_KEY"):
# create_config("openai")
# credentials = LettaCredentials(
# openai_key=os.getenv("OPENAI_API_KEY"),
# )
# else: # hosted
# create_config("letta_hosted")
# credentials = LettaCredentials()
#
# config = LettaConfig.load()
#
# # set to use postgres
# config.archival_storage_uri = db_url
# config.recall_storage_uri = db_url
# config.metadata_storage_uri = db_url
# config.archival_storage_type = "postgres"
# config.recall_storage_type = "postgres"
# config.metadata_storage_type = "postgres"
#
# config.save()
# credentials.save()
# print("_reset_config :: ", config.config_path)
#
#
# def run_server():
#
# load_dotenv()
#
# _reset_config()
#
# from letta.server.rest_api.server import start_server
#
# print("Starting server...")
# start_server(debug=True)
#
#
## Fixture to create clients with different configurations
# @pytest.fixture(
# params=[ # whether to use REST API server
# {"server": True},
# # {"server": False} # TODO: add when implemented
# ],
# scope="module",
# )
# def admin_client(request):
# if request.param["server"]:
# # get URL from enviornment
# server_url = os.getenv("MEMGPT_SERVER_URL")
# if server_url is None:
# # run server in thread
# # NOTE: must set MEMGPT_SERVER_PASS enviornment variable
# server_url = "http://localhost:8283"
# print("Starting server thread")
# thread = threading.Thread(target=run_server, daemon=True)
# thread.start()
# time.sleep(5)
# print("Running client tests with server:", server_url)
# # create user via admin client
# admin = Admin(server_url, test_server_token)
# response = admin.create_user(test_user_id) # Adjust as per your client's method
#
# yield admin
#
#
# def test_concurrent_messages(admin_client):
# # test concurrent messages
#
# # create three
#
# results = []
#
# def _send_message():
# try:
# print("START SEND MESSAGE")
# response = admin_client.create_user()
# token = response.api_key
# client = create_client(base_url=admin_client.base_url, token=token)
# agent = client.create_agent()
#
# print("Agent created", agent.id)
#
# st = time.time()
# message = "Hello, how are you?"
# response = client.send_message(agent_id=agent.id, message=message, role="user")
# et = time.time()
# print(f"Message sent from {st} to {et}")
# print(response.messages)
# results.append((st, et))
# except Exception as e:
# print("ERROR", e)
#
# threads = []
# print("Starting threads...")
# for i in range(5):
# thread = threading.Thread(target=_send_message)
# threads.append(thread)
# thread.start()
# print("CREATED THREAD")
#
# print("waiting for threads to finish...")
# for thread in threads:
# print(thread.join())
#
# # make sure runtime are overlapping
# assert (results[0][0] < results[1][0] and results[0][1] > results[1][0]) or (
# results[1][0] < results[0][0] and results[1][1] > results[0][0]
# ), f"Threads should have overlapping runtimes {results}"
#

121
tests/test_different_embedding_size.py
View File

@@ -1,121 +0,0 @@
# TODO: add back once tests are cleaned up
# import os
# import uuid
#
# from letta import create_client
# from letta.agent_store.storage import StorageConnector, TableType
# from letta.schemas.passage import Passage
# from letta.embeddings import embedding_model
# from tests import TEST_MEMGPT_CONFIG
#
# from .utils import create_config, wipe_config
#
# test_agent_name = f"test_client_{str(uuid.uuid4())}"
# test_agent_state = None
# client = None
#
# test_agent_state_post_message = None
# test_user_id = uuid.uuid4()
#
#
# def generate_passages(user, agent):
# # Note: the database will filter out rows that do not correspond to agent1 and test_user by default.
# texts = [
# "This is a test passage",
# "This is another test passage",
# "Cinderella wept",
# ]
# embed_model = embedding_model(agent.embedding_config)
# orig_embeddings = []
# passages = []
# for text in texts:
# embedding = embed_model.get_text_embedding(text)
# orig_embeddings.append(list(embedding))
# passages.append(
# Passage(
# user_id=user.id,
# agent_id=agent.id,
# text=text,
# embedding=embedding,
# embedding_dim=agent.embedding_config.embedding_dim,
# embedding_model=agent.embedding_config.embedding_model,
# )
# )
# return passages, orig_embeddings
#
#
# def test_create_user():
# if not os.getenv("OPENAI_API_KEY"):
# print("Skipping test, missing OPENAI_API_KEY")
# return
#
# wipe_config()
#
# # create client
# create_config("openai")
# client = create_client()
#
# # openai: create agent
# openai_agent = client.create_agent(
# name="openai_agent",
# )
# assert (
# openai_agent.embedding_config.embedding_endpoint_type == "openai"
# ), f"openai_agent.embedding_config.embedding_endpoint_type={openai_agent.embedding_config.embedding_endpoint_type}"
#
# # openai: add passages
# passages, openai_embeddings = generate_passages(client.user, openai_agent)
# openai_agent_run = client.server.load_agent(user_id=client.user.id, agent_id=openai_agent.id)
# openai_agent_run.persistence_manager.archival_memory.storage.insert_many(passages)
#
# # create client
# create_config("letta_hosted")
# client = create_client()
#
# # hosted: create agent
# hosted_agent = client.create_agent(
# name="hosted_agent",
# )
# # check to make sure endpoint overriden
# assert (
# hosted_agent.embedding_config.embedding_endpoint_type == "hugging-face"
# ), f"hosted_agent.embedding_config.embedding_endpoint_type={hosted_agent.embedding_config.embedding_endpoint_type}"
#
# # hosted: add passages
# passages, hosted_embeddings = generate_passages(client.user, hosted_agent)
# hosted_agent_run = client.server.load_agent(user_id=client.user.id, agent_id=hosted_agent.id)
# hosted_agent_run.persistence_manager.archival_memory.storage.insert_many(passages)
#
# # test passage dimentionality
# storage = StorageConnector.get_storage_connector(TableType.PASSAGES, TEST_MEMGPT_CONFIG, client.user.id)
# storage.filters = {} # clear filters to be able to get all passages
# passages = storage.get_all()
# for passage in passages:
# if passage.agent_id == hosted_agent.id:
# assert (
# passage.embedding_dim == hosted_agent.embedding_config.embedding_dim
# ), f"passage.embedding_dim={passage.embedding_dim} != hosted_agent.embedding_config.embedding_dim={hosted_agent.embedding_config.embedding_dim}"
#
# # ensure was in original embeddings
# embedding = passage.embedding[: passage.embedding_dim]
# assert embedding in hosted_embeddings, f"embedding={embedding} not in hosted_embeddings={hosted_embeddings}"
#
# # make sure all zeros
# assert not any(
# passage.embedding[passage.embedding_dim :]
# ), f"passage.embedding[passage.embedding_dim:]={passage.embedding[passage.embedding_dim:]}"
# elif passage.agent_id == openai_agent.id:
# assert (
# passage.embedding_dim == openai_agent.embedding_config.embedding_dim
# ), f"passage.embedding_dim={passage.embedding_dim} != openai_agent.embedding_config.embedding_dim={openai_agent.embedding_config.embedding_dim}"
#
# # ensure was in original embeddings
# embedding = passage.embedding[: passage.embedding_dim]
# assert embedding in openai_embeddings, f"embedding={embedding} not in openai_embeddings={openai_embeddings}"
#
# # make sure all zeros
# assert not any(
# passage.embedding[passage.embedding_dim :]
# ), f"passage.embedding[passage.embedding_dim:]={passage.embedding[passage.embedding_dim:]}"
#

48
tests/test_function_parser.py
View File

@@ -1,48 +0,0 @@
import letta.system as system
from letta.local_llm.function_parser import patch_function
from letta.utils import json_dumps
EXAMPLE_FUNCTION_CALL_SEND_MESSAGE = {
"message_history": [
{"role": "user", "content": system.package_user_message("hello")},
],
# "new_message": {
# "role": "function",
# "name": "send_message",
# "content": system.package_function_response(was_success=True, response_string="None"),
# },
"new_message": {
"role": "assistant",
"content": "I'll send a message.",
"function_call": {
"name": "send_message",
"arguments": "null",
},
},
}
EXAMPLE_FUNCTION_CALL_CORE_MEMORY_APPEND_MISSING = {
"message_history": [
{"role": "user", "content": system.package_user_message("hello")},
],
"new_message": {
"role": "assistant",
"content": "I'll append to memory.",
"function_call": {
"name": "core_memory_append",
"arguments": json_dumps({"content": "new_stuff"}),
},
},
}
def test_function_parsers():
"""Try various broken JSON and check that the parsers can fix it"""
og_message = EXAMPLE_FUNCTION_CALL_SEND_MESSAGE["new_message"]
corrected_message = patch_function(**EXAMPLE_FUNCTION_CALL_SEND_MESSAGE)
assert corrected_message == og_message, f"Uncorrected:\n{og_message}\nCorrected:\n{corrected_message}"
og_message = EXAMPLE_FUNCTION_CALL_CORE_MEMORY_APPEND_MISSING["new_message"].copy()
corrected_message = patch_function(**EXAMPLE_FUNCTION_CALL_CORE_MEMORY_APPEND_MISSING)
assert corrected_message != og_message, f"Uncorrected:\n{og_message}\nCorrected:\n{corrected_message}"

99
tests/test_json_parsers.py
View File

@@ -1,99 +0,0 @@
import letta.local_llm.json_parser as json_parser
from letta.utils import json_loads
EXAMPLE_ESCAPED_UNDERSCORES = """{
"function":"send\_message",
"params": {
"inner\_thoughts": "User is asking for information about themselves. Retrieving data from core memory.",
"message": "I know that you are Chad. Is there something specific you would like to know or talk about regarding yourself?"
"""
EXAMPLE_MISSING_CLOSING_BRACE = """{
"function": "send_message",
"params": {
"inner_thoughts": "Oops, I got their name wrong! I should apologize and correct myself.",
"message": "Sorry about that! I assumed you were Chad. Welcome, Brad! "
}
"""
EXAMPLE_BAD_TOKEN_END = """{
"function": "send_message",
"params": {
"inner_thoughts": "Oops, I got their name wrong! I should apologize and correct myself.",
"message": "Sorry about that! I assumed you were Chad. Welcome, Brad! "
}
}<|>"""
EXAMPLE_DOUBLE_JSON = """{
"function": "core_memory_append",
"params": {
"name": "human",
"content": "Brad, 42 years old, from Germany."
}
}
{
"function": "send_message",
"params": {
"message": "Got it! Your age and nationality are now saved in my memory."
}
}
"""
EXAMPLE_HARD_LINE_FEEDS = """{
"function": "send_message",
"params": {
"message": "Let's create a list:
- First, we can do X
- Then, we can do Y!
- Lastly, we can do Z :)"
}
}
"""
# Situation where beginning of send_message call is fine (and thus can be extracted)
# but has a long training garbage string that comes after
EXAMPLE_SEND_MESSAGE_PREFIX_OK_REST_BAD = """{
"function": "send_message",
"params": {
"inner_thoughts": "User request for debug assistance",
"message": "Of course, Chad. Please check the system log file for 'assistant.json' and send me the JSON output you're getting. Armed with that data, I'll assist you in debugging the issue.",
GARBAGEGARBAGEGARBAGEGARBAGE
GARBAGEGARBAGEGARBAGEGARBAGE
GARBAGEGARBAGEGARBAGEGARBAGE
"""
EXAMPLE_ARCHIVAL_SEARCH = """
{
"function": "archival_memory_search",
"params": {
"inner_thoughts": "Looking for WaitingForAction.",
"query": "WaitingForAction",
"""
def test_json_parsers():
"""Try various broken JSON and check that the parsers can fix it"""
test_strings = [
EXAMPLE_ESCAPED_UNDERSCORES,
EXAMPLE_MISSING_CLOSING_BRACE,
EXAMPLE_BAD_TOKEN_END,
EXAMPLE_DOUBLE_JSON,
EXAMPLE_HARD_LINE_FEEDS,
EXAMPLE_SEND_MESSAGE_PREFIX_OK_REST_BAD,
EXAMPLE_ARCHIVAL_SEARCH,
]
for string in test_strings:
try:
json_loads(string)
assert False, f"Test JSON string should have failed basic JSON parsing:\n{string}"
except:
print("String failed (expectedly)")
try:
json_parser.clean_json(string)
except:
f"Failed to repair test JSON string:\n{string}"
raise

19
tests/test_local_client.py
View File

@@ -4,7 +4,7 @@ import pytest
from letta import create_client
from letta.client.client import LocalClient
from letta.schemas.agent import PersistedAgentState
from letta.schemas.agent import AgentState
from letta.schemas.embedding_config import EmbeddingConfig
from letta.schemas.llm_config import LLMConfig
from letta.schemas.memory import BasicBlockMemory, ChatMemory, Memory
@@ -13,6 +13,7 @@ from letta.schemas.memory import BasicBlockMemory, ChatMemory, Memory
@pytest.fixture(scope="module")
def client():
client = create_client()
# client.set_default_llm_config(LLMConfig.default_config("gpt-4o-mini"))
client.set_default_llm_config(LLMConfig.default_config("gpt-4o-mini"))
client.set_default_embedding_config(EmbeddingConfig.default_config(provider="openai"))
@@ -29,7 +30,6 @@ def agent(client):
yield agent_state
client.delete_agent(agent_state.id)
assert client.get_agent(agent_state.id) is None, f"Failed to properly delete agent {agent_state.id}"
def test_agent(client: LocalClient):
@@ -80,16 +80,15 @@ def test_agent(client: LocalClient):
assert isinstance(agent_state.memory, Memory)
# update agent: tools
tool_to_delete = "send_message"
assert tool_to_delete in agent_state.tool_names
new_agent_tools = [t_name for t_name in agent_state.tool_names if t_name != tool_to_delete]
client.update_agent(agent_state_test.id, tools=new_agent_tools)
assert client.get_agent(agent_state_test.id).tool_names == new_agent_tools
assert tool_to_delete in [t.name for t in agent_state.tools]
new_agent_tool_ids = [t.id for t in agent_state.tools if t.name != tool_to_delete]
client.update_agent(agent_state_test.id, tool_ids=new_agent_tool_ids)
assert sorted([t.id for t in client.get_agent(agent_state_test.id).tools]) == sorted(new_agent_tool_ids)
assert isinstance(agent_state.memory, Memory)
# update agent: memory
new_human = "My name is Mr Test, 100 percent human."
new_persona = "I am an all-knowing AI."
new_memory = ChatMemory(human=new_human, persona=new_persona)
assert agent_state.memory.get_block("human").value != new_human
assert agent_state.memory.get_block("persona").value != new_persona
@@ -216,7 +215,7 @@ def test_agent_with_shared_blocks(client: LocalClient):
client.delete_agent(second_agent_state_test.id)
def test_memory(client: LocalClient, agent: PersistedAgentState):
def test_memory(client: LocalClient, agent: AgentState):
# get agent memory
original_memory = client.get_in_context_memory(agent.id)
assert original_memory is not None
@@ -229,7 +228,7 @@ def test_memory(client: LocalClient, agent: PersistedAgentState):
assert updated_memory.get_block("human").value != original_memory_value # check if the memory has been updated
def test_archival_memory(client: LocalClient, agent: PersistedAgentState):
def test_archival_memory(client: LocalClient, agent: AgentState):
"""Test functions for interacting with archival memory store"""
# add archival memory
@@ -244,7 +243,7 @@ def test_archival_memory(client: LocalClient, agent: PersistedAgentState):
client.delete_archival_memory(agent.id, passage.id)
def test_recall_memory(client: LocalClient, agent: PersistedAgentState):
def test_recall_memory(client: LocalClient, agent: AgentState):
"""Test functions for interacting with recall memory store"""
# send message to the agent

897
tests/test_managers.py
View File

File diff suppressed because it is too large Load Diff

126
tests/test_new_cli.py
View File

@@ -1,126 +0,0 @@
# TODO: fix later
# import os
# import random
# import string
# import unittest.mock
#
# import pytest
#
# from letta.cli.cli_config import add, delete, list
# from letta.config import LettaConfig
# from letta.credentials import LettaCredentials
# from tests.utils import create_config
#
#
# def _reset_config():
#
# if os.getenv("OPENAI_API_KEY"):
# create_config("openai")
# credentials = LettaCredentials(
# openai_key=os.getenv("OPENAI_API_KEY"),
# )
# else: # hosted
# create_config("letta_hosted")
# credentials = LettaCredentials()
#
# config = LettaConfig.load()
# config.save()
# credentials.save()
# print("_reset_config :: ", config.config_path)
#
#
# @pytest.mark.skip(reason="This is a helper function.")
# def generate_random_string(length):
# characters = string.ascii_letters + string.digits
# random_string = "".join(random.choices(characters, k=length))
# return random_string
#
#
# @pytest.mark.skip(reason="Ensures LocalClient is used during testing.")
# def unset_env_variables():
# server_url = os.environ.pop("MEMGPT_BASE_URL", None)
# token = os.environ.pop("MEMGPT_SERVER_PASS", None)
# return server_url, token
#
#
# @pytest.mark.skip(reason="Set env variables back to values before test.")
# def reset_env_variables(server_url, token):
# if server_url is not None:
# os.environ["MEMGPT_BASE_URL"] = server_url
# if token is not None:
# os.environ["MEMGPT_SERVER_PASS"] = token
#
#
# def test_crud_human(capsys):
# _reset_config()
#
# server_url, token = unset_env_variables()
#
# # Initialize values that won't interfere with existing ones
# human_1 = generate_random_string(16)
# text_1 = generate_random_string(32)
# human_2 = generate_random_string(16)
# text_2 = generate_random_string(32)
# text_3 = generate_random_string(32)
#
# # Add inital human
# add("human", human_1, text_1)
#
# # Expect inital human to be listed
# list("humans")
# captured = capsys.readouterr()
# output = captured.out[captured.out.find(human_1) :]
#
# assert human_1 in output
# assert text_1 in output
#
# # Add second human
# add("human", human_2, text_2)
#
# # Expect to see second human
# list("humans")
# captured = capsys.readouterr()
# output = captured.out[captured.out.find(human_1) :]
#
# assert human_1 in output
# assert text_1 in output
# assert human_2 in output
# assert text_2 in output
#
# with unittest.mock.patch("questionary.confirm") as mock_confirm:
# mock_confirm.return_value.ask.return_value = True
#
# # Update second human
# add("human", human_2, text_3)
#
# # Expect to see update text
# list("humans")
# captured = capsys.readouterr()
# output = captured.out[captured.out.find(human_1) :]
#
# assert human_1 in output
# assert text_1 in output
# assert human_2 in output
# assert output.count(human_2) == 1
# assert text_3 in output
# assert text_2 not in output
#
# # Delete second human
# delete("human", human_2)
#
# # Expect second human to be deleted
# list("humans")
# captured = capsys.readouterr()
# output = captured.out[captured.out.find(human_1) :]
#
# assert human_1 in output
# assert text_1 in output
# assert human_2 not in output
# assert text_2 not in output
#
# # Clean up
# delete("human", human_1)
#
# reset_env_variables(server_url, token)
#

93
tests/test_openai_client.py
View File

@@ -1,93 +0,0 @@
from logging import getLogger
from openai import APIConnectionError, OpenAI
logger = getLogger(__name__)
def test_openai_assistant():
client = OpenAI(base_url="http://127.0.0.1:8080/v1")
# create assistant
try:
assistant = client.beta.assistants.create(
name="Math Tutor",
instructions="You are a personal math tutor. Write and run code to answer math questions.",
# tools=[{"type": "code_interpreter"}],
model="gpt-4-turbo-preview",
)
except APIConnectionError as e:
logger.error("Connection issue with localhost openai stub: %s", e)
return
# create thread
thread = client.beta.threads.create()
message = client.beta.threads.messages.create(
thread_id=thread.id, role="user", content="I need to solve the equation `3x + 11 = 14`. Can you help me?"
)
run = client.beta.threads.runs.create(
thread_id=thread.id, assistant_id=assistant.id, instructions="Please address the user as Jane Doe. The user has a premium account."
)
# run = client.beta.threads.runs.create(
# thread_id=thread.id,
# assistant_id=assistant.id,
# model="gpt-4-turbo-preview",
# instructions="New instructions that override the Assistant instructions",
# tools=[{"type": "code_interpreter"}, {"type": "retrieval"}]
# )
# Store the run ID
run_id = run.id
print(run_id)
# NOTE: Letta does not support polling yet, so run status is always "completed"
# Retrieve all messages from the thread
messages = client.beta.threads.messages.list(thread_id=thread.id)
# Print all messages from the thread
for msg in messages.messages:
role = msg["role"]
content = msg["content"][0]
print(f"{role.capitalize()}: {content}")
# TODO: add once polling works
## Polling for the run status
# while True:
# # Retrieve the run status
# run_status = client.beta.threads.runs.retrieve(
# thread_id=thread.id,
# run_id=run_id
# )
# # Check and print the step details
# run_steps = client.beta.threads.runs.steps.list(
# thread_id=thread.id,
# run_id=run_id
# )
# for step in run_steps.data:
# if step.type == 'tool_calls':
# print(f"Tool {step.type} invoked.")
# # If step involves code execution, print the code
# if step.type == 'code_interpreter':
# print(f"Python Code Executed: {step.step_details['code_interpreter']['input']}")
# if run_status.status == 'completed':
# # Retrieve all messages from the thread
# messages = client.beta.threads.messages.list(
# thread_id=thread.id
# )
# # Print all messages from the thread
# for msg in messages.data:
# role = msg.role
# content = msg.content[0].text.value
# print(f"{role.capitalize()}: {content}")
# break # Exit the polling loop since the run is complete
# elif run_status.status in ['queued', 'in_progress']:
# print(f'{run_status.status.capitalize()}... Please wait.')
# time.sleep(1.5) # Wait before checking again
# else:
# print(f"Run status: {run_status.status}")
# break # Exit the polling loop if the status is neither 'in_progress' nor 'completed'

52
tests/test_persistence.py
View File

@@ -1,52 +0,0 @@
# test state saving between client session
# TODO: update this test with correct imports
# def test_save_load(client):
# """Test that state is being persisted correctly after an /exit
#
# Create a new agent, and request a message
#
# Then trigger
# """
# assert client is not None, "Run create_agent test first"
# assert test_agent_state is not None, "Run create_agent test first"
# assert test_agent_state_post_message is not None, "Run test_user_message test first"
#
# # Create a new client (not thread safe), and load the same agent
# # The agent state inside should correspond to the initial state pre-message
# if os.getenv("OPENAI_API_KEY"):
# client2 = Letta(quickstart="openai", user_id=test_user_id)
# else:
# client2 = Letta(quickstart="letta_hosted", user_id=test_user_id)
# print(f"\n\n[3] CREATING CLIENT2, LOADING AGENT {test_agent_state.id}!")
# client2_agent_obj = client2.server.load_agent(user_id=test_user_id, agent_id=test_agent_state.id)
# client2_agent_state = client2_agent_obj.update_state()
# print(f"[3] LOADED AGENT! AGENT {client2_agent_state.id}\n\tmessages={client2_agent_state.state['messages']}")
#
# # assert test_agent_state == client2_agent_state, f"{vars(test_agent_state)}\n{vars(client2_agent_state)}"
# def check_state_equivalence(state_1, state_2):
# """Helper function that checks the equivalence of two AgentState objects"""
# assert state_1.keys() == state_2.keys(), f"{state_1.keys()}\n{state_2.keys}"
# for k, v1 in state_1.items():
# v2 = state_2[k]
# if isinstance(v1, LLMConfig) or isinstance(v1, EmbeddingConfig):
# assert vars(v1) == vars(v2), f"{vars(v1)}\n{vars(v2)}"
# else:
# assert v1 == v2, f"{v1}\n{v2}"
#
# check_state_equivalence(vars(test_agent_state), vars(client2_agent_state))
#
# # Now, write out the save from the original client
# # This should persist the test message into the agent state
# client.save()
#
# if os.getenv("OPENAI_API_KEY"):
# client3 = Letta(quickstart="openai", user_id=test_user_id)
# else:
# client3 = Letta(quickstart="letta_hosted", user_id=test_user_id)
# client3_agent_obj = client3.server.load_agent(user_id=test_user_id, agent_id=test_agent_state.id)
# client3_agent_state = client3_agent_obj.update_state()
#
# check_state_equivalence(vars(test_agent_state_post_message), vars(client3_agent_state))
#

62
tests/test_schema_generator.py
View File

@@ -1,62 +0,0 @@
from letta.functions.schema_generator import generate_schema
def send_message(self, message: str):
"""
Sends a message to the human user.
Args:
message (str): Message contents. All unicode (including emojis) are supported.
Returns:
Optional[str]: None is always returned as this function does not produce a response.
"""
return None
def send_message_missing_types(self, message):
"""
Sends a message to the human user.
Args:
message (str): Message contents. All unicode (including emojis) are supported.
Returns:
Optional[str]: None is always returned as this function does not produce a response.
"""
return None
def send_message_missing_docstring(self, message: str):
return None
def test_schema_generator():
# Check that a basic function schema converts correctly
correct_schema = {
"name": "send_message",
"description": "Sends a message to the human user.",
"parameters": {
"type": "object",
"properties": {"message": {"type": "string", "description": "Message contents. All unicode (including emojis) are supported."}},
"required": ["message"],
},
}
generated_schema = generate_schema(send_message)
print(f"\n\nreference_schema={correct_schema}")
print(f"\n\ngenerated_schema={generated_schema}")
assert correct_schema == generated_schema
# Check that missing types results in an error
try:
_ = generate_schema(send_message_missing_types)
assert False
except:
pass
# Check that missing docstring results in an error
try:
_ = generate_schema(send_message_missing_docstring)
assert False
except:
pass

153
tests/test_server.py
View File

@@ -19,8 +19,6 @@ from letta.schemas.letta_message import (
)
from letta.schemas.user import User
from .test_managers import DEFAULT_EMBEDDING_CONFIG
utils.DEBUG = True
from letta.config import LettaConfig
from letta.schemas.agent import CreateAgent
@@ -266,6 +264,7 @@ Lise, young Bolkónski's wife, this very evening, and perhaps the
thing can be arranged. It shall be on your family's behalf that I'll
start my apprenticeship as old maid."""
@pytest.fixture(scope="module")
def server():
config = LettaConfig.load()
@@ -302,42 +301,66 @@ def user_id(server, org_id):
@pytest.fixture(scope="module")
def agent_id(server, user_id):
def base_tools(server, user_id):
actor = server.user_manager.get_user_or_default(user_id)
tools = []
for tool_name in BASE_TOOLS:
tools.append(server.tool_manager.get_tool_by_name(tool_name=tool_name, actor=actor))
yield tools
@pytest.fixture(scope="module")
def base_memory_tools(server, user_id):
actor = server.user_manager.get_user_or_default(user_id)
tools = []
for tool_name in BASE_MEMORY_TOOLS:
tools.append(server.tool_manager.get_tool_by_name(tool_name=tool_name, actor=actor))
yield tools
@pytest.fixture(scope="module")
def agent_id(server, user_id, base_tools):
# create agent
actor = server.user_manager.get_user_or_default(user_id)
agent_state = server.create_agent(
request=CreateAgent(
name="test_agent",
tools=BASE_TOOLS,
tool_ids=[t.id for t in base_tools],
memory_blocks=[],
llm_config=LLMConfig.default_config("gpt-4"),
embedding_config=EmbeddingConfig.default_config(provider="openai"),
),
actor=server.get_user_or_default(user_id),
actor=actor,
)
print(f"Created agent\n{agent_state}")
yield agent_state.id
# cleanup
server.delete_agent(user_id, agent_state.id)
server.agent_manager.delete_agent(agent_state.id, actor=actor)
@pytest.fixture(scope="module")
def other_agent_id(server, user_id):
def other_agent_id(server, user_id, base_tools):
# create agent
actor = server.user_manager.get_user_or_default(user_id)
agent_state = server.create_agent(
request=CreateAgent(
name="test_agent_other",
tools=BASE_TOOLS,
tool_ids=[t.id for t in base_tools],
memory_blocks=[],
llm_config=LLMConfig.default_config("gpt-4"),
embedding_config=EmbeddingConfig.default_config(provider="openai"),
),
actor=server.get_user_or_default(user_id),
actor=actor,
)
print(f"Created agent\n{agent_state}")
yield agent_state.id
# cleanup
server.delete_agent(user_id, agent_state.id)
server.agent_manager.delete_agent(agent_state.id, actor=actor)
def test_error_on_nonexistent_agent(server, user_id, agent_id):
try:
@@ -416,6 +439,7 @@ def test_user_message(server, user_id, agent_id):
@pytest.mark.order(5)
def test_get_recall_memory(server, org_id, user_id, agent_id):
# test recall memory cursor pagination
actor = server.user_manager.get_user_or_default(user_id=user_id)
messages_1 = server.get_agent_recall_cursor(user_id=user_id, agent_id=agent_id, limit=2)
cursor1 = messages_1[-1].id
messages_2 = server.get_agent_recall_cursor(user_id=user_id, agent_id=agent_id, after=cursor1, limit=1000)
@@ -427,7 +451,9 @@ def test_get_recall_memory(server, org_id, user_id, agent_id):
assert len(messages_4) == 1
# test in-context message ids
in_context_ids = server.get_in_context_message_ids(agent_id=agent_id)
# in_context_ids = server.get_in_context_message_ids(agent_id=agent_id)
in_context_ids = server.agent_manager.get_agent_by_id(agent_id=agent_id, actor=actor).message_ids
message_ids = [m.id for m in messages_3]
for message_id in in_context_ids:
assert message_id in message_ids, f"{message_id} not in {message_ids}"
@@ -437,10 +463,13 @@ def test_get_recall_memory(server, org_id, user_id, agent_id):
def test_get_archival_memory(server, user_id, agent_id):
# test archival memory cursor pagination
user = server.user_manager.get_user_by_id(user_id=user_id)
# List latest 2 passages
passages_1 = server.passage_manager.list_passages(
actor=user, agent_id=agent_id, ascending=False, limit=2,
actor=user,
agent_id=agent_id,
ascending=False,
limit=2,
)
assert len(passages_1) == 2, f"Returned {[p.text for p in passages_1]}, not equal to 2"
@@ -483,12 +512,13 @@ def test_agent_rethink_rewrite_retry(server, user_id, agent_id):
- "rewrite" replaces the text of the last assistant message
- "retry" retries the last assistant message
"""
actor = server.user_manager.get_user_or_default(user_id)
# Send an initial message
server.user_message(user_id=user_id, agent_id=agent_id, message="Hello?")
# Grab the raw Agent object
letta_agent = server.load_agent(agent_id=agent_id)
letta_agent = server.load_agent(agent_id=agent_id, actor=actor)
assert letta_agent._messages[-1].role == MessageRole.tool
assert letta_agent._messages[-2].role == MessageRole.assistant
last_agent_message = letta_agent._messages[-2]
@@ -496,10 +526,10 @@ def test_agent_rethink_rewrite_retry(server, user_id, agent_id):
# Try "rethink"
new_thought = "I am thinking about the meaning of life, the universe, and everything. Bananas?"
assert last_agent_message.text is not None and last_agent_message.text != new_thought
server.rethink_agent_message(agent_id=agent_id, new_thought=new_thought)
server.rethink_agent_message(agent_id=agent_id, new_thought=new_thought, actor=actor)
# Grab the agent object again (make sure it's live)
letta_agent = server.load_agent(agent_id=agent_id)
letta_agent = server.load_agent(agent_id=agent_id, actor=actor)
assert letta_agent._messages[-1].role == MessageRole.tool
assert letta_agent._messages[-2].role == MessageRole.assistant
last_agent_message = letta_agent._messages[-2]
@@ -513,10 +543,10 @@ def test_agent_rethink_rewrite_retry(server, user_id, agent_id):
assert "message" in args_json and args_json["message"] is not None and args_json["message"] != ""
new_text = "Why hello there my good friend! Is 42 what you're looking for? Bananas?"
server.rewrite_agent_message(agent_id=agent_id, new_text=new_text)
server.rewrite_agent_message(agent_id=agent_id, new_text=new_text, actor=actor)
# Grab the agent object again (make sure it's live)
letta_agent = server.load_agent(agent_id=agent_id)
letta_agent = server.load_agent(agent_id=agent_id, actor=actor)
assert letta_agent._messages[-1].role == MessageRole.tool
assert letta_agent._messages[-2].role == MessageRole.assistant
last_agent_message = letta_agent._messages[-2]
@@ -524,10 +554,10 @@ def test_agent_rethink_rewrite_retry(server, user_id, agent_id):
assert "message" in args_json and args_json["message"] is not None and args_json["message"] == new_text
# Try retry
server.retry_agent_message(agent_id=agent_id)
server.retry_agent_message(agent_id=agent_id, actor=actor)
# Grab the agent object again (make sure it's live)
letta_agent = server.load_agent(agent_id=agent_id)
letta_agent = server.load_agent(agent_id=agent_id, actor=actor)
assert letta_agent._messages[-1].role == MessageRole.tool
assert letta_agent._messages[-2].role == MessageRole.assistant
last_agent_message = letta_agent._messages[-2]
@@ -581,33 +611,6 @@ def test_get_context_window_overview(server: SyncServer, user_id: str, agent_id:
)
def test_load_agent_with_nonexistent_tool_names_does_not_error(server: SyncServer, user_id: str):
fake_tool_name = "blahblahblah"
tools = BASE_TOOLS + [fake_tool_name]
agent_state = server.create_agent(
request=CreateAgent(
name="nonexistent_tools_agent",
tools=tools,
memory_blocks=[],
llm_config=LLMConfig.default_config("gpt-4"),
embedding_config=EmbeddingConfig.default_config(provider="openai"),
),
actor=server.get_user_or_default(user_id),
)
# Check that the tools in agent_state do NOT include the fake name
assert fake_tool_name not in agent_state.tool_names
assert set(BASE_TOOLS).issubset(set(agent_state.tool_names))
# Load the agent from the database and check that it doesn't error / tools are correct
saved_tools = server.get_tools_from_agent(agent_id=agent_state.id, user_id=user_id)
assert fake_tool_name not in agent_state.tool_names
assert set(BASE_TOOLS).issubset(set(agent_state.tool_names))
# cleanup
server.delete_agent(user_id, agent_state.id)
def test_delete_agent_same_org(server: SyncServer, org_id: str, user_id: str):
agent_state = server.create_agent(
request=CreateAgent(
@@ -616,14 +619,14 @@ def test_delete_agent_same_org(server: SyncServer, org_id: str, user_id: str):
llm_config=LLMConfig.default_config("gpt-4"),
embedding_config=EmbeddingConfig.default_config(provider="openai"),
),
actor=server.get_user_or_default(user_id),
actor=server.user_manager.get_user_or_default(user_id),
)
# create another user in the same org
another_user = server.user_manager.create_user(User(organization_id=org_id, name="another"))
# test that another user in the same org can delete the agent
server.delete_agent(another_user.id, agent_state.id)
server.agent_manager.delete_agent(agent_state.id, actor=another_user)
def _test_get_messages_letta_format(
@@ -887,14 +890,14 @@ def test_composio_client_simple(server):
assert len(actions) > 0
def test_memory_rebuild_count(server, user_id, mock_e2b_api_key_none):
def test_memory_rebuild_count(server, user_id, mock_e2b_api_key_none, base_tools, base_memory_tools):
"""Test that the memory rebuild is generating the correct number of role=system messages"""
actor = server.user_manager.get_user_or_default(user_id)
# create agent
agent_state = server.create_agent(
request=CreateAgent(
name="memory_rebuild_test_agent",
tools=BASE_TOOLS + BASE_MEMORY_TOOLS,
tool_ids=[t.id for t in base_tools + base_memory_tools],
memory_blocks=[
CreateBlock(label="human", value="The human's name is Bob."),
CreateBlock(label="persona", value="My name is Alice."),
@@ -902,7 +905,7 @@ def test_memory_rebuild_count(server, user_id, mock_e2b_api_key_none):
llm_config=LLMConfig.default_config("gpt-4"),
embedding_config=EmbeddingConfig.default_config(provider="openai"),
),
actor=server.get_user_or_default(user_id),
actor=actor,
)
print(f"Created agent\n{agent_state}")
@@ -929,31 +932,28 @@ def test_memory_rebuild_count(server, user_id, mock_e2b_api_key_none):
try:
# At this stage, there should only be 1 system message inside of recall storage
num_system_messages, all_messages = count_system_messages_in_recall()
# assert num_system_messages == 1, (num_system_messages, all_messages)
assert num_system_messages == 2, (num_system_messages, all_messages)
assert num_system_messages == 1, (num_system_messages, all_messages)
# Assuming core memory append actually ran correctly, at this point there should be 2 messages
server.user_message(user_id=user_id, agent_id=agent_state.id, message="Append 'banana' to your core memory")
# At this stage, there should only be 1 system message inside of recall storage
# At this stage, there should be 2 system message inside of recall storage
num_system_messages, all_messages = count_system_messages_in_recall()
# assert num_system_messages == 2, (num_system_messages, all_messages)
assert num_system_messages == 3, (num_system_messages, all_messages)
assert num_system_messages == 2, (num_system_messages, all_messages)
# Run server.load_agent, and make sure that the number of system messages is still 2
server.load_agent(agent_id=agent_state.id)
server.load_agent(agent_id=agent_state.id, actor=actor)
num_system_messages, all_messages = count_system_messages_in_recall()
# assert num_system_messages == 2, (num_system_messages, all_messages)
assert num_system_messages == 3, (num_system_messages, all_messages)
assert num_system_messages == 2, (num_system_messages, all_messages)
finally:
# cleanup
server.delete_agent(user_id, agent_state.id)
server.agent_manager.delete_agent(agent_state.id, actor=actor)
def test_load_file_to_source(server: SyncServer, user_id: str, agent_id: str, other_agent_id: str, tmp_path):
user = server.get_user_or_default(user_id)
actor = server.user_manager.get_user_or_default(user_id)
# Create a source
source = server.source_manager.create_source(
@@ -962,7 +962,7 @@ def test_load_file_to_source(server: SyncServer, user_id: str, agent_id: str, ot
embedding_config=EmbeddingConfig.default_config(provider="openai"),
created_by_id=user_id,
),
actor=user
actor=actor,
)
# Create a test file with some content
@@ -971,11 +971,10 @@ def test_load_file_to_source(server: SyncServer, user_id: str, agent_id: str, ot
test_file.write_text(test_content)
# Attach source to agent first
agent = server.load_agent(agent_id=agent_id)
agent.attach_source(user=user, source_id=source.id, source_manager=server.source_manager, ms=server.ms)
server.agent_manager.attach_source(agent_id=agent_id, source_id=source.id, actor=actor)
# Get initial passage count
initial_passage_count = server.passage_manager.size(actor=user, agent_id=agent_id, source_id=source.id)
initial_passage_count = server.passage_manager.size(actor=actor, agent_id=agent_id, source_id=source.id)
assert initial_passage_count == 0
# Create a job for loading the first file
@@ -984,7 +983,7 @@ def test_load_file_to_source(server: SyncServer, user_id: str, agent_id: str, ot
user_id=user_id,
metadata_={"type": "embedding", "filename": test_file.name, "source_id": source.id},
),
actor=user
actor=actor,
)
# Load the first file to source
@@ -992,17 +991,17 @@ def test_load_file_to_source(server: SyncServer, user_id: str, agent_id: str, ot
source_id=source.id,
file_path=str(test_file),
job_id=job.id,
actor=user,
actor=actor,
)
# Verify job completed successfully
job = server.job_manager.get_job_by_id(job_id=job.id, actor=user)
job = server.job_manager.get_job_by_id(job_id=job.id, actor=actor)
assert job.status == "completed"
assert job.metadata_["num_passages"] == 1
assert job.metadata_["num_passages"] == 1
assert job.metadata_["num_documents"] == 1
# Verify passages were added
first_file_passage_count = server.passage_manager.size(actor=user, agent_id=agent_id, source_id=source.id)
first_file_passage_count = server.passage_manager.size(actor=actor, agent_id=agent_id, source_id=source.id)
assert first_file_passage_count > initial_passage_count
# Create a second test file with different content
@@ -1015,7 +1014,7 @@ def test_load_file_to_source(server: SyncServer, user_id: str, agent_id: str, ot
user_id=user_id,
metadata_={"type": "embedding", "filename": test_file2.name, "source_id": source.id},
),
actor=user
actor=actor,
)
# Load the second file to source
@@ -1023,22 +1022,22 @@ def test_load_file_to_source(server: SyncServer, user_id: str, agent_id: str, ot
source_id=source.id,
file_path=str(test_file2),
job_id=job2.id,
actor=user,
actor=actor,
)
# Verify second job completed successfully
job2 = server.job_manager.get_job_by_id(job_id=job2.id, actor=user)
job2 = server.job_manager.get_job_by_id(job_id=job2.id, actor=actor)
assert job2.status == "completed"
assert job2.metadata_["num_passages"] >= 10
assert job2.metadata_["num_documents"] == 1
# Verify passages were appended (not replaced)
final_passage_count = server.passage_manager.size(actor=user, agent_id=agent_id, source_id=source.id)
final_passage_count = server.passage_manager.size(actor=actor, agent_id=agent_id, source_id=source.id)
assert final_passage_count > first_file_passage_count
# Verify both old and new content is searchable
passages = server.passage_manager.list_passages(
actor=user,
actor=actor,
agent_id=agent_id,
source_id=source.id,
query_text="what does Timber like to eat",

6
tests/test_summarize.py
View File

@@ -33,7 +33,7 @@ def create_test_agent():
)
global agent_obj
agent_obj = client.server.load_agent(agent_id=agent_state.id)
agent_obj = client.server.load_agent(agent_id=agent_state.id, actor=client.user)
def test_summarize_messages_inplace(mock_e2b_api_key_none):
@@ -74,7 +74,7 @@ def test_summarize_messages_inplace(mock_e2b_api_key_none):
print(f"test_summarize: response={response}")
# reload agent object
agent_obj = client.server.load_agent(agent_id=agent_obj.agent_state.id)
agent_obj = client.server.load_agent(agent_id=agent_obj.agent_state.id, actor=client.user)
agent_obj.summarize_messages_inplace()
print(f"Summarization succeeded: messages[1] = \n{agent_obj.messages[1]}")
@@ -121,7 +121,7 @@ def test_auto_summarize(mock_e2b_api_key_none):
# check if the summarize message is inside the messages
assert isinstance(client, LocalClient), "Test only works with LocalClient"
agent_obj = client.server.load_agent(agent_id=agent_state.id)
agent_obj = client.server.load_agent(agent_id=agent_state.id, actor=client.user)
print("SUMMARY", summarize_message_exists(agent_obj._messages))
if summarize_message_exists(agent_obj._messages):
break

14
tests/test_v1_routes.py
View File

@@ -169,7 +169,7 @@ def configure_mock_sync_server(mock_sync_server):
mock_sync_server.sandbox_config_manager.list_sandbox_env_vars_by_key.return_value = [mock_api_key]
# Mock user retrieval
mock_sync_server.get_user_or_default.return_value = Mock() # Provide additional attributes if needed
mock_sync_server.user_manager.get_user_or_default.return_value = Mock() # Provide additional attributes if needed
# ======================================================================================================================
@@ -182,7 +182,7 @@ def test_delete_tool(client, mock_sync_server, add_integers_tool):
assert response.status_code == 200
mock_sync_server.tool_manager.delete_tool_by_id.assert_called_once_with(
tool_id=add_integers_tool.id, actor=mock_sync_server.get_user_or_default.return_value
tool_id=add_integers_tool.id, actor=mock_sync_server.user_manager.get_user_or_default.return_value
)
@@ -195,7 +195,7 @@ def test_get_tool(client, mock_sync_server, add_integers_tool):
assert response.json()["id"] == add_integers_tool.id
assert response.json()["source_code"] == add_integers_tool.source_code
mock_sync_server.tool_manager.get_tool_by_id.assert_called_once_with(
tool_id=add_integers_tool.id, actor=mock_sync_server.get_user_or_default.return_value
tool_id=add_integers_tool.id, actor=mock_sync_server.user_manager.get_user_or_default.return_value
)
@@ -216,7 +216,7 @@ def test_get_tool_id(client, mock_sync_server, add_integers_tool):
assert response.status_code == 200
assert response.json() == add_integers_tool.id
mock_sync_server.tool_manager.get_tool_by_name.assert_called_once_with(
tool_name=add_integers_tool.name, actor=mock_sync_server.get_user_or_default.return_value
tool_name=add_integers_tool.name, actor=mock_sync_server.user_manager.get_user_or_default.return_value
)
@@ -268,7 +268,7 @@ def test_update_tool(client, mock_sync_server, update_integers_tool, add_integer
assert response.status_code == 200
assert response.json()["id"] == add_integers_tool.id
mock_sync_server.tool_manager.update_tool_by_id.assert_called_once_with(
tool_id=add_integers_tool.id, tool_update=update_integers_tool, actor=mock_sync_server.get_user_or_default.return_value
tool_id=add_integers_tool.id, tool_update=update_integers_tool, actor=mock_sync_server.user_manager.get_user_or_default.return_value
)
@@ -280,7 +280,9 @@ def test_add_base_tools(client, mock_sync_server, add_integers_tool):
assert response.status_code == 200
assert len(response.json()) == 1
assert response.json()[0]["id"] == add_integers_tool.id
mock_sync_server.tool_manager.add_base_tools.assert_called_once_with(actor=mock_sync_server.get_user_or_default.return_value)
mock_sync_server.tool_manager.add_base_tools.assert_called_once_with(
actor=mock_sync_server.user_manager.get_user_or_default.return_value
)
def test_list_composio_apps(client, mock_sync_server, composio_apps):

21
tests/test_vector_embeddings.py
View File

@@ -1,42 +1,39 @@
import numpy as np
import sqlite3
import base64
from numpy.testing import assert_array_almost_equal
import pytest
from letta.orm.sqlalchemy_base import adapt_array
from letta.orm.sqlite_functions import convert_array, verify_embedding_dimension
from letta.orm.sqlalchemy_base import adapt_array, convert_array
from letta.orm.sqlite_functions import verify_embedding_dimension
def test_vector_conversions():
"""Test the vector conversion functions"""
# Create test data
original = np.random.random(4096).astype(np.float32)
print(f"Original shape: {original.shape}")
# Test full conversion cycle
encoded = adapt_array(original)
print(f"Encoded type: {type(encoded)}")
print(f"Encoded length: {len(encoded)}")
decoded = convert_array(encoded)
print(f"Decoded shape: {decoded.shape}")
print(f"Dimension verification: {verify_embedding_dimension(decoded)}")
# Verify data integrity
np.testing.assert_array_almost_equal(original, decoded)
print("✓ Data integrity verified")
# Test with a list
list_data = original.tolist()
encoded_list = adapt_array(list_data)
decoded_list = convert_array(encoded_list)
np.testing.assert_array_almost_equal(original, decoded_list)
print("✓ List conversion verified")
# Test None handling
assert adapt_array(None) is None
assert convert_array(None) is None
print("✓ None handling verified")
# Run the tests
# Run the tests