bedrock-agentcore-starter-toolkit 0.1.14__py3-none-any.whl → 0.1.15__py3-none-any.whl

bedrock_agentcore_starter_toolkit/operations/memory/constants.py

@@ -1,7 +1,7 @@
 """Constants for Bedrock AgentCore Memory SDK."""
 
 from enum import Enum
-from typing import Dict, List, Optional
+from typing import Optional
 
 
 class StrategyType(Enum):
@@ -88,11 +88,3 @@ class MemoryStrategyStatus(Enum):
     ACTIVE = "ACTIVE"
     DELETING = "DELETING"
     FAILED = "FAILED"
-
-
-# Default namespaces for each strategy type
-DEFAULT_NAMESPACES: Dict[StrategyType, List[str]] = {
-    StrategyType.SEMANTIC: ["/actor/{actorId}/strategy/{strategyId}/{sessionId}"],
-    StrategyType.SUMMARY: ["/actor/{actorId}/strategy/{strategyId}/{sessionId}"],
-    StrategyType.USER_PREFERENCE: ["/actor/{actorId}/strategy/{strategyId}"],
-}

bedrock_agentcore_starter_toolkit/operations/memory/manager.py

@@ -4,16 +4,19 @@ import copy
 import logging
 import time
 import uuid
-import warnings
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 import boto3
+from botocore.config import Config as BotocoreConfig
 from botocore.exceptions import ClientError
 
-from .constants import DEFAULT_NAMESPACES, MemoryStatus, MemoryStrategyStatus, OverrideType, StrategyType
+from .constants import MemoryStatus, MemoryStrategyStatus, OverrideType, StrategyType
+from .models import convert_strategies_to_dicts
 from .models.Memory import Memory
 from .models.MemoryStrategy import MemoryStrategy
 from .models.MemorySummary import MemorySummary
+from .models.strategies import BaseStrategy
+from .strategy_validator import validate_existing_memory_strategies
 
 logger = logging.getLogger(__name__)
 
@@ -24,14 +27,52 @@ class MemoryManager:
     This class handles all CONTROL PLANE CRUD operations.
     """
 
-    def __init__(self, region_name: str):
+    def __init__(
+        self,
+        region_name: Optional[str] = None,
+        boto3_session: Optional[boto3.Session] = None,
+        boto_client_config: Optional[BotocoreConfig] = None,
+    ):
         """Initialize MemoryManager with AWS region.
 
         Args:
-            region_name: AWS region name for the bedrock-agentcore-control client.
+            region_name: AWS region for the bedrock-agentcore-control client. If not provided,
+                   will use the region from boto3_session or default session.
+            boto3_session: Optional boto3 Session to use. If provided and region_name
+                          parameter is also specified, validation will ensure they match.
+            boto_client_config: Optional boto3 client configuration. If provided, will be
+                              merged with default configuration including user agent.
+
+        Raises:
+            ValueError: If region_name parameter conflicts with boto3_session region.
         """
+        session = boto3_session or boto3.Session()
+        session_region = session.region_name
+
+        # Validate region consistency if both are provided
+        if region_name and boto3_session and session_region and region_name != session_region:
+            raise ValueError(
+                f"Region mismatch: provided region_name '{region_name}' does not match "
+                f"boto3_session region '{session_region}'. Please ensure both "
+                f"parameters specify the same region or omit the region_name parameter "
+                f"to use the session's region."
+            )
+
+        # Configure boto3 client with merged configuration
+        if boto_client_config:
+            existing_user_agent = getattr(boto_client_config, "user_agent_extra", None)
+            if existing_user_agent:
+                new_user_agent = f"{existing_user_agent} bedrock-agentcore-starter-toolkit"
+            else:
+                new_user_agent = "bedrock-agentcore-starter-toolkit"
+            client_config = boto_client_config.merge(BotocoreConfig(user_agent_extra=new_user_agent))
+        else:
+            client_config = BotocoreConfig(user_agent_extra="bedrock-agentcore-starter-toolkit")
+
         self.region_name = region_name
-        self._control_plane_client = boto3.client("bedrock-agentcore-control", region_name=region_name)
+        self._control_plane_client = session.client(
+            "bedrock-agentcore-control", region_name=self.region_name, config=client_config
+        )
 
         # AgentCore Memory control plane methods
         self._ALLOWED_CONTROL_PLANE_METHODS = {
@@ -153,6 +194,7 @@ class MemoryManager:
         description: Optional[str] = None,
         event_expiry_days: int = 90,
         memory_execution_role_arn: Optional[str] = None,
+        encryption_key_arn: Optional[str] = None,
     ) -> Memory:
         """Create a memory resource and return the raw response.
 
