feat(server): add json format to docling extraction

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

Author: aleskalfas

apps/agentstack-server/src/agentstack_server/api/routes/files.py

@@ -14,7 +14,7 @@
     RequiresContextPermissions,
 )
 from agentstack_server.api.schema.common import EntityModel
-from agentstack_server.api.schema.files import FileListQuery
+from agentstack_server.api.schema.files import FileListQuery, TextExtractionRequest
 from agentstack_server.domain.models.common import PaginatedResult
 from agentstack_server.domain.models.file import AsyncFile, ExtractionStatus, File, TextExtraction
 from agentstack_server.domain.models.permissions import AuthorizedUser
@@ -114,6 +114,7 @@ async def create_text_extraction(
     file_id: UUID,
     file_service: FileServiceDependency,
     user: Annotated[AuthorizedUser, Depends(RequiresContextPermissions(files={"write", "extract"}))],
+    request: TextExtractionRequest | None = None,
 ) -> EntityModel[TextExtraction]:
     """Create or return text extraction for a file.
 
@@ -122,8 +123,13 @@ async def create_text_extraction(
     - If extraction is pending/in-progress, returns current status
     - If no extraction exists, creates a new one
     """
+    if request is None:
+        request = TextExtractionRequest()
+
     return EntityModel(
-        await file_service.create_extraction(file_id=file_id, user=user.user, context_id=user.context_id)
+        await file_service.create_extraction(
+            file_id=file_id, user=user.user, context_id=user.context_id, extras=request.extras
+        )
     )
 
 

apps/agentstack-server/src/agentstack_server/api/schema/files.py

@@ -36,3 +36,12 @@ class FileListQuery(PaginationQuery):
         description="Case-insensitive partial match search on filename (e.g., 'doc' matches 'my_document.pdf')",
     )
     order_by: str = Field(default_factory=lambda: "created_at", pattern="^created_at|filename|file_size_bytes$")
+
+
+class TextExtractionRequest(BaseModel):
+    """Request schema for text extraction."""
+
+    extras: dict | None = Field(
+        default=None,
+        description="Additional options for text extraction (e.g., {'json_format': True})",
+    )

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

@@ -60,6 +60,7 @@ class TextExtraction(BaseModel):
     job_id: str | None = None
     error_message: str | None = None
     extraction_metadata: ExtractionMetadata | None = None
+    extras: dict | None = None
     started_at: AwareDatetime | None = None
     finished_at: AwareDatetime | None = None
     created_at: AwareDatetime = Field(default_factory=utc_now)

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

@@ -70,5 +70,11 @@ async def get_file_metadata(self, *, file_id: UUID) -> FileMetadata: ...
 @runtime_checkable
 class ITextExtractionBackend(Protocol):
     @asynccontextmanager
-    async def extract_text(self, *, file_url: AnyUrl, timeout: timedelta | None = None) -> AsyncIterator[AsyncFile]:  # noqa: ASYNC109
+    async def extract_text(
+        self,
+        *,
+        file_url: AnyUrl,
+        timeout: timedelta | None = None,  # noqa: ASYNC109
+        extras: dict | None = None,
+    ) -> AsyncIterator[AsyncFile]:
         yield ...  # pyright: ignore [reportReturnType]

apps/agentstack-server/src/agentstack_server/infrastructure/persistence/migrations/alembic/versions/90f8cd7d6bef_.py

