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

spatial_memory/__init__.py

@@ -1,97 +1,97 @@
-"""Spatial Memory MCP Server - Vector-based semantic memory for LLMs."""
-
-__version__ = "1.0.3"
-__author__ = "arman-tech"
-
-# Re-export core components for convenience
-# Adapters
-from spatial_memory.adapters.lancedb_repository import LanceDBMemoryRepository
-from spatial_memory.config import Settings, get_settings
-from spatial_memory.core import (
-    ClusterInfo,
-    ClusteringError,
-    ConfigurationError,
-    # Core services
-    Database,
-    EmbeddingError,
-    EmbeddingService,
-    Filter,
-    FilterGroup,
-    FilterOperator,
-    JourneyStep,
-    # Models
-    Memory,
-    MemoryNotFoundError,
-    MemoryResult,
-    MemorySource,
-    NamespaceNotFoundError,
-    # Errors
-    SpatialMemoryError,
-    StorageError,
-    ValidationError,
-    VisualizationCluster,
-    VisualizationData,
-    VisualizationEdge,
-    VisualizationError,
-    VisualizationNode,
-)
-
-# Server
-from spatial_memory.server import SpatialMemoryServer, create_server
-
-# Services
-from spatial_memory.services.memory import (
-    ForgetResult,
-    MemoryService,
-    NearbyResult,
-    RecallResult,
-    RememberBatchResult,
-    RememberResult,
-)
-
-__all__ = [
-    # Version info
-    "__version__",
-    "__author__",
-    # Configuration
-    "Settings",
-    "get_settings",
-    # Errors
-    "SpatialMemoryError",
-    "MemoryNotFoundError",
-    "NamespaceNotFoundError",
-    "EmbeddingError",
-    "StorageError",
-    "ValidationError",
-    "ConfigurationError",
-    "ClusteringError",
-    "VisualizationError",
-    # Models
-    "Memory",
-    "MemorySource",
-    "MemoryResult",
-    "ClusterInfo",
-    "JourneyStep",
-    "VisualizationNode",
-    "VisualizationEdge",
-    "VisualizationCluster",
-    "VisualizationData",
-    "Filter",
-    "FilterOperator",
-    "FilterGroup",
-    # Core services
-    "Database",
-    "EmbeddingService",
-    # Services
-    "MemoryService",
-    "RememberResult",
-    "RememberBatchResult",
-    "RecallResult",
-    "NearbyResult",
-    "ForgetResult",
-    # Adapters
-    "LanceDBMemoryRepository",
-    # Server
-    "SpatialMemoryServer",
-    "create_server",
-]
+"""Spatial Memory MCP Server - Vector-based semantic memory for LLMs."""
+
+__version__ = "1.6.0"
+__author__ = "arman-tech"
+
+# Re-export core components for convenience
+# Adapters
+from spatial_memory.adapters.lancedb_repository import LanceDBMemoryRepository
+from spatial_memory.config import Settings, get_settings
+from spatial_memory.core import (
+    ClusterInfo,
+    ClusteringError,
+    ConfigurationError,
+    # Core services
+    Database,
+    EmbeddingError,
+    EmbeddingService,
+    Filter,
+    FilterGroup,
+    FilterOperator,
+    JourneyStep,
+    # Models
+    Memory,
+    MemoryNotFoundError,
+    MemoryResult,
+    MemorySource,
+    NamespaceNotFoundError,
+    # Errors
+    SpatialMemoryError,
+    StorageError,
+    ValidationError,
+    VisualizationCluster,
+    VisualizationData,
+    VisualizationEdge,
+    VisualizationError,
+    VisualizationNode,
+)
+
+# Server
+from spatial_memory.server import SpatialMemoryServer, create_server
+
+# Services
+from spatial_memory.services.memory import (
+    ForgetResult,
+    MemoryService,
+    NearbyResult,
+    RecallResult,
+    RememberBatchResult,
+    RememberResult,
+)
+
+__all__ = [
+    # Version info
+    "__version__",
+    "__author__",
+    # Configuration
+    "Settings",
+    "get_settings",
+    # Errors
+    "SpatialMemoryError",
+    "MemoryNotFoundError",
+    "NamespaceNotFoundError",
+    "EmbeddingError",
+    "StorageError",
+    "ValidationError",
+    "ConfigurationError",
+    "ClusteringError",
+    "VisualizationError",
+    # Models
+    "Memory",
+    "MemorySource",
+    "MemoryResult",
+    "ClusterInfo",
+    "JourneyStep",
+    "VisualizationNode",
+    "VisualizationEdge",
+    "VisualizationCluster",
+    "VisualizationData",
+    "Filter",
+    "FilterOperator",
+    "FilterGroup",
+    # Core services
+    "Database",
+    "EmbeddingService",
+    # Services
+    "MemoryService",
+    "RememberResult",
+    "RememberBatchResult",
+    "RecallResult",
+    "NearbyResult",
+    "ForgetResult",
+    # Adapters
+    "LanceDBMemoryRepository",
+    # Server
+    "SpatialMemoryServer",
+    "create_server",
+]

