Check embedding size consistency when opening existing databases

test/test_embedding_consistency.py

@@ -0,0 +1,148 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Test embedding consistency checks between database and settings."""
+
+import pytest
+import tempfile
+import os
+from typeagent import create_conversation
+from typeagent.transcripts.transcript import TranscriptMessage, TranscriptMessageMeta
+from typeagent.knowpro.convsettings import ConversationSettings
+from typeagent.aitools.embeddings import AsyncEmbeddingModel
+from typeagent.storage.sqlite import SqliteStorageProvider
+
+
+@pytest.mark.asyncio
+async def test_embedding_size_mismatch_in_message_index():
+    """Test that opening a DB with mismatched embedding size raises an error."""
+    # Create a temporary database file
+    with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as tmp:
+        db_path = tmp.name
+
+    try:
+        # Create a conversation with test model (embedding size 3)
+        settings1 = ConversationSettings(
+            model=AsyncEmbeddingModel(embedding_size=3, model_name="test")
+        )
+        # Disable LLM knowledge extraction to avoid API key requirement
+        settings1.semantic_ref_index_settings.auto_extract_knowledge = False
+        conv1 = await create_conversation(
+            db_path, TranscriptMessage, settings=settings1
+        )
+
+        # Add some messages to populate the index
+        messages = [
+            TranscriptMessage(
+                text_chunks=["Hello world"],
+                metadata=TranscriptMessageMeta(speaker="Alice"),
+            )
+        ]
+        await conv1.add_messages_with_indexing(messages)
+        await conv1.storage_provider.close()
+
+        # Now try to open the same database with a different embedding size
+        # This should raise an error
+        settings2 = ConversationSettings(
+            model=AsyncEmbeddingModel(embedding_size=5, model_name="test")
+        )
+
+        with pytest.raises(ValueError, match="embedding size mismatch"):
+            provider = SqliteStorageProvider(
+                db_path=db_path,
+                message_type=TranscriptMessage,
+                message_text_index_settings=settings2.message_text_index_settings,
+                related_term_index_settings=settings2.related_term_index_settings,
+            )
+            await provider.close()
+
+    finally:
+        # Clean up the temporary database
+        if os.path.exists(db_path):
+            os.unlink(db_path)
+
+
+@pytest.mark.asyncio
+async def test_embedding_size_mismatch_in_related_terms():
+    """Test that opening a DB with mismatched embedding size in related terms raises an error."""
+    # Create a temporary database file
+    with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as tmp:
+        db_path = tmp.name
+
+    try:
+        # Create a conversation with default embedding size
+        settings1 = ConversationSettings(
+            model=AsyncEmbeddingModel(embedding_size=3, model_name="test")
+        )
+        # Disable LLM knowledge extraction to avoid API key requirement
+        settings1.semantic_ref_index_settings.auto_extract_knowledge = False
+        conv1 = await create_conversation(
+            db_path, TranscriptMessage, settings=settings1
+        )
+
+        # Add some messages to populate the related terms index
+        messages = [
+            TranscriptMessage(
+                text_chunks=["Apple is a fruit"],
+                metadata=TranscriptMessageMeta(speaker="Alice"),
+            )
+        ]
+        await conv1.add_messages_with_indexing(messages)
+        await conv1.storage_provider.close()
+
+        # Now try to open the same database with a different embedding size
+        # This should raise an error
+        settings2 = ConversationSettings(
+            model=AsyncEmbeddingModel(embedding_size=5, model_name="test")
+        )
+
+        with pytest.raises(ValueError, match="embedding size mismatch"):
+            provider = SqliteStorageProvider(
+                db_path=db_path,
+                message_type=TranscriptMessage,
+                message_text_index_settings=settings2.message_text_index_settings,
+                related_term_index_settings=settings2.related_term_index_settings,
+            )
+            await provider.close()
+
+    finally:
+        # Clean up the temporary database
+        if os.path.exists(db_path):
+            os.unlink(db_path)
+
+
+@pytest.mark.asyncio
+async def test_empty_db_no_error():
+    """Test that opening an empty DB doesn't raise an error regardless of embedding size."""
+    # Create a temporary database file
+    with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as tmp:
+        db_path = tmp.name
+
+    try:
+        # Create an empty database
+        settings1 = ConversationSettings(
+            model=AsyncEmbeddingModel(embedding_size=3, model_name="test")
+        )
+        # Disable LLM knowledge extraction to avoid API key requirement
+        settings1.semantic_ref_index_settings.auto_extract_knowledge = False
+        conv1 = await create_conversation(
+            db_path, TranscriptMessage, settings=settings1
+        )
+        await conv1.storage_provider.close()
+
+        # Open with different embedding size should work since DB is empty
+        settings2 = ConversationSettings(
+            model=AsyncEmbeddingModel(embedding_size=5, model_name="test")
+        )
+        provider = SqliteStorageProvider(
+            db_path=db_path,
+            message_type=TranscriptMessage,
+            message_text_index_settings=settings2.message_text_index_settings,
+            related_term_index_settings=settings2.related_term_index_settings,
+        )
+        await provider.close()
+
+    finally:
+        # Clean up the temporary database
+        if os.path.exists(db_path):
+            os.unlink(db_path)

