Doc streaming from librarian

trustgraph-flow/trustgraph/librarian/blob_store.py

@@ -81,6 +81,28 @@ class BlobStore:
 
         return resp.read()
 
+    async def get_range(self, object_id, offset: int, length: int) -> bytes:
+        """Fetch a specific byte range from an object."""
+        resp = self.client.get_object(
+            bucket_name=self.bucket_name,
+            object_name="doc/" + str(object_id),
+            offset=offset,
+            length=length,
+        )
+        try:
+            return resp.read()
+        finally:
+            resp.close()
+            resp.release_conn()
+
+    async def get_size(self, object_id) -> int:
+        """Get the size of an object without downloading it."""
+        stat = self.client.stat_object(
+            bucket_name=self.bucket_name,
+            object_name="doc/" + str(object_id),
+        )
+        return stat.size
+
     def get_stream(self, object_id, chunk_size: int = 1024 * 1024) -> Iterator[bytes]:
         """
         Stream document content in chunks.

trustgraph-flow/trustgraph/librarian/librarian.py

@@ -657,19 +657,20 @@ class Librarian:
         """
         logger.debug(f"Streaming document {request.document_id}, chunk {request.chunk_index}")
 
+        # Server-enforced limits
+        MAX_CHUNK_SIZE = 10 * 1024 * 1024  # 10MB max
+        DEFAULT_CHUNK_SIZE = 1024 * 1024   # 1MB default
+
+        chunk_size = request.chunk_size if request.chunk_size > 0 else DEFAULT_CHUNK_SIZE
+        chunk_size = min(chunk_size, MAX_CHUNK_SIZE)
+
         object_id = await self.table_store.get_document_object_id(
             request.user,
             request.document_id
         )
 
-        # Default chunk size of 1MB
-        chunk_size = request.chunk_size if request.chunk_size > 0 else 1024 * 1024
-
-        # Get the full content and slice out the requested chunk
-        # Note: This is a simple implementation. For true streaming, we'd need
-        # range requests on the object storage.
-        content = await self.blob_store.get(object_id)
-        total_size = len(content)
+        # Get size via stat (no content download)
+        total_size = await self.blob_store.get_size(object_id)
         total_chunks = math.ceil(total_size / chunk_size)
 
         if request.chunk_index >= total_chunks:
@@ -678,12 +679,15 @@ class Librarian:
                 f"document has {total_chunks} chunks"
             )
 
-        start = request.chunk_index * chunk_size
-        end = min(start + chunk_size, total_size)
-        chunk_content = content[start:end]
+        # Calculate byte range
+        offset = request.chunk_index * chunk_size
+        length = min(chunk_size, total_size - offset)
+
+        # Fetch only the requested range
+        chunk_content = await self.blob_store.get_range(object_id, offset, length)
 
         logger.debug(f"Returning chunk {request.chunk_index}/{total_chunks}, "
-                    f"bytes {start}-{end} of {total_size}")
+                    f"bytes {offset}-{offset + length} of {total_size}")
 
         return LibrarianResponse(
             error=None,
@@ -691,7 +695,7 @@ class Librarian:
             chunk_index=request.chunk_index,
             chunks_received=1,  # Using as "current chunk" indicator
             total_chunks=total_chunks,
-            bytes_received=end,
+            bytes_received=offset + length,
             total_bytes=total_size,
         )