fix: make extract_text streamable again

Signed-off-by: Aleš Kalfas <kalfas.ales@gmail.com>

Author: aleskalfas

apps/agentstack-server/src/agentstack_server/infrastructure/text_extraction/docling.py

@@ -12,7 +12,7 @@
 from agentstack_server.configuration import DoclingExtractionConfiguration
 from agentstack_server.domain.models.file import AsyncFile, ExtractionFormat, TextExtractionSettings
 from agentstack_server.domain.repositories.file import ITextExtractionBackend
-from agentstack_server.utils.utils import extract_multiple_values_stream
+from agentstack_server.utils.utils import StreamingExtractor, build_stream_reader
 
 logger = logging.getLogger(__name__)
 
@@ -58,6 +58,34 @@ async def extract_text(
         timeout: timedelta | None = None,  # noqa: ASYNC109
         settings: TextExtractionSettings | None = None,
     ) -> AsyncIterator[list[tuple[AsyncFile, ExtractionFormat]]]:
+        """
+        Extract text from a file using the Docling service.
+
+        IMPORTANT: This method MUST be used as an async context manager. The returned AsyncFile
+        objects stream data directly from the HTTP response and must be consumed INSIDE the
+        context block:
+
+            async with backend.extract_text(file_url) as files:
+                for async_file, format in files:
+                    # Read and process files HERE, inside the context
+                    data = await async_file.read()
+            # Files are no longer readable after context exits
+
+        Do NOT store AsyncFile references outside the context - they will fail to read with
+        "StreamingExtractor is closed" error.
+
+        Args:
+            file_url: URL of the file to extract text from
+            timeout: Maximum time to wait for extraction (default: 5 minutes)
+            settings: Optional extraction settings (formats to extract)
+
+        Yields:
+            List of (AsyncFile, ExtractionFormat) tuples containing extracted content
+
+        Raises:
+            RuntimeError: If Docling extraction backend is not enabled
+            HTTPError: If Docling service returns an error
+        """
         if not self._enabled:
             raise RuntimeError(
                 "Docling extraction backend is not enabled, please check the documentation how to enable it"
@@ -92,34 +120,19 @@ async def extract_text(
             # Build a dict of keys to extract with their types
             keys_to_extract = {info.field_key: "object" if info.is_object else "string" for _, info in format_pairs}
 
-            # Extract all requested formats in a single pass through the stream
-            extracted_contents = await extract_multiple_values_stream(response.aiter_text, keys_to_extract)
-
-            # Create (AsyncFile, ExtractionFormat) tuples for each format
-            files = []
-            for extraction_format, info in format_pairs:
-                content = extracted_contents[info.field_key]
-
-                # Create a read function that returns the buffered content
-                def create_read_fn(buffered_content: bytes):
-                    position = 0
-
-                    async def read(chunk_size: int = 1024) -> bytes:
-                        nonlocal position
-                        if position >= len(buffered_content):
-                            return b""
-                        chunk = buffered_content[position : position + chunk_size]
-                        position += len(chunk)
-                        return chunk
-
-                    return read
-
-                async_file = AsyncFile(
-                    filename=f"extracted_response.{info.format_value}",
-                    content_type=info.content_type,
-                    read=create_read_fn(content),
-                    size=len(content),
-                )
-                files.append((async_file, extraction_format))
-
-            yield files
+            # Create streaming extractor with context manager for proper cleanup
+            async with StreamingExtractor(response.aiter_text, keys_to_extract) as extractor:
+                iterators = extractor.get_iterators()
+
+                # Create (AsyncFile, ExtractionFormat) tuples for each format
+                files = []
+                for extraction_format, info in format_pairs:
+                    async_file = AsyncFile(
+                        filename=f"extracted_response.{info.format_value}",
+                        content_type=info.content_type,
+                        read=build_stream_reader(iterators[info.field_key]),
+                        size=None,  # Size unknown when streaming
+                    )
+                    files.append((async_file, extraction_format))
+
+                yield files

apps/agentstack-server/src/agentstack_server/service_layer/services/files.py

@@ -51,6 +51,20 @@ def __init__(
         self._extraction_backend = extraction_backend
 
     async def extract_text(self, file_id: UUID, job_id: str):
+        """
+        Extract text from a file using the configured extraction backend.
+
+        This method coordinates the entire extraction process: fetching the file,
+        extracting text, and uploading the extracted content.
+
+        Args:
+            file_id: ID of the file to extract text from
+            job_id: Background job ID for tracking
+
+        Raises:
+            CancelledError: If the job is cancelled
+            Exception: For any extraction or upload errors
+        """
         error_log = []
         async with self._uow() as uow:
             extraction = await uow.files.get_extraction_by_file_id(file_id=file_id)
@@ -67,6 +81,7 @@ async def extract_text(self, file_id: UUID, job_id: str):
         try:
             file_url = await self._object_storage.get_file_url(file_id=file_id)
             error_log.append(f"file url: {file_url}")
+
             async with self._extraction_backend.extract_text(
                 file_url=file_url,
                 settings=extraction.extraction_metadata.settings if extraction.extraction_metadata else None,
@@ -81,6 +96,7 @@ async def extract_text(self, file_id: UUID, job_id: str):
                         parent_file_id=file_id,
                     )
                     extracted_files.append(ExtractedFileInfo(file_id=extracted_db_file.id, format=extraction_format))
+
             extraction.set_completed(extracted_files=extracted_files)
             async with self._uow() as uow:
                 await uow.files.update_extraction(extraction=extraction)