spatial_memory/__main__.py

@@ -1,14 +1,253 @@
-"""Entry point for running the Spatial Memory MCP Server."""
+"""Entry point for running the Spatial Memory MCP Server and CLI commands."""
 
+from __future__ import annotations
+
+import argparse
 import asyncio
+import logging
+import sys
+from typing import NoReturn
+
+logger = logging.getLogger(__name__)
 
 
-def main() -> None:
+def run_server() -> None:
     """Run the Spatial Memory MCP Server."""
     from spatial_memory.server import main as server_main
 
     asyncio.run(server_main())
 
 
+def run_migrate(args: argparse.Namespace) -> int:
+    """Run database migrations.
+
+    Args:
+        args: Parsed command line arguments.
+
+    Returns:
+        Exit code (0 for success, 1 for error).
+    """
+    from spatial_memory.config import get_settings
+    from spatial_memory.core.database import Database
+    from spatial_memory.core.db_migrations import (
+        CURRENT_SCHEMA_VERSION,
+        MigrationManager,
+    )
+    from spatial_memory.core.embeddings import EmbeddingService
+
+    settings = get_settings()
+
+    # Set up logging based on verbosity
+    log_level = logging.DEBUG if args.verbose else logging.INFO
+    logging.basicConfig(
+        level=log_level,
+        format="%(levelname)s: %(message)s",
+    )
+
+    print(f"Spatial Memory Migration Tool")
+    print(f"Target schema version: {CURRENT_SCHEMA_VERSION}")
+    print(f"Database path: {settings.memory_path}")
+    print()
+
+    try:
+        # Create embedding service if needed for migrations
+        embeddings = None
+        if not args.dry_run:
+            # Only load embeddings for actual migrations (some may need re-embedding)
+            print("Loading embedding service...")
+            embeddings = EmbeddingService(
+                model_name=settings.embedding_model,
+                openai_api_key=settings.openai_api_key,
+                backend=settings.embedding_backend,
+            )
+
+        # Connect to database
+        print("Connecting to database...")
+        db = Database(
+            storage_path=settings.memory_path,
+            embedding_dim=embeddings.dimensions if embeddings else 384,
+            auto_create_indexes=settings.auto_create_indexes,
+        )
+        db.connect()
+
+        # Create migration manager
+        manager = MigrationManager(db, embeddings)
+        manager.register_builtin_migrations()
+
+        current_version = manager.get_current_version()
+        print(f"Current schema version: {current_version}")
+
+        if args.status:
+            # Just show status, don't run migrations
+            pending = manager.get_pending_migrations()
+            if pending:
+                print(f"\nPending migrations ({len(pending)}):")
+                for m in pending:
+                    print(f"  - {m.version}: {m.description}")
+            else:
+                print("\nNo pending migrations. Database is up to date.")
+
+            applied = manager.get_applied_migrations()
+            if applied:
+                print(f"\nApplied migrations ({len(applied)}):")
+                for m in applied:
+                    print(f"  - {m.version}: {m.description} (applied: {m.applied_at})")
+
+            db.close()
+            return 0
+
+        if args.rollback:
+            # Rollback to specified version
+            print(f"\nRolling back to version {args.rollback}...")
+            result = manager.rollback(args.rollback)
+
+            if result.errors:
+                print("\nRollback failed with errors:")
+                for error in result.errors:
+                    print(f"  - {error}")
+                db.close()
+                return 1
+
+            if result.migrations_applied:
+                print(f"\nRolled back migrations:")
+                for v in result.migrations_applied:
+                    print(f"  - {v}")
+                print(f"\nCurrent version: {result.current_version}")
+            else:
+                print("\nNo migrations to rollback.")
+
+            db.close()
+            return 0
+
+        # Run pending migrations
+        pending = manager.get_pending_migrations()
+        if not pending:
+            print("\nNo pending migrations. Database is up to date.")
+            db.close()
+            return 0
+
+        print(f"\nPending migrations ({len(pending)}):")
+        for m in pending:
+            print(f"  - {m.version}: {m.description}")
+
+        if args.dry_run:
+            print("\n[DRY RUN] Would apply the above migrations.")
+            print("Run without --dry-run to apply.")
+            db.close()
+            return 0
+
+        # Confirm before applying
+        if not args.yes:
+            print()
+            response = input("Apply these migrations? [y/N] ").strip().lower()
+            if response not in ("y", "yes"):
+                print("Aborted.")
+                db.close()
+                return 0
+
+        print("\nApplying migrations...")
+        result = manager.run_pending(dry_run=False)
+
+        if result.errors:
+            print("\nMigration failed with errors:")
+            for error in result.errors:
+                print(f"  - {error}")
+            print("\nSome migrations may have been applied. Check database state.")
+            db.close()
+            return 1
+
+        print(f"\nSuccessfully applied {len(result.migrations_applied)} migration(s):")
+        for v in result.migrations_applied:
+            print(f"  - {v}")
+        print(f"\nCurrent version: {result.current_version}")
+
+        db.close()
+        return 0
+
+    except Exception as e:
+        logger.error(f"Migration failed: {e}", exc_info=args.verbose)
+        print(f"\nError: {e}")
+        return 1
+
+
+def run_version() -> None:
+    """Print version information."""
+    from spatial_memory import __version__
+
+    print(f"spatial-memory {__version__}")
+
+
+def main() -> NoReturn:
+    """Main entry point with subcommand support."""
+    parser = argparse.ArgumentParser(
+        prog="spatial-memory",
+        description="Spatial Memory MCP Server and CLI tools",
+    )
+    parser.add_argument(
+        "--version", "-V",
+        action="store_true",
+        help="Show version and exit",
+    )
+
+    subparsers = parser.add_subparsers(
+        dest="command",
+        title="commands",
+        description="Available commands",
+    )
+
+    # Server command (default)
+    server_parser = subparsers.add_parser(
+        "serve",
+        help="Start the MCP server (default if no command given)",
+    )
+
+    # Migrate command
+    migrate_parser = subparsers.add_parser(
+        "migrate",
+        help="Run database migrations",
+    )
+    migrate_parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Preview migrations without applying",
+    )
+    migrate_parser.add_argument(
+        "--status",
+        action="store_true",
+        help="Show migration status and exit",
+    )
+    migrate_parser.add_argument(
+        "--rollback",
+        metavar="VERSION",
+        help="Rollback to specified version (e.g., 1.0.0)",
+    )
+    migrate_parser.add_argument(
+        "-y", "--yes",
+        action="store_true",
+        help="Skip confirmation prompt",
+    )
+    migrate_parser.add_argument(
+        "-v", "--verbose",
+        action="store_true",
+        help="Enable verbose output",
+    )
+
+    args = parser.parse_args()
+
+    if args.version:
+        run_version()
+        sys.exit(0)
+
+    if args.command == "migrate":
+        sys.exit(run_migrate(args))
+    elif args.command == "serve" or args.command is None:
+        # Default to running the server
+        run_server()
+        sys.exit(0)
+    else:
+        parser.print_help()
+        sys.exit(1)
+
+
 if __name__ == "__main__":
     main()