typeagent/storage/sqlite/provider.py

@@ -32,7 +32,12 @@
 class SqliteStorageProvider[TMessage: interfaces.IMessage](
     interfaces.IStorageProvider[TMessage]
 ):
-    """SQLite-backed storage provider implementation."""
+    """SQLite-backed storage provider implementation.
+
+    This provider performs consistency checks on database initialization to ensure
+    that existing embeddings match the configured embedding_size. If a mismatch is
+    detected, a ValueError is raised with a descriptive error message.
+    """
 
     def __init__(
         self,
@@ -80,6 +85,9 @@ def __init__(
         # Initialize schema
         init_db_schema(self.db)
 
+        # Check embedding consistency before initializing indexes
+        self._check_embedding_consistency()
+
         # Initialize collections
         # Initialize message collection first
         self._message_collection = SqliteMessageCollection(self.db, self.message_type)
@@ -102,6 +110,53 @@ def __init__(
         # Connect message collection to message text index for automatic indexing
         self._message_collection.set_message_text_index(self._message_text_index)
 
+    def _check_embedding_consistency(self) -> None:
+        """Check that existing embeddings in the database match the expected embedding size.
+
+        This method is called during initialization to ensure that embeddings stored in the
+        database match the embedding_size specified in ConversationSettings. This prevents
+        runtime errors when trying to use embeddings of incompatible sizes.
+
+        Raises:
+            ValueError: If embeddings in the database don't match the expected size.
+        """
+        from .schema import deserialize_embedding
+
+        cursor = self.db.cursor()
+        expected_size = (
+            self.message_text_index_settings.embedding_index_settings.embedding_size
+        )
+
+        # Check message text index embeddings
+        cursor.execute("SELECT embedding FROM MessageTextIndex LIMIT 1")
+        row = cursor.fetchone()
+        if row and row[0]:
+            embedding = deserialize_embedding(row[0])
+            actual_size = len(embedding)
+            if actual_size != expected_size:
+                raise ValueError(
+                    f"Message text index embedding size mismatch: "
+                    f"database contains embeddings of size {actual_size}, "
+                    f"but ConversationSettings specifies embedding_size={expected_size}. "
+                    f"The database was likely created with a different embedding model. "
+                    f"Please use the same embedding model or create a new database."
+                )
+
+        # Check related terms fuzzy index embeddings
+        cursor.execute("SELECT term_embedding FROM RelatedTermsFuzzy LIMIT 1")
+        row = cursor.fetchone()
+        if row and row[0]:
+            embedding = deserialize_embedding(row[0])
+            actual_size = len(embedding)
+            if actual_size != expected_size:
+                raise ValueError(
+                    f"Related terms index embedding size mismatch: "
+                    f"database contains embeddings of size {actual_size}, "
+                    f"but ConversationSettings specifies embedding_size={expected_size}. "
+                    f"The database was likely created with a different embedding model. "
+                    f"Please use the same embedding model or create a new database."
+                )
+
     async def __aenter__(self) -> "SqliteStorageProvider[TMessage]":
         """Enter transaction context."""
         self.db.execute("BEGIN IMMEDIATE")