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

spatial_memory/__init__.py

@@ -1,6 +1,6 @@
 """Spatial Memory MCP Server - Vector-based semantic memory for LLMs."""
 
-__version__ = "1.5.3"
+__version__ = "1.6.0"
 __author__ = "arman-tech"
 
 # Re-export core components for convenience

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

spatial_memory/config.py

@@ -1,7 +1,7 @@
 """Configuration system for Spatial Memory MCP Server."""
 
 from pathlib import Path
-from typing import Any
+from typing import Any, Literal
 
 from pydantic import Field, SecretStr
 from pydantic_settings import BaseSettings
@@ -27,6 +27,14 @@ class Settings(BaseSettings):
         default=Path("./.spatial-memory"),
         description="Path to LanceDB storage directory",
     )
+    acknowledge_network_filesystem_risk: bool = Field(
+        default=False,
+        description=(
+            "Set to True to suppress warnings about network filesystem usage. "
+            "File-based locking does not work reliably on NFS/SMB/CIFS. "
+            "Only set this if you are certain only one instance will access the storage."
+        ),
+    )
 
     # Embedding Model
     embedding_model: str = Field(
@@ -133,7 +141,7 @@ class Settings(BaseSettings):
         ge=1,
         description="Re-rank top (refine_factor * limit) candidates for accuracy",
     )
-    index_type: str = Field(
+    index_type: Literal["IVF_PQ", "IVF_FLAT", "HNSW_SQ"] = Field(
         default="IVF_PQ",
         description="Vector index type: IVF_PQ, IVF_FLAT, or HNSW_SQ",
     )

spatial_memory/core/__init__.py

@@ -6,6 +6,10 @@ from spatial_memory.core.circuit_breaker import (
     CircuitState,
 )
 from spatial_memory.core.database import Database
+from spatial_memory.core.db_idempotency import IdempotencyManager, IdempotencyRecord
+from spatial_memory.core.db_indexes import IndexManager
+from spatial_memory.core.db_search import SearchManager
+from spatial_memory.core.db_versioning import VersionManager
 from spatial_memory.core.embeddings import EmbeddingService
 from spatial_memory.core.rate_limiter import RateLimiter
 from spatial_memory.core.errors import (
@@ -87,6 +91,11 @@ __all__ = [
     "FilterGroup",
     # Core services
     "Database",
+    "VersionManager",
+    "IndexManager",
+    "SearchManager",
+    "IdempotencyManager",
+    "IdempotencyRecord",
     "EmbeddingService",
     "RateLimiter",
     # Utilities

spatial_memory/core/connection_pool.py

@@ -62,6 +62,7 @@ class ConnectionPool:
         self,
         uri: str,
         read_consistency_interval_ms: int = 0,
+        validate_health: bool = True,
         **kwargs: Any,
     ) -> DBConnection:
         """Get existing connection or create new one.
@@ -69,6 +70,7 @@ class ConnectionPool:
         Args:
             uri: Database URI/path.
             read_consistency_interval_ms: Read consistency interval.
+            validate_health: Whether to validate cached connection health (default: True).
             **kwargs: Additional args for lancedb.connect().
 
         Returns:
@@ -77,9 +79,21 @@ class ConnectionPool:
         with self._lock:
             # Check if exists
             if uri in self._connections:
-                # Move to end (most recently used)
-                self._connections.move_to_end(uri)
-                return self._connections[uri]
+                conn = self._connections[uri]
+
+                # Optionally validate health of cached connection
+                if validate_health and not self._validate_connection(conn, uri):
+                    # Connection is stale, remove and create new
+                    logger.info(f"Stale connection detected for {uri}, recreating")
+                    self._connections.pop(uri)
+                    try:
+                        conn.close()
+                    except Exception:
+                        pass
+                else:
+                    # Connection is healthy, move to end (most recently used)
+                    self._connections.move_to_end(uri)
+                    return conn
 
             # Evict oldest if at capacity
             while len(self._connections) >= self._max_size:
@@ -98,6 +112,30 @@ class ConnectionPool:
             logger.debug(f"Created new connection for {uri} (pool size: {len(self._connections)})")
             return conn
 
+    def _validate_connection(self, conn: DBConnection, uri: str) -> bool:
+        """Validate that a cached connection is still healthy.
+
+        Performs a lightweight operation to verify the connection
+        is still usable.
+
+        Args:
+            conn: The connection to validate.
+            uri: URI for logging purposes.
+
+        Returns:
+            True if connection is healthy, False if stale.
+        """
+        try:
+            # List tables is a lightweight operation that validates connection
+            conn.table_names()
+            return True
+        except Exception as e:
+            if self.is_stale_connection_error(e):
+                logger.debug(f"Connection health check failed for {uri}: {e}")
+                return False
+            # Other errors might be transient, consider healthy
+            return True
+
     def _evict_oldest(self) -> None:
         """Evict the oldest (least recently used) connection."""
         if self._connections: