spatial-memory-mcp 1.0.3__py3-none-any.whl → 1.6.0__py3-none-any.whl

spatial_memory/services/lifecycle.py

@@ -16,9 +16,10 @@ from __future__ import annotations
 
 import logging
 from dataclasses import dataclass
-from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Literal
 
+import numpy as np
+
 from spatial_memory.core.errors import (
     ConsolidationError,
     DecayError,
@@ -26,6 +27,10 @@ from spatial_memory.core.errors import (
     ReinforcementError,
     ValidationError,
 )
+from spatial_memory.core.consolidation_strategies import (
+    ConsolidationAction,
+    get_strategy,
+)
 from spatial_memory.core.lifecycle_ops import (
     apply_decay,
     calculate_decay_factor,
@@ -34,9 +39,6 @@ from spatial_memory.core.lifecycle_ops import (
     extract_candidates,
     find_duplicate_groups,
     jaccard_similarity,
-    merge_memory_content,
-    merge_memory_metadata,
-    select_representative,
 )
 from spatial_memory.core.models import (
     ConsolidateResult,
@@ -53,7 +55,7 @@ from spatial_memory.core.models import (
 
 # Alias for backward compatibility
 ConsolidationGroupResult = ConsolidationGroup
-from spatial_memory.core.utils import to_naive_utc, utc_now_naive
+from spatial_memory.core.utils import to_naive_utc, utc_now, utc_now_naive
 from spatial_memory.core.validation import validate_namespace
 
 logger = logging.getLogger(__name__)
@@ -110,6 +112,7 @@ class LifecycleConfig:
     consolidate_min_threshold: float = 0.7
     consolidate_content_weight: float = 0.3
     consolidate_max_batch: int = 1000
+    consolidate_chunk_size: int = 200  # Process in smaller chunks for memory efficiency
 
 
 # =============================================================================
@@ -271,16 +274,31 @@ class LifecycleService:
                 total_decay_factor / len(all_memories) if all_memories else 1.0
             )
 
-            # Apply updates if not dry run
+            # Apply updates if not dry run - use batch update for efficiency
             failed_updates: list[str] = []
             if not dry_run and memories_to_update:
                 logger.info(f"Applying decay to {len(memories_to_update)} memories")
-                for memory_id, new_importance in memories_to_update:
-                    try:
-                        self._repo.update(memory_id, {"importance": new_importance})
-                    except Exception as e:
-                        logger.warning(f"Failed to update {memory_id}: {e}")
-                        failed_updates.append(memory_id)
+                # Convert to batch update format: list of (memory_id, updates_dict)
+                batch_updates = [
+                    (memory_id, {"importance": new_importance})
+                    for memory_id, new_importance in memories_to_update
+                ]
+                try:
+                    success_count, failed_ids = self._repo.update_batch(batch_updates)
+                    failed_updates = failed_ids
+                    logger.debug(
+                        f"Batch decay update: {success_count} succeeded, "
+                        f"{len(failed_ids)} failed"
+                    )
+                except Exception as e:
+                    logger.warning(f"Batch decay update failed: {e}")
+                    # Fall back to individual updates on batch failure
+                    for memory_id, new_importance in memories_to_update:
+                        try:
+                            self._repo.update(memory_id, {"importance": new_importance})
+                        except Exception as update_err:
+                            logger.warning(f"Failed to update {memory_id}: {update_err}")
+                            failed_updates.append(memory_id)
 
             return DecayResult(
                 memories_analyzed=len(all_memories),
@@ -345,14 +363,30 @@ class LifecycleService:
             failed_updates: list[str] = []
             total_boost = 0.0
 
-            for memory_id in memory_ids:
-                memory = self._repo.get(memory_id)
+            # Batch fetch all memories in a single query
+            memory_map = self._repo.get_batch(memory_ids)
 
-                if memory is None:
+            # Track which memories were not found
+            for memory_id in memory_ids:
+                if memory_id not in memory_map:
                     not_found.append(memory_id)
                     logger.warning(f"Memory not found for reinforcement: {memory_id}")
-                    continue
 
+            if not memory_map:
+                return ReinforceResult(
+                    memories_reinforced=0,
+                    avg_boost=0.0,
+                    reinforced=[],
+                    not_found=not_found,
+                    failed_updates=[],
+                )
+
+            # Calculate reinforcement for all found memories
+            now = utc_now()
+            batch_updates: list[tuple[str, dict[str, Any]]] = []
+            reinforcement_info: list[tuple[str, Memory, float, float]] = []  # id, memory, new_imp, boost
+
+            for memory_id, memory in memory_map.items():
                 # Calculate new importance
                 new_importance, actual_boost = calculate_reinforcement(
                     current_importance=memory.importance,
@@ -364,12 +398,36 @@ class LifecycleService:
                 # Prepare update
                 updates: dict[str, Any] = {"importance": new_importance}
                 if update_access:
-                    updates["last_accessed"] = datetime.now(timezone.utc)
+                    updates["last_accessed"] = now
                     updates["access_count"] = memory.access_count + 1
 
-                # Apply update
-                try:
-                    self._repo.update(memory_id, updates)
+                batch_updates.append((memory_id, updates))
+                reinforcement_info.append((memory_id, memory, new_importance, actual_boost))
+
+            # Apply all updates in a single batch operation
+            try:
+                success_count, batch_failed_ids = self._repo.update_batch(batch_updates)
+                failed_updates = batch_failed_ids
+                logger.debug(
+                    f"Batch reinforce update: {success_count} succeeded, "
+                    f"{len(batch_failed_ids)} failed"
+                )
+            except Exception as e:
+                logger.warning(f"Batch reinforce update failed: {e}, falling back to individual updates")
+                # Fall back to individual updates on batch failure
+                batch_failed_ids = []
+                for memory_id, updates in batch_updates:
+                    try:
+                        self._repo.update(memory_id, updates)
+                    except Exception as update_err:
+                        logger.warning(f"Failed to reinforce {memory_id}: {update_err}")
+                        batch_failed_ids.append(memory_id)
+                failed_updates = batch_failed_ids
+
+            # Build reinforced results for successful updates
+            failed_set = set(failed_updates)
+            for memory_id, memory, new_importance, actual_boost in reinforcement_info:
+                if memory_id not in failed_set:
                     reinforced.append(
                         ReinforcedMemory(
                             id=memory_id,
@@ -382,9 +440,6 @@ class LifecycleService:
                         )
                     )
                     total_boost += actual_boost
-                except Exception as e:
-                    logger.warning(f"Failed to reinforce {memory_id}: {e}")
-                    failed_updates.append(memory_id)
 
             avg_boost = total_boost / len(reinforced) if reinforced else 0.0
 
@@ -462,11 +517,17 @@ class LifecycleService:
                     extractions=[],
                 )
 
+            # Generate embeddings for all candidates in a single batch
+            # This is much more efficient than generating one at a time
+            candidate_texts = [c.content for c in candidates]
+            candidate_vectors = self._embeddings.embed_batch(candidate_texts)
+            logger.debug(f"Generated {len(candidate_vectors)} embeddings in batch")
+
             extractions: list[ExtractedMemory] = []
             memories_created = 0
             deduplicated_count = 0
 
-            for candidate in candidates:
+            for candidate, vector in zip(candidates, candidate_vectors):
                 extraction = ExtractedMemory(
                     content=candidate.content,
                     confidence=candidate.confidence,
@@ -477,22 +538,24 @@ class LifecycleService:
                     memory_id=None,
                 )
 
-                # Check for duplicates if requested
+                # Check for duplicates if requested (use pre-computed vector)
                 if deduplicate:
-                    is_duplicate = self._check_duplicate(
-                        candidate.content,
-                        effective_namespace,
-                        dedup_threshold,
+                    is_duplicate = self._check_duplicate_with_vector(
+                        content=candidate.content,
+                        vector=vector,
+                        namespace=effective_namespace,
+                        threshold=dedup_threshold,
                     )
                     if is_duplicate:
                         deduplicated_count += 1
                         extractions.append(extraction)
                         continue
 
-                # Store the extracted memory
+                # Store the extracted memory (use pre-computed vector)
                 try:
-                    memory_id = self._store_extracted_memory(
+                    memory_id = self._store_extracted_memory_with_vector(
                         content=candidate.content,
+                        vector=vector,
                         namespace=effective_namespace,
                         confidence=candidate.confidence,
                         pattern_type=candidate.pattern_type,
@@ -557,25 +620,20 @@ class LifecycleService:
                 "similarity_threshold must be between 0.7 and 0.99"
             )
 
-        if strategy not in (
-            "keep_newest",
-            "keep_oldest",
-            "keep_highest_importance",
-            "merge_content",
-        ):
-            raise ValidationError(f"Invalid strategy: {strategy}")
+        # Validate strategy using the strategy registry
+        try:
+            strategy_impl = get_strategy(strategy)
+        except ValueError as e:
+            raise ValidationError(str(e)) from e
 
         if max_groups < 1:
             raise ValidationError("max_groups must be at least 1")
 
         try:
-            # Fetch all memories in namespace with vectors
-            all_memories = self._repo.get_all(
-                namespace=namespace,
-                limit=self._config.consolidate_max_batch,
-            )
+            # Get total count to decide processing strategy
+            total_count = self._repo.count(namespace=namespace)
 
-            if len(all_memories) < 2:
+            if total_count < 2:
                 logger.info("Not enough memories for consolidation")
                 return ConsolidateResult(
                     groups_found=0,
@@ -585,9 +643,34 @@ class LifecycleService:
                     dry_run=dry_run,
                 )
 
-            # Build lookup structures
-            import numpy as np
+            # Use chunked processing for large namespaces to reduce memory usage
+            chunk_size = min(
+                self._config.consolidate_chunk_size,
+                self._config.consolidate_max_batch,
+            )
+            use_chunked = total_count > chunk_size
 
+            if use_chunked:
+                logger.info(
+                    f"Using chunked consolidation: {total_count} memories in "
+                    f"chunks of {chunk_size}"
+                )
+                return self._consolidate_chunked(
+                    namespace=namespace,
+                    similarity_threshold=similarity_threshold,
+                    strategy=strategy,
+                    dry_run=dry_run,
+                    max_groups=max_groups,
+                    chunk_size=chunk_size,
+                )
+
+            # Standard single-pass processing for smaller namespaces
+            all_memories = self._repo.get_all(
+                namespace=namespace,
+                limit=self._config.consolidate_max_batch,
+            )
+
+            # Build lookup structures
             memories = [m for m, _ in all_memories]
             vectors_list = [v for _, v in all_memories]
             vectors_array = np.array(vectors_list, dtype=np.float32)
@@ -636,12 +719,6 @@ class LifecycleService:
                 group_member_dicts = [memory_dicts[i] for i in member_indices]
                 group_member_ids = [str(d["id"]) for d in group_member_dicts]
 
-                # Select representative
-                rep_idx = select_representative(group_member_dicts, strategy)
-                rep_id = str(group_member_dicts[rep_idx]["id"])
-
-                action = "preview" if dry_run else "merged"
-
                 # Calculate average similarity for the group
                 total_sim = 0.0
                 pair_count = 0
@@ -665,81 +742,32 @@ class LifecycleService:
                         pair_count += 1
                 avg_similarity = total_sim / pair_count if pair_count > 0 else 0.0
 
-                if not dry_run:
-                    try:
-                        if strategy == "merge_content":
-                            # Create merged content in memory first (no DB write yet)
-                            group_contents = [str(d["content"]) for d in group_member_dicts]
-                            merged_content = merge_memory_content(group_contents)
-                            merged_meta = merge_memory_metadata(group_member_dicts)
-
-                            # Generate new embedding before any DB changes
-                            new_vector = self._embeddings.embed(merged_content)
-
-                            # Prepare merged memory object (not persisted yet)
-                            merged_memory = Memory(
-                                id="",  # Will be assigned
-                                content=merged_content,
-                                namespace=namespace,
-                                tags=merged_meta.get("tags", []),
-                                importance=merged_meta.get("importance", 0.5),
-                                source=MemorySource.CONSOLIDATED,
-                                metadata=merged_meta.get("metadata", {}),
-                            )
-
-                            # DELETE FIRST pattern: remove originals before adding merge
-                            # This prevents duplicates if add fails after delete
-                            deleted_ids: list[str] = []
-                            try:
-                                for mid in group_member_ids:
-                                    self._repo.delete(mid)
-                                    deleted_ids.append(mid)
-                                    memories_deleted += 1
-                            except Exception as del_err:
-                                # Partial delete - log for manual recovery
-                                logger.critical(
-                                    f"Partial consolidation failure: deleted {deleted_ids}, "
-                                    f"failed on {mid}: {del_err}. "
-                                    f"Remaining members may need manual cleanup: "
-                                    f"{[m for m in group_member_ids if m not in deleted_ids]}"
-                                )
-                                raise
-
-                            # Now add the merged memory
-                            try:
-                                new_id = self._repo.add(merged_memory, new_vector)
-                            except Exception as add_err:
-                                # CRITICAL: Originals deleted but merge failed
-                                # Log for manual recovery - data is in merged_content
-                                logger.critical(
-                                    f"Consolidation add failed after deleting originals. "
-                                    f"Deleted IDs: {deleted_ids}. "
-                                    f"Merged content (save for recovery): {merged_content[:500]}... "
-                                    f"Error: {add_err}"
-                                )
-                                raise
-
-                            rep_id = new_id
-                            memories_merged += 1
-                            action = "merged"
-                        else:
-                            # Keep representative, delete others
-                            for mid in group_member_ids:
-                                if mid != rep_id:
-                                    self._repo.delete(mid)
-                                    memories_deleted += 1
-                            memories_merged += 1
-                            action = "kept_representative"
-                    except Exception as e:
-                        logger.warning(f"Failed to consolidate group: {e}")
-                        action = "failed"
+                # Apply consolidation strategy
+                try:
+                    action_result: ConsolidationAction = strategy_impl.apply(
+                        members=group_member_dicts,
+                        member_ids=group_member_ids,
+                        namespace=namespace,
+                        repository=self._repo,
+                        embeddings=self._embeddings,
+                        dry_run=dry_run,
+                    )
+                    memories_merged += action_result.memories_merged
+                    memories_deleted += action_result.memories_deleted
+                except Exception as e:
+                    logger.warning(f"Failed to consolidate group: {e}")
+                    action_result = ConsolidationAction(
+                        representative_id=group_member_ids[0],
+                        deleted_ids=[],
+                        action="failed",
+                    )
 
                 result_groups.append(
                     ConsolidationGroupResult(
-                        representative_id=rep_id,
+                        representative_id=action_result.representative_id,
                         member_ids=group_member_ids,
                         avg_similarity=avg_similarity,
-                        action_taken=action,
+                        action_taken=action_result.action,
                     )
                 )
 
@@ -760,6 +788,171 @@ class LifecycleService:
     # Helper Methods
     # =========================================================================
 
+    def _consolidate_chunked(
+        self,
+        namespace: str,
+        similarity_threshold: float,
+        strategy: Literal[
+            "keep_newest", "keep_oldest", "keep_highest_importance", "merge_content"
+        ],
+        dry_run: bool,
+        max_groups: int,
+        chunk_size: int,
+    ) -> ConsolidateResult:
+        """Process consolidation in memory-efficient chunks.
+
+        Processes memories in smaller chunks to reduce peak memory usage.
+        Note: This may miss duplicates that span chunk boundaries.
+
+        Args:
+            namespace: Namespace to consolidate.
+            similarity_threshold: Minimum similarity for duplicates.
+            strategy: How to handle duplicates.
+            dry_run: Preview without changes.
+            max_groups: Maximum groups to process total.
+            chunk_size: Memories per chunk.
+
+        Returns:
+            Aggregated ConsolidateResult from all chunks.
+        """
+        all_groups: list[ConsolidationGroupResult] = []
+        total_merged = 0
+        total_deleted = 0
+        offset = 0
+        groups_remaining = max_groups
+
+        while groups_remaining > 0:
+            # Fetch chunk of memories
+            chunk_memories = self._repo.get_all(
+                namespace=namespace,
+                limit=chunk_size,
+            )
+
+            # Skip already processed memories by filtering by offset
+            # Note: This is a simplified approach - in production, you'd want
+            # to track processed IDs or use cursor-based pagination
+            if offset > 0:
+                # Re-fetch with offset simulation (get more and skip)
+                all_chunk = self._repo.get_all(
+                    namespace=namespace,
+                    limit=offset + chunk_size,
+                )
+                if len(all_chunk) <= offset:
+                    # No more memories to process
+                    break
+                chunk_memories = all_chunk[offset:offset + chunk_size]
+
+            if len(chunk_memories) < 2:
+                break
+
+            # Build lookup structures for this chunk
+            memories = [m for m, _ in chunk_memories]
+            vectors_list = [v for _, v in chunk_memories]
+            vectors_array = np.array(vectors_list, dtype=np.float32)
+            memory_ids = [m.id for m in memories]
+            contents = [m.content for m in memories]
+            memory_dicts: list[dict[str, Any]] = [
+                {
+                    "id": m.id,
+                    "content": m.content,
+                    "created_at": m.created_at,
+                    "last_accessed": m.last_accessed,
+                    "access_count": m.access_count,
+                    "importance": m.importance,
+                    "tags": list(m.tags),
+                }
+                for m in memories
+            ]
+
+            # Find duplicate groups in this chunk
+            group_indices = find_duplicate_groups(
+                memory_ids=memory_ids,
+                vectors=vectors_array,
+                contents=contents,
+                threshold=similarity_threshold,
+                content_weight=self._config.consolidate_content_weight,
+            )
+
+            # Limit groups for this chunk
+            group_indices = group_indices[:groups_remaining]
+
+            if not group_indices:
+                offset += len(chunk_memories)
+                continue
+
+            # Get strategy implementation
+            strategy_impl = get_strategy(strategy)
+
+            # Process groups in this chunk
+            for member_indices in group_indices:
+                group_member_dicts = [memory_dicts[i] for i in member_indices]
+                group_member_ids = [str(d["id"]) for d in group_member_dicts]
+
+                # Calculate average similarity
+                total_sim = 0.0
+                pair_count = 0
+                for i_idx, i in enumerate(member_indices):
+                    for j in member_indices[i_idx + 1:]:
+                        v1, v2 = vectors_array[i], vectors_array[j]
+                        dot = float(np.dot(v1, v2))
+                        norm1 = float(np.linalg.norm(v1))
+                        norm2 = float(np.linalg.norm(v2))
+                        if norm1 > 1e-10 and norm2 > 1e-10:
+                            v_sim = dot / (norm1 * norm2)
+                        else:
+                            v_sim = 0.0
+                        c_sim = jaccard_similarity(contents[i], contents[j])
+                        combined = combined_similarity(
+                            v_sim, c_sim, self._config.consolidate_content_weight
+                        )
+                        total_sim += combined
+                        pair_count += 1
+                avg_similarity = total_sim / pair_count if pair_count > 0 else 0.0
+
+                # Apply consolidation strategy
+                try:
+                    action_result: ConsolidationAction = strategy_impl.apply(
+                        members=group_member_dicts,
+                        member_ids=group_member_ids,
+                        namespace=namespace,
+                        repository=self._repo,
+                        embeddings=self._embeddings,
+                        dry_run=dry_run,
+                    )
+                    total_merged += action_result.memories_merged
+                    total_deleted += action_result.memories_deleted
+                except Exception as e:
+                    logger.warning(f"Failed to consolidate group: {e}")
+                    action_result = ConsolidationAction(
+                        representative_id=group_member_ids[0],
+                        deleted_ids=[],
+                        action="failed",
+                    )
+
+                all_groups.append(
+                    ConsolidationGroupResult(
+                        representative_id=action_result.representative_id,
+                        member_ids=group_member_ids,
+                        avg_similarity=avg_similarity,
+                        action_taken=action_result.action,
+                    )
+                )
+                groups_remaining -= 1
+
+            offset += len(chunk_memories)
+            logger.debug(
+                f"Processed chunk at offset {offset - len(chunk_memories)}, "
+                f"found {len(group_indices)} groups"
+            )
+
+        return ConsolidateResult(
+            groups_found=len(all_groups),
+            memories_merged=total_merged,
+            memories_deleted=total_deleted,
+            groups=all_groups,
+            dry_run=dry_run,
+        )
+
     def _check_duplicate(
         self,
         content: str,
@@ -843,3 +1036,85 @@ class LifecycleService:
         )
 
         return self._repo.add(memory, vector)
+
+    def _check_duplicate_with_vector(
+        self,
+        content: str,
+        vector: np.ndarray,
+        namespace: str,
+        threshold: float,
+    ) -> bool:
+        """Check if similar content already exists using pre-computed vector.
+
+        Args:
+            content: Content to check.
+            vector: Pre-computed embedding vector.
+            namespace: Namespace to search.
+            threshold: Similarity threshold.
+
+        Returns:
+            True if a similar memory exists.
+        """
+        try:
+            # Search for similar memories using pre-computed vector
+            results = self._repo.search(vector, limit=5, namespace=namespace)
+
+            for result in results:
+                # Check vector similarity
+                if result.similarity >= threshold:
+                    return True
+
+                # Also check content overlap
+                content_sim = jaccard_similarity(content, result.content)
+                combined = combined_similarity(
+                    result.similarity,
+                    content_sim,
+                    self._config.consolidate_content_weight,
+                )
+                if combined >= threshold:
+                    return True
+
+            return False
+
+        except Exception as e:
+            logger.warning(f"Duplicate check failed: {e}")
+            return False
+
+    def _store_extracted_memory_with_vector(
+        self,
+        content: str,
+        vector: np.ndarray,
+        namespace: str,
+        confidence: float,
+        pattern_type: str,
+    ) -> str:
+        """Store an extracted memory using pre-computed vector.
+
+        Args:
+            content: Memory content.
+            vector: Pre-computed embedding vector.
+            namespace: Target namespace.
+            confidence: Extraction confidence.
+            pattern_type: Type of pattern matched.
+
+        Returns:
+            The new memory's ID.
+        """
+        # Scale importance by confidence but keep lower than manual memories
+        importance = self._config.extract_default_importance * confidence
+
+        # Create memory
+        memory = Memory(
+            id="",  # Will be assigned
+            content=content,
+            namespace=namespace,
+            tags=[f"extracted-{pattern_type}"],
+            importance=importance,
+            source=MemorySource.EXTRACTED,
+            metadata={
+                "extraction_confidence": confidence,
+                "extraction_pattern": pattern_type,
+            },
+        )
+
+        return self._repo.add(memory, vector)