fix: Add additional testing for anthropic token counting (#2619)

This commit is contained in:
Matthew Zhou
2025-06-03 20:56:39 -07:00
committed by GitHub
parent 3d8704395b
commit ebccd8176a
4 changed files with 65 additions and 19 deletions

View File

@@ -706,11 +706,24 @@ async def test_create_get_list_agent(server: SyncServer, comprehensive_test_agen
assert len(list_agents) == 0
@pytest.fixture(params=["", "PRODUCTION"])
def set_letta_environment(request):
original = os.environ.get("LETTA_ENVIRONMENT")
os.environ["LETTA_ENVIRONMENT"] = request.param
yield request.param
# Restore original environment variable
if original is not None:
os.environ["LETTA_ENVIRONMENT"] = original
else:
os.environ.pop("LETTA_ENVIRONMENT", None)
@pytest.mark.asyncio
async def test_get_context_window_basic(server: SyncServer, comprehensive_test_agent_fixture, default_user, default_file, event_loop):
async def test_get_context_window_basic(
server: SyncServer, comprehensive_test_agent_fixture, default_user, default_file, event_loop, set_letta_environment
):
# Test agent creation
created_agent, create_agent_request = comprehensive_test_agent_fixture
comprehensive_agent_checks(created_agent, create_agent_request, actor=default_user)
# Attach a file
assoc = await server.file_agent_manager.attach_file(
@@ -723,7 +736,7 @@ async def test_get_context_window_basic(server: SyncServer, comprehensive_test_a
# Get context window and check for basic appearances
context_window_overview = await server.agent_manager.get_context_window(agent_id=created_agent.id, actor=default_user)
validate_context_window_overview(context_window_overview, assoc)
validate_context_window_overview(created_agent, context_window_overview, assoc)
# Test deleting the agent
server.agent_manager.delete_agent(created_agent.id, default_user)
@@ -731,6 +744,24 @@ async def test_get_context_window_basic(server: SyncServer, comprehensive_test_a
assert len(list_agents) == 0
@pytest.mark.asyncio
async def test_get_context_window_composio_tool(
server: SyncServer, comprehensive_test_agent_fixture, default_user, default_file, event_loop, set_letta_environment
):
# Test agent creation
created_agent, create_agent_request = comprehensive_test_agent_fixture
# Attach a composio tool
tool_create = ToolCreate.from_composio(action_name="GITHUB_GET_EMOJIS")
tool = server.tool_manager.create_or_update_composio_tool(tool_create=tool_create, actor=default_user)
created_agent = server.agent_manager.attach_tool(agent_id=created_agent.id, tool_id=tool.id, actor=default_user)
# Get context window and check for basic appearances
context_window_overview = await server.agent_manager.get_context_window(agent_id=created_agent.id, actor=default_user)
validate_context_window_overview(created_agent, context_window_overview)
@pytest.mark.asyncio
async def test_create_agent_passed_in_initial_messages(server: SyncServer, default_user, default_block, event_loop):
memory_blocks = [CreateBlock(label="human", value="BananaBoy"), CreateBlock(label="persona", value="I am a helpful assistant")]