test(integration): bill OCR annotation pages at annotation_cost_per_page (Pylon #7958)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

tests/integration/pricing/test_ocr_page_pricing.py

@@ -0,0 +1,59 @@
+import uuid
+from typing import Final
+
+import pytest
+
+from tests.integration._support.client import Gateway, eventually, string_value
+from tests.integration._support.database import read_rows
+from tests.integration._support.upstream import delete_scenario, register_scenario
+from tests.integration.cost_calculation.cost_tracking_case import JsonResponse
+
+
+@pytest.mark.covers("pricing.ocr.annotation_pages_billed_at_annotation_rate")
+def test_ocr_annotation_pages_are_billed_at_annotation_cost_per_page(gateway: Gateway) -> None:
+    with gateway.scenario() as scenario:
+        scenario_id: Final = f"ocr-annotation-{uuid.uuid4().hex[:12]}"
+        handle: Final = register_scenario(
+            scenario_id,
+            JsonResponse(
+                content_type="application/json",
+                body={
+                    "pages": [{"index": index, "markdown": f"page {index}"} for index in range(3)],
+                    "model": "integration-ocr",
+                    "document_annotation": '{"title": "annotated"}',
+                    "usage_info": {"pages_processed": 3, "pages_processed_annotation": 2, "doc_size_bytes": 4096},
+                },
+            ),
+        )
+        scenario.cleanups.callback(delete_scenario, handle)
+        model: Final = scenario.model(
+            model=f"mistral/integration-ocr-{scenario_id}",
+            api_base=f"{handle.api_base()}/v1",
+            ocr_cost_per_page=0.002,
+            annotation_cost_per_page=0.01,
+        )
+        response: Final = gateway.request(
+            "POST",
+            "/v1/ocr",
+            {
+                "model": model,
+                "document": {"type": "document_url", "document_url": "https://example.com/annotated.pdf"},
+                "document_annotation_format": {"type": "json_schema", "json_schema": {"name": "title"}},
+            },
+        )
+        assert response.status_code == 200, response.text
+        assert response.json()["usage_info"] == {
+            "pages_processed": 3,
+            "pages_processed_annotation": 2,
+            "credits": None,
+            "doc_size_bytes": 4096,
+        }, response.text
+        expected: Final = 3 * 0.002 + 2 * 0.01
+        assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected), response.text
+        request_id: Final = string_value(response.headers["x-litellm-call-id"])
+        rows: Final = eventually(
+            lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,)),
+            lambda values: len(values) == 1,
+            seconds=70,
+        )
+        assert float(rows[0]["spend"]) == pytest.approx(expected)