diff --git a/src/typeagent/storage/memory/related_terms_index.py b/src/typeagent/storage/memory/related_terms_index.py index ec074bf9..6e800fe4 100644 --- a/src/typeagent/storage/memory/related_terms_index.py +++ b/src/typeagent/storage/memory/related_terms_index.py @@ -79,19 +79,18 @@ async def serialize(self) -> TermToRelatedTermsData: return TermToRelatedTermsData(relatedTerms=related_terms) async def deserialize(self, data: TermToRelatedTermsData | None) -> None: + # Parse everything before clearing, so malformed data leaves the + # existing aliases in place. + items: list[tuple[str, list[Term]]] = [] + if data is not None: + for item in data.get("relatedTerms") or []: + related_terms = [ + Term(term_data["text"], weight=term_data.get("weight")) + for term_data in item["relatedTerms"] + ] + items.append((item["termText"], related_terms)) self.map.clear() - if data is None: - return - related_terms_data = data.get("relatedTerms") - if related_terms_data is None: - return - for item in related_terms_data: - term_text = item["termText"] - related_terms_data = item["relatedTerms"] - related_terms: list[Term] = [ - Term(term_data["text"], weight=term_data.get("weight")) - for term_data in related_terms_data - ] + for term_text, related_terms in items: await self.add_related_term(term_text, related_terms) diff --git a/src/typeagent/storage/sqlite/related_terms_index.py b/src/typeagent/storage/sqlite/related_terms_index.py index d918fd6c..3a642419 100644 --- a/src/typeagent/storage/sqlite/related_terms_index.py +++ b/src/typeagent/storage/sqlite/related_terms_index.py @@ -62,10 +62,10 @@ async def set_related_terms(self, term: str, related_terms: list[str]) -> None: cursor = self.db.cursor() # Clear existing aliases for this term cursor.execute("DELETE FROM RelatedTermsAliases WHERE term = ?", (term,)) - # Add new aliases + # Add new aliases (use INSERT OR IGNORE to skip duplicates) for alias in related_terms: cursor.execute( - "INSERT INTO RelatedTermsAliases (term, alias) VALUES (?, ?)", + "INSERT OR IGNORE INTO RelatedTermsAliases (term, alias) VALUES (?, ?)", (term, alias), ) @@ -109,33 +109,25 @@ async def serialize(self) -> TermToRelatedTermsData: return TermToRelatedTermsData(relatedTerms=items) async def deserialize(self, data: TermToRelatedTermsData | None) -> None: - """Deserialize alias data.""" - cursor = self.db.cursor() - - # Clear existing data - cursor.execute("DELETE FROM RelatedTermsAliases") - - if data is None: - return - - related_terms = data.get("relatedTerms", []) - - if related_terms: - # Prepare all insertion data for bulk operation - insertion_data = [] - for item in related_terms: + """Deserialize alias data, replacing the existing aliases.""" + # Parse everything before deleting, so malformed data leaves the + # existing aliases in place. + insertion_data: list[tuple[str, str]] = [] + if data is not None: + for item in data.get("relatedTerms") or []: if item and item.get("termText") and item.get("relatedTerms"): term = item["termText"] for term_data in item["relatedTerms"]: - alias = term_data["text"] - insertion_data.append((term, alias)) - - # Bulk insert all the data - if insertion_data: - cursor.executemany( - "INSERT INTO RelatedTermsAliases (term, alias) VALUES (?, ?)", - insertion_data, - ) + insertion_data.append((term, term_data["text"])) + + cursor = self.db.cursor() + cursor.execute("DELETE FROM RelatedTermsAliases") + if insertion_data: + # INSERT OR IGNORE: the same alias may appear more than once. + cursor.executemany( + "INSERT OR IGNORE INTO RelatedTermsAliases (term, alias) VALUES (?, ?)", + insertion_data, + ) class SqliteRelatedTermsFuzzy(ITermToRelatedTermsFuzzy): diff --git a/tests/test_related_terms_index.py b/tests/test_related_terms_index.py index 1a6c4d29..0ac211c0 100644 --- a/tests/test_related_terms_index.py +++ b/tests/test_related_terms_index.py @@ -14,7 +14,12 @@ MessageTextIndexSettings, RelatedTermIndexSettings, ) -from typeagent.knowpro.interfaces import IMessage, ITermToRelatedTermsIndex, Term +from typeagent.knowpro.interfaces import ( + IMessage, + ITermToRelatedTermsIndex, + Term, + TermToRelatedTermsData, +) from typeagent.knowpro.knowledge_schema import KnowledgeResponse from typeagent.knowpro.query import CompiledSearchTerm, CompiledTermGroup from typeagent.storage import SqliteStorageProvider @@ -119,6 +124,38 @@ async def test_serialize_and_deserialize( assert any(term.text == "programming" for term in result) +@pytest.mark.asyncio +async def test_deserialize_accepts_duplicate_aliases( + related_terms_index: ITermToRelatedTermsIndex, +) -> None: + aliases = related_terms_index.aliases + data: TermToRelatedTermsData = { + "relatedTerms": [ + {"termText": "car", "relatedTerms": [{"text": "auto"}, {"text": "auto"}]}, + {"termText": "car", "relatedTerms": [{"text": "auto"}]}, + ] + } + await aliases.deserialize(data) + result = await aliases.lookup_term("car") + assert result is not None + assert [t.text for t in result] == ["auto"] + assert await aliases.size() == 1 + + +@pytest.mark.asyncio +async def test_deserialize_malformed_data_keeps_existing_aliases( + related_terms_index: ITermToRelatedTermsIndex, +) -> None: + aliases = related_terms_index.aliases + await aliases.add_related_term("keep", Term(text="kept")) + malformed = {"relatedTerms": [{"termText": "car", "relatedTerms": [{"weight": 1}]}]} + with pytest.raises(KeyError): + await aliases.deserialize(malformed) # type: ignore[arg-type] + result = await aliases.lookup_term("keep") + assert result is not None + assert [t.text for t in result] == ["kept"] + + @pytest.mark.asyncio async def test_related_terms_index_basic(really_needs_auth: None) -> None: settings = RelatedTermIndexSettings(TextEmbeddingIndexSettings()) diff --git a/tests/test_sqlite_indexes.py b/tests/test_sqlite_indexes.py index c7c9af01..264de303 100644 --- a/tests/test_sqlite_indexes.py +++ b/tests/test_sqlite_indexes.py @@ -294,6 +294,16 @@ async def test_serialize_deserialize(self, sqlite_db: sqlite3.Connection): assert not await index.is_empty() assert await index.size() == 2 + @pytest.mark.asyncio + async def test_set_related_terms_accepts_duplicates( + self, sqlite_db: sqlite3.Connection + ): + index = SqliteRelatedTermsAliases(sqlite_db) + await index.set_related_terms("car", ["auto", "auto", "vehicle"]) + results = await index.lookup_term("car") + assert results is not None + assert sorted(t.text for t in results) == ["auto", "vehicle"] + class TestSqliteRelatedTermsFuzzy: """Test SqliteRelatedTermsFuzzy with real embeddings."""