diff --git a/src/bedrock_agentcore/memory/client.py b/src/bedrock_agentcore/memory/client.py index 8c8649ac..55932a16 100644 --- a/src/bedrock_agentcore/memory/client.py +++ b/src/bedrock_agentcore/memory/client.py @@ -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]: @@ -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} } @@ -1961,13 +1975,16 @@ 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 @@ -1975,11 +1992,14 @@ def _wrap_configuration( 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 diff --git a/tests/bedrock_agentcore/memory/test_client.py b/tests/bedrock_agentcore/memory/test_client.py index f6d0dc5f..84d46c80 100644 --- a/tests/bedrock_agentcore/memory/test_client.py +++ b/tests/bedrock_agentcore/memory/test_client.py @@ -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"):