From f4f79ad51ae209f083208959d2243c80f39a51d8 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Tue, 9 Jun 2026 04:56:14 -0500 Subject: [PATCH 1/7] docs: remove outdated success property reference from ChunkProcessingResult docstring The success property was removed in a previous commit but the docstring still referenced it, causing confusion. --- src/typeagent/knowpro/add_messages.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/typeagent/knowpro/add_messages.py b/src/typeagent/knowpro/add_messages.py index 4418b972..a9a02a73 100644 --- a/src/typeagent/knowpro/add_messages.py +++ b/src/typeagent/knowpro/add_messages.py @@ -218,9 +218,6 @@ class ChunkProcessingResult[TMessage: IMessage]: related_terms: Lowercased, 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. - - The ``success`` property is True only when extraction succeeded, chunk embedding was - generated, related-term embeddings were generated, and no error occurred. """ chunk_id: TextLocation From 95804a15c806dddcd6e2c5386ae41e0b998a635c Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Tue, 9 Jun 2026 05:01:53 -0500 Subject: [PATCH 2/7] feat: add shutdown_event support to _dispatcher_task Propagate shutdown_event to dispatcher for coordinated pipeline shutdown. When shutdown is requested, skip remaining chunks instead of processing them after the producer has stopped. --- src/typeagent/knowpro/add_messages.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/src/typeagent/knowpro/add_messages.py b/src/typeagent/knowpro/add_messages.py index a9a02a73..bcf4873b 100644 --- a/src/typeagent/knowpro/add_messages.py +++ b/src/typeagent/knowpro/add_messages.py @@ -138,6 +138,7 @@ async def _dispatcher_task[TMessage: IMessage]( embedding_model: IEmbeddingModel, concurrency: int, skip_failed_messages: bool, + shutdown_event: asyncio.Event | None = None, ) -> None: """Dispatch chunk work items to bounded per-item worker tasks. @@ -153,13 +154,18 @@ async def _dispatcher_task[TMessage: IMessage]( Args: skip_failed_messages: If True, don't halt producer on extraction/embedding failures; continue processing. If False, halt on first failure. + shutdown_event: If set, stop processing new chunks and let the pipeline drain. """ sem = asyncio.Semaphore(concurrency) async def _process_one(work_item: ChunkWorkItem[TMessage]) -> None: try: stop_at = stop_state.stop_at_message_id - if work_item.chunk_id.message_ordinal >= stop_at: + if ( + work_item.chunk_id.message_ordinal >= stop_at + or shutdown_event is not None + and shutdown_event.is_set() + ): result: "ChunkProcessingResult[TMessage]" = ChunkProcessingResult( chunk_id=work_item.chunk_id, chunk_count=work_item.chunk_count, @@ -195,7 +201,7 @@ async def _process_one(work_item: ChunkWorkItem[TMessage]) -> None: await result_queue.put(result) async with asyncio.TaskGroup() as tg: - while True: + while not (shutdown_event is not None and shutdown_event.is_set()): item = await chunk_queue.get() if item is None: break @@ -631,6 +637,7 @@ async def _commit_batch( embedding_model, concurrency=sem_ref_settings.concurrency, skip_failed_messages=skip_failed_messages, + shutdown_event=shutdown_event, ) ) reassembler_task = tg.create_task( From 683b4e8d2084f10669895b5897241650e8e91dcd Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Tue, 9 Jun 2026 05:23:11 -0500 Subject: [PATCH 3/7] fix: prevent producer deadlock when shutdown fires with a full chunk_queue --- src/typeagent/knowpro/add_messages.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/typeagent/knowpro/add_messages.py b/src/typeagent/knowpro/add_messages.py index bcf4873b..e50dfec9 100644 --- a/src/typeagent/knowpro/add_messages.py +++ b/src/typeagent/knowpro/add_messages.py @@ -112,6 +112,8 @@ async def _producer_task[TMessage: IMessage]( for chunk_ordinal, chunk_text in enumerate(message.text_chunks): if message_id >= stop_state.stop_at_message_id: break + if shutdown_event is not None and shutdown_event.is_set(): + break await chunk_queue.put( ChunkWorkItem[TMessage]( chunk_id=TextLocation(message_id, chunk_ordinal), @@ -207,6 +209,13 @@ async def _process_one(work_item: ChunkWorkItem[TMessage]) -> None: break await sem.acquire() tg.create_task(_process_one(item)) + else: + # Shutdown was set: drain remaining items so the producer's put() + # calls can unblock and it can send the None sentinel. + while True: + item = await chunk_queue.get() + if item is None: + break await result_queue.put(None) From 6233478e79ca3bbbcf16dedc374143cc00a90c03 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Tue, 9 Jun 2026 05:27:07 -0500 Subject: [PATCH 4/7] perf: store first error message on MessageAssembly to avoid chunk scan on skip --- src/typeagent/knowpro/add_messages.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/src/typeagent/knowpro/add_messages.py b/src/typeagent/knowpro/add_messages.py index e50dfec9..6288c882 100644 --- a/src/typeagent/knowpro/add_messages.py +++ b/src/typeagent/knowpro/add_messages.py @@ -357,6 +357,7 @@ class MessageAssembly[TMessage: IMessage]: message: TMessage chunks: dict[ChunkOrdinal, ChunkProcessingResult[TMessage]] has_error: bool = False + first_error_msg: str = "Unknown error" def is_complete(self) -> bool: return len(self.chunks) == self.chunk_count @@ -434,16 +435,9 @@ async def _drain_consecutive_complete(force: bool = False) -> None: return if assembly.has_error: if skip_failed_messages: - # Skip this failed message and continue - # Find the error from one of the chunks for logging - error_msg = "Unknown error" - for chunk_result in assembly.chunks.values(): - if chunk_result.error is not None: - error_msg = str(chunk_result.error) - break print( f"Skipping message {state.first_uncommitted_ordinal} " - f"due to chunk processing error: {error_msg}" + f"due to chunk processing error: {assembly.first_error_msg}" ) del assemblies[state.first_uncommitted_ordinal] state.first_uncommitted_ordinal += 1 @@ -530,6 +524,8 @@ async def _drain_consecutive_complete(force: bool = False) -> None: assembly.chunks[chunk_ordinal] = item if item.error is not None: + if not assembly.has_error: + assembly.first_error_msg = str(item.error) assembly.has_error = True state.chunk_failures += 1 if not skip_failed_messages: From 1da5fff717d4f8836d94a943f4f5dee8301bbb18 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Tue, 9 Jun 2026 06:02:29 -0500 Subject: [PATCH 5/7] fix: remap chunk results after skipped messages --- src/typeagent/knowpro/add_messages.py | 35 +++++--------- src/typeagent/knowpro/conversation_base.py | 47 +++++++++++++++--- tests/test_add_messages_pipeline.py | 56 ++++++++++++++++++++++ tests/test_add_messages_streaming.py | 30 ++++++++++++ 4 files changed, 140 insertions(+), 28 deletions(-) diff --git a/src/typeagent/knowpro/add_messages.py b/src/typeagent/knowpro/add_messages.py index 6288c882..4418b972 100644 --- a/src/typeagent/knowpro/add_messages.py +++ b/src/typeagent/knowpro/add_messages.py @@ -112,8 +112,6 @@ async def _producer_task[TMessage: IMessage]( for chunk_ordinal, chunk_text in enumerate(message.text_chunks): if message_id >= stop_state.stop_at_message_id: break - if shutdown_event is not None and shutdown_event.is_set(): - break await chunk_queue.put( ChunkWorkItem[TMessage]( chunk_id=TextLocation(message_id, chunk_ordinal), @@ -140,7 +138,6 @@ async def _dispatcher_task[TMessage: IMessage]( embedding_model: IEmbeddingModel, concurrency: int, skip_failed_messages: bool, - shutdown_event: asyncio.Event | None = None, ) -> None: """Dispatch chunk work items to bounded per-item worker tasks. @@ -156,18 +153,13 @@ async def _dispatcher_task[TMessage: IMessage]( Args: skip_failed_messages: If True, don't halt producer on extraction/embedding failures; continue processing. If False, halt on first failure. - shutdown_event: If set, stop processing new chunks and let the pipeline drain. """ sem = asyncio.Semaphore(concurrency) async def _process_one(work_item: ChunkWorkItem[TMessage]) -> None: try: stop_at = stop_state.stop_at_message_id - if ( - work_item.chunk_id.message_ordinal >= stop_at - or shutdown_event is not None - and shutdown_event.is_set() - ): + if work_item.chunk_id.message_ordinal >= stop_at: result: "ChunkProcessingResult[TMessage]" = ChunkProcessingResult( chunk_id=work_item.chunk_id, chunk_count=work_item.chunk_count, @@ -203,19 +195,12 @@ async def _process_one(work_item: ChunkWorkItem[TMessage]) -> None: await result_queue.put(result) async with asyncio.TaskGroup() as tg: - while not (shutdown_event is not None and shutdown_event.is_set()): + while True: item = await chunk_queue.get() if item is None: break await sem.acquire() tg.create_task(_process_one(item)) - else: - # Shutdown was set: drain remaining items so the producer's put() - # calls can unblock and it can send the None sentinel. - while True: - item = await chunk_queue.get() - if item is None: - break await result_queue.put(None) @@ -233,6 +218,9 @@ class ChunkProcessingResult[TMessage: IMessage]: related_terms: Lowercased, 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. + + The ``success`` property is True only when extraction succeeded, chunk embedding was + generated, related-term embeddings were generated, and no error occurred. """ chunk_id: TextLocation @@ -357,7 +345,6 @@ class MessageAssembly[TMessage: IMessage]: message: TMessage chunks: dict[ChunkOrdinal, ChunkProcessingResult[TMessage]] has_error: bool = False - first_error_msg: str = "Unknown error" def is_complete(self) -> bool: return len(self.chunks) == self.chunk_count @@ -435,9 +422,16 @@ async def _drain_consecutive_complete(force: bool = False) -> None: return if assembly.has_error: if skip_failed_messages: + # Skip this failed message and continue + # Find the error from one of the chunks for logging + error_msg = "Unknown error" + for chunk_result in assembly.chunks.values(): + if chunk_result.error is not None: + error_msg = str(chunk_result.error) + break print( f"Skipping message {state.first_uncommitted_ordinal} " - f"due to chunk processing error: {assembly.first_error_msg}" + f"due to chunk processing error: {error_msg}" ) del assemblies[state.first_uncommitted_ordinal] state.first_uncommitted_ordinal += 1 @@ -524,8 +518,6 @@ async def _drain_consecutive_complete(force: bool = False) -> None: assembly.chunks[chunk_ordinal] = item if item.error is not None: - if not assembly.has_error: - assembly.first_error_msg = str(item.error) assembly.has_error = True state.chunk_failures += 1 if not skip_failed_messages: @@ -642,7 +634,6 @@ async def _commit_batch( embedding_model, concurrency=sem_ref_settings.concurrency, skip_failed_messages=skip_failed_messages, - shutdown_event=shutdown_event, ) ) reassembler_task = tg.create_task( 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..9308d5b0 100644 --- a/tests/test_add_messages_streaming.py +++ b/tests/test_add_messages_streaming.py @@ -200,6 +200,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.""" From c3f65ddc94f315320d621ce50df7dc676ecfc9d7 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Tue, 22 Sep 2026 04:30:29 -0400 Subject: [PATCH 6/7] test: cover repair-and-replay reimport after a skipped message Adds the fail -> repair -> replay -> replay coverage requested in review, to confirm a repaired reimport does not duplicate messages. - ControlledExtractor gains fail_on_text: chunks are extracted concurrently, so call-index-based failure cannot deterministically target a specific message. The set is mutable so a test can "repair" the extractor. - _replay mirrors what importers do (tools/ingest_email.py): pre-filter already-ingested source IDs via are_sources_ingested, then stream the rest. - test_streaming_replay_after_repair_ingests_each_source_once: the skipped message's source is never marked ingested, so a repaired replay picks up exactly that message and a second replay is a no-op. - test_streaming_replay_without_repair_does_not_duplicate: replaying while the failure persists re-skips rather than duplicating. --- tests/test_add_messages_streaming.py | 119 ++++++++++++++++++++++++++- 1 file changed, 118 insertions(+), 1 deletion(-) diff --git a/tests/test_add_messages_streaming.py b/tests/test_add_messages_streaming.py index 9308d5b0..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) @@ -736,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() From eccb297bc96ba5d25ea056a517e30d095838587c Mon Sep 17 00:00:00 2001 From: Bernhard Merkle Date: Wed, 23 Sep 2026 11:11:13 +0200 Subject: [PATCH 7/7] Update src/typeagent/knowpro/conversation_base.py --- src/typeagent/knowpro/conversation_base.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/src/typeagent/knowpro/conversation_base.py b/src/typeagent/knowpro/conversation_base.py index 879fa753..a1acb889 100644 --- a/src/typeagent/knowpro/conversation_base.py +++ b/src/typeagent/knowpro/conversation_base.py @@ -341,14 +341,22 @@ async def _commit_batch_from_chunk_results( 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, + remapped_knowledge_items: list[ + tuple[MessageOrdinal, int, kplib.KnowledgeResponse] + ] = [] + for source_message_id, chunk_ordinal, knowledge in knowledge_items: + storage_message_id = source_to_storage_message_id.get( + source_message_id ) - for source_message_id, chunk_ordinal, knowledge in knowledge_items - ] + if storage_message_id is None: + raise ValueError( + "No storage message id for chunk result: " + f"source_message={source_message_id}, chunk={chunk_ordinal}" + ) + remapped_knowledge_items.append( + (storage_message_id, chunk_ordinal, knowledge) + ) + knowledge_items = remapped_knowledge_items # Use precomputed embeddings to avoid redundant embedding work await self.messages.extend(