@@ -162,12 +204,10 @@ class MemoryManager:
             strategies = []
 
         try:
-            processed_strategies = self._add_default_namespaces(strategies)
-
             params = {
                 "name": name,
                 "eventExpiryDuration": event_expiry_days,
-                "memoryStrategies": processed_strategies,
+                "memoryStrategies": strategies,
                 "clientToken": str(uuid.uuid4()),
             }
 
@@ -177,6 +217,9 @@ class MemoryManager:
             if memory_execution_role_arn is not None:
                 params["memoryExecutionRoleArn"] = memory_execution_role_arn
 
+            if encryption_key_arn is not None:
+                params["encryptionKeyArn"] = encryption_key_arn
+
             response = self._control_plane_client.create_memory(**params)
 
             memory = response["memory"]
@@ -193,12 +236,13 @@ class MemoryManager:
     def _create_memory_and_wait(
         self,
         name: str,
-        strategies: List[Dict[str, Any]],
+        strategies: Optional[List[Dict[str, Any]]],
         description: Optional[str] = None,
         event_expiry_days: int = 90,
         memory_execution_role_arn: Optional[str] = None,
         max_wait: int = 300,
         poll_interval: int = 10,
+        encryption_key_arn: Optional[str] = None,
     ) -> Memory:
         """Create a memory and wait for it to become ACTIVE.
 
@@ -213,6 +257,7 @@ class MemoryManager:
             memory_execution_role_arn: IAM role ARN for memory execution
             max_wait: Maximum seconds to wait (default: 300)
             poll_interval: Seconds between status checks (default: 10)
+            encryption_key_arn: kms key ARN for encryption
 
         Returns:
             Created memory object in ACTIVE status
@@ -228,6 +273,7 @@ class MemoryManager:
             description=description,
             event_expiry_days=event_expiry_days,
             memory_execution_role_arn=memory_execution_role_arn,
+            encryption_key_arn=encryption_key_arn,
         )
 
         memory_id = memory.id
@@ -264,20 +310,64 @@ class MemoryManager:
     def create_memory_and_wait(
         self,
         name: str,
-        strategies: List[Dict[str, Any]],
+        strategies: Optional[List[Union[BaseStrategy, Dict[str, Any]]]] = None,
         description: Optional[str] = None,
         event_expiry_days: int = 90,
         memory_execution_role_arn: Optional[str] = None,
+        encryption_key_arn: Optional[str] = None,
         max_wait: int = 300,
         poll_interval: int = 10,
     ) -> Memory:
-        """Create a memory and wait for it to become ACTIVE - public method."""
+        """Create a memory and wait for it to become ACTIVE - public method.
+
+        Args:
+            name: Name for the memory resource
+            strategies: List of typed strategy objects or dictionary configurations
+            description: Optional description
+            event_expiry_days: How long to retain events (default: 90 days)
+            memory_execution_role_arn: IAM role ARN for memory execution
+            max_wait: Maximum seconds to wait (default: 300)
+            poll_interval: Seconds between status checks (default: 10)
+            encryption_key_arn: kms key ARN for encryption
+
+        Returns:
+            Created memory object in ACTIVE status
+
+        Example:
+            from bedrock_agentcore_starter_toolkit.operations.memory.models import (
+                SemanticStrategy, CustomSemanticStrategy, ExtractionConfig, ConsolidationConfig
+            )
+
+            # Create typed strategies
+            semantic = SemanticStrategy(name="MySemanticStrategy")
+            custom = CustomSemanticStrategy(
+                name="MyCustomStrategy",
+                extraction_config=ExtractionConfig(
+                    append_to_prompt="Extract insights",
+                    model_id="anthropic.claude-3-sonnet-20240229-v1:0"
+                ),
+                consolidation_config=ConsolidationConfig(
+                    append_to_prompt="Consolidate insights",
+                    model_id="anthropic.claude-3-haiku-20240307-v1:0"
+                )
+            )
+
+            # Create memory with typed strategies
+            memory = manager.create_memory_and_wait(
+                name="TypedMemory",
+                strategies=[semantic, custom]
+            )
+        """
+        # Convert typed strategies to dicts for internal processing
+        dict_strategies = convert_strategies_to_dicts(strategies) if strategies else None
+
         return self._create_memory_and_wait(
             name=name,
-            strategies=strategies,
+            strategies=dict_strategies,
             description=description,
             event_expiry_days=event_expiry_days,
             memory_execution_role_arn=memory_execution_role_arn,
+            encryption_key_arn=encryption_key_arn,
             max_wait=max_wait,
             poll_interval=poll_interval,
         )
@@ -285,15 +375,37 @@ class MemoryManager:
     def get_or_create_memory(
         self,
         name: str,
-        strategies: Optional[List[Dict[str, Any]]] = None,
+        strategies: Optional[List[Union[BaseStrategy, Dict[str, Any]]]] = None,
         description: Optional[str] = None,
         event_expiry_days: int = 90,
         memory_execution_role_arn: Optional[str] = None,
+        encryption_key_arn: Optional[str] = None,
     ) -> Memory:
         """Fetch an existing memory resource or create the memory.
 
