|
34 | 34 | FILES_TOOLS, |
35 | 35 | LETTA_TOOL_EXECUTION_DIR, |
36 | 36 | LETTA_TOOL_SET, |
| 37 | + LOCAL_ONLY_MULTI_AGENT_TOOLS, |
37 | 38 | MCP_TOOL_TAG_NAME_PREFIX, |
38 | 39 | MULTI_AGENT_TOOLS, |
39 | 40 | ) |
|
86 | 87 | from letta.server.db import db_registry |
87 | 88 | from letta.server.server import SyncServer |
88 | 89 | from letta.services.block_manager import BlockManager |
89 | | -from letta.services.helpers.agent_manager_helper import calculate_base_tools |
| 90 | +from letta.services.helpers.agent_manager_helper import calculate_base_tools, calculate_multi_agent_tools |
90 | 91 | from letta.services.step_manager import FeedbackType |
91 | 92 | from letta.settings import tool_settings |
92 | 93 | from tests.helpers.utils import comprehensive_agent_checks, validate_context_window_overview |
@@ -747,6 +748,71 @@ async def test_create_agent_include_base_tools(server: SyncServer, default_user, |
747 | 748 | assert sorted(tool_names) == sorted(expected_tools) |
748 | 749 |
|
749 | 750 |
|
| 751 | +def test_calculate_multi_agent_tools(set_letta_environment): |
| 752 | + """Test that calculate_multi_agent_tools excludes local-only tools in production.""" |
| 753 | + result = calculate_multi_agent_tools() |
| 754 | + |
| 755 | + if set_letta_environment == "PRODUCTION": |
| 756 | + # Production environment should exclude local-only tools |
| 757 | + expected_tools = set(MULTI_AGENT_TOOLS) - set(LOCAL_ONLY_MULTI_AGENT_TOOLS) |
| 758 | + assert result == expected_tools, "Production should exclude local-only multi-agent tools" |
| 759 | + assert not set(LOCAL_ONLY_MULTI_AGENT_TOOLS).intersection(result), "Production should not include local-only tools" |
| 760 | + |
| 761 | + # Verify specific tools |
| 762 | + assert "send_message_to_agent_and_wait_for_reply" in result, "Standard multi-agent tools should be in production" |
| 763 | + assert "send_message_to_agents_matching_tags" in result, "Standard multi-agent tools should be in production" |
| 764 | + assert "send_message_to_agent_async" not in result, "Local-only tools should not be in production" |
| 765 | + else: |
| 766 | + # Non-production environment should include all multi-agent tools |
| 767 | + assert result == set(MULTI_AGENT_TOOLS), "Non-production should include all multi-agent tools" |
| 768 | + assert set(LOCAL_ONLY_MULTI_AGENT_TOOLS).issubset(result), "Non-production should include local-only tools" |
| 769 | + |
| 770 | + # Verify specific tools |
| 771 | + assert "send_message_to_agent_and_wait_for_reply" in result, "All multi-agent tools should be in non-production" |
| 772 | + assert "send_message_to_agents_matching_tags" in result, "All multi-agent tools should be in non-production" |
| 773 | + assert "send_message_to_agent_async" in result, "Local-only tools should be in non-production" |
| 774 | + |
| 775 | + |
| 776 | +async def test_upsert_base_tools_excludes_local_only_in_production(server: SyncServer, default_user, set_letta_environment, event_loop): |
| 777 | + """Test that upsert_base_tools excludes local-only multi-agent tools in production.""" |
| 778 | + # Upsert all base tools |
| 779 | + tools = await server.tool_manager.upsert_base_tools_async(actor=default_user) |
| 780 | + tool_names = {tool.name for tool in tools} |
| 781 | + |
| 782 | + if set_letta_environment == "PRODUCTION": |
| 783 | + # Production environment should exclude local-only multi-agent tools |
| 784 | + for local_only_tool in LOCAL_ONLY_MULTI_AGENT_TOOLS: |
| 785 | + assert local_only_tool not in tool_names, f"Local-only tool '{local_only_tool}' should not be upserted in production" |
| 786 | + |
| 787 | + # But should include standard multi-agent tools |
| 788 | + standard_multi_agent_tools = set(MULTI_AGENT_TOOLS) - set(LOCAL_ONLY_MULTI_AGENT_TOOLS) |
| 789 | + for standard_tool in standard_multi_agent_tools: |
| 790 | + assert standard_tool in tool_names, f"Standard multi-agent tool '{standard_tool}' should be upserted in production" |
| 791 | + else: |
| 792 | + # Non-production environment should include all multi-agent tools |
| 793 | + for tool in MULTI_AGENT_TOOLS: |
| 794 | + assert tool in tool_names, f"Multi-agent tool '{tool}' should be upserted in non-production" |
| 795 | + |
| 796 | + |
| 797 | +async def test_upsert_multi_agent_tools_only(server: SyncServer, default_user, set_letta_environment, event_loop): |
| 798 | + """Test that upserting only multi-agent tools respects production filtering.""" |
| 799 | + from letta.orm.enums import ToolType |
| 800 | + |
| 801 | + # Upsert only multi-agent tools |
| 802 | + tools = await server.tool_manager.upsert_base_tools_async(actor=default_user, allowed_types={ToolType.LETTA_MULTI_AGENT_CORE}) |
| 803 | + tool_names = {tool.name for tool in tools} |
| 804 | + |
| 805 | + if set_letta_environment == "PRODUCTION": |
| 806 | + # Should only have non-local multi-agent tools |
| 807 | + expected_tools = set(MULTI_AGENT_TOOLS) - set(LOCAL_ONLY_MULTI_AGENT_TOOLS) |
| 808 | + assert tool_names == expected_tools, "Production multi-agent upsert should exclude local-only tools" |
| 809 | + assert "send_message_to_agent_async" not in tool_names, "Local-only async tool should not be upserted in production" |
| 810 | + else: |
| 811 | + # Should have all multi-agent tools |
| 812 | + assert tool_names == set(MULTI_AGENT_TOOLS), "Non-production multi-agent upsert should include all tools" |
| 813 | + assert "send_message_to_agent_async" in tool_names, "Local-only async tool should be upserted in non-production" |
| 814 | + |
| 815 | + |
750 | 816 | @pytest.mark.asyncio |
751 | 817 | async def test_create_agent_with_default_source(server: SyncServer, default_user, print_tool, default_block, event_loop): |
752 | 818 | """Test agent creation with include_default_source=True""" |
|
0 commit comments