docling-project / docling-project/docling
Preset/engine VLM runtime (ApiVlmEngine) drops token usage from VlmPrediction
- Dominant language
- Python
- Stars
- 66.4k
- Forks
- 4.8k
- Avg merge
- 2d 21h
- Merged PRs (30d)
- 84
Description
### Bug
When converting a document via the new pluggable VLM runtime system (`VlmConvertOptions.from_preset(...)` / `ApiVlmEngine`), the token usage returned by the API is silently dropped and never reaches `VlmPrediction.usage`.
`ApiVlmEngine.predict_batch()` (`docling/models/inference_engines/vlm/api_openai_compatible_engine.py`) correctly extracts `usage` from the raw API response and stores it on `VlmEngineOutput.metadata["usage"]`:
```python
return VlmEngineOutput(
text=generated_text,
stop_reason=stop_reason,
metadata={
"generation_time": generation_time,
"num_tokens": num_tokens,
"usage": api_response.usage,
},
)
```
(`usage` here is the full raw usage object from the API response, e.g. `{"prompt_tokens": ..., "completion_tokens": ..., "total_tokens": ...}`; `num_tokens` is just the derived `total_tokens` — see `_extract_total_tokens` in `docling/utils/api_image_request.py`. Both are computed from the same API response and both are lost the same way below.)
However, `_prediction_from_engine_output` (`docling/models/stages/vlm_convert/vlm_convert_model.py`, line 32), the function that converts a `VlmEngineOutput` into the `VlmPrediction` exposed to callers, never reads `output.metadata`:
```python
def _prediction_from_engine_output(output: VlmEngineOutput) -> VlmPrediction:
stop_reason = VlmStopReason.UNSPECIFIED
if output.stop_reason in _VLM_STOP_REASON_VALUES:
stop_reason = VlmStopReason(output.stop_reason)
return VlmPrediction(text=output.text, stop_reason=stop_reason)
```
`output.metadata["usage"]` and `output.metadata["num_tokens"]` are computed, then discarded. Every caller using a preset (`VlmConvertOptions.from_preset(...)`), including all first-party presets shipped with Docling (`granite_docling`, `smoldocling`, `qwen`, etc.), gets `VlmPrediction.usage = None` and `VlmPrediction.num_tokens = None`, even though the underlying API response did include usage data.
By contrast, the older "legacy" VLM path (`ApiVlmOptions` → `ApiVlmModel`, `docling/models/vlm_pipeline_models/api_vlm_model.py`) does forward this correctly:
```python
return VlmPrediction(
text=page_tags,
num_tokens=num_tokens,
usage=api_response.usage,
stop_reason=stop_reason,
input_prompt=input_prompt,
)
```
This means the new preset/engine system, the one actively being developed and the one the legacy path's own deprecation warning steers users toward (`_initialize_legacy_vlm_models` in `docling/pipeline/vlm_pipeline.py`), is currently unable to forward token usage correctly compared to the legacy path.
### Suggested fix
`_prediction_from_engine_output` should also copy `usage`/`num_tokens` from `output.metadata`, mirroring what `ApiVlmModel` already does for the legacy path:
```python
def _prediction_from_engine_output(output: VlmEngineOutput) -> VlmPrediction:
stop_reason = VlmStopReason.UNSPECIFIED
if output.stop_reason in _VLM_STOP_REASON_VALUES:
stop_reason = VlmStopReason(output.stop_reason)
return VlmPrediction(
text=output.text,
stop_reason=stop_reason,
num_tokens=output.metadata.get("num_tokens"),
usage=output.metadata.get("usage"),
)
```
Since `VlmEngineOutput.metadata` is a generic `dict[str, Any]` shared across all inference engines (not just `ApiVlmEngine`), this fix is safe for engines that don't populate `usage`/`num_tokens` — `dict.get(...)` returns `None` for those, which is the existing "no usage data available" contract `VlmPrediction` already supports (its `usage`/`num_tokens` fields both default to `None`).
Happy to open a PR with this change if that's helpful since it looks like a small, self-contained fix.
### Steps to reproduce
1. Serve any OpenAI-compatible VLM endpoint (e.g. via vLLM) that returns a `usage` object in its `chat/completions` response.
2. Run a conversion using a preset, e.g.:
```python
from docling.datamodel.pipeline_options import VlmConvertOptions, VlmPipelineOptions
from docling.datamodel.pipeline_options_vlm_model import ApiVlmEngineOptions
from docling.document_converter import DocumentConverter, PdfFormatOption
from docling.pipeline.vlm_pipeline import VlmPipeline
from docling.datamodel.base_models import InputFormat
from pydantic import AnyUrl
vlm_options = VlmConvertOptions.from_preset(
"granite_docling",
engine_options=ApiVlmEngineOptions(url=AnyUrl("http://localhost:8000/v1/chat/completions")),
)
pipeline_options = VlmPipelineOptions(vlm_options=vlm_options, enable_remote_services=True)
converter = DocumentConverter(
format_options={InputFormat.PDF: PdfFormatOption(pipeline_cls=VlmPipeline, pipeline_options=pipeline_options)}
)
result = converter.convert("some.pdf")
for page in result.pages:
print(page.predictions.vlm_response.usage, page.predictions.vlm_response.num_tokens)
```
3. Observe `usage` and `num_tokens` are always `None`, even though the vLLM server's raw HTTP response includes a populated `usage` object (confirmed via server-side logging of the raw response).
4. For comparison, running the same conversion through the legacy path (`ApiVlmOptions` directly, instead of `VlmConvertOptions.from_preset(...)`) against the same endpoint correctly populates `usage`/`num_tokens` on every page.
### Docling version
```
Docling version: 2.113.0
```
### Python version
```
Python 3.12.13
```
Contributor guide
Research direction
Start in docling/models/stages/vlm_convert/vlm_convert_model.py at _prediction_from_engine_output, then inspect ApiVlmEngine.predict_batch() in docling/models/inference_engines/vlm/api_openai_compatible_engine.py and the legacy forwarding in docling/models/vlm_pipeline_models/api_vlm_model.py. Reproduce a preset conversion against an OpenAI-compatible endpoint and verify that VlmPrediction.usage and num_tokens contain the API response values instead of None.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- ai, api
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 82/100