+        Args:
+            name: Memory name
+            strategies: Optional List of typed strategy objects or dictionary configurations
+            description: Optional description
+            event_expiry_days: How long to retain events (default: 90 days)
+            memory_execution_role_arn: IAM role ARN for memory execution
+            encryption_key_arn: kms key ARN for encryption
+
         Returns:
             Memory object, either newly created or existing
+
+        Raises:
+            ValueError: If strategies are provided but existing memory has different strategies
+
+        Example:
+            from bedrock_agentcore_starter_toolkit.operations.memory.models import SemanticStrategy
+
+            # Create with typed strategy
+            semantic = SemanticStrategy(name="MyStrategy")
+            memory = manager.get_or_create_memory(
+                name="MyMemory",
+                strategies=[semantic]
+            )
         """
         memory: Memory = None
         try:
@@ -302,16 +414,27 @@ class MemoryManager:
 
             # Create Memory if it doesn't exist
             if memory_summary is None:
+                # Convert typed strategies to dicts for internal processing
+                dict_strategies = convert_strategies_to_dicts(strategies) if strategies else None
+
                 memory = self._create_memory_and_wait(
                     name=name,
-                    strategies=strategies,
+                    strategies=dict_strategies,
                     description=description,
                     event_expiry_days=event_expiry_days,
                     memory_execution_role_arn=memory_execution_role_arn,
+                    encryption_key_arn=encryption_key_arn,
                 )
             else:
                 logger.info("Memory already exists. Using existing memory ID: %s", memory_summary.id)
                 memory = self.get_memory(memory_summary.id)
+
+                # Validate strategies if provided using deep comparison
+                if strategies is not None:
+                    existing_strategies = memory.get("strategies", memory.get("memoryStrategies", []))
+                    memory_name = memory.get("name")
+                    validate_existing_memory_strategies(existing_strategies, strategies, memory_name)
+
             return memory
         except ClientError as e:
             # Failed to create memory
@@ -470,7 +593,7 @@ class MemoryManager:
         if namespaces:
             strategy[StrategyType.SEMANTIC.value]["namespaces"] = namespaces
 
-        return self._add_strategy(memory_id, strategy)
+        return self.add_strategy(memory_id, strategy)
 
     def add_semantic_strategy_and_wait(
         self,
@@ -514,7 +637,7 @@ class MemoryManager:
         if namespaces:
             strategy[StrategyType.SUMMARY.value]["namespaces"] = namespaces
 
-        return self._add_strategy(memory_id, strategy)
+        return self.add_strategy(memory_id, strategy)
 
     def add_summary_strategy_and_wait(
         self,
@@ -551,7 +674,7 @@ class MemoryManager:
         if namespaces:
             strategy[StrategyType.USER_PREFERENCE.value]["namespaces"] = namespaces
 
-        return self._add_strategy(memory_id, strategy)
+        return self.add_strategy(memory_id, strategy)
 
     def add_user_preference_strategy_and_wait(
         self,
@@ -610,7 +733,7 @@ class MemoryManager:
         if namespaces:
             strategy[StrategyType.CUSTOM.value]["namespaces"] = namespaces
 
-        return self._add_strategy(memory_id, strategy)
+        return self.add_strategy(memory_id, strategy)
 
     def add_custom_semantic_strategy_and_wait(
         self,
@@ -656,17 +779,38 @@ class MemoryManager:
     def update_memory_strategies(
         self,
         memory_id: str,
-        add_strategies: Optional[List[Dict[str, Any]]] = None,
+        add_strategies: Optional[List[Union[BaseStrategy, Dict[str, Any]]]] = None,
         modify_strategies: Optional[List[Dict[str, Any]]] = None,
         delete_strategy_ids: Optional[List[str]] = None,
     ) -> Memory:
-        """Update memory strategies - add, modify, or delete."""
+        """Update memory strategies - add, modify, or delete.
+
+        Args:
+            memory_id: Memory resource ID
+            add_strategies: List of typed strategy objects or dictionaries to add
+            modify_strategies: List of strategy modification dictionaries
+            delete_strategy_ids: List of strategy IDs to delete
+
+        Returns:
+            Updated Memory object
+
+        Example:
+            from bedrock_agentcore_starter_toolkit.operations.memory.models import SemanticStrategy
+
+            # Add typed strategy
+            semantic = SemanticStrategy(name="NewStrategy")
+            memory = manager.update_memory_strategies(
+                memory_id="mem-123",
+                add_strategies=[semantic]
+            )
+        """
         try:
             memory_strategies = {}
 
             if add_strategies:
-                processed_add = self._add_default_namespaces(add_strategies)
-                memory_strategies["addMemoryStrategies"] = processed_add
+                # Convert typed strategies to dicts for internal processing
+                dict_strategies = convert_strategies_to_dicts(add_strategies)
+                memory_strategies["addMemoryStrategies"] = dict_strategies
 
             if modify_strategies:
                 current_strategies = self.get_memory_strategies(memory_id)
@@ -722,7 +866,7 @@ class MemoryManager:
     def update_memory_strategies_and_wait(
         self,
         memory_id: str,
-        add_strategies: Optional[List[Dict[str, Any]]] = None,
+        add_strategies: Optional[List[Union[BaseStrategy, Dict[str, Any]]]] = None,
         modify_strategies: Optional[List[Dict[str, Any]]] = None,
         delete_strategy_ids: Optional[List[str]] = None,
         max_wait: int = 300,
@@ -732,6 +876,27 @@ class MemoryManager:
 
         This method handles the temporary CREATING state that occurs when
         updating strategies, preventing subsequent update errors.
+
+        Args:
+            memory_id: Memory resource ID
+            add_strategies: List of typed strategy objects or dictionaries to add
+            modify_strategies: List of strategy modification dictionaries
+            delete_strategy_ids: List of strategy IDs to delete
+            max_wait: Maximum seconds to wait (default: 300)
+            poll_interval: Seconds between checks (default: 10)
+
+        Returns:
+            Updated Memory object in ACTIVE state
+
+        Example:
+            from bedrock_agentcore_starter_toolkit.operations.memory.models import SummaryStrategy
+
+            # Add typed strategy and wait
+            summary = SummaryStrategy(name="NewSummaryStrategy")
+            memory = manager.update_memory_strategies_and_wait(
+                memory_id="mem-123",
+                add_strategies=[summary]
+            )
         """
         # Update strategies
         self.update_memory_strategies(memory_id, add_strategies, modify_strategies, delete_strategy_ids)
