diff --git a/src/extract_bench/inference/providers/extract/llamaextract_v2_api.py b/src/extract_bench/inference/providers/extract/llamaextract_v2_api.py index 588a9c17..8c5c2c9f 100644 --- a/src/extract_bench/inference/providers/extract/llamaextract_v2_api.py +++ b/src/extract_bench/inference/providers/extract/llamaextract_v2_api.py @@ -472,7 +472,7 @@ def _poll_job(self, client: httpx.Client, job_id: str, params: dict[str, str]) - start = time.monotonic() poll_started_at = datetime.now().isoformat() - # Request both expandable blocks on every poll. The V2 extract API + # Request expandable metadata and telemetry on every poll. The V2 extract API # strips them from the GET response unless the caller opts in. # # ``extract_metadata``: without it citations have no ``bounding_boxes`` @@ -482,7 +482,10 @@ def _poll_job(self, client: httpx.Client, job_id: str, params: dict[str, str]) - # ``usage.num_pages_billed``. It comes back null unless expanded, and # without it there is no billable page count, so cost falls back to a # schema estimate (or to nothing at all on non-Agentic-Plus tiers). - poll_params: dict[str, Any] = {**params, "expand": ["metadata", "extract_metadata"]} + poll_params: dict[str, Any] = { + **params, + "expand": ["metadata", "extract_metadata", "latency", "metadata_state_transitions"], + } poll_history: list[dict[str, Any]] = [] last_recorded_status: str | None = None @@ -543,6 +546,8 @@ def _poll_job(self, client: httpx.Client, job_id: str, params: dict[str, str]) - # estimate, so it has to survive into the raw result. "metadata": data.get("metadata"), "status": status, + "latency": data.get("latency"), + "metadata_state_transitions": data.get("metadata_state_transitions"), "poll_history": poll_history, "poll_started_at": poll_started_at, "poll_completed_at": datetime.now().isoformat(), diff --git a/tests/extract_bench/inference/providers/extract/test_llamaextract_v2_telemetry.py b/tests/extract_bench/inference/providers/extract/test_llamaextract_v2_telemetry.py new file mode 100644 index 00000000..12afe3b3 --- /dev/null +++ b/tests/extract_bench/inference/providers/extract/test_llamaextract_v2_telemetry.py @@ -0,0 +1,71 @@ +from datetime import datetime + +import pytest + +httpx = pytest.importorskip("httpx") + +from extract_bench.inference.providers.extract import llamaextract_v2_api as module # noqa: E402 +from extract_bench.schemas.pipeline import PipelineSpec # noqa: E402 +from extract_bench.schemas.pipeline_io import InferenceRequest, RawInferenceResult # noqa: E402 +from extract_bench.schemas.product import ProductType # noqa: E402 + + +@pytest.mark.parametrize( + "telemetry", + [ + {"latency": {"parse_ms": 1000}, "metadata_state_transitions": [{"status": "COMPLETED"}]}, + {}, + {"latency": None, "metadata_state_transitions": None}, + ], +) +def test_poll_retains_optional_service_telemetry(monkeypatch, telemetry): + provider = module.LlamaExtractV2Provider("llamaextract_v2", {"api_key": "offline"}) + calls = [] + monkeypatch.setattr(module.time, "sleep", lambda seconds: None) + + def handler(request): + calls.append(request) + assert request.url.params["project_id"] == "project" + assert set(request.url.params.get_list("expand")) == { + "metadata", + "extract_metadata", + "latency", + "metadata_state_transitions", + } + status = "RUNNING" if len(calls) == 1 else "COMPLETED" + return httpx.Response( + 200, + json={ + "status": status, + "extract_result": {"amount": 12}, + "metadata": {"usage": {"num_pages_extracted": 1}}, + "extract_metadata": {"parse_job_id": "parse-1"}, + **telemetry, + }, + ) + + with httpx.Client(base_url="https://example.test", transport=httpx.MockTransport(handler)) as client: + result = provider._poll_job(client, "job-1", {"project_id": "project"}) + assert len(calls) == 2 + assert result["latency"] == telemetry.get("latency") + assert result["metadata_state_transitions"] == telemetry.get("metadata_state_transitions") + assert result["metadata"]["usage"]["num_pages_extracted"] == 1 + assert result["spawned_parse_job_id"] == "parse-1" + assert result["data"] == {"amount": 12} + pipeline = PipelineSpec( + pipeline_name="fixture", provider_name="llamaextract_v2", product_type=ProductType.EXTRACT, config={} + ) + raw = RawInferenceResult( + request=InferenceRequest(example_id="one", source_file_path="one.pdf", product_type=ProductType.EXTRACT), + pipeline=pipeline, + pipeline_name="fixture", + product_type=ProductType.EXTRACT, + raw_output=result, + started_at=datetime(2026, 1, 1), + completed_at=datetime(2026, 1, 1), + latency_in_ms=0, + ) + normalized = provider.normalize(raw) + saved = normalized.model_dump(mode="json") + assert saved["raw_output"]["latency"] == telemetry.get("latency") + assert saved["raw_output"]["metadata_state_transitions"] == telemetry.get("metadata_state_transitions")