Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 28 additions & 0 deletions src/bedrock_agentcore/_utils/namespace.py
Original file line number Diff line number Diff line change
@@ -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}
72 changes: 42 additions & 30 deletions src/bedrock_agentcore/memory/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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:
Expand All @@ -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)
Expand Down Expand Up @@ -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.

Expand All @@ -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)
Expand Down Expand Up @@ -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))

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down
4 changes: 2 additions & 2 deletions src/bedrock_agentcore/memory/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -913,7 +913,7 @@ def search_long_term_memories(
params = {
"memoryId": self._memory_id,
"searchCriteria": search_criteria,
"namespace": namespace,
"namespacePath": namespace,
"maxResults": max_results,
}

Expand All @@ -940,7 +940,7 @@ def list_long_term_memory_records(

params = {
"memoryId": self._memory_id,
"namespace": namespace_prefix,
"namespacePath": namespace_prefix,
}

if strategy_id:
Expand Down
64 changes: 64 additions & 0 deletions tests/bedrock_agentcore/memory/test_client.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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"):
Expand Down
2 changes: 1 addition & 1 deletion tests/bedrock_agentcore/memory/test_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
27 changes: 27 additions & 0 deletions tests/unit/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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/")
4 changes: 2 additions & 2 deletions tests_integ/memory/test_memory_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
Loading