Skip to content

Commit cfbc478

Browse files
committed
Fix test to allow search_memory to include message type in enum
Update test_all_tool_schemas_exclude_message_type to only check creation/editing tools (add_memory_to_working_memory, create_long_term_memory, edit_long_term_memory) for message type exclusion. Search tools like search_memory are allowed to include message type for filtering existing memories.
1 parent ea5f617 commit cfbc478

File tree

1 file changed

+20
-7
lines changed

1 file changed

+20
-7
lines changed

tests/test_client_tool_calls.py

Lines changed: 20 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -563,10 +563,21 @@ def test_add_memory_tool_schema_excludes_message_type(self):
563563
assert "message" not in memory_type_prop["enum"]
564564

565565
def test_all_tool_schemas_exclude_message_type(self):
566-
"""Test that all tool schemas with memory_type exclude 'message'."""
566+
"""Test that creation/editing tool schemas exclude 'message' type.
567+
568+
Note: search_memory CAN include 'message' in its filter enum since it's for
569+
searching/reading existing memories, not creating new ones.
570+
"""
567571
# Get all schemas
568572
all_schemas = MemoryAPIClient.get_all_memory_tool_schemas()
569573

574+
# Tools that should NOT expose message type (creation/editing tools)
575+
restricted_tools = {
576+
"add_memory_to_working_memory",
577+
"create_long_term_memory",
578+
"edit_long_term_memory",
579+
}
580+
570581
# Check each schema that has memory_type parameter
571582
for schema in all_schemas:
572583
function_name = schema["function"]["name"]
@@ -575,18 +586,20 @@ def test_all_tool_schemas_exclude_message_type(self):
575586
# Check if this schema has memory_type in properties
576587
if "memory_type" in params["properties"]:
577588
memory_type_prop = params["properties"]["memory_type"]
578-
assert "message" not in memory_type_prop.get(
579-
"enum", []
580-
), f"Tool {function_name} should not expose 'message' memory type"
589+
if function_name in restricted_tools:
590+
assert (
591+
"message" not in memory_type_prop.get("enum", [])
592+
), f"Creation/editing tool {function_name} should not expose 'message' memory type"
581593

582594
# Check nested properties (like in create_long_term_memory)
583595
if "memories" in params["properties"]:
584596
items = params["properties"]["memories"].get("items", {})
585597
if "properties" in items and "memory_type" in items["properties"]:
586598
memory_type_prop = items["properties"]["memory_type"]
587-
assert (
588-
"message" not in memory_type_prop.get("enum", [])
589-
), f"Tool {function_name} should not expose 'message' memory type in nested properties"
599+
if function_name in restricted_tools:
600+
assert (
601+
"message" not in memory_type_prop.get("enum", [])
602+
), f"Creation/editing tool {function_name} should not expose 'message' memory type in nested properties"
590603

591604

592605
class TestToolCallErrorHandling:

0 commit comments

Comments
 (0)