diff --git a/trustgraph-base/trustgraph/messaging/translators/library.py b/trustgraph-base/trustgraph/messaging/translators/library.py index 8d7cc930..62350e9f 100644 --- a/trustgraph-base/trustgraph/messaging/translators/library.py +++ b/trustgraph-base/trustgraph/messaging/translators/library.py @@ -173,7 +173,14 @@ class LibraryResponseTranslator(MessageTranslator): return result def from_response_with_completion(self, obj: LibrarianResponse) -> Tuple[Dict[str, Any], bool]: - """Returns (response_dict, is_final)""" - # For streaming responses, check end_of_stream to determine if this is the final message - is_final = getattr(obj, 'end_of_stream', True) + """Returns (response_dict, is_final) + + For chunked streaming responses (total_chunks > 0), completion is + determined by whether we've reached the final chunk. + For non-streaming responses (total_chunks = 0), always final. + """ + if obj.total_chunks > 0: + is_final = (obj.chunk_index >= obj.total_chunks - 1) + else: + is_final = True return self.from_pulsar(obj), is_final diff --git a/trustgraph-base/trustgraph/schema/services/library.py b/trustgraph-base/trustgraph/schema/services/library.py index 4ec206ca..6dcdee1a 100644 --- a/trustgraph-base/trustgraph/schema/services/library.py +++ b/trustgraph-base/trustgraph/schema/services/library.py @@ -212,10 +212,6 @@ class LibrarianResponse: # list-uploads response upload_sessions: list[UploadSession] = field(default_factory=list) - # stream-document response - indicates final chunk in stream - # Default True so non-streaming operations are treated as final - end_of_stream: bool = True - # FIXME: Is this right? Using persistence on librarian so that # message chunking works diff --git a/trustgraph-flow/trustgraph/librarian/librarian.py b/trustgraph-flow/trustgraph/librarian/librarian.py index aecc1289..f3e7eabc 100644 --- a/trustgraph-flow/trustgraph/librarian/librarian.py +++ b/trustgraph-flow/trustgraph/librarian/librarian.py @@ -656,9 +656,8 @@ class Librarian: This is an async generator that yields document content in smaller chunks, allowing memory-efficient processing of large documents. Each yielded - response includes chunk information and an end_of_stream flag. - - The final chunk will have end_of_stream=True. + response includes chunk_index and total_chunks for tracking progress. + Completion is determined by chunk_index reaching total_chunks - 1. """ logger.debug(f"Streaming document {request.document_id}") @@ -688,11 +687,8 @@ class Librarian: # Fetch only the requested range chunk_content = await self.blob_store.get_range(object_id, offset, length) - is_last_chunk = (chunk_index == total_chunks - 1) - - logger.debug(f"Streaming chunk {chunk_index}/{total_chunks}, " - f"bytes {offset}-{offset + length} of {total_size}, " - f"end_of_stream={is_last_chunk}") + logger.debug(f"Streaming chunk {chunk_index + 1}/{total_chunks}, " + f"bytes {offset}-{offset + length} of {total_size}") yield LibrarianResponse( error=None, @@ -702,6 +698,5 @@ class Librarian: total_chunks=total_chunks, bytes_received=offset + length, total_bytes=total_size, - end_of_stream=is_last_chunk, ) diff --git a/trustgraph-flow/trustgraph/librarian/service.py b/trustgraph-flow/trustgraph/librarian/service.py index 9f24bbc7..009d23bd 100755 --- a/trustgraph-flow/trustgraph/librarian/service.py +++ b/trustgraph-flow/trustgraph/librarian/service.py @@ -527,7 +527,6 @@ class Processor(AsyncProcessor): type = "request-error", message = str(e), ), - end_of_stream = True, ) await self.librarian_response_producer.send( @@ -541,7 +540,6 @@ class Processor(AsyncProcessor): type = "unexpected-error", message = str(e), ), - end_of_stream = True, ) await self.librarian_response_producer.send(