remove streaming parser

Signed-off-by: Radek Ježek <radek.jezek@ibm.com>

Author: jezekra1

apps/agentstack-server/pyproject.toml

@@ -47,6 +47,7 @@ dependencies = [
     "mcp>=1.13.1",
     "opentelemetry-instrumentation-httpx>=0.59b0",
     "opentelemetry-instrumentation-fastapi>=0.59b0",
+    "ijson>=3.4.0.post0",
 ]
 
 [dependency-groups]

apps/agentstack-server/src/agentstack_server/domain/models/file.py

@@ -1,8 +1,9 @@
 # Copyright 2025 © BeeAI a Series of LF Projects, LLC
 # SPDX-License-Identifier: Apache-2.0
 
-from collections.abc import Awaitable, Callable
+from collections.abc import AsyncIterator, Awaitable, Callable
 from enum import StrEnum
+from typing import Self
 from uuid import UUID, uuid4
 
 from pydantic import AwareDatetime, BaseModel, Field
@@ -55,6 +56,36 @@ class AsyncFile(BaseModel):
     read: Callable[[int], Awaitable[bytes]]
     size: int | None = None
 
+    @classmethod
+    def from_async_iterator(cls, iterator: AsyncIterator[bytes], filename: str, content_type: str) -> Self:
+        buffer = b""
+
+        async def read(size: int = 8192) -> bytes:
+            nonlocal buffer
+            while len(buffer) < size:
+                try:
+                    buffer += await anext(iterator)
+                except StopAsyncIteration:
+                    break
+
+            result = buffer[:size]
+            buffer = buffer[size:]
+            return result
+
+        return cls(filename=filename, content_type=content_type, read=read)
+
+    @classmethod
+    def from_bytes(cls, content: bytes, filename: str, content_type: str) -> Self:
+        pos = 0
+
+        async def read(size: int = 8192) -> bytes:
+            nonlocal pos
+            result = content[pos : pos + size]
+            pos += len(result)
+            return result
+
+        return cls(filename=filename, content_type=content_type, read=read, size=len(content))
+
 
 class File(BaseModel):
     id: UUID = Field(default_factory=uuid4)

apps/agentstack-server/src/agentstack_server/domain/repositories/file.py

@@ -89,5 +89,5 @@ async def extract_text(
         file_url: AnyUrl,
         timeout: timedelta | None = None,  # noqa: ASYNC109
         settings: TextExtractionSettings | None = None,
-    ) -> AsyncIterator[list[tuple[AsyncFile, ExtractionFormat]]]:
+    ) -> AsyncIterator[AsyncIterator[tuple[AsyncFile, ExtractionFormat]]]:
         yield ...  # pyright: ignore [reportReturnType]

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

@@ -5,42 +5,66 @@
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from datetime import timedelta
+from decimal import Decimal
+from typing import NamedTuple, cast
 
-from httpx import AsyncClient
+import ijson
+import orjson
+from httpx import AsyncClient as HttpxAsyncClient
 from pydantic import AnyUrl
 
 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 StreamingExtractor, build_stream_reader
 
 logger = logging.getLogger(__name__)
 
 
-class DoclingFormatInfo:
-    """Docling-specific format information."""
+class DoclingFormatInfo(NamedTuple):
+    api_option: str
+    file_extension: str
+    response_field_key: str
+    content_type: str
 
-    def __init__(self, format_value: str, field_key: str, content_type: str):
-        self.format_value = format_value  # The format value for the API (e.g., "md", "json")
-        self.field_key = field_key  # The JSON response field key (e.g., "md_content", "json_content")
-        self.content_type = content_type  # The MIME type for the file
-        self.is_object = content_type == "application/json"  # Whether the field contains a JSON object
 
+_DOCLING_FORMAT_INFO: dict[ExtractionFormat, DoclingFormatInfo] = {
+    ExtractionFormat.MARKDOWN: DoclingFormatInfo("md", "md", "md_content", "text/markdown"),
+    ExtractionFormat.VENDOR_SPECIFIC_JSON: DoclingFormatInfo("json", "json", "json_content", "application/json"),
+}
 
