powermem 0.1.0__py3-none-any.whl

powermem/storage/adapter.py

@@ -0,0 +1,896 @@
+"""
+Storage adapter for Memory class
+
+This module provides an adapter that bridges the VectorStoreBase interface
+with the interface expected by the Memory class.
+"""
+
+import logging
+import uuid
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+from powermem.storage.base import VectorStoreBase
+from powermem.utils.utils import serialize_datetime
+
+logger = logging.getLogger(__name__)
+
+
+class StorageAdapter:
+    """Adapter that bridges VectorStoreBase interface with Memory class expectations."""
+    
+    def __init__(self, vector_store: VectorStoreBase, embedding_service=None):
+        """Initialize the adapter with a vector store and embedding service."""
+        self.vector_store = vector_store
+        self.embedding_service = embedding_service
+        # get collection name from vector store attribute collection_name
+        self.collection_name = getattr(vector_store, 'collection_name', 'memories')
+
+        # Sub stores support (optional, for multi-table routing)
+        self.sub_stores: Dict[str, 'SubStoreConfig'] = {}
+        self.migration_manager = None
+
+        # Ensure collection exists (will be created with actual vector size when first vector is added)
+        # self.vector_store.create_col(self.collection_name, vector_size=1536, distance="cosine")
+    
+    def add_memory(self, memory_data: Dict[str, Any]) -> int:
+        """Add a memory to the store."""
+        # ID will be generated using Snowflake algorithm before insertion
+        
+        # Create vector from content using embedding service
+        content = memory_data.get("content", "")
+        metadata = memory_data.get("metadata", {})
+
+        # Route to target store (main or sub store)
+        target_store = self._route_to_store(metadata)
+
+        # Check if embedding is already provided (preferred way)
+        vector = memory_data.get("embedding")
+
+        if vector is None:
+            # No embedding provided, generate using embedding service
+            if self.embedding_service:
+                try:
+                    vector = self.embedding_service.embed(content, memory_action="add")
+                except Exception as e:
+                    logger.warning(f"Failed to generate embedding, using mock vector: {e}")
+                    vector = [0.1] * 1536  # Use 1536 dimensions for OceanBase compatibility
+            else:
+                # No embedding service available, use mock vector
+                vector = [0.1] * 1536
+
+        # Create collection with actual vector size if not exists
+        collection_name = getattr(target_store, 'collection_name', self.collection_name)
+        if not hasattr(self, '_collection_created'):
+            target_store.create_col(collection_name, vector_size=len(vector), distance="cosine")
+            self._collection_created = True
+        
+        # Store the memory data as payload - unified format based on OceanBase
+        payload = {
+            "data": content,  # Unified field name for text content
+            "user_id": memory_data.get("user_id", ""),
+            "agent_id": memory_data.get("agent_id", ""),
+            "run_id": memory_data.get("run_id", ""),
+            "actor_id": memory_data.get("actor_id", ""),
+            "hash": memory_data.get("hash", ""),
+            "created_at": serialize_datetime(memory_data.get("created_at", "")),
+            "updated_at": serialize_datetime(memory_data.get("updated_at", "")),
+            "category": memory_data.get("category", ""),
+            "fulltext_content": content,  # For full-text search
+        }
+        
+        # Add only user-defined metadata (not system fields)
+        user_metadata = memory_data.get("metadata", {})
+        payload["metadata"] = serialize_datetime(user_metadata) if user_metadata else {}
+        
+        # Add any extra fields (excluding system fields and embedding)
+        excluded_fields = ["id", "content", "data", "user_id", "agent_id", "run_id", "metadata", "filters", 
+                          "created_at", "updated_at", "actor_id", "hash", "category", "embedding"]
+        for key, value in memory_data.items():
+            if key not in excluded_fields:
+                payload[key] = serialize_datetime(value)
+        
+        # Insert and get generated Snowflake ID
+        generated_ids = target_store.insert([vector], [payload])
+        if not generated_ids:
+            raise ValueError("Failed to insert memory: no ID returned from vector store")
+        memory_id = generated_ids[0]  # Get the first (and only) generated Snowflake ID
+        return memory_id
+    
+    def search_memories(
+        self,
+        query_embedding: List[float],
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        run_id: Optional[str] = None,
+        filters: Optional[Dict[str, Any]] = None,
+        limit: int = 30,
+        query: Optional[str] = None,
+    ) -> List[Dict[str, Any]]:
+        """Search for memories."""
+        # Use the provided query embedding or generate one
+        if query_embedding:
+            query_vector = query_embedding
+        else:
+            # If no query embedding provided, we can't search meaningfully
+            logger.warning("No query embedding provided for search")
+            return []
+        
+        # Merge user_id/agent_id/run_id into filters to ensure consistency
+        # This ensures filters are applied at the database level, avoiding redundant filtering
+        effective_filters = filters.copy() if filters else {}
+        if user_id is not None:
+            effective_filters["user_id"] = user_id
+        if agent_id is not None:
+            effective_filters["agent_id"] = agent_id
+        if run_id is not None:
+            effective_filters["run_id"] = run_id
+        
+        # Route to target store (main or sub store)
+        target_store = self._route_to_store(effective_filters)
+
+        # Unified search method - try OceanBase format first, fallback to SQLite
+        # Pass query text to enable hybrid search (vector + full-text search)
+        try:
+            # Try OceanBase format first - pass query text for hybrid search
+            search_query = query if query else ""
+            results = target_store.search(search_query, vectors=query_vector, limit=limit, filters=effective_filters)
+        except TypeError:
+            # Fallback to SQLite format (doesn't support query text parameter)
+            # Pass filters to ensure filtering works correctly
+            results = target_store.search(search_query if query else "", vectors=[query_vector], limit=limit, filters=effective_filters)
+        
+        # Convert results to unified format
+        memories = []
+        for result in results:
+            # Handle different result formats
+            if hasattr(result, 'payload') and result.payload:
+                # Result with payload attribute
+                payload = result.payload
+                memory_id = result.id
+                # Extract score - use 0.0 as default instead of 1.0 to avoid false high scores
+                # Score should always exist from vector search, but handle None case gracefully
+                score = getattr(result, 'score', None)
+                if score is None:
+                    logger.warning(f"Result {memory_id} missing score, using 0.0")
+                    score = 0.0
+            elif hasattr(result, 'payload') and isinstance(result.payload, dict):
+                # Result with dict payload
+                payload = result.payload
+                memory_id = result.id
+                # Extract score - use 0.0 as default instead of 1.0
+                score = getattr(result, 'score', None)
+                if score is None:
+                    logger.warning(f"Result {memory_id} missing score, using 0.0")
+                    score = 0.0
+            elif isinstance(result, dict):
+                # Direct dict result
+                payload = result
+                memory_id = result.get("id")
+                # Extract score - use 0.0 as default instead of 1.0
+                score = result.get("score")
+                if score is None:
+                    logger.warning(f"Result {memory_id} missing score, using 0.0")
+                    score = 0.0
+            else:
+                continue
+            
+            # Extract unified fields
+            # Core and promoted keys that should not be in metadata
+            promoted_payload_keys = ["user_id", "agent_id", "run_id", "actor_id", "role"]
+            core_and_promoted_keys = {"data", "hash", "created_at", "updated_at", "id", "metadata", *promoted_payload_keys}
+            
+            # Extract core fields
+            content = payload.get("data", "")
+            created_at = payload.get("created_at")
+            updated_at = payload.get("updated_at")
+            
+            # Extract promoted fields
+            promoted_fields = {}
+            for key in promoted_payload_keys:
+                if key in payload:
+                    promoted_fields[key] = payload[key]
+            
+            # Extract user metadata from payload
+            # If payload contains "metadata" field (nested user metadata), use it directly
+            # Otherwise, extract additional metadata from other fields
+            if "metadata" in payload:
+                user_metadata = payload["metadata"].copy() if payload["metadata"] else {}
+            else:
+                # Extract additional metadata (all fields not in core_and_promoted_keys)
+                user_metadata = {k: v for k, v in payload.items() if k not in core_and_promoted_keys}
+            
+            # Merge any user-defined fields from payload top-level into metadata
+            # These fields (like "category") were extracted from metadata for filtering purposes
+            # but should still be visible in the returned metadata
+            for key, value in payload.items():
+                if key not in core_and_promoted_keys and key not in user_metadata and value:
+                    # Only include non-empty values that aren't already in metadata
+                    user_metadata[key] = value
+            
+            memory = {
+                "id": memory_id,
+                "memory": content, 
+                "created_at": created_at,
+                "updated_at": updated_at,
+                "score": score,
+                **promoted_fields,  # Add promoted fields at top level
+                "metadata": user_metadata if user_metadata else {},  # Add user metadata
+            }
+            
+            # No need to apply filters here - filters are already applied at the database level
+            # in vector_store.search(), so all returned results should already match the filters
+            memories.append(memory)
+        
+        # Vector store already applied limit, no need to slice again
+        return memories
+    
+    def get_memory(
+        self,
+        memory_id: int,
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+    ) -> Optional[Dict[str, Any]]:
+        """Get a specific memory by ID."""
+        result = self.vector_store.get(memory_id)
+        
+        if result and result.payload:
+            memory = {
+                "id": result.id,
+                "content": result.payload.get("content", ""),
+                "user_id": result.payload.get("user_id"),
+                "agent_id": result.payload.get("agent_id"),
+                "run_id": result.payload.get("run_id"),
+                "metadata": result.payload.get("metadata", {}),
+                "created_at": result.payload.get("created_at"),
+                "updated_at": result.payload.get("updated_at"),
+            }
+            
+            # Check access control
+            if user_id and memory.get("user_id") != user_id:
+                return None
+            if agent_id and memory.get("agent_id") != agent_id:
+                return None
+            
+            return memory
+        
+        # If not found in main store and sub stores exist, search sub stores
+        if self.sub_stores:
+            for sub_config in self.sub_stores.values():
+                try:
+                    result = sub_config.vector_store.get(memory_id)
+                    if result and result.payload:
+                        memory = {
+                            "id": result.id,
+                            "content": result.payload.get("content", ""),
+                            "user_id": result.payload.get("user_id"),
+                            "agent_id": result.payload.get("agent_id"),
+                            "run_id": result.payload.get("run_id"),
+                            "metadata": result.payload.get("metadata", {}),
+                            "created_at": result.payload.get("created_at"),
+                            "updated_at": result.payload.get("updated_at"),
+                        }
+
+                        # Check access control
+                        if user_id and memory.get("user_id") != user_id:
+                            continue
+                        if agent_id and memory.get("agent_id") != agent_id:
+                            continue
+
+                        return memory
+                except Exception as e:
+                    logger.debug(f"Error searching in sub store {sub_config.name}: {e}")
+                    continue
+
+        return None
+    
+    def update_memory(
+        self,
+        memory_id: int,
+        update_data: Dict[str, Any],
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+    ) -> Optional[Dict[str, Any]]:
+        """Update a memory."""
+        # First check if memory exists and user has access (get_memory returns dict)
+        existing_memory_dict = self.get_memory(memory_id, user_id, agent_id)
+        if not existing_memory_dict:
+            logger.warning(f"Memory {memory_id} not found or access denied")
+            return None
+        
+        # Get raw OutputData object from vector store to access payload
+        existing_result = self.vector_store.get(memory_id)
+        target_store = self.vector_store
+
+        # If not found in main store, search sub stores
+        if (not existing_result or not existing_result.payload) and self.sub_stores:
+            for sub_config in self.sub_stores.values():
+                try:
+                    sub_result = sub_config.vector_store.get(memory_id)
+                    if sub_result and sub_result.payload:
+                        # Verify access control matches
+                        sub_payload = sub_result.payload
+                        if user_id and sub_payload.get("user_id") != user_id:
+                            continue
+                        if agent_id and sub_payload.get("agent_id") != agent_id:
+                            continue
+                        existing_result = sub_result
+                        target_store = sub_config.vector_store
+                        break
+                except Exception as e:
+                    logger.debug(f"Error searching in sub store {sub_config.name}: {e}")
+                    continue
+
+        if not existing_result or not existing_result.payload:
+            logger.warning(f"Memory {memory_id} not found in vector store")
+            return None
+        
+        # Get existing payload
+        existing_payload = existing_result.payload
+        
+        # Merge update_data into payload
+        updated_payload = existing_payload.copy()
+        
+        # Handle content field - map to "data" in payload
+        if "content" in update_data:
+            updated_payload["data"] = update_data["content"]
+            updated_payload["fulltext_content"] = update_data["content"]
+            # Remove content from update_data to avoid confusion
+            update_data = update_data.copy()
+            del update_data["content"]
+        
+        # Serialize datetime objects in update_data before merging
+        serialized_update_data = serialize_datetime(update_data)
+        
+        # Update other fields
+        updated_payload.update(serialized_update_data)
+        
+        # Ensure datetime fields are serialized as ISO format strings
+        if "updated_at" in updated_payload:
+            updated_at = updated_payload["updated_at"]
+            if isinstance(updated_at, datetime):
+                updated_payload["updated_at"] = updated_at.isoformat()
+        if "created_at" in updated_payload:
+            created_at = updated_payload["created_at"]
+            if isinstance(created_at, datetime):
+                updated_payload["created_at"] = created_at.isoformat()
+
+        # Update updated_at if not provided
+        if "updated_at" not in updated_payload:
+            updated_payload["updated_at"] = datetime.utcnow().isoformat()
+        
+        # Update in vector store with proper payload
+        target_store.update(memory_id, vector=update_data.get("embedding"), payload=updated_payload)
+        
+        return updated_payload
+    
+    def delete_memory(
+        self,
+        memory_id: int,
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+    ) -> bool:
+        """Delete a memory."""
+        # Check if memory exists and user has access
+        existing = self.get_memory(memory_id, user_id, agent_id)
+        if not existing:
+            return False
+
+        # Try to delete from main store
+        try:
+            self.vector_store.delete(memory_id)
+            return True
+        except Exception as e:
+            logger.debug(f"Memory {memory_id} not in main store: {e}")
+
+        # If not in main store, try sub stores
+        if self.sub_stores:
+            for sub_config in self.sub_stores.values():
+                try:
+                    sub_config.vector_store.delete(memory_id)
+                    logger.debug(f"Deleted memory {memory_id} from sub store {sub_config.name}")
+                    return True
+                except Exception as e:
+                    logger.debug(f"Memory {memory_id} not in sub store {sub_config.name}: {e}")
+                    continue
+
+        # If we reach here, memory existed in get_memory but couldn't be deleted
+        logger.warning(f"Failed to delete memory {memory_id}")
+        return False
+    
+    def get_all_memories(
+        self,
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        run_id: Optional[str] = None,
+        limit: int = 100,
+        offset: int = 0,
+    ) -> List[Dict[str, Any]]:
+        """Get all memories with optional filtering."""
+        # Build filters for database-level filtering
+        filters = {}
+        if user_id:
+            filters["user_id"] = user_id
+        if agent_id:
+            filters["agent_id"] = agent_id
+        if run_id:
+            filters["run_id"] = run_id
+        
+        # Get memories from vector store with filters (if supported)
+        if filters and hasattr(self.vector_store, 'list'):
+            # Pass filters to vector store's list method for database-level filtering
+            # Request more records to support offset
+            results = self.vector_store.list(filters=filters, limit=limit + offset)
+        else:
+            # Fallback: get all and filter in memory
+            results = self.vector_store.list(limit=limit + offset)
+        
+        # OceanBase returns [memories], SQLite/PGVector return memories directly
+        if results and isinstance(results[0], list):
+            raw_results = results[0]
+        else:
+            raw_results = results
+        
+        # Convert to expected format and apply filters
+        memories = []
+        for result in raw_results:
+            # Handle different result formats
+            if hasattr(result, 'payload') and result.payload:
+                # Result with payload attribute (e.g., from OceanBase OutputData)
+                payload = result.payload
+                memory_id = result.id
+            elif isinstance(result, dict):
+                # Direct dict result (e.g., from SQLite)
+                payload = result
+                memory_id = result.get("id")
+            else:
+                continue
+            
+            # Convert datetime objects to ISO format strings
+            created_at = payload.get("created_at")
+            if created_at is not None:
+                from datetime import datetime
+                if isinstance(created_at, datetime):
+                    created_at = created_at.isoformat()
+            
+            updated_at = payload.get("updated_at")
+            if updated_at is not None:
+                from datetime import datetime
+                if isinstance(updated_at, datetime):
+                    updated_at = updated_at.isoformat()
+            
+            memory = {
+                "id": memory_id,
+                "memory": payload.get("data", ""),  # Unified field name to match search_memories format
+                "user_id": payload.get("user_id"),
+                "agent_id": payload.get("agent_id"),
+                "run_id": payload.get("run_id"),
+                "metadata": payload.get("metadata", {}),
+                "created_at": created_at,
+                "updated_at": updated_at,
+            }
+            
+            # Apply filters (as double-check if database didn't filter)
+            # Note: If filters were applied at database level, these will all pass
+            if user_id and memory.get("user_id") != user_id:
+                continue
+            if agent_id and memory.get("agent_id") != agent_id:
+                continue
+            if run_id and memory.get("run_id") != run_id:
+                continue
+            
+            memories.append(memory)
+        
+        # Apply offset and limit
+        return memories[offset:offset + limit]
+    
+    def clear_memories(
+        self,
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        run_id: Optional[str] = None,
+    ) -> bool:
+        """Clear all memories for a user or agent or run."""
+        # Build filters for database query
+        filters = {}
+        if user_id:
+            filters["user_id"] = user_id
+        if agent_id:
+            filters["agent_id"] = agent_id
+        if run_id:
+            filters["run_id"] = run_id
+        
+        # Use batch processing to avoid timeout
+        batch_size = 1000
+        deleted_count = 0
+        
+        while True:
+            # Get a batch of memories with filtering
+            batch = self.get_all_memories(user_id, agent_id, run_id, limit=batch_size, offset=deleted_count)
+            
+            # If no more records, we're done
+            if not batch:
+                break
+            
+            # Delete each memory in the batch
+            for memory in batch:
+                try:
+                    self.vector_store.delete(memory["id"])
+                except Exception as e:
+                    logger.warning(f"Failed to delete memory {memory.get('id')}: {e}")
+            
+            deleted_count += len(batch)
+            
+            # If we got fewer records than batch_size, we've reached the end
+            if len(batch) < batch_size:
+                break
+        
+        logger.info(f"Deleted {deleted_count} memories with filters: {filters}")
+        return True
+    
+    async def get_all_memories_async(
+        self,
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        run_id: Optional[str] = None,
+        limit: int = 100,
+        offset: int = 0,
+    ) -> List[Dict[str, Any]]:
+        """Get all memories with optional filtering asynchronously."""
+        import asyncio
+        return await asyncio.to_thread(self.get_all_memories, user_id, agent_id, run_id, limit, offset)
+    
+    async def clear_memories_async(
+        self,
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        run_id: Optional[str] = None,
+    ) -> bool:
+        """Clear all memories for a user or agent or run asynchronously."""
+        import asyncio
+        return await asyncio.to_thread(self.clear_memories, user_id, agent_id, run_id)
+    
+    async def initialize_async(self):
+        """Initialize storage asynchronously."""
+        # No-op for now
+        pass
+    
+    async def add_memory_async(self, memory_data: Dict[str, Any]) -> int:
+        """Add a memory to the store asynchronously."""
+        import asyncio
+        return await asyncio.to_thread(self.add_memory, memory_data)
+    
+    async def search_memories_async(
+        self,
+        query_embedding: List[float],
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        run_id: Optional[str] = None,
+        filters: Optional[Dict[str, Any]] = None,
+        limit: int = 30,
+        query: Optional[str] = None,
+    ) -> List[Dict[str, Any]]:
+        """Search for memories asynchronously."""
+        import asyncio
+        return await asyncio.to_thread(self.search_memories, query_embedding, user_id, agent_id, run_id, filters, limit, query)
+    
+    async def get_memory_async(
+        self,
+        memory_id: int,
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+    ) -> Optional[Dict[str, Any]]:
+        """Get a specific memory by ID asynchronously."""
+        import asyncio
+        return await asyncio.to_thread(self.get_memory, memory_id, user_id, agent_id)
+    
+    async def delete_memory_async(
+        self,
+        memory_id: int,
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+    ) -> bool:
+        """Delete a memory asynchronously."""
+        import asyncio
+        return await asyncio.to_thread(self.delete_memory, memory_id, user_id, agent_id)
+    
+    async def update_memory_async(
+        self,
+        memory_id: int,
+        update_data: Dict[str, Any],
+        user_id: Optional[str] = None,
+        agent_id: Optional[str] = None,
+    ) -> Optional[Dict[str, Any]]:
+        """Update a memory asynchronously."""
+        import asyncio
+        return await asyncio.to_thread(self.update_memory, memory_id, update_data, user_id, agent_id)
+
+    # ==================== Routing Support Methods ====================
+
+    def _route_to_store(self, filters_or_metadata: Optional[Dict] = None) -> VectorStoreBase:
+        """
+        Route to correct storage instance (main or sub store).
+
+        Args:
+            filters_or_metadata: Query conditions or memory metadata
+
+        Returns:
+            Target VectorStoreBase instance
+        """
+        # If no sub stores configured, always use main store
+        if not self.sub_stores:
+            return self.vector_store
+
+        # Try to find matching sub store
+        if filters_or_metadata:
+            for sub_config in self.sub_stores.values():
+                # Check if sub store is ready (only for query operations)
+                if self.migration_manager and not self.migration_manager.is_ready(sub_config.name):
+                    continue
+
+                # Check if filters_or_metadata matches routing rules
+                routing_filter = sub_config.routing_filter
+                if all(
+                    key in filters_or_metadata and filters_or_metadata[key] == value
+                    for key, value in routing_filter.items()
+                ):
+                    logger.debug(f"Routing to sub store: {sub_config.name}")
+                    return sub_config.vector_store
+
+        # Default to main store
+        logger.debug("Routing to main store")
+        return self.vector_store
+
+    def get_target_store_name(self, filters_or_metadata: Optional[Dict] = None) -> str:
+        """
+        Get target store name for given filters/metadata.
+
+        Args:
+            filters_or_metadata: Query conditions or memory metadata
+
+        Returns:
+            Target storage name
+        """
+        target_store = self._route_to_store(filters_or_metadata)
+        return getattr(target_store, 'collection_name', self.collection_name)
+
+    def is_sub_store_ready(self, store_name: str) -> bool:
+        """
+        Check if sub store is ready (migration completed).
+
+        Args:
+            store_name: Sub store name
+
+        Returns:
+            True if ready, False otherwise (or if no migration manager)
+        """
+        if self.migration_manager:
+            return self.migration_manager.is_ready(store_name)
+        return False
+
+
+# ==================== Sub Store Configuration ====================
+
+class SubStoreConfig:
+    """Configuration for a sub store."""
+    def __init__(self, name: str, routing_filter: Dict, vector_store: VectorStoreBase, embedding_service=None):
+        self.name = name
+        self.routing_filter = routing_filter
+        self.vector_store = vector_store
+        self.embedding_service = embedding_service
+
+
+class SubStorageAdapter(StorageAdapter):
+    """
+    Extended storage adapter with sub-store management capabilities.
+
+    This adapter extends the basic StorageAdapter to support multiple sub-stores
+    with intelligent routing based on metadata filters. All CRUD operations are
+    inherited from the parent class and automatically support routing.
+
+    This class only contains sub-store management methods.
+    """
+
+    def __init__(self, vector_store: VectorStoreBase, embedding_service=None):
+        """
+        Initialize the sub-storage adapter.
+
+        Args:
+            vector_store: The main vector store instance
+            embedding_service: Optional embedding service for generating vectors
+        """
+        # Initialize parent class
+        super().__init__(vector_store, embedding_service)
+
+        # Initialize migration status management (database-backed)
+        from powermem.storage.migration_manager import SubStoreMigrationManager
+        self.migration_manager = SubStoreMigrationManager(vector_store, self.collection_name)
+
+    # ==================== Sub Store Management Methods ====================
+
+    def register_sub_store(
+        self,
+        store_name: str,
+        routing_filter: Dict,
+        vector_store: VectorStoreBase,
+        embedding_service=None,
+    ):
+        """
+        Register a sub store for routing.
+
+        Args:
+            store_name: Name of the sub store
+            routing_filter: Dictionary of metadata conditions for routing
+            vector_store: Vector store instance for the sub store
+            embedding_service: Optional embedding service for this sub store (for migration)
+        """
+        sub_config = SubStoreConfig(
+            name=store_name,
+            routing_filter=routing_filter,
+            vector_store=vector_store,
+            embedding_service=embedding_service
+        )
+        self.sub_stores[store_name] = sub_config
+
+        # Register in migration manager
+        if self.migration_manager:
+            self.migration_manager.register_sub_store(
+                sub_store_name=store_name,
+                routing_filter=routing_filter
+            )
+
+        logger.info(f"Registered sub store: {store_name} with filter: {routing_filter}")
+
+    def migrate_to_sub_store(
+        self,
+        store_name: str,
+        delete_source: bool = False,
+        batch_size: int = 100
+    ) -> int:
+        """
+        Migrate data from main store to sub store based on routing filter.
+
+        Args:
+            store_name: Name of the sub store to migrate to
+            delete_source: Whether to delete source data after migration
+            batch_size: Number of records to process in each batch
+
+        Returns:
+            Number of records migrated
+
+        Raises:
+            ValueError: If sub store not found or not registered
+        """
+        if store_name not in self.sub_stores:
+            raise ValueError(f"Sub store '{store_name}' not found. Please register it first.")
+
+        sub_config = self.sub_stores[store_name]
+        routing_filter = sub_config.routing_filter
+        target_store = sub_config.vector_store
+        sub_embedding_service = sub_config.embedding_service
+
+        # Validate that embedding service is provided
+        if not sub_embedding_service:
+            raise ValueError(f"Sub store '{store_name}' does not have an embedding service configured. "
+                           "Cannot migrate without re-embedding the data.")
+
+        # Mark migration as started
+        if self.migration_manager:
+            self.migration_manager.mark_migrating(store_name, 0)
+
+        try:
+            # Query all matching records from main store
+            migrated_count = 0
+
+            # Get all memories that match the routing filter
+            from powermem.storage.oceanbase.oceanbase import OceanBaseVectorStore
+
+            if isinstance(self.vector_store, OceanBaseVectorStore):
+                # Use OceanBase specific query
+                # Build SQL query to find matching records (only need ID and content fields)
+                filter_conditions = " AND ".join([
+                    f"JSON_EXTRACT(metadata, '$.{key}') = '{value}'"
+                    for key, value in routing_filter.items()
+                ])
+
+                # Query only IDs first for efficiency
+                id_query_sql = f"""
+                SELECT id
+                FROM {self.collection_name}
+                WHERE {filter_conditions}
+                LIMIT {batch_size}
+                """
+
+                while True:
+                    id_results = self.vector_store.execute_sql(id_query_sql)
+                    if not id_results:
+                        break
+
+                    for id_record in id_results:
+                        record_id = id_record['id']
+
+                        # Use get() method to retrieve the full record
+                        result = self.vector_store.get(record_id)
+                        if not result or not result.payload:
+                            logger.warning(f"Record {record_id} not found, skipping")
+                            continue
+
+                        # Use payload from result
+                        payload = result.payload.copy()
+                        payload['id'] = record_id
+
+                        # Extract content for re-embedding
+                        content = payload.get('data', '')
+                        if not content:
+                            logger.warning(f"Record {record_id} has no content, skipping")
+                            continue
+
+                        # Re-generate vector using sub store's embedding service
+                        try:
+                            vector = sub_embedding_service.embed(content, memory_action="add")
+                            logger.debug(f"Re-embedded record {record_id} with dimension {len(vector)}")
+                        except Exception as embed_error:
+                            logger.error(f"Failed to re-embed record {record_id}: {embed_error}")
+                            continue
+
+                        try:
+                            target_store.insert([vector], [payload])
+                            migrated_count += 1
+
+                            # Delete from source if requested
+                            if delete_source:
+                                self.vector_store.delete(record_id)
+
+                            # Update progress
+                            if self.migration_manager and migrated_count % 10 == 0:
+                                self.migration_manager.update_progress(
+                                    store_name,
+                                    migrated_count,
+                                    migrated_count  # Total is unknown in this approach
+                                )
+                        except Exception as e:
+                            logger.error(f"Error migrating record {record_id}: {e}")
+                            continue
+
+                    # If we got fewer results than batch_size, we're done
+                    if len(id_results) < batch_size:
+                        break
+            else:
+                logger.warning(f"Migration not fully supported for {type(self.vector_store).__name__}")
+
+            # Mark migration as completed
+            if self.migration_manager:
+                self.migration_manager.mark_completed(store_name, migrated_count)
+
+            logger.info(f"Migration completed: {migrated_count} records migrated to {store_name}")
+            return migrated_count
+
+        except Exception as e:
+            # Mark migration as failed
+            if self.migration_manager:
+                self.migration_manager.mark_failed(store_name, str(e))
+            logger.error(f"Migration failed: {e}")
+            raise
+
+    def get_migration_status(self, store_name: str) -> Optional[Dict[str, Any]]:
+        """
+        Get migration status for a sub store.
+
+        Args:
+            store_name: Sub store name
+
+        Returns:
+            Migration status dict or None if not found
+        """
+        if self.migration_manager:
+            return self.migration_manager.get_status(store_name)
+        return None
+
+    def list_sub_stores(self) -> List[str]:
+        """
+        List all registered sub stores.
+
+        Returns:
+            List of sub store names
+        """
+        return list(self.sub_stores.keys())