fix: migrate mental_models.embedding dimension alongside memory_units (#526)

nicoloboschi · web-flow · commit 7accac94b20e · 2026-03-09T12:23:50.000+01:00
ensure_embedding_dimension() now also checks and migrates mental_models.embedding, fixing silent failures when changing embedding model dimensions. Extracted shared per-table logic into _migrate_table_embedding_dimension() to avoid duplication. Adds test coverage for the mental_models dimension migration path. Fixes #523
diff --git a/hindsight-api/hindsight_api/migrations.py b/hindsight-api/hindsight_api/migrations.py
@@ -25,7 +25,7 @@
 from alembic import command
 from alembic.config import Config
 from alembic.script.revision import ResolutionError
-from sqlalchemy import create_engine, text
+from sqlalchemy import Connection, create_engine, text
 
 from .utils import mask_network_location
 
@@ -471,17 +471,135 @@ def check_migration_status(
         return None, None
 
 
+def _migrate_table_embedding_dimension(
+    conn: Connection,
+    schema_name: str,
+    table_name: str,
+    required_dimension: int,
+    vector_ext: str,
+) -> None:
+    """
+    Migrate the embedding column of a single table to the required dimension.
+
+    - If dimensions match: no action needed
+    - If dimensions differ and table is empty: ALTER COLUMN to new dimension
+    - If dimensions differ and table has data: raise error with migration guidance
+    """
+    current_dim = conn.execute(
+        text("""
+            SELECT atttypmod
+            FROM pg_attribute a
+            JOIN pg_class c ON a.attrelid = c.oid
+            JOIN pg_namespace n ON c.relnamespace = n.oid
+            WHERE n.nspname = :schema
+              AND c.relname = :table
+              AND a.attname = 'embedding'
+        """),
+        {"schema": schema_name, "table": table_name},
+    ).scalar()
+
+    if current_dim is None:
+        logger.debug(f"No embedding column found on {table_name}, skipping")
+        return
+
+    if current_dim == required_dimension:
+        logger.debug(f"Embedding dimension OK for {table_name}: {current_dim}")
+        return
+
+    logger.info(
+        f"Embedding dimension mismatch on {table_name}: database has {current_dim}, model requires {required_dimension}"
+    )
+
+    row_count = conn.execute(
+        text(f"SELECT COUNT(*) FROM {schema_name}.{table_name} WHERE embedding IS NOT NULL")
+    ).scalar()
+
+    if row_count > 0:
+        raise RuntimeError(
+            f"Cannot change embedding dimension from {current_dim} to {required_dimension}: "
+            f"{table_name} table contains {row_count} rows with embeddings. "
+            f"To change dimensions, you must either:\n"
+            f"  1. Re-embed all data: DELETE FROM {schema_name}.{table_name}; then restart\n"
+            f"  2. Use a model with {current_dim}-dimensional embeddings"
+        )
+
+    logger.info(f"Altering {table_name}.embedding column dimension from {current_dim} to {required_dimension}")
+
+    # Drop existing vector index (works for both HNSW and vchordrq)
+    conn.execute(
+        text(f"""
+            DO $$
+            DECLARE idx_name TEXT;
+            BEGIN
+                FOR idx_name IN
+                    SELECT indexname FROM pg_indexes
+                    WHERE schemaname = '{schema_name}'
+                      AND tablename = '{table_name}'
+                      AND (indexdef LIKE '%hnsw%' OR indexdef LIKE '%vchordrq%' OR indexdef LIKE '%diskann%')
+                      AND indexdef LIKE '%embedding%'
+                LOOP
+                    EXECUTE 'DROP INDEX IF EXISTS {schema_name}.' || idx_name;
+                END LOOP;
+            END $$;
+        """)
+    )
+
+    conn.execute(
+        text(f"ALTER TABLE {schema_name}.{table_name} ALTER COLUMN embedding TYPE vector({required_dimension})")
+    )
+    conn.commit()
+
+    # Recreate index with appropriate type based on detected extension
+    if vector_ext == "pgvectorscale":
+        conn.execute(
+            text(f"""
+                CREATE INDEX IF NOT EXISTS idx_{table_name}_embedding_diskann
+                ON {schema_name}.{table_name}
+                USING diskann (embedding vector_cosine_ops)
+                WITH (num_neighbors = 50)
+            """)
+        )
+        logger.info(f"Created DiskANN index on {table_name} for {required_dimension}-dimensional embeddings")
+    elif vector_ext == "vchord":
+        conn.execute(
+            text(f"""
+                CREATE INDEX IF NOT EXISTS idx_{table_name}_embedding_vchordrq
+                ON {schema_name}.{table_name}
+                USING vchordrq (embedding vector_l2_ops)
+            """)
+        )
+        logger.info(f"Created vchordrq index on {table_name} for {required_dimension}-dimensional embeddings")
+    else:  # pgvector
+        if required_dimension > 2000:
+            raise RuntimeError(
+                f"Embedding dimension {required_dimension} exceeds pgvector HNSW index limit of 2000. "
+                f"Use an embedding model with <= 2000 dimensions, or switch to a vector extension "
+                f"that supports higher dimensions (e.g., pgvectorscale/DiskANN)."
+            )
+        conn.execute(
+            text(f"""
+                CREATE INDEX IF NOT EXISTS idx_{table_name}_embedding_hnsw
+                ON {schema_name}.{table_name}
+                USING hnsw (embedding vector_cosine_ops)
+                WITH (m = 16, ef_construction = 64)
+            """)
+        )
+        logger.info(f"Created HNSW index on {table_name} for {required_dimension}-dimensional embeddings")
+    conn.commit()
+
+    logger.info(f"Successfully changed {table_name}.embedding dimension to {required_dimension}")
+
+
 def ensure_embedding_dimension(
     database_url: str,
     required_dimension: int,
     schema: str | None = None,
     vector_extension: str = "pgvector",
 ) -> None:
     """
-    Ensure the embedding column dimension matches the model's dimension.
+    Ensure the embedding column dimension matches the model's dimension for all tables.
 
-    This function checks the current vector column dimension in the database
-    and adjusts it if necessary:
+    Checks and adjusts memory_units.embedding and mental_models.embedding:
     - If dimensions match: no action needed
     - If dimensions differ and table is empty: ALTER COLUMN to new dimension
     - If dimensions differ and table has data: raise error with migration guidance
@@ -499,7 +617,7 @@ def ensure_embedding_dimension(
 
     engine = create_engine(database_url)
     with engine.connect() as conn:
-        # Check if memory_units table exists
+        # Check if memory_units table exists (proxy for schema being initialized)
         table_exists = conn.execute(
             text("""
                 SELECT EXISTS (
@@ -518,117 +636,8 @@ def ensure_embedding_dimension(
         vector_ext = _detect_vector_extension(conn, vector_extension)
         logger.info(f"Using vector extension: {vector_ext}")
 
-        # Get current column dimension from pg_attribute
-        # pgvector stores dimension in atttypmod
-        current_dim = conn.execute(
-            text("""
-                SELECT atttypmod
-                FROM pg_attribute a
-                JOIN pg_class c ON a.attrelid = c.oid
-                JOIN pg_namespace n ON c.relnamespace = n.oid
-                WHERE n.nspname = :schema
-                  AND c.relname = 'memory_units'
-                  AND a.attname = 'embedding'
-            """),
-            {"schema": schema_name},
-        ).scalar()
-
-        if current_dim is None:
-            logger.warning("Could not determine current embedding dimension, skipping check")
-            return
-
-        # pgvector stores dimension directly in atttypmod (no offset like other types)
-        current_dimension = current_dim
-
-        if current_dimension == required_dimension:
-            logger.debug(f"Embedding dimension OK: {current_dimension}")
-            return
-
-        logger.info(
-            f"Embedding dimension mismatch: database has {current_dimension}, model requires {required_dimension}"
-        )
-
-        # Check if table has data
-        row_count = conn.execute(
-            text(f"SELECT COUNT(*) FROM {schema_name}.memory_units WHERE embedding IS NOT NULL")
-        ).scalar()
-
-        if row_count > 0:
-            raise RuntimeError(
-                f"Cannot change embedding dimension from {current_dimension} to {required_dimension}: "
-                f"memory_units table contains {row_count} rows with embeddings. "
-                f"To change dimensions, you must either:\n"
-                f"  1. Re-embed all data: DELETE FROM {schema_name}.memory_units; then restart\n"
-                f"  2. Use a model with {current_dimension}-dimensional embeddings"
-            )
-
-        # Table is empty, safe to alter column
-        logger.info(f"Altering embedding column dimension from {current_dimension} to {required_dimension}")
-
-        # Drop existing vector index (works for both HNSW and vchordrq)
-        conn.execute(
-            text(f"""
-                DO $$
-                DECLARE idx_name TEXT;
-                BEGIN
-                    FOR idx_name IN
-                        SELECT indexname FROM pg_indexes
-                        WHERE schemaname = '{schema_name}'
-                          AND tablename = 'memory_units'
-                          AND (indexdef LIKE '%hnsw%' OR indexdef LIKE '%vchordrq%')
-                          AND indexdef LIKE '%embedding%'
-                    LOOP
-                        EXECUTE 'DROP INDEX IF EXISTS {schema_name}.' || idx_name;
-                    END LOOP;
-                END $$;
-            """)
-        )
-
-        # Alter the column type
-        conn.execute(
-            text(f"ALTER TABLE {schema_name}.memory_units ALTER COLUMN embedding TYPE vector({required_dimension})")
-        )
-        conn.commit()
-
-        # Recreate index with appropriate type based on detected extension
-        if vector_ext == "pgvectorscale":
-            conn.execute(
-                text(f"""
-                    CREATE INDEX IF NOT EXISTS idx_memory_units_embedding_diskann
-                    ON {schema_name}.memory_units
-                    USING diskann (embedding vector_cosine_ops)
-                    WITH (num_neighbors = 50)
-                """)
-            )
-            logger.info(f"Created DiskANN index for {required_dimension}-dimensional embeddings")
-        elif vector_ext == "vchord":
-            conn.execute(
-                text(f"""
-                    CREATE INDEX IF NOT EXISTS idx_memory_units_embedding_vchordrq
-                    ON {schema_name}.memory_units
-                    USING vchordrq (embedding vector_l2_ops)
-                """)
-            )
-            logger.info(f"Created vchordrq index for {required_dimension}-dimensional embeddings")
-        else:  # pgvector
-            if required_dimension > 2000:
-                raise RuntimeError(
-                    f"Embedding dimension {required_dimension} exceeds pgvector HNSW index limit of 2000. "
-                    f"Use an embedding model with <= 2000 dimensions, or switch to a vector extension "
-                    f"that supports higher dimensions (e.g., pgvectorscale/DiskANN)."
-                )
-            conn.execute(
-                text(f"""
-                    CREATE INDEX IF NOT EXISTS idx_memory_units_embedding_hnsw
-                    ON {schema_name}.memory_units
-                    USING hnsw (embedding vector_cosine_ops)
-                    WITH (m = 16, ef_construction = 64)
-                """)
-            )
-            logger.info(f"Created HNSW index for {required_dimension}-dimensional embeddings")
-        conn.commit()
-
-        logger.info(f"Successfully changed embedding dimension to {required_dimension}")
+        _migrate_table_embedding_dimension(conn, schema_name, "memory_units", required_dimension, vector_ext)
+        _migrate_table_embedding_dimension(conn, schema_name, "mental_models", required_dimension, vector_ext)
 
 
 def ensure_vector_extension(
diff --git a/hindsight-api/tests/test_custom_embedding_dimension.py b/hindsight-api/tests/test_custom_embedding_dimension.py
@@ -75,7 +75,7 @@ def drop_schema(db_url: str, schema_name: str):
         conn.commit()
 
 
-def get_column_dimension(db_url: str, schema: str = "public") -> int | None:
+def get_column_dimension(db_url: str, schema: str = "public", table: str = "memory_units") -> int | None:
     """Get the current embedding column dimension from the database."""
     engine = create_engine(db_url)
     with engine.connect() as conn:
@@ -86,10 +86,10 @@ def get_column_dimension(db_url: str, schema: str = "public") -> int | None:
                 JOIN pg_class c ON a.attrelid = c.oid
                 JOIN pg_namespace n ON c.relnamespace = n.oid
                 WHERE n.nspname = :schema
-                  AND c.relname = 'memory_units'
+                  AND c.relname = :table
                   AND a.attname = 'embedding'
             """),
-            {"schema": schema},
+            {"schema": schema, "table": table},
         ).scalar()
         return result
 
@@ -125,6 +125,38 @@ def clear_embeddings(db_url: str, schema: str):
         conn.commit()
 
 
+def insert_test_mental_model_embedding(db_url: str, schema: str, dimension: int):
+    """Insert a test mental model row with a dummy embedding."""
+    engine = create_engine(db_url)
+    embedding = [0.1] * dimension
+    embedding_str = "[" + ",".join(str(x) for x in embedding) + "]"
+
+    with engine.connect() as conn:
+        # Ensure test bank exists
+        conn.execute(
+            text(f"""
+                INSERT INTO {schema}.banks (bank_id, name)
+                VALUES ('test-bank-mm', 'Test Bank')
+                ON CONFLICT (bank_id) DO NOTHING
+            """)
+        )
+        conn.execute(
+            text(f"""
+                INSERT INTO {schema}.mental_models (bank_id, name, source_query, content, embedding)
+                VALUES ('test-bank-mm', 'test model', 'test query', 'test content', '{embedding_str}'::vector)
+            """)
+        )
+        conn.commit()
+
+
+def clear_mental_model_embeddings(db_url: str, schema: str):
+    """Clear all rows from mental_models."""
+    engine = create_engine(db_url)
+    with engine.connect() as conn:
+        conn.execute(text(f"DELETE FROM {schema}.mental_models"))
+        conn.commit()
+
+
 # =============================================================================
 # Embedding Dimension Tests (Local Embeddings)
 # =============================================================================
@@ -199,6 +231,49 @@ def test_dimension_change_blocked_with_data(self, dimension_test_schema):
         # Cleanup
         clear_embeddings(db_url, schema)
 
+    def test_mental_models_dimension_matches_no_change(self, dimension_test_schema):
+        """When mental_models dimension matches, no changes should be made."""
+        db_url, schema = dimension_test_schema
+
+        initial_dim = get_column_dimension(db_url, schema, table="mental_models")
+        assert initial_dim == 384, f"Expected 384, got {initial_dim}"
+
+        ensure_embedding_dimension(db_url, 384, schema=schema)
+
+        assert get_column_dimension(db_url, schema, table="mental_models") == 384
+
+    def test_mental_models_dimension_change_empty_table(self, dimension_test_schema):
+        """When mental_models is empty, dimension can be changed."""
+        db_url, schema = dimension_test_schema
+
+        clear_mental_model_embeddings(db_url, schema)
+
+        ensure_embedding_dimension(db_url, 768, schema=schema)
+
+        assert get_column_dimension(db_url, schema, table="mental_models") == 768
+
+        # Change back for other tests
+        ensure_embedding_dimension(db_url, 384, schema=schema)
+        assert get_column_dimension(db_url, schema, table="mental_models") == 384
+
+    def test_mental_models_dimension_change_blocked_with_data(self, dimension_test_schema):
+        """When mental_models has data, dimension change should be blocked."""
+        db_url, schema = dimension_test_schema
+
+        clear_mental_model_embeddings(db_url, schema)
+        insert_test_mental_model_embedding(db_url, schema, 384)
+
+        with pytest.raises(RuntimeError) as exc_info:
+            ensure_embedding_dimension(db_url, 768, schema=schema)
+
+        assert "Cannot change embedding dimension" in str(exc_info.value)
+        assert "mental_models" in str(exc_info.value)
+
+        assert get_column_dimension(db_url, schema, table="mental_models") == 384
+
+        # Cleanup
+        clear_mental_model_embeddings(db_url, schema)
+
     def test_local_embeddings_dimension_detection(self, embeddings):
         """Test that LocalSTEmbeddings correctly detects dimension."""
         # Initialize embeddings if not already done