-def _get_docling_format_info(format: ExtractionFormat) -> DoclingFormatInfo:
-    """Get docling-specific format information for an ExtractionFormat."""
-    mapping = {
-        ExtractionFormat.MARKDOWN: DoclingFormatInfo(
-            format_value="md", field_key="md_content", content_type="text/markdown"
-        ),
-        ExtractionFormat.VENDOR_SPECIFIC_JSON: DoclingFormatInfo(
-            format_value="json", field_key="json_content", content_type="application/json"
-        ),
-    }
-    return mapping[format]
+
+async def _process_docling_stream(
+    async_file: AsyncFile, formats: list[ExtractionFormat]
+) -> AsyncIterator[tuple[AsyncFile, ExtractionFormat]]:
+    key_map = {info.response_field_key: (fmt, info) for fmt, info in _DOCLING_FORMAT_INFO.items() if fmt in formats}
+
+    def serialize(obj):
+        if isinstance(obj, Decimal):
+            return float(obj)
+        raise TypeError
+
+    async for k, v in ijson.kvitems_async(async_file, "document", use_float=False):  # pyright: ignore[reportAny]
+        if k in key_map:
+            fmt, info = key_map[k]
+
+            content = v.encode("utf-8") if isinstance(v, str) else cast(bytes, orjson.dumps(v, default=serialize))  # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue]
+
+            async_file = AsyncFile.from_bytes(
+                filename=f"extracted_response.{info.file_extension}",
+                content_type=info.content_type,
+                content=content,
+            )
+
+            yield (async_file, fmt)
 
 
 class DoclingTextExtractionBackend(ITextExtractionBackend):
+    """
+    # TODO: [DISCLAIMER] this is loading entire field value to memory (which in case of docling is entire document)
+    # Implementing a streaming parser has been deemed too complex, instead the text extraction worker should be
+    # scaled independently with memory bounded by concurrency and max_single_file_size configuration
+    """
+
     def __init__(self, config: DoclingExtractionConfiguration):
         self._config = config
         self._enabled = config.enabled
@@ -57,82 +81,35 @@ async def extract_text(
         file_url: AnyUrl,
         timeout: timedelta | None = None,  # noqa: ASYNC109
         settings: TextExtractionSettings | None = None,
-    ) -> AsyncIterator[list[tuple[AsyncFile, ExtractionFormat]]]:
+    ) -> AsyncIterator[AsyncIterator[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
+        Streams the response and yields files as they are parsed.
         """
         if not self._enabled:
             raise RuntimeError(
                 "Docling extraction backend is not enabled, please check the documentation how to enable it"
             )
 
-        # Exports to both markdown and vendor-specific JSON by default
         formats = settings.formats if settings else [ExtractionFormat.MARKDOWN, ExtractionFormat.VENDOR_SPECIFIC_JSON]
 
-        # Get docling-specific format information for all requested formats
-        # Keep pairs of (format_info, extraction_format) to maintain association
-        format_pairs = [(fmt, _get_docling_format_info(fmt)) for fmt in formats]
-        to_formats = [info.format_value for _, info in format_pairs]
-
         timeout = timeout or timedelta(minutes=5)
+
         async with (
-            AsyncClient(base_url=str(self._config.docling_service_url), timeout=timeout.seconds) as client,
+            HttpxAsyncClient(base_url=str(self._config.docling_service_url), timeout=timeout.seconds) as client,
             client.stream(
                 "POST",
                 "/v1/convert/source",
                 json={
                     "options": {
-                        "to_formats": to_formats,
+                        "to_formats": [_DOCLING_FORMAT_INFO[fmt].api_option for fmt in formats],
                         "document_timeout": timeout.total_seconds(),
                         "image_export_mode": "placeholder",
                     },
                     "sources": [{"kind": "http", "url": str(file_url)}],
                 },
             ) as response,
         ):
-            response.raise_for_status()
-
-            # 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}
-
-            # 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
+            response.raise_for_status()  # pyright: ignore[reportUnusedCallResult]
+            async_file = AsyncFile.from_async_iterator(response.aiter_bytes(), "tmp", "application/json")
+            yield _process_docling_stream(async_file, formats)

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

@@ -91,9 +91,9 @@ async def extract_text(self, file_id: UUID, job_id: str):
             async with self._extraction_backend.extract_text(
                 file_url=file_url,
                 settings=extraction.extraction_metadata.settings if extraction.extraction_metadata else None,
-            ) as extracted_files_with_format:
+            ) as extracted_files_iterator:
                 extracted_files = []
-                for async_file, extraction_format in extracted_files_with_format:
+                async for async_file, extraction_format in extracted_files_iterator:
                     extracted_db_file = await self.upload_file(
                         file=async_file,
                         user=user,