From 0e6e6bf48f49da7a45e151580a598fa359dcdd40 Mon Sep 17 00:00:00 2001 From: Nicolas Borges Date: Fri, 1 May 2026 11:32:37 -0400 Subject: [PATCH] chore: update DP APIs to support namespace re-design --- src/bedrock_agentcore/_utils/namespace.py | 28 ++++++++ src/bedrock_agentcore/memory/client.py | 72 +++++++++++-------- .../integrations/strands/session_manager.py | 2 +- src/bedrock_agentcore/memory/session.py | 4 +- tests/bedrock_agentcore/memory/test_client.py | 64 +++++++++++++++++ .../bedrock_agentcore/memory/test_session.py | 2 +- tests/unit/test_utils.py | 27 +++++++ tests_integ/memory/test_memory_client.py | 4 +- 8 files changed, 167 insertions(+), 36 deletions(-) create mode 100644 src/bedrock_agentcore/_utils/namespace.py diff --git a/src/bedrock_agentcore/_utils/namespace.py b/src/bedrock_agentcore/_utils/namespace.py new file mode 100644 index 00000000..b32ce4d8 --- /dev/null +++ b/src/bedrock_agentcore/_utils/namespace.py @@ -0,0 +1,28 @@ +"""Namespace utilities for data plane API calls.""" + +from typing import Dict, Optional + + +def build_namespace_params(namespace: Optional[str] = None, namespace_path: Optional[str] = None) -> Dict[str, str]: + """Build the namespace kwargs for a data plane API call. + + Exactly one of ``namespace`` (exact match) or ``namespace_path`` + (hierarchical path prefix) must be provided. Wildcards (``*``) are not + supported in either field. + + Raises: + ValueError: if both arguments are provided, neither is provided, or + the provided value contains a wildcard. + """ + if namespace is not None and namespace_path is not None: + raise ValueError("'namespace' and 'namespace_path' are mutually exclusive.") + if namespace is None and namespace_path is None: + raise ValueError("At least one of 'namespace' or 'namespace_path' must be provided.") + + value = namespace if namespace is not None else namespace_path + if "*" in value: + raise ValueError("Wildcards (*) are not supported in namespaces.") + + if namespace is not None: + return {"namespace": namespace} + return {"namespacePath": namespace_path} diff --git a/src/bedrock_agentcore/memory/client.py b/src/bedrock_agentcore/memory/client.py index b74e55c9..55bfc936 100644 --- a/src/bedrock_agentcore/memory/client.py +++ b/src/bedrock_agentcore/memory/client.py @@ -20,6 +20,7 @@ from botocore.config import Config from botocore.exceptions import ClientError +from bedrock_agentcore._utils.namespace import build_namespace_params from bedrock_agentcore._utils.snake_case import accept_snake_case_kwargs from bedrock_agentcore._utils.user_agent import build_user_agent_suffix @@ -126,8 +127,11 @@ def __getattr__(self, name: str): # Access any boto3 method directly client = MemoryClient() - # These calls are forwarded to the appropriate boto3 client - response = client.list_memory_records(memoryId="mem-123", namespace="test/") + # These calls are forwarded to the appropriate boto3 client. + # Use `namespace` for exact match, or `namespace_path` for + # hierarchical path-prefix retrieval. + response = client.list_memory_records(memoryId="mem-123", namespace="/actor/Jane/") + response = client.list_memory_records(memoryId="mem-123", namespace_path="/org/MyOrg/") metadata = client.get_memory_metadata(memoryId="mem-123") """ if name in self._ALLOWED_GMDP_METHODS and hasattr(self.gmdp_client, name): @@ -307,46 +311,48 @@ def create_memory_and_wait( raise TimeoutError("Memory %s did not become ACTIVE within %d seconds" % (memory_id, max_wait)) def retrieve_memories( - self, memory_id: str, namespace: str, query: str, actor_id: Optional[str] = None, top_k: int = 3 + self, + memory_id: str, + namespace: Optional[str] = None, + query: str = None, + actor_id: Optional[str] = None, + top_k: int = 3, + namespace_path: Optional[str] = None, ) -> List[Dict[str, Any]]: - """Retrieve relevant memories from a namespace. + """Retrieve relevant memories using exact match or hierarchical path prefix. - Note: Wildcards (*) are NOT supported in namespaces. You must provide the - exact namespace path with all variables resolved. + Exactly one of ``namespace`` or ``namespace_path`` must be provided. Args: memory_id: Memory resource ID - namespace: Exact namespace path (no wildcards) - query: Search query + namespace: Exact namespace to match (e.g., "/actor/Jane/") + query: Search query (required) actor_id: Optional actor ID (deprecated, use namespace) top_k: Number of results to return + namespace_path: Hierarchical path prefix (e.g., "/org/team/") Returns: - List of memory records - - Example: - # Correct - exact namespace - memories = client.retrieve_memories( - memory_id="mem-123", - namespace="support/facts/session-456/", - query="customer preferences" - ) - - # Incorrect - wildcards not supported - # memories = client.retrieve_memories(..., namespace="support/facts/*/", ...) + List of memory records. Returns an empty list if the namespace + arguments are invalid (both provided, neither provided, or contain + wildcards) or if the service call fails. """ - if "*" in namespace: - logger.error("Wildcards are not supported in namespaces. Please provide exact namespace.") + if query is None: + raise TypeError("retrieve_memories() missing required argument: 'query'") + + try: + ns_params = build_namespace_params(namespace, namespace_path) + except ValueError as e: + logger.error(str(e)) return [] + ns_value = namespace or namespace_path + try: - # Let service handle all namespace validation response = self.gmdp_client.retrieve_memory_records( - memoryId=memory_id, namespace=namespace, searchCriteria={"searchQuery": query, "topK": top_k} + memoryId=memory_id, searchCriteria={"searchQuery": query, "topK": top_k}, **ns_params ) - memories = response.get("memoryRecordSummaries", []) - logger.info("Retrieved %d memories from namespace: %s", len(memories), namespace) + logger.info("Retrieved %d memories from namespace: %s", len(memories), ns_value) return memories except ClientError as e: @@ -357,7 +363,7 @@ def retrieve_memories( logger.warning( "Memory or namespace not found. Ensure memory %s exists and namespace '%s' is configured", memory_id, - namespace, + ns_value, ) elif error_code == "ValidationException": logger.warning("Invalid search parameters: %s", error_msg) @@ -662,6 +668,7 @@ def process_turn_with_llm( retrieval_query: Optional[str] = None, top_k: int = 3, event_timestamp: Optional[datetime] = None, + retrieval_namespace_path: Optional[str] = None, ) -> Tuple[List[Dict[str, Any]], str, Dict[str, Any]]: r"""Complete conversation turn with LLM callback integration. @@ -676,10 +683,11 @@ def process_turn_with_llm( llm_callback: Function that takes (user_input, memories) and returns agent_response The callback receives the user input and retrieved memories, and should return the agent's response string - retrieval_namespace: Namespace to search for memories (optional) + retrieval_namespace: Namespace for exact match retrieval (optional) retrieval_query: Custom search query (defaults to user_input) top_k: Number of memories to retrieve event_timestamp: Optional timestamp for the event + retrieval_namespace_path: Namespace path for hierarchical prefix retrieval (optional) Returns: Tuple of (retrieved_memories, agent_response, created_event) @@ -709,10 +717,14 @@ def my_llm(user_input: str, memories: List[Dict]) -> str: """ # Step 1: Retrieve relevant memories retrieved_memories = [] - if retrieval_namespace: + if retrieval_namespace or retrieval_namespace_path: search_query = retrieval_query or user_input retrieved_memories = self.retrieve_memories( - memory_id=memory_id, namespace=retrieval_namespace, query=search_query, top_k=top_k + memory_id=memory_id, + namespace=retrieval_namespace, + namespace_path=retrieval_namespace_path, + query=search_query, + top_k=top_k, ) logger.info("Retrieved %d memories for LLM context", len(retrieved_memories)) diff --git a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py index 622fa736..9203c519 100644 --- a/src/bedrock_agentcore/memory/integrations/strands/session_manager.py +++ b/src/bedrock_agentcore/memory/integrations/strands/session_manager.py @@ -848,7 +848,7 @@ def retrieve_for_namespace(namespace: str, retrieval_config: RetrievalConfig): memories = self.memory_client.retrieve_memories( memory_id=self.config.memory_id, - namespace=resolved_namespace, + namespace_path=resolved_namespace, query=user_query, top_k=retrieval_config.top_k, ) diff --git a/src/bedrock_agentcore/memory/session.py b/src/bedrock_agentcore/memory/session.py index 279f29c5..d8c31eb8 100644 --- a/src/bedrock_agentcore/memory/session.py +++ b/src/bedrock_agentcore/memory/session.py @@ -913,7 +913,7 @@ def search_long_term_memories( params = { "memoryId": self._memory_id, "searchCriteria": search_criteria, - "namespace": namespace, + "namespacePath": namespace, "maxResults": max_results, } @@ -940,7 +940,7 @@ def list_long_term_memory_records( params = { "memoryId": self._memory_id, - "namespace": namespace_prefix, + "namespacePath": namespace_prefix, } if strategy_id: diff --git a/tests/bedrock_agentcore/memory/test_client.py b/tests/bedrock_agentcore/memory/test_client.py index d5333fd9..73822491 100644 --- a/tests/bedrock_agentcore/memory/test_client.py +++ b/tests/bedrock_agentcore/memory/test_client.py @@ -1,11 +1,13 @@ """Unit tests for Memory Client - no external connections.""" +import logging import time import uuid import warnings from datetime import datetime from unittest.mock import MagicMock, patch +import pytest from botocore.exceptions import ClientError from bedrock_agentcore.memory import MemoryClient @@ -1627,6 +1629,68 @@ def test_retrieve_memories_wildcard_namespace(): assert not mock_gmdp.retrieve_memory_records.called +def test_retrieve_memories_with_namespace_path(): + """Test retrieve_memories uses namespacePath for hierarchical retrieval.""" + with patch("boto3.Session"): + client = MemoryClient() + + mock_gmdp = MagicMock() + mock_gmdp.retrieve_memory_records.return_value = {"memoryRecordSummaries": [{"memoryRecordId": "rec-1"}]} + client.gmdp_client = mock_gmdp + + result = client.retrieve_memories(memory_id="mem-123", namespace_path="/org/team/", query="test", top_k=3) + + assert len(result) == 1 + call_kwargs = mock_gmdp.retrieve_memory_records.call_args[1] + assert "namespacePath" in call_kwargs + assert call_kwargs["namespacePath"] == "/org/team/" + assert "namespace" not in call_kwargs + + +def test_retrieve_memories_mutual_exclusivity(caplog): + """Test retrieve_memories soft-fails when both namespace and namespace_path are passed.""" + with patch("boto3.Session"): + client = MemoryClient() + + mock_gmdp = MagicMock() + client.gmdp_client = mock_gmdp + + with caplog.at_level(logging.ERROR): + result = client.retrieve_memories(memory_id="mem-123", namespace="/a/", namespace_path="/b/", query="test") + + # Should log the error and return [] without calling the service + assert result == [] + assert not mock_gmdp.retrieve_memory_records.called + assert any( + "mutually exclusive" in record.message and record.levelno == logging.ERROR for record in caplog.records + ) + + +def test_retrieve_memories_missing_namespace_and_path(caplog): + """Test retrieve_memories soft-fails when neither namespace nor namespace_path is passed.""" + with patch("boto3.Session"): + client = MemoryClient() + + mock_gmdp = MagicMock() + client.gmdp_client = mock_gmdp + + with caplog.at_level(logging.ERROR): + result = client.retrieve_memories(memory_id="mem-123", query="test") + + assert result == [] + assert not mock_gmdp.retrieve_memory_records.called + assert any("At least one" in record.message and record.levelno == logging.ERROR for record in caplog.records) + + +def test_retrieve_memories_missing_query_raises(): + """Test retrieve_memories raises TypeError when query is omitted.""" + with patch("boto3.Session"): + client = MemoryClient() + + with pytest.raises(TypeError, match="query"): + client.retrieve_memories(memory_id="mem-123", namespace="/actor/Jane/") + + def test_add_semantic_strategy_and_wait(): """Test add_semantic_strategy_and_wait functionality.""" with patch("boto3.Session"): diff --git a/tests/bedrock_agentcore/memory/test_session.py b/tests/bedrock_agentcore/memory/test_session.py index 56a12da1..b286e46b 100644 --- a/tests/bedrock_agentcore/memory/test_session.py +++ b/tests/bedrock_agentcore/memory/test_session.py @@ -1197,7 +1197,7 @@ def test_search_long_term_memories_success(self): assert call_args["memoryId"] == "testMemory-1234567890" assert call_args["searchCriteria"]["searchQuery"] == "test query" assert call_args["searchCriteria"]["topK"] == 5 - assert call_args["namespace"] == "test/namespace/" + assert call_args["namespacePath"] == "test/namespace/" def test_search_long_term_memories_with_strategy(self): """Test search_long_term_memories with strategy_id.""" diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 2b40f8b1..50e6ec54 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -5,6 +5,7 @@ import pytest from botocore.exceptions import ClientError +from bedrock_agentcore._utils.namespace import build_namespace_params from bedrock_agentcore._utils.polling import wait_until, wait_until_deleted @@ -112,3 +113,29 @@ def test_timeout(self, _mock_time, _mock_sleep): poll_fn = Mock(return_value={"status": "DELETING"}) with pytest.raises(TimeoutError): wait_until_deleted(poll_fn) + + +class TestBuildNamespaceParams: + """Tests for build_namespace_params utility.""" + + def test_namespace_only(self): + assert build_namespace_params(namespace="/actor/Jane/") == {"namespace": "/actor/Jane/"} + + def test_namespace_path_only(self): + assert build_namespace_params(namespace_path="/org/team/") == {"namespacePath": "/org/team/"} + + def test_both_raises(self): + with pytest.raises(ValueError, match="mutually exclusive"): + build_namespace_params(namespace="/a/", namespace_path="/b/") + + def test_neither_raises(self): + with pytest.raises(ValueError, match="At least one"): + build_namespace_params() + + def test_wildcard_in_namespace_raises(self): + with pytest.raises(ValueError, match="[Ww]ildcard"): + build_namespace_params(namespace="/actor/*/") + + def test_wildcard_in_namespace_path_raises(self): + with pytest.raises(ValueError, match="[Ww]ildcard"): + build_namespace_params(namespace_path="/org/*/team/") diff --git a/tests_integ/memory/test_memory_client.py b/tests_integ/memory/test_memory_client.py index 278696c0..33875b25 100644 --- a/tests_integ/memory/test_memory_client.py +++ b/tests_integ/memory/test_memory_client.py @@ -371,11 +371,11 @@ def test_retrieve_summary_memories(self): assert len(results) > 0, "Expected summary memories in /summaries/ namespace" @pytest.mark.order(14) - # retrieve_memories with a prefix namespace matches broader results + # retrieve_memories with a path prefix matches all descendants under /facts/ def test_retrieve_memories_prefix_namespace(self): results = self.client.retrieve_memories( memory_id=self.prepopulated_memory_id, - namespace="/facts/", + namespace_path="/facts/", query="developer preferences", ) assert len(results) > 0, "Expected prefix namespace to match"