@@ -0,0 +1,35 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+"""add extras column to text_extractions table
+
+Revision ID: 90f8cd7d6bef
+Revises: 214ed3790c6d
+Create Date: 2025-11-26 17:44:56.490088
+
+"""
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "90f8cd7d6bef"
+down_revision: str | None = "214ed3790c6d"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("text_extractions", sa.Column("extras", sa.JSON(), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("text_extractions", "extras")
+    # ### end Alembic commands ###

apps/agentstack-server/src/agentstack_server/infrastructure/persistence/repositories/file.py

@@ -54,6 +54,7 @@
     Column("job_id", String(255), nullable=True),
     Column("error_message", Text, nullable=True),
     Column("extraction_metadata", JSON, nullable=True),
+    Column("extras", JSON, nullable=True),
     Column("started_at", DateTime(timezone=True), nullable=True),
     Column("finished_at", DateTime(timezone=True), nullable=True),
     Column("created_at", DateTime(timezone=True), nullable=False),
@@ -211,6 +212,7 @@ def _to_text_extraction(self, row: Row) -> TextExtraction:
                 "job_id": row.job_id,
                 "error_message": row.error_message,
                 "extraction_metadata": row.extraction_metadata,
+                "extras": row.extras,
                 "started_at": row.started_at,
                 "finished_at": row.finished_at,
                 "created_at": row.created_at,
@@ -227,6 +229,7 @@ async def create_extraction(self, *, extraction: TextExtraction) -> None:
             job_id=extraction.job_id,
             error_message=extraction.error_message,
             extraction_metadata=extraction_metadata and extraction_metadata.model_dump(mode="json"),
+            extras=extraction.extras,
             started_at=extraction.started_at,
             finished_at=extraction.finished_at,
             created_at=extraction.created_at,

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

@@ -1,6 +1,7 @@
 # Copyright 2025 © BeeAI a Series of LF Projects, LLC
 # SPDX-License-Identifier: Apache-2.0
 
+import logging
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from datetime import timedelta
@@ -11,7 +12,9 @@
 from agentstack_server.configuration import DoclingExtractionConfiguration
 from agentstack_server.domain.models.file import AsyncFile
 from agentstack_server.domain.repositories.file import ITextExtractionBackend
-from agentstack_server.utils.utils import extract_string_value_stream
+from agentstack_server.utils.utils import extract_object_value_stream, extract_string_value_stream
+
+logger = logging.getLogger(__name__)
 
 
 class DoclingTextExtractionBackend(ITextExtractionBackend):
@@ -20,12 +23,22 @@ def __init__(self, config: DoclingExtractionConfiguration):
         self._enabled = config.enabled
 
     @asynccontextmanager
-    async def extract_text(self, *, file_url: AnyUrl, timeout: timedelta | None = None) -> AsyncIterator[AsyncFile]:  # noqa: ASYNC109
+    async def extract_text(
+        self,
+        *,
+        file_url: AnyUrl,
+        timeout: timedelta | None = None,  # noqa: ASYNC109
+        extras: dict | None = None,
+    ) -> AsyncIterator[AsyncFile]:
         if not self._enabled:
             raise RuntimeError(
                 "Docling extraction backend is not enabled, please check the documentation how to enable it"
             )
 
+        # Switch to json output formats if specified in extras
+        is_json_format = extras and extras.get("json_format") is True
+        to_formats = ["json"] if is_json_format else ["md"]
+
         timeout = timeout or timedelta(minutes=5)
         async with (
             AsyncClient(base_url=str(self._config.docling_service_url), timeout=timeout.seconds) as client,
@@ -34,7 +47,7 @@ async def extract_text(self, *, file_url: AnyUrl, timeout: timedelta | None = No
                 "/v1/convert/source",
                 json={
                     "options": {
-                        "to_formats": ["md"],
+                        "to_formats": to_formats,
                         "document_timeout": timeout.total_seconds(),
                         "image_export_mode": "placeholder",
                     },
@@ -44,19 +57,22 @@ async def extract_text(self, *, file_url: AnyUrl, timeout: timedelta | None = No
         ):
             response.raise_for_status()
 
-            md_stream = None
+            resp_stream = None
 
             async def read(chunk_size: int = 1024) -> bytes:
-                nonlocal md_stream
-                if not md_stream:
-                    md_stream = extract_string_value_stream(response.aiter_text, "md_content", chunk_size)
-                async for text_chunk in md_stream:
+                nonlocal resp_stream
+                if not resp_stream:
+                    if is_json_format:
+                        resp_stream = extract_object_value_stream(response.aiter_text, "json_content", chunk_size)
+                    else:
+                        resp_stream = extract_string_value_stream(response.aiter_text, "md_content", chunk_size)
+                async for text_chunk in resp_stream:
                     return text_chunk.encode("utf-8")
                 return b""
 
             yield AsyncFile(
-                filename="extracted_text.md",
-                content_type="text/markdown",
+                filename="extracted_response.json" if is_json_format else "extracted_response.md",
+                content_type="application/json" if is_json_format else "text/markdown",
                 read=read,
                 size=None,  # size is unknown beforehand
             )

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

@@ -60,7 +60,9 @@ 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) as extracted_file:
+            async with self._extraction_backend.extract_text(
+                file_url=file_url, extras=extraction.extras
+            ) as extracted_file:
                 extracted_db_file = await self.upload_file(
                     file=extracted_file,
                     user=user,
@@ -157,7 +159,9 @@ async def delete(self, *, file_id: UUID, user: User, context_id: UUID | None = N
                 await self._object_storage.delete_files(file_ids=[file_id])
                 await uow.commit()
 
-    async def create_extraction(self, *, file_id: UUID, user: User, context_id: UUID | None = None) -> TextExtraction:
+    async def create_extraction(
+        self, *, file_id: UUID, user: User, context_id: UUID | None = None, extras: dict | None = None
+    ) -> TextExtraction:
         async with self._uow() as uow:
             # Check user permissions
             await uow.files.get(file_id=file_id, user_id=user.id, context_id=context_id, file_type=FileType.USER_UPLOAD)
@@ -174,7 +178,7 @@ async def create_extraction(self, *, file_id: UUID, user: User, context_id: UUID
                         raise TypeError(f"Unknown extraction status: {extraction.status}")
             except EntityNotFoundError:
                 file_metadata = await self._object_storage.get_file_metadata(file_id=file_id)
-                extraction = TextExtraction(file_id=file_id)
+                extraction = TextExtraction(file_id=file_id, extras=extras)
                 if file_metadata.content_type in {"text/plain", "text/markdown"}:
                     extraction.set_completed(
                         extracted_file_id=file_id,  # Point to itself since it's already text

apps/agentstack-server/src/agentstack_server/utils/utils.py

@@ -106,3 +106,106 @@ async def extract_string_value_stream(
         raise EOFError("Unterminated string value in JSON input")
     else:
         raise KeyError(f"Key {key} not found in JSON input")
+
+
+async def extract_object_value_stream(
+    async_stream: Callable[[int], AsyncIterable[str]], key: str, chunk_size: int = 1024
+) -> AsyncIterable[str]:
+    """
+    Extract a JSON object or array value from a streaming JSON response.
+
+    Parses a JSON stream incrementally to find a specific key and extract its
+    associated object or array value. Yields the content in chunks as it's read,
+    without waiting for the entire stream to complete. Properly handles nested
+    objects, arrays, and escaped characters within string values.
+
+    Args:
+        async_stream: Async function that yields text chunks
+        key: The JSON key to extract the object from
+        chunk_size: Size of chunks to read from the stream
+
+    Yields:
+        String chunks containing the JSON object content
+
+    Raises:
+        KeyError: If the key is not found in the JSON
+        EOFError: If the JSON object is unterminated
+    """
+    buffer = ""
+    max_buffer_size = len(key) * 4 + 20  # Enough to store key pattern like: "key" : {
+    state = "outside"
+    brace_depth = 0
+    bracket_depth = 0
+    in_string = False
+    escape_next = False
+    processed_idx = 0  # Track how much of the buffer we've already yielded
+
+    async for chunk in async_stream(chunk_size):
+        buffer += chunk
+
+        if state == "outside":
+            # Look for the key followed by : and either { or [
+            if match := re.search(rf'"{key}" *: *([{{[])', buffer):
+                buffer = buffer[match.end() - 1 :]  # Start from the opening brace/bracket
+                state = "inside"
+                processed_idx = 0
+            elif len(buffer) > max_buffer_size:
+                # Only trim buffer if it's getting too large
+                buffer = buffer[-max_buffer_size:]
+
+        if state == "inside":
+            # Process only the new characters in the buffer
+            chars_to_yield = []
+            idx = processed_idx
+
+            while idx < len(buffer):
+                char = buffer[idx]
+
+                if escape_next:
+                    escape_next = False
+                    chars_to_yield.append(char)
+                    idx += 1
+                    continue
+
+                if char == "\\":
+                    escape_next = True
+                    chars_to_yield.append(char)
+                    idx += 1
+                    continue
+
+                if char == '"':
+                    in_string = not in_string
+                    chars_to_yield.append(char)
+                    idx += 1
+                    continue
+
+                if not in_string:
+                    if char == "{":
+                        brace_depth += 1
+                    elif char == "}":
+                        brace_depth -= 1
+                    elif char == "[":
+                        bracket_depth += 1
+                    elif char == "]":
+                        bracket_depth -= 1
+
+                chars_to_yield.append(char)
+
+                # Check if we've reached the end of the object
+                if not in_string and brace_depth == 0 and bracket_depth == 0:
+                    # Yield any remaining characters and return
+                    if chars_to_yield:
+                        yield "".join(chars_to_yield)
+                    return
+
+                idx += 1
+
+            # Yield only the newly processed characters
+            if chars_to_yield:
+                yield "".join(chars_to_yield)
+                processed_idx = idx
+
+    if state == "inside":
+        raise EOFError("Unterminated JSON object in input")
+    else:
+        raise KeyError(f"Key {key} not found in JSON input")