spatial_memory/adapters/lancedb_repository.py

@@ -178,7 +178,7 @@ class LanceDBMemoryRepository:
             logger.error(f"Unexpected error in delete: {e}")
             raise StorageError(f"Failed to delete memory: {e}") from e
 
-    def delete_batch(self, memory_ids: list[str]) -> int:
+    def delete_batch(self, memory_ids: list[str]) -> tuple[int, list[str]]:
         """Delete multiple memories atomically.
 
         Delegates to Database.delete_batch for proper encapsulation.
@@ -187,7 +187,9 @@ class LanceDBMemoryRepository:
             memory_ids: List of memory UUIDs to delete.
 
         Returns:
-            Number of memories actually deleted.
+            Tuple of (count_deleted, list_of_deleted_ids) where:
+                - count_deleted: Number of memories actually deleted
+                - list_of_deleted_ids: IDs that were actually deleted
 
         Raises:
             ValidationError: If any memory_id is invalid.
@@ -206,6 +208,7 @@ class LanceDBMemoryRepository:
         query_vector: np.ndarray,
         limit: int = 5,
         namespace: str | None = None,
+        include_vector: bool = False,
     ) -> list[MemoryResult]:
         """Search for similar memories by vector.
 
@@ -213,16 +216,24 @@ class LanceDBMemoryRepository:
             query_vector: Query embedding vector.
             limit: Maximum number of results.
             namespace: Filter to specific namespace.
