diff --git a/src/typeagent/storage/sqlite/related_terms_index.py b/src/typeagent/storage/sqlite/related_terms_index.py index d918fd6c..81550457 100644 --- a/src/typeagent/storage/sqlite/related_terms_index.py +++ b/src/typeagent/storage/sqlite/related_terms_index.py @@ -159,6 +159,7 @@ def __init__(self, db: sqlite3.Connection, settings: TextEmbeddingIndexSettings) for term, blob in rows: assert blob is not None, term self._terms_list.append(term) + self._added_terms.add(term) embeddings.append(deserialize_embedding(blob)) # Bulk add embeddings to VectorBase self._vector_base.add_embeddings(None, np.array(embeddings)) @@ -202,6 +203,9 @@ async def remove_term(self, term: str) -> None: async def clear(self) -> None: cursor = self.db.cursor() cursor.execute("DELETE FROM RelatedTermsFuzzy") + self._vector_base.clear() + self._terms_list.clear() + self._added_terms.clear() async def size(self) -> int: return self._size() @@ -292,6 +296,7 @@ async def deserialize(self, data: TextEmbeddingIndexData) -> None: cursor.execute("DELETE FROM RelatedTermsFuzzy") # Clear local mappings + self._vector_base.clear() self._terms_list.clear() self._added_terms.clear() diff --git a/tests/test_sqlite_indexes.py b/tests/test_sqlite_indexes.py index c7c9af01..30875859 100644 --- a/tests/test_sqlite_indexes.py +++ b/tests/test_sqlite_indexes.py @@ -442,6 +442,65 @@ async def test_fuzzy_serialize( embeddings = data["embeddings"] assert embeddings is not None + @pytest.mark.asyncio + async def test_fuzzy_clear_resets_lookup_state( + self, + sqlite_db: sqlite3.Connection, + embedding_settings: TextEmbeddingIndexSettings, + ): + """clear() forgets cleared terms, so re-adding them stores them again.""" + index = SqliteRelatedTermsFuzzy(sqlite_db, embedding_settings) + await index.add_terms(["apple"]) + + await index.clear() + assert await index.lookup_term("apple", min_score=0.0) == [] + assert index.serialize()["textItems"] == [] + + await index.add_terms(["apple"]) + assert await index.get_terms() == ["apple"] + assert [t.text for t in await index.lookup_term("apple", min_score=0.0)] == [ + "apple" + ] + + @pytest.mark.asyncio + async def test_fuzzy_empty_deserialize_resets_lookup_state( + self, + sqlite_db: sqlite3.Connection, + embedding_settings: TextEmbeddingIndexSettings, + ): + """Terms added after deserializing empty data map to the right embeddings.""" + index = SqliteRelatedTermsFuzzy(sqlite_db, embedding_settings) + await index.add_terms(["pear"]) + + await index.deserialize({"textItems": [], "embeddings": None}) + await index.add_terms(["kiwi"]) + + results = await index.lookup_term("kiwi", max_hits=1, min_score=0.0) + assert [t.text for t in results] == ["kiwi"] + assert await index.get_terms() == ["kiwi"] + + @pytest.mark.asyncio + async def test_fuzzy_reopen_does_not_duplicate_existing_terms( + self, + sqlite_db: sqlite3.Connection, + embedding_settings: TextEmbeddingIndexSettings, + temp_db_path: str, + ): + """Terms loaded from an existing database count as already added.""" + index = SqliteRelatedTermsFuzzy(sqlite_db, embedding_settings) + await index.add_terms(["plum"]) + sqlite_db.commit() + + reopened_db = sqlite3.connect(temp_db_path) + try: + reopened = SqliteRelatedTermsFuzzy(reopened_db, embedding_settings) + await reopened.add_terms(["plum"]) + results = await reopened.lookup_term("plum", min_score=0.0) + assert [t.text for t in results] == ["plum"] + assert reopened.serialize()["textItems"] == ["plum"] + finally: + reopened_db.close() + class TestSqliteRelatedTermsIndex: """Test SqliteRelatedTermsIndex combined functionality."""