@@ -739,31 +904,77 @@ class MemoryManager:
         # Wait for memory to return to ACTIVE
         return self._wait_for_memory_active(memory_id, max_wait, poll_interval)
 
-    def add_strategy(self, memory_id: str, strategy: Dict[str, Any]) -> Memory:
+    def add_strategy(self, memory_id: str, strategy: Union[BaseStrategy, Dict[str, Any]]) -> Memory:
         """Add a strategy to a memory (without waiting).
 
         WARNING: After adding a strategy, the memory enters CREATING state temporarily.
-        Use add_*_strategy_and_wait() methods instead to avoid errors.
+        Use add_strategy_and_wait() method instead to avoid errors.
 
         Args:
             memory_id: Memory resource ID
-            strategy: Strategy configuration dictionary
+            strategy: Typed strategy object or dictionary configuration
 
         Returns:
             Updated memory response
-        """
-        warnings.warn(
-            "add_strategy() may leave memory in CREATING state. "
-            "Use add_*_strategy_and_wait() methods to avoid subsequent errors.",
-            UserWarning,
-            stacklevel=2,
-        )
-        return self._add_strategy(memory_id, strategy)
 
-    def _add_strategy(self, memory_id: str, strategy: Dict[str, Any]) -> Memory:
-        """Internal method to add a single strategy."""
+        Example:
+            from bedrock_agentcore_starter_toolkit.operations.memory.models.semantic import SemanticStrategy
+
+            # Using typed strategy (recommended)
+            semantic = SemanticStrategy(name="MyStrategy", description="Test")
+            memory = manager.add_strategy(memory_id="mem-123", strategy=semantic)
+
+            # Using dictionary (legacy support)
+            strategy_dict = {"semanticMemoryStrategy": {"name": "MyStrategy"}}
+            memory = manager.add_strategy(memory_id="mem-123", strategy=strategy_dict)
+        """
         return self.update_memory_strategies(memory_id=memory_id, add_strategies=[strategy])
 
+    def add_strategy_and_wait(
+        self,
+        memory_id: str,
+        strategy: Union[BaseStrategy, Dict[str, Any]],
+        max_wait: int = 300,
+        poll_interval: int = 10,
+    ) -> Memory:
+        """Add a strategy to a memory and wait for it to return to ACTIVE state.
+
+        Args:
+            memory_id: Memory resource ID
+            strategy: Typed strategy object or dictionary configuration
+            max_wait: Maximum seconds to wait (default: 300)
+            poll_interval: Seconds between status checks (default: 10)
+
+        Returns:
+            Updated memory response in ACTIVE state
+
+        Example:
+            from bedrock_agentcore_starter_toolkit.operations.memory.models.strategies import (
+                SemanticStrategy, CustomSemanticStrategy, ExtractionConfig, ConsolidationConfig
+            )
+
+            # Using typed strategy (recommended)
+            semantic = SemanticStrategy(name="MyStrategy", description="Test")
+            memory = manager.add_strategy_and_wait(memory_id="mem-123", strategy=semantic)
+
+            # Using custom strategy with configurations
+            custom = CustomSemanticStrategy(
+                name="CustomStrategy",
+                extraction_config=ExtractionConfig(
+                    append_to_prompt="Extract insights",
+                    model_id="anthropic.claude-3-sonnet-20240229-v1:0"
+                ),
+                consolidation_config=ConsolidationConfig(
+                    append_to_prompt="Consolidate insights",
+                    model_id="anthropic.claude-3-haiku-20240307-v1:0"
+                )
+            )
+            memory = manager.add_strategy_and_wait(memory_id="mem-123", strategy=custom)
+        """
+        return self.update_memory_strategies_and_wait(
+            memory_id=memory_id, add_strategies=[strategy], max_wait=max_wait, poll_interval=poll_interval
+        )
+
     def _check_strategies_terminal_state(self, strategies: List[Dict[str, Any]]) -> tuple[bool, List[str], List[str]]:
         """Check if all strategies are in terminal states.
 