+            include_vector: Whether to include embedding vectors in results.
+                Defaults to False to reduce response size.
 
         Returns:
             List of MemoryResult objects with similarity scores.
+            If include_vector=True, each result includes its embedding vector.
 
         Raises:
             ValidationError: If input validation fails.
             StorageError: If database operation fails.
         """
         try:
-            results = self._db.vector_search(query_vector, limit=limit, namespace=namespace)
+            results = self._db.vector_search(
+                query_vector,
+                limit=limit,
+                namespace=namespace,
+                include_vector=include_vector,
+            )
             return [self._record_to_memory_result(r) for r in results]
         except (ValidationError, StorageError):
             raise
@@ -292,6 +303,53 @@ class LanceDBMemoryRepository:
             logger.error(f"Unexpected error in update: {e}")
             raise StorageError(f"Failed to update memory: {e}") from e
 
+    def get_batch(self, memory_ids: list[str]) -> dict[str, Memory]:
+        """Get multiple memories by ID in a single query.
+
+        Args:
+            memory_ids: List of memory UUIDs to retrieve.
+
+        Returns:
+            Dict mapping memory_id to Memory object. Missing IDs are not included.
+
+        Raises:
+            ValidationError: If any memory_id format is invalid.
+            StorageError: If database operation fails.
+        """
+        try:
+            raw_results = self._db.get_batch(memory_ids)
+            result: dict[str, Memory] = {}
+            for memory_id, record in raw_results.items():
+                result[memory_id] = self._record_to_memory(record)
+            return result
+        except (ValidationError, StorageError):
+            raise
+        except Exception as e:
+            logger.error(f"Unexpected error in get_batch: {e}")
+            raise StorageError(f"Failed to batch get memories: {e}") from e
+
+    def update_batch(
+        self, updates: list[tuple[str, dict[str, Any]]]
+    ) -> tuple[int, list[str]]:
+        """Update multiple memories in a single batch operation.
+
+        Args:
+            updates: List of (memory_id, updates_dict) tuples.
+
+        Returns:
+            Tuple of (success_count, list of failed memory_ids).
+
+        Raises:
+            StorageError: If database operation fails completely.
+        """
+        try:
+            return self._db.update_batch(updates)
+        except StorageError:
+            raise
+        except Exception as e:
+            logger.error(f"Unexpected error in update_batch: {e}")
+            raise StorageError(f"Failed to batch update memories: {e}") from e
+
     def count(self, namespace: str | None = None) -> int:
         """Count memories.
 
@@ -533,6 +591,13 @@ class LanceDBMemoryRepository:
         similarity = record.get("similarity", 0.0)
         similarity = max(0.0, min(1.0, similarity))
 
+        # Include vector if present in record (when include_vector=True in search)
+        vector = None
+        if "vector" in record and record["vector"] is not None:
+            # Convert to list for JSON serialization
+            vec = record["vector"]
+            vector = vec.tolist() if hasattr(vec, "tolist") else list(vec)
+
         return MemoryResult(
             id=record["id"],
             content=record["content"],
@@ -542,6 +607,7 @@ class LanceDBMemoryRepository:
             importance=record["importance"],
             created_at=record["created_at"],
             metadata=record.get("metadata", {}),
+            vector=vector,
         )
 
     # ========================================================================
@@ -586,6 +652,7 @@ class LanceDBMemoryRepository:
         query_vectors: list[np.ndarray],
         limit_per_query: int = 3,
         namespace: str | None = None,
+        include_vector: bool = False,
     ) -> list[list[dict[str, Any]]]:
         """Search for memories near multiple query points.
 
@@ -597,10 +664,13 @@ class LanceDBMemoryRepository:
             query_vectors: List of query embedding vectors.
             limit_per_query: Maximum results per query vector.
             namespace: Filter to specific namespace.
+            include_vector: Whether to include embedding vectors in results.
+                Defaults to False to reduce response size.
 
         Returns:
             List of result lists (one per query vector). Each result
             is a dict containing memory fields and similarity score.
+            If include_vector=True, each dict includes the 'vector' field.
 
         Raises:
             ValidationError: If input validation fails.
@@ -611,8 +681,7 @@ class LanceDBMemoryRepository:
                 query_vectors=query_vectors,
                 limit_per_query=limit_per_query,
                 namespace=namespace,
-                parallel=len(query_vectors) > 3,  # Use parallel for larger batches
-                max_workers=4,
+                include_vector=include_vector,
             )
         except (ValidationError, StorageError):
             raise