diff --git a/src/typeagent/knowpro/conversation_base.py b/src/typeagent/knowpro/conversation_base.py index ee3b0295..879fa753 100644 --- a/src/typeagent/knowpro/conversation_base.py +++ b/src/typeagent/knowpro/conversation_base.py @@ -243,6 +243,16 @@ async def _commit_batch_from_chunk_results( if not messages_batch: return AddMessagesResult() + # Chunk locations use source-stream ordinals. Those ordinals can have + # gaps when skip_failed_messages omits a message, so preserve their + # order for remapping to consecutive storage ordinals below. + source_message_ids: list[MessageOrdinal] = [] + source_message_id: MessageOrdinal | None = None + for result in chunk_results: + if result.chunk_id.message_ordinal != source_message_id: + source_message_ids.append(result.chunk_id.message_ordinal) + source_message_id = result.chunk_id.message_ordinal + # Process chunk results first to collect embeddings and knowledge items knowledge_items: list[tuple[MessageOrdinal, int, kplib.KnowledgeResponse]] = [] fuzzy_terms: list[str] = [] @@ -290,7 +300,7 @@ async def _commit_batch_from_chunk_results( ) fuzzy_terms.extend(result.related_terms) fuzzy_term_embeddings.extend(result.related_term_embeddings) - # Store embedding for later retrieval in correct message/chunk order + # Store embedding by source message/chunk location until commit remaps it. chunk_embedding_map[ (result.chunk_id.message_ordinal, result.chunk_id.chunk_ordinal) ] = result.chunk_embedding @@ -301,20 +311,45 @@ async def _commit_batch_from_chunk_results( semref_count=await self.semantic_refs.size(), ) - # Build chunk_embeddings in the correct order (matching message/chunk iteration) + source_to_storage_message_id: dict[MessageOrdinal, MessageOrdinal] = {} chunk_embeddings: list[NormalizedEmbedding] = [] - for msg_ord, message in enumerate( + result_group_index = 0 + for storage_message_id, message in enumerate( messages_batch, start_points.message_count ): - for chunk_ord in range(len(message.text_chunks)): - embedding = chunk_embedding_map.get((msg_ord, chunk_ord)) + if not message.text_chunks: + continue + if result_group_index >= len(source_message_ids): + raise ValueError( + "Missing chunk results for staged message: " + f"message={storage_message_id}" + ) + source_message_id = source_message_ids[result_group_index] + result_group_index += 1 + source_to_storage_message_id[source_message_id] = storage_message_id + for chunk_ordinal in range(len(message.text_chunks)): + embedding = chunk_embedding_map.get( + (source_message_id, chunk_ordinal) + ) if embedding is None: raise ValueError( "Missing chunk embedding for staged message chunk: " - f"message={msg_ord}, chunk={chunk_ord}" + f"message={storage_message_id}, chunk={chunk_ordinal}" ) chunk_embeddings.append(embedding) + if result_group_index != len(source_message_ids): + raise ValueError("Chunk results exceed staged messages with chunks") + + knowledge_items = [ + ( + source_to_storage_message_id[source_message_id], + chunk_ordinal, + knowledge, + ) + for source_message_id, chunk_ordinal, knowledge in knowledge_items + ] + # Use precomputed embeddings to avoid redundant embedding work await self.messages.extend( messages_batch, chunk_embeddings=chunk_embeddings diff --git a/tests/test_add_messages_pipeline.py b/tests/test_add_messages_pipeline.py index 41be8d5c..6a26edf5 100644 --- a/tests/test_add_messages_pipeline.py +++ b/tests/test_add_messages_pipeline.py @@ -445,6 +445,62 @@ async def test_dispatcher_stops_on_sentinel_and_emits_result_sentinel() -> None: assert items[0].error is None +@pytest.mark.asyncio +async def test_shutdown_finishes_message_already_admitted_to_full_queue() -> None: + chunk_queue: asyncio.Queue[ChunkWorkItem[_Message] | None] = asyncio.Queue( + maxsize=1 + ) + result_queue = asyncio.Queue() + stop_state = PipelineStopState() + producer_state = ProducerState(next_message_id=0) + shutdown_event = asyncio.Event() + first_message = _Message(["first chunk", "second chunk"]) + second_message = _Message(["not admitted"]) + + async def _iter_messages() -> AsyncIterator[_Message]: + yield first_message + yield second_message + + producer_task = asyncio.create_task( + _producer_task( + _iter_messages(), + chunk_queue, + stop_state, + producer_state, + result_queue, + shutdown_event, + ) + ) + await asyncio.sleep(0) + assert chunk_queue.full() + shutdown_event.set() + + await _dispatcher_task( + chunk_queue, + result_queue, + stop_state, + _SequenceExtractor( + [ + typechat.Success(_empty_knowledge()), + typechat.Success(_empty_knowledge()), + ] + ), + _StubEmbeddingModel(), + concurrency=1, + skip_failed_messages=False, + ) + await asyncio.wait_for(producer_task, timeout=1) + + items = await _drain_result_queue(result_queue) + results = [item for item in items if item is not None] + assert [result.chunk_id for result in results] == [ + TextLocation(0, 0), + TextLocation(0, 1), + ] + assert producer_state.produced_messages == 1 + assert producer_state.produced_chunks == 2 + + @pytest.mark.asyncio async def test_dispatcher_extraction_failure_lowers_stop() -> None: """A Failure from the extractor sets error and lowers stop_at_message_id.""" diff --git a/tests/test_add_messages_streaming.py b/tests/test_add_messages_streaming.py index cd1f15e2..ce75a5fb 100644 --- a/tests/test_add_messages_streaming.py +++ b/tests/test_add_messages_streaming.py @@ -15,7 +15,7 @@ from typeagent.knowpro import knowledge_schema as kplib from typeagent.knowpro.add_messages import add_messages_streaming from typeagent.knowpro.convsettings import ConversationSettings -from typeagent.knowpro.interfaces_core import IKnowledgeExtractor +from typeagent.knowpro.interfaces_core import AddMessagesResult, IKnowledgeExtractor from typeagent.storage.sqlite.provider import SqliteStorageProvider from typeagent.transcripts.transcript import ( Transcript, @@ -97,6 +97,10 @@ class ControlledExtractor: ``fail_on`` is a set of 0-based call indices for which the extractor returns a Failure instead of a Success. ``raise_on`` is a set of call indices that raise an exception. + ``fail_on_text`` is a set of chunk texts that always fail, regardless of + call order -- chunks are extracted concurrently, so call indices are not + deterministic when the failing chunk must be a specific one. It is a + mutable set so a test can "repair" the extractor between runs. """ def __init__( @@ -104,9 +108,11 @@ def __init__( *, fail_on: set[int] | None = None, raise_on: set[int] | None = None, + fail_on_text: set[str] | None = None, ) -> None: self.fail_on = fail_on or set() self.raise_on = raise_on or set() + self.fail_on_text = fail_on_text or set() self.call_count = 0 async def extract(self, message: str) -> typechat.Result[kplib.KnowledgeResponse]: @@ -116,6 +122,8 @@ async def extract(self, message: str) -> typechat.Result[kplib.KnowledgeResponse raise RuntimeError(f"Systemic failure at call {idx}") if idx in self.fail_on: return typechat.Failure(f"Extraction failed for call {idx}") + if message in self.fail_on_text: + return typechat.Failure(f"Extraction failed for chunk {message!r}") return typechat.Success(_EMPTY_RESPONSE) @@ -200,6 +208,36 @@ async def test_streaming_extraction_failure_stops_at_failing_message() -> None: await storage.close() +@pytest.mark.asyncio +async def test_streaming_skips_failed_message_and_commits_next_message() -> None: + """Skipped messages do not leave gaps in persisted message ordinals.""" + with tempfile.TemporaryDirectory() as tmpdir: + db_path = os.path.join(tmpdir, "test.db") + extractor = ControlledExtractor(fail_on={0}) + transcript, storage = await _create_transcript( + db_path, auto_extract=True, knowledge_extractor=extractor + ) + messages = [ + _make_message("failed", source_id="failed-source"), + _make_message("succeeds", source_id="successful-source"), + ] + + result = await add_messages_streaming( + transcript, + _async_iter(messages), + skip_failed_messages=True, + ) + + assert result.messages_added == 1 + assert result.messages_skipped == 1 + assert result.chunks_added == 1 + assert await transcript.messages.get_slice(0, 1) == [messages[1]] + assert not await storage.is_source_ingested("failed-source") + assert await storage.is_source_ingested("successful-source") + + await storage.close() + + @pytest.mark.asyncio async def test_streaming_exception_stops_run() -> None: """A raised exception stops processing; committed batches survive.""" @@ -706,3 +744,112 @@ async def test_streaming_extraction_returns_none_for_empty_chunks() -> None: assert extractor.call_count == 0 await storage.close() + + +# --------------------------------------------------------------------------- +# Repair-and-reimport (replay) behavior +# --------------------------------------------------------------------------- + + +async def _message_texts(transcript: Transcript) -> list[str]: + """Return every stored chunk text, in storage order.""" + size = await transcript.messages.size() + messages = await transcript.messages.get_slice(0, size) + return [chunk for message in messages for chunk in message.text_chunks] + + +async def _replay( + transcript: Transcript, + storage: SqliteStorageProvider, + messages: list[TranscriptMessage], +) -> AddMessagesResult: + """Re-submit ``messages``, pre-filtering sources the DB already has. + + This mirrors what importers do (see ``tools/ingest_email.py``): ask the + storage provider which source IDs are already ingested and only stream the + remainder. Re-running it must be idempotent. + """ + source_ids = [m.source_id for m in messages if m.source_id is not None] + already_ingested = await storage.are_sources_ingested(source_ids) + pending = [m for m in messages if m.source_id not in already_ingested] + return await add_messages_streaming( + transcript, + _async_iter(pending), + skip_failed_messages=True, + ) + + +@pytest.mark.asyncio +async def test_streaming_replay_after_repair_ingests_each_source_once() -> None: + """fail -> repair -> replay -> replay stores each source exactly once. + + The first run skips a message whose extraction fails, so its source is + never marked ingested. After the extractor is repaired, a replay picks up + only that message, and a second replay is a no-op -- no duplicates. + """ + with tempfile.TemporaryDirectory() as tmpdir: + db_path = os.path.join(tmpdir, "test.db") + extractor = ControlledExtractor(fail_on_text={"msg-1"}) + transcript, storage = await _create_transcript( + db_path, auto_extract=True, knowledge_extractor=extractor + ) + messages = [_make_message(f"msg-{i}", source_id=f"s-{i}") for i in range(3)] + + # 1. Fail: msg-1 fails extraction and is skipped; the others commit. + first = await _replay(transcript, storage, messages) + assert first.messages_added == 2 + assert first.messages_skipped == 1 + assert await _message_texts(transcript) == ["msg-0", "msg-2"] + assert not await storage.is_source_ingested("s-1") + assert _ingested_count(storage) == 2 + + # 2. Repair: the extractor no longer fails on that chunk. + extractor.fail_on_text.clear() + + # 3. Replay: only the previously failed source is re-submitted. + second = await _replay(transcript, storage, messages) + assert second.messages_added == 1 + assert second.messages_skipped == 0 + assert await storage.is_source_ingested("s-1") + + # 4. Replay again: everything is ingested, so nothing is streamed. + third = await _replay(transcript, storage, messages) + assert third.messages_added == 0 + assert third.messages_skipped == 0 + assert third.chunks_added == 0 + + texts = await _message_texts(transcript) + assert len(texts) == len(set(texts)) # no duplicated messages + assert sorted(texts) == ["msg-0", "msg-1", "msg-2"] + assert await transcript.messages.size() == 3 + assert _ingested_count(storage) == 3 + + await storage.close() + + +@pytest.mark.asyncio +async def test_streaming_replay_without_repair_does_not_duplicate() -> None: + """Replaying while the failure persists re-skips instead of duplicating.""" + with tempfile.TemporaryDirectory() as tmpdir: + db_path = os.path.join(tmpdir, "test.db") + extractor = ControlledExtractor(fail_on_text={"msg-1"}) + transcript, storage = await _create_transcript( + db_path, auto_extract=True, knowledge_extractor=extractor + ) + messages = [_make_message(f"msg-{i}", source_id=f"s-{i}") for i in range(3)] + + first = await _replay(transcript, storage, messages) + assert first.messages_added == 2 + assert first.messages_skipped == 1 + + # Replay with the failure still in place: the two good sources are + # filtered out, and msg-1 fails and is skipped again. + second = await _replay(transcript, storage, messages) + assert second.messages_added == 0 + assert second.messages_skipped == 1 + + texts = await _message_texts(transcript) + assert texts == ["msg-0", "msg-2"] + assert _ingested_count(storage) == 2 + + await storage.close()