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
32 changes: 26 additions & 6 deletions src/bedrock_agentcore/memory/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -1920,6 +1920,20 @@ def _validate_strategy_config(self, strategy: Dict[str, Any], strategy_type: str
for namespace in namespaces:
self._validate_namespace(namespace)

def _try_get_override_type(self, override_type: Optional[str]) -> Optional[OverrideType]:
"""Safely convert override_type string to OverrideType enum.

Returns None if override_type is None or not a valid OverrideType value
(e.g., 'SELF_MANAGED' which is a valid configuration type but not in the enum).
"""
if override_type is None:
return None
try:
return OverrideType(override_type)
except ValueError:
# Unknown override type (e.g., SELF_MANAGED), return None
return None

def _wrap_configuration(
self, config: Dict[str, Any], strategy_type: str, override_type: Optional[str] = None
) -> Dict[str, Any]:
Expand All @@ -1932,8 +1946,8 @@ def _wrap_configuration(
builtin_config_keys = ["triggerEveryNMessages", "historicalContextWindowSize"]

if strategy_type == "CUSTOM" and override_type:
override_enum = OverrideType(override_type)
if override_enum in CUSTOM_EXTRACTION_WRAPPER_KEYS:
override_enum = self._try_get_override_type(override_type)
if override_enum and override_enum in CUSTOM_EXTRACTION_WRAPPER_KEYS:
wrapped_config["extraction"] = {
"customExtractionConfiguration": {CUSTOM_EXTRACTION_WRAPPER_KEYS[override_enum]: extraction}
}
Expand Down Expand Up @@ -1961,25 +1975,31 @@ def _wrap_configuration(
}
}
elif strategy_type == "CUSTOM" and override_type:
override_enum = OverrideType(override_type)
if override_enum in CUSTOM_CONSOLIDATION_WRAPPER_KEYS:
override_enum = self._try_get_override_type(override_type)
if override_enum and override_enum in CUSTOM_CONSOLIDATION_WRAPPER_KEYS:
wrapped_config["consolidation"] = {
"customConsolidationConfiguration": {
CUSTOM_CONSOLIDATION_WRAPPER_KEYS[override_enum]: consolidation
}
}
else:
# Unknown override type (e.g., SELF_MANAGED), pass through as-is
wrapped_config["consolidation"] = consolidation
else:
wrapped_config["consolidation"] = consolidation

if "reflection" in config:
reflection = config["reflection"]

if strategy_type == "CUSTOM" and override_type:
override_enum = OverrideType(override_type)
if override_enum in CUSTOM_REFLECTION_WRAPPER_KEYS:
override_enum = self._try_get_override_type(override_type)
if override_enum and override_enum in CUSTOM_REFLECTION_WRAPPER_KEYS:
wrapped_config["reflection"] = {
"customReflectionConfiguration": {CUSTOM_REFLECTION_WRAPPER_KEYS[override_enum]: reflection}
}
else:
# Unknown override type (e.g., SELF_MANAGED), pass through as-is
wrapped_config["reflection"] = reflection
else:
wrapped_config["reflection"] = reflection

Expand Down
68 changes: 68 additions & 0 deletions tests/bedrock_agentcore/memory/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -3283,6 +3283,74 @@ def test_wrap_configuration_custom_episodic_override():
)


def test_wrap_configuration_custom_self_managed():
"""Test _wrap_configuration with CUSTOM strategy and SELF_MANAGED type.

SELF_MANAGED is a valid configuration type but not in the OverrideType enum.
The method should handle this gracefully by passing configuration through as-is.
See: https://github.com/aws/bedrock-agentcore-sdk-python/issues/212
"""
with patch("boto3.client"):
client = MemoryClient()

config = {
"extraction": {
"historicalContextWindowSize": 10,
"triggerEveryNMessages": 5,
},
"consolidation": {
"appendToPrompt": "Consolidate data",
"modelId": "consolidation-model",
},
"reflection": {
"appendToPrompt": "Reflect on data",
"modelId": "reflection-model",
},
}

# Should NOT raise ValueError for SELF_MANAGED
wrapped = client._wrap_configuration(config, "CUSTOM", "SELF_MANAGED")

# Configuration should be passed through as-is (no special wrapping)
assert "extraction" in wrapped
assert wrapped["extraction"]["historicalContextWindowSize"] == 10
assert wrapped["extraction"]["triggerEveryNMessages"] == 5

# Consolidation should also be passed through (no wrapping for unknown override types)
assert "consolidation" in wrapped
assert wrapped["consolidation"]["appendToPrompt"] == "Consolidate data"
assert wrapped["consolidation"]["modelId"] == "consolidation-model"

# Reflection should also be passed through
assert "reflection" in wrapped
assert wrapped["reflection"]["appendToPrompt"] == "Reflect on data"
assert wrapped["reflection"]["modelId"] == "reflection-model"


def test_try_get_override_type_valid():
"""Test _try_get_override_type returns enum for valid override types."""
with patch("boto3.client"):
from bedrock_agentcore.memory.constants import OverrideType

client = MemoryClient()

assert client._try_get_override_type("SEMANTIC_OVERRIDE") == OverrideType.SEMANTIC_OVERRIDE
assert client._try_get_override_type("EPISODIC_OVERRIDE") == OverrideType.EPISODIC_OVERRIDE
assert client._try_get_override_type("SUMMARY_OVERRIDE") == OverrideType.SUMMARY_OVERRIDE
assert client._try_get_override_type("USER_PREFERENCE_OVERRIDE") == OverrideType.USER_PREFERENCE_OVERRIDE


def test_try_get_override_type_invalid():
"""Test _try_get_override_type returns None for invalid override types."""
with patch("boto3.client"):
client = MemoryClient()

assert client._try_get_override_type(None) is None
assert client._try_get_override_type("SELF_MANAGED") is None
assert client._try_get_override_type("UNKNOWN_TYPE") is None
assert client._try_get_override_type("") is None


def test_get_last_k_turns_auto_pagination():
"""Test get_last_k_turns automatically paginates until k turns are found."""
with patch("boto3.client"):
Expand Down