diff --git a/CHANGELOG.md b/CHANGELOG.md index d29f2d0a..c7a3e1b7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -66,6 +66,16 @@ loads. Other dataclasses using this decorator also serialize previously unaliased multi-word fields as camelCase when `by_alias=True`; their snake_case input names remain accepted unless explicitly overridden. +- `add_messages_with_indexing` and `add_messages_streaming` now add the same + terms to the fuzzy related-terms index: every term added to the semantic ref + index (#346), so related-term expansion no longer depends on the ingestion + path. `add_messages_with_indexing` used to add only entity names, topics and + each action verb separately; it now embeds more terms per message, and + multi-word verbs are embedded as one term (e.g. `"looks up"` rather than + `"looks"` and `"up"`), matching the term index. `add_messages_streaming` + used to skip terms from message metadata (e.g. speakers and recipients) and + now embeds them at commit time; it also no longer embeds the types and + facets of entities without a name, which the term index skips. ### 0.5.0 (May 26) diff --git a/src/typeagent/knowpro/add_messages.py b/src/typeagent/knowpro/add_messages.py index 9e244928..8eb8312d 100644 --- a/src/typeagent/knowpro/add_messages.py +++ b/src/typeagent/knowpro/add_messages.py @@ -13,13 +13,13 @@ from ..aitools.embeddings import IEmbeddingModel, NormalizedEmbedding from ..storage.memory import semantic_ref_index -from .common import normalize_term from .interfaces import ( AddMessagesResult, IKnowledgeExtractor, IMessage, MessageOrdinal, TextLocation, + Topic, ) from .knowledge_schema import KnowledgeResponse @@ -221,7 +221,7 @@ class ChunkProcessingResult[TMessage: IMessage]: message: Original message object containing this chunk. extracted_knowledge: Extracted KnowledgeResponse if extraction succeeded, else None. chunk_embedding: Normalized embedding vector for the message chunk, or None if extraction or embedding failed. - related_terms: Lowercased, deduplicated related-term texts extracted from knowledge. + related_terms: Normalized, deduplicated related-term texts extracted from knowledge. related_term_embeddings: Embeddings for related_terms in the same order, or [] when there are no related terms. error: Exception from the first failing operation, or None if extraction and embedding succeeded. @@ -244,32 +244,18 @@ def _collect_related_terms_for_fuzzy_index( ) -> list[str]: """Collect canonical related-term texts for the fuzzy related-terms index. - These terms are derived from the same knowledge that feeds semantic refs. - We normalize each term with `normalize_term` (strip, NFC, collapse - whitespace, lowercase), drop empty results, and deduplicate while - preserving order to match index behavior. + These are the normalized, deduplicated terms the knowledge adds to the + semantic ref index, in the order they are first added. Terms from message + metadata are added at commit time. """ - seen: set[str] = set() - related_terms: list[str] = [] - - def _add_term(term: str) -> None: - canonical = normalize_term(term) - if canonical and canonical not in seen: - seen.add(canonical) - related_terms.append(canonical) - - for entity in knowledge.entities: - for term in semantic_ref_index.collect_entity_terms(entity): - _add_term(term) - - for action in chain(knowledge.actions, knowledge.inverse_actions): - for term in semantic_ref_index.collect_action_terms(action): - _add_term(term) - - for topic in knowledge.topics: - _add_term(topic) - - return related_terms + return semantic_ref_index.collect_related_terms_for_fuzzy_index( + chain( + knowledge.entities, + knowledge.actions, + knowledge.inverse_actions, + (Topic(text=topic) for topic in knowledge.topics), + ) + ) # "Public", imported by tests diff --git a/src/typeagent/knowpro/conversation_base.py b/src/typeagent/knowpro/conversation_base.py index b2bcfecb..421b5350 100644 --- a/src/typeagent/knowpro/conversation_base.py +++ b/src/typeagent/knowpro/conversation_base.py @@ -4,7 +4,7 @@ """Base class for conversations with incremental indexing support.""" import asyncio -from collections.abc import AsyncIterable, Callable, Sequence +from collections.abc import AsyncIterable, Callable, Collection, Sequence from dataclasses import dataclass from datetime import datetime, timezone from typing import Generic, Protocol, Self, TypeVar @@ -22,7 +22,6 @@ from ..aitools import model_adapters, utils from ..aitools.embeddings import NormalizedEmbedding from ..storage.memory import semantic_ref_index -from .common import normalize_term from .convsettings import ConversationSettings from .interfaces import ( AddMessagesResult, @@ -36,9 +35,8 @@ ITermToSemanticRefIndex, MessageOrdinal, TextLocation, - Topic, ) -from .knowledge_schema import Action, ConcreteEntity, KnowledgeResponse +from .knowledge_schema import KnowledgeResponse from .messageutils import get_all_message_chunk_locations TMessage = TypeVar("TMessage", bound=IMessage) @@ -439,6 +437,11 @@ async def _update_secondary_indexes_incremental_with_embeddings( related_terms, related_term_embeddings, ) + # Metadata knowledge is only indexed at commit time, so its terms + # have no precomputed embeddings; embed them now. + await self._update_related_terms_incremental( + start_points.semref_count, exclude=set(related_terms) + ) async def _add_metadata_knowledge_incremental( self, @@ -534,8 +537,12 @@ async def _add_timestamps_for_messages( async def _update_related_terms_incremental( self, start_from_semref_ordinal: int, + exclude: Collection[str] = (), ) -> None: - """Update related terms index with new semantic refs.""" + """Update related terms index with new semantic refs. + + Terms in ``exclude`` (already added by the caller) are skipped. + """ if ( self.secondary_indexes is None or self.secondary_indexes.term_to_related_terms_index is None @@ -549,20 +556,15 @@ async def _update_related_terms_incremental( fuzzy_index = self.secondary_indexes.term_to_related_terms_index.fuzzy_index if fuzzy_index is not None and new_semrefs: - new_terms: set[str] = set() - for semref in new_semrefs: - knowledge = semref.knowledge - if isinstance(knowledge, ConcreteEntity): - new_terms.add(normalize_term(knowledge.name)) - elif isinstance(knowledge, Topic): - new_terms.add(normalize_term(knowledge.text)) - elif isinstance(knowledge, Action): - for verb in knowledge.verbs: - new_terms.add(normalize_term(verb)) - new_terms.discard("") # Whitespace-only terms must not be embedded - + new_terms = [ + term + for term in semantic_ref_index.collect_related_terms_for_fuzzy_index( + semref.knowledge for semref in new_semrefs + ) + if term not in exclude + ] if new_terms: - await fuzzy_index.add_terms(list(new_terms)) + await fuzzy_index.add_terms(new_terms) async def _update_message_index_incremental( self, diff --git a/src/typeagent/storage/memory/semantic_ref_index.py b/src/typeagent/storage/memory/semantic_ref_index.py index 4cf45183..ff8e96e6 100644 --- a/src/typeagent/storage/memory/semantic_ref_index.py +++ b/src/typeagent/storage/memory/semantic_ref_index.py @@ -3,7 +3,7 @@ from __future__ import annotations # TODO: Avoid -from collections.abc import AsyncIterable, Callable, Sequence +from collections.abc import AsyncIterable, Callable, Iterable, Sequence from typechat import Failure @@ -585,6 +585,38 @@ def collect_action_terms(action: Action) -> list[str]: return terms +def collect_knowledge_terms(knowledge: Knowledge) -> list[str]: + """Collect all terms a knowledge item would add to the semantic ref index.""" + if isinstance(knowledge, ConcreteEntity): + return collect_entity_terms(knowledge) if validate_entity(knowledge) else [] + if isinstance(knowledge, Action): + return collect_action_terms(knowledge) + if isinstance(knowledge, Topic): + return [knowledge.text] + return [] # Tags are not added to the semantic ref index + + +def collect_related_terms_for_fuzzy_index( + knowledge_items: Iterable[Knowledge], +) -> list[str]: + """Collect the terms to add to the fuzzy related-terms index. + + These are the terms the knowledge items add to the semantic ref index, + so the same knowledge yields the same fuzzy-index terms on every + ingestion path. Each term is normalized with `normalize_term`; empty + results are dropped and duplicates removed, keeping the first occurrence. + """ + seen: set[str] = set() + related_terms: list[str] = [] + for knowledge in knowledge_items: + for term in collect_knowledge_terms(knowledge): + canonical = normalize_term(term) + if canonical and canonical not in seen: + seen.add(canonical) + related_terms.append(canonical) + return related_terms + + async def add_metadata_to_index_from_list[TMessage: IMessage]( messages: list[TMessage], semantic_refs: ISemanticRefCollection, diff --git a/tests/test_add_messages_pipeline.py b/tests/test_add_messages_pipeline.py index 0849168c..bf82c714 100644 --- a/tests/test_add_messages_pipeline.py +++ b/tests/test_add_messages_pipeline.py @@ -24,6 +24,7 @@ process_chunk_with_extraction_and_embeddings, ProducerState, ) +from typeagent.knowpro.common import normalize_term from typeagent.knowpro.core_types import ( DeletionInfo, IMessageMetadata, @@ -31,10 +32,12 @@ ) from typeagent.knowpro.knowledge_schema import ( Action, + ActionParam, ConcreteEntity, Facet, KnowledgeResponse, ) +from typeagent.storage.memory import semantic_ref_index @dataclass @@ -201,6 +204,55 @@ async def test_collect_related_terms_lowercases_dedupes_and_preserves_order() -> assert len(terms) == len(set(terms)) +def test_fuzzy_terms_match_semantic_ref_index_terms() -> None: + """Both ingestion paths collect exactly the term index terms for knowledge.""" + knowledge = KnowledgeResponse( + entities=[ + ConcreteEntity( + name="Space Needle", + type=["Tower"], + facets=[Facet(name="Height", value=184)], + ), + ConcreteEntity(name="", type=["Ghost"]), # Not indexed + ], + actions=[ + Action( + verbs=["Looks", "Up"], + verb_tense="past", + subject_entity_name="Alice", + object_entity_name="Space Needle", + params=["Slowly", ActionParam(name="Time", value="Noon")], + subject_entity_facet=Facet(name="Mood", value="Happy"), + ) + ], + inverse_actions=[ + Action( + verbs=["Seen"], + verb_tense="past", + subject_entity_name="Space Needle", + object_entity_name="none", + ) + ], + topics=["Seattle"], + ) + refs, index_terms = semantic_ref_index._collect_knowledge_refs_and_terms( + 0, 0, 0, knowledge + ) + expected = list(dict.fromkeys(normalize_term(term) for term, _ in index_terms)) + + # Path used by add_messages_streaming (extracted knowledge). + assert _collect_related_terms_for_fuzzy_index(knowledge) == expected + # Path used by add_messages_with_indexing (new semantic refs). + assert ( + semantic_ref_index.collect_related_terms_for_fuzzy_index( + ref.knowledge for ref in refs + ) + == expected + ) + assert "looks up" in expected and "looks" not in expected + assert "ghost" not in expected + + @pytest.mark.asyncio async def test_process_chunk_success_with_related_terms() -> None: extractor = _SequenceExtractor([typechat.Success(_knowledge_with_terms())]) diff --git a/tests/test_add_messages_streaming.py b/tests/test_add_messages_streaming.py index d4f509e4..2046005a 100644 --- a/tests/test_add_messages_streaming.py +++ b/tests/test_add_messages_streaming.py @@ -22,8 +22,14 @@ IKnowledgeExtractor, TextLocation, ) -from typeagent.knowpro.knowledge_schema import ConcreteEntity, KnowledgeResponse +from typeagent.knowpro.knowledge_schema import ( + Action, + ConcreteEntity, + Facet, + KnowledgeResponse, +) from typeagent.storage.sqlite.provider import SqliteStorageProvider +from typeagent.storage.sqlite.related_terms_index import SqliteRelatedTermsFuzzy from typeagent.transcripts.transcript import ( Transcript, TranscriptMessage, @@ -986,3 +992,61 @@ async def test_streaming_replay_without_repair_does_not_duplicate() -> None: assert _ingested_count(storage) == 2 await storage.close() + + +# --------------------------------------------------------------------------- +# Fuzzy related-terms index contents (#346) +# --------------------------------------------------------------------------- + + +class _RichExtractor: + """Returns knowledge that uses every kind of indexed term.""" + + async def extract(self, message: str) -> typechat.Result[KnowledgeResponse]: + return typechat.Success( + KnowledgeResponse( + entities=[ + ConcreteEntity( + name="Space Needle", + type=["Tower"], + facets=[Facet(name="Height", value=184)], + ) + ], + actions=[ + Action( + verbs=["Looks", "Up"], + verb_tense="past", + subject_entity_name="Carol", + object_entity_name="Space Needle", + ) + ], + inverse_actions=[], + topics=["Seattle"], + ) + ) + + +@pytest.mark.asyncio +async def test_streaming_fuzzy_terms_match_term_index() -> None: + """Streaming embeds every semantic ref index term, metadata included.""" + with tempfile.TemporaryDirectory() as tmpdir: + transcript, storage = await _create_transcript( + os.path.join(tmpdir, "test.db"), + auto_extract=True, + knowledge_extractor=_RichExtractor(), + ) + message = TranscriptMessage( + text_chunks=["Look up at the Space Needle"], + metadata=TranscriptMessageMeta(speaker="Carol", recipients=["Bob"]), + tags=["test"], + ) + await add_messages_streaming(transcript, _async_iter([message])) + + fuzzy_index = storage.related_terms_index.fuzzy_index + assert isinstance(fuzzy_index, SqliteRelatedTermsFuzzy) + term_index_terms = await storage.semantic_ref_index.get_terms() + # Metadata terms (speaker, recipient) and extracted terms. + assert {"carol", "bob", "person", "looks up", "height"} <= set(term_index_terms) + assert sorted(await fuzzy_index.get_terms()) == sorted(term_index_terms) + + await storage.close() diff --git a/tests/test_add_messages_with_indexing.py b/tests/test_add_messages_with_indexing.py index d3df2c4d..78455069 100644 --- a/tests/test_add_messages_with_indexing.py +++ b/tests/test_add_messages_with_indexing.py @@ -8,9 +8,18 @@ import pytest +import typechat + from typeagent.aitools.model_adapters import create_test_embedding_model from typeagent.knowpro.convsettings import ConversationSettings +from typeagent.knowpro.knowledge_schema import ( + Action, + ConcreteEntity, + Facet, + KnowledgeResponse, +) from typeagent.storage.sqlite.provider import SqliteStorageProvider +from typeagent.storage.sqlite.related_terms_index import SqliteRelatedTermsFuzzy from typeagent.transcripts.transcript import ( Transcript, TranscriptMessage, @@ -162,3 +171,72 @@ async def test_transaction_rollback_on_error(): pass await storage.close() + + +class _RichExtractor: + """Returns knowledge that uses every kind of indexed term.""" + + async def extract(self, message: str) -> typechat.Result[KnowledgeResponse]: + return typechat.Success( + KnowledgeResponse( + entities=[ + ConcreteEntity( + name="Space Needle", + type=["Tower"], + facets=[Facet(name="Height", value=184)], + ) + ], + actions=[ + Action( + verbs=["Looks", "Up"], + verb_tense="past", + subject_entity_name="Alice", + object_entity_name="Space Needle", + ) + ], + inverse_actions=[], + topics=["Seattle"], + ) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("auto_extract", [False, True]) +async def test_add_messages_with_indexing_fuzzy_terms_match_term_index( + auto_extract: bool, +): + """The fuzzy related-terms index gets every semantic ref index term (#346).""" + with tempfile.TemporaryDirectory() as tmpdir: + db_path = os.path.join(tmpdir, "test.db") + + test_model = create_test_embedding_model() + settings = ConversationSettings(model=test_model) + settings.semantic_ref_index_settings.auto_extract_knowledge = auto_extract + settings.semantic_ref_index_settings.knowledge_extractor = _RichExtractor() + + storage = SqliteStorageProvider( + db_path, + message_type=TranscriptMessage, + message_text_index_settings=settings.message_text_index_settings, + related_term_index_settings=settings.related_term_index_settings, + ) + settings.storage_provider = storage + transcript = await Transcript.create(settings, name="test") + + messages = [ + TranscriptMessage( + text_chunks=["Hello world"], + metadata=TranscriptMessageMeta(speaker="Alice", recipients=["Bob"]), + ), + ] + await transcript.add_messages_with_indexing(messages) + + fuzzy_index = storage.related_terms_index.fuzzy_index + assert isinstance(fuzzy_index, SqliteRelatedTermsFuzzy) + term_index_terms = await storage.semantic_ref_index.get_terms() + assert {"alice", "bob", "person"} <= set(term_index_terms) + if auto_extract: + assert {"looks up", "height", "seattle"} <= set(term_index_terms) + assert sorted(await fuzzy_index.get_terms()) == sorted(term_index_terms) + + await storage.close()