@@ -851,26 +1062,6 @@ class MemoryManager:
             % (memory_id, max_wait)
         )
 
-    def _add_default_namespaces(self, strategies: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
-        """Add default namespaces to strategies that don't have them."""
-        processed = []
-
-        for strategy in strategies:
-            strategy_copy = copy.deepcopy(strategy)
-
-            strategy_type_key = list(strategy.keys())[0]
-            strategy_config = strategy_copy[strategy_type_key]
-
-            if "namespaces" not in strategy_config:
-                strategy_type = StrategyType(strategy_type_key)
-                strategy_config["namespaces"] = DEFAULT_NAMESPACES.get(strategy_type, ["custom/{actorId}/{sessionId}"])
-
-            self._validate_strategy_config(strategy_copy, strategy_type_key)
-
-            processed.append(strategy_copy)
-
-        return processed
-
     def _validate_namespace(self, namespace: str) -> bool:
         """Validate namespace format - basic check only."""
         # Only check for template variables in namespace definition

bedrock_agentcore_starter_toolkit/operations/memory/models/__init__.py

@@ -0,0 +1,106 @@
+"""Bedrock AgentCore Memory Models.
+
+This module provides strongly typed models for memory operations,
+including memory resources, strategies, and configurations.
+
+Example:
+    # Import strategy classes
+    from bedrock_agentcore_starter_toolkit.operations.memory.models import (
+        SemanticStrategy,
+        SummaryStrategy,
+        CustomSemanticStrategy,
+        ExtractionConfig,
+        ConsolidationConfig
+    )
+
+    # Create typed strategies
+    semantic_strategy = SemanticStrategy(
+        name="ConversationSemantics",
+        description="Extract semantic information",
+        namespaces=["semantics/{actorId}/{sessionId}"]
+    )
+
+    custom_strategy = CustomSemanticStrategy(
+        name="CustomExtraction",
+        extraction_config=ExtractionConfig(
+            append_to_prompt="Extract key insights",
+            model_id="anthropic.claude-3-sonnet-20240229-v1:0"
+        ),
+        consolidation_config=ConsolidationConfig(
+            append_to_prompt="Consolidate insights",
+            model_id="anthropic.claude-3-haiku-20240307-v1:0"
+        )
+    )
+"""
+
+# Memory resource models
+from typing import Any, Dict, List
+
+from .Memory import Memory
+from .MemoryStrategy import MemoryStrategy
+from .MemorySummary import MemorySummary
+
+# Strategy models
+from .strategies import (
+    BaseStrategy,
+    ConsolidationConfig,
+    CustomSemanticStrategy,
+    ExtractionConfig,
+    SemanticStrategy,
+    StrategyType,
+    SummaryStrategy,
+    UserPreferenceStrategy,
+)
+
+
+def convert_strategies_to_dicts(strategies: List[StrategyType]) -> List[Dict[str, Any]]:
+    """Convert mixed strategy types to dictionary format for API calls.
+
+    This function handles both new typed strategies and legacy dictionary
+    strategies, ensuring backward compatibility.
+
+    Args:
+        strategies: List of strategy objects (typed or dict)
+
+    Returns:
+        List of strategy dictionaries compatible with the API
+
+    Raises:
+        ValueError: If an invalid strategy type is provided
+
+    Example:
+        strategies = [
+            SemanticStrategy(name="Test"),
+            {"semanticMemoryStrategy": {"name": "Legacy"}}
+        ]
+        dicts = convert_strategies_to_dicts(strategies)
+    """
+    result = []
+    for strategy in strategies:
+        if isinstance(strategy, BaseStrategy):
+            result.append(strategy.to_dict())
+        elif isinstance(strategy, dict):
+            result.append(strategy)  # Backward compatibility
+        else:
+            raise ValueError(f"Invalid strategy type: {type(strategy)}. Expected BaseStrategy or dict.")
+    return result
+
+
+__all__ = [
+    # Memory models
+    "Memory",
+    "MemorySummary",
+    "MemoryStrategy",
+    # Strategy base classes and types
+    "BaseStrategy",
+    "StrategyType",
+    "ExtractionConfig",
+    "ConsolidationConfig",
+    # Strategy models
+    "SemanticStrategy",
+    "SummaryStrategy",
+    "UserPreferenceStrategy",
+    "CustomSemanticStrategy",
+    # Utility functions
+    "convert_strategies_to_dicts",
+]

bedrock_agentcore_starter_toolkit/operations/memory/models/strategies/__init__.py

@@ -0,0 +1,52 @@
+"""Memory strategy models and configurations.
+
+This module provides strongly typed strategy classes for creating
+memory strategies with full type safety and IDE support.
+
+Example:
+    from bedrock_agentcore_starter_toolkit.operations.memory.models.strategies import (
+        SemanticStrategy,
+        CustomSemanticStrategy,
+        ExtractionConfig,
+        ConsolidationConfig
+    )
+
+    # Create typed strategies manually
+    semantic = SemanticStrategy(
+        name="MySemanticStrategy",
+        description="Extract key information"
+    )
+
+    custom = CustomSemanticStrategy(
+        name="MyCustomStrategy",
+        extraction_config=ExtractionConfig(
+            append_to_prompt="Extract insights",
+            model_id="anthropic.claude-3-sonnet-20240229-v1:0"
+        ),
+        consolidation_config=ConsolidationConfig(
+            append_to_prompt="Consolidate insights",
+            model_id="anthropic.claude-3-haiku-20240307-v1:0"
+        )
+    )
+"""
+
+from .base import BaseStrategy, ConsolidationConfig, ExtractionConfig, StrategyType
+from .custom import CustomSemanticStrategy, CustomSummaryStrategy, CustomUserPreferenceStrategy
+from .semantic import SemanticStrategy
+from .summary import SummaryStrategy
+from .user_preference import UserPreferenceStrategy
+
+__all__ = [
+    # Base classes and types
+    "BaseStrategy",
+    "StrategyType",
+    "ExtractionConfig",
+    "ConsolidationConfig",
+    # Concrete strategy classes
+    "SemanticStrategy",
+    "SummaryStrategy",
+    "UserPreferenceStrategy",
+    "CustomSemanticStrategy",
+    "CustomSummaryStrategy",
+    "CustomUserPreferenceStrategy",
+]