From 6c10a93fdeaf59d3a2f1564a2d171ccdcfd7f8c1 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 10 Oct 2026 11:54:11 -0700 Subject: [PATCH] test: pin OCR parse bodies and billed cost, group new e2e cases into classes --- .../test_messages_beta_headers_e2e.py | 90 ++--- ...ages_prompt_caching_and_tool_search_e2e.py | 356 +++++++++--------- .../test_aws_secret_manager_e2e.py | 16 +- tests/integration/sdk/test_ocr_matrix_sdk.py | 101 +++-- 4 files changed, 300 insertions(+), 263 deletions(-) diff --git a/tests/e2e/llm_translation/test_messages_beta_headers_e2e.py b/tests/e2e/llm_translation/test_messages_beta_headers_e2e.py index a688517c77e..91a31201b0f 100644 --- a/tests/e2e/llm_translation/test_messages_beta_headers_e2e.py +++ b/tests/e2e/llm_translation/test_messages_beta_headers_e2e.py @@ -37,50 +37,52 @@ def _deployment(backend: str) -> LiteLLMParamsBody: ) -@pytest.mark.parametrize( - ("backend", "beta_provider", "tools"), - [ - pytest.param(_ANTHROPIC_MODEL, "anthropic", [_CODE_EXECUTION], id="anthropic"), - pytest.param(f"bedrock/invoke/{_BEDROCK_MODEL_ID}", "bedrock", [], id="bedrock-invoke"), - pytest.param(f"bedrock/converse/{_BEDROCK_MODEL_ID}", "bedrock_converse", [], id="bedrock-converse"), - ], -) -@pytest.mark.covers( - "llm.messages.anthropic.basic.nonstream.works", - "llm.messages.bedrock_invoke.basic.nonstream.works", - "llm.messages.bedrock_converse.basic.nonstream.works", -) -@meta( - Subject( - domain=Domain.LLM_TRANSLATION, - route=Route.MESSAGES, - providers=(Provider.ANTHROPIC, Provider.BEDROCK), - models=(_ANTHROPIC_MODEL, _BEDROCK_MODEL_ID), - mode=Mode.NONSTREAM, +class TestAnthropicBetaHeaders: + @pytest.mark.parametrize( + ("backend", "beta_provider", "tools"), + [ + pytest.param(_ANTHROPIC_MODEL, "anthropic", [_CODE_EXECUTION], id="anthropic"), + pytest.param(f"bedrock/invoke/{_BEDROCK_MODEL_ID}", "bedrock", [], id="bedrock-invoke"), + pytest.param(f"bedrock/converse/{_BEDROCK_MODEL_ID}", "bedrock_converse", [], id="bedrock-converse"), + ], ) -) -def test_every_mapped_beta_header_is_accepted_by_the_provider( - proxy: ProxyClient, - resources: ResourceManager, - sdk: SdkClients, - backend: str, - beta_provider: str, - tools: list[ToolUnionParam], -) -> None: - model: Final = f"e2e-messages-beta-{unique_marker()}" - model_id: Final = proxy.create_model(model, _deployment(backend)) - resources.defer(lambda: proxy.delete_model(model_id)) - - message: Final = sdk.anthropic(resources.key()).messages.create( - model=model, - max_tokens=64, - messages=[{"role": "user", "content": "Say 'hello' and nothing else"}], - tools=tools, - extra_headers={"anthropic-beta": _mapped_beta_headers(beta_provider)}, - extra_body=NO_PROXY_CACHE, + @pytest.mark.covers( + "llm.messages.anthropic.basic.nonstream.works", + "llm.messages.bedrock_invoke.basic.nonstream.works", + "llm.messages.bedrock_converse.basic.nonstream.works", ) + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.MESSAGES, + providers=(Provider.ANTHROPIC, Provider.BEDROCK), + models=(_ANTHROPIC_MODEL, _BEDROCK_MODEL_ID), + mode=Mode.NONSTREAM, + ) + ) + def test_every_mapped_beta_header_is_accepted_by_the_provider( + self, + proxy: ProxyClient, + resources: ResourceManager, + sdk: SdkClients, + backend: str, + beta_provider: str, + tools: list[ToolUnionParam], + ) -> None: + model: Final = f"e2e-messages-beta-{unique_marker()}" + model_id: Final = proxy.create_model(model, _deployment(backend)) + resources.defer(lambda: proxy.delete_model(model_id)) - text: Final = "".join(block.text for block in message.content if isinstance(block, TextBlock)) - assert message.role == "assistant" - assert "hello" in text.lower(), message.content - assert message.usage.input_tokens > 0 and message.usage.output_tokens > 0, message.usage + message: Final = sdk.anthropic(resources.key()).messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": "Say 'hello' and nothing else"}], + tools=tools, + extra_headers={"anthropic-beta": _mapped_beta_headers(beta_provider)}, + extra_body=NO_PROXY_CACHE, + ) + + text: Final = "".join(block.text for block in message.content if isinstance(block, TextBlock)) + assert message.role == "assistant" + assert "hello" in text.lower(), message.content + assert message.usage.input_tokens > 0 and message.usage.output_tokens > 0, message.usage diff --git a/tests/e2e/llm_translation/test_messages_prompt_caching_and_tool_search_e2e.py b/tests/e2e/llm_translation/test_messages_prompt_caching_and_tool_search_e2e.py index 1b7b44b5504..400420b0073 100644 --- a/tests/e2e/llm_translation/test_messages_prompt_caching_and_tool_search_e2e.py +++ b/tests/e2e/llm_translation/test_messages_prompt_caching_and_tool_search_e2e.py @@ -19,6 +19,8 @@ TOOL_SEARCH_BACKENDS: Final = ( pytest.param(ANTHROPIC_MODEL, id="anthropic"), *BEDROCK_BACKENDS, ) + + def _tool_search_tools() -> list[ToolUnionParam]: return [ cast( @@ -52,6 +54,8 @@ def _tool_search_tools() -> list[ToolUnionParam]: }, ), ] + + CACHED_DOCUMENT: Final = "This agreement describes payment terms, renewal dates, and service obligations. " * 120 pytestmark = pytest.mark.e2e @@ -103,195 +107,197 @@ def _message_text(message: Message) -> str: return "".join(block.text for block in message.content if isinstance(block, TextBlock)) -@pytest.mark.parametrize("backend", BEDROCK_BACKENDS) -@pytest.mark.covers( - "llm.messages.bedrock_converse.prompt_cache_5m.nonstream.works", - "llm.messages.bedrock_invoke.prompt_cache_5m.nonstream.works", -) -@meta( - Subject( - domain=Domain.LLM_TRANSLATION, - route=Route.MESSAGES, - providers=(Provider.BEDROCK,), - models=(BEDROCK_MODEL_ID,), - mode=Mode.NONSTREAM, +class TestBedrockMessagesPromptCaching: + @pytest.mark.parametrize("backend", BEDROCK_BACKENDS) + @pytest.mark.covers( + "llm.messages.bedrock_converse.prompt_cache_5m.nonstream.works", + "llm.messages.bedrock_invoke.prompt_cache_5m.nonstream.works", ) -) -def test_bedrock_messages_prompt_caching_creates_cache( - proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, backend: str -) -> None: - model, key = _register_deployment(proxy, resources, _deployment_params(backend), "e2e-prompt-cache-create") - client = sdk.anthropic(key) - message = client.messages.create( - model=model, - max_tokens=100, - messages=_cached_messages(unique_marker()), - extra_body=NO_PROXY_CACHE, + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.MESSAGES, + providers=(Provider.BEDROCK,), + models=(BEDROCK_MODEL_ID,), + mode=Mode.NONSTREAM, + ) ) + def test_bedrock_messages_prompt_caching_creates_cache( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, backend: str + ) -> None: + model, key = _register_deployment(proxy, resources, _deployment_params(backend), "e2e-prompt-cache-create") + client = sdk.anthropic(key) + message = client.messages.create( + model=model, + max_tokens=100, + messages=_cached_messages(unique_marker()), + extra_body=NO_PROXY_CACHE, + ) - assert message.role == "assistant", f"unexpected role: {message.role!r}" - assert (message.usage.cache_creation_input_tokens or 0) > 0 or (message.usage.cache_read_input_tokens or 0) > 0, ( - f"prompt cache was not created or read: {message.usage!r}" + assert message.role == "assistant", f"unexpected role: {message.role!r}" + assert (message.usage.cache_creation_input_tokens or 0) > 0 or ( + message.usage.cache_read_input_tokens or 0 + ) > 0, f"prompt cache was not created or read: {message.usage!r}" + + @pytest.mark.parametrize("backend", BEDROCK_BACKENDS) + @pytest.mark.covers( + "llm.messages.bedrock_converse.prompt_cache_5m.nonstream.works", + "llm.messages.bedrock_invoke.prompt_cache_5m.nonstream.works", ) - - -@pytest.mark.parametrize("backend", BEDROCK_BACKENDS) -@pytest.mark.covers( - "llm.messages.bedrock_converse.prompt_cache_5m.nonstream.works", - "llm.messages.bedrock_invoke.prompt_cache_5m.nonstream.works", -) -@meta( - Subject( - domain=Domain.LLM_TRANSLATION, - route=Route.MESSAGES, - providers=(Provider.BEDROCK,), - models=(BEDROCK_MODEL_ID,), - mode=Mode.NONSTREAM, + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.MESSAGES, + providers=(Provider.BEDROCK,), + models=(BEDROCK_MODEL_ID,), + mode=Mode.NONSTREAM, + ) ) -) -def test_bedrock_messages_prompt_caching_reads_cache_on_second_call( - proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, backend: str -) -> None: - model, key = _register_deployment(proxy, resources, _deployment_params(backend), "e2e-prompt-cache-read") - client = sdk.anthropic(key) - messages: Final = _cached_messages(unique_marker()) + def test_bedrock_messages_prompt_caching_reads_cache_on_second_call( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, backend: str + ) -> None: + model, key = _register_deployment(proxy, resources, _deployment_params(backend), "e2e-prompt-cache-read") + client = sdk.anthropic(key) + messages: Final = _cached_messages(unique_marker()) - client.messages.create(model=model, max_tokens=100, messages=messages, extra_body=NO_PROXY_CACHE) - message = client.messages.create(model=model, max_tokens=100, messages=messages, extra_body=NO_PROXY_CACHE) + client.messages.create(model=model, max_tokens=100, messages=messages, extra_body=NO_PROXY_CACHE) + message = client.messages.create(model=model, max_tokens=100, messages=messages, extra_body=NO_PROXY_CACHE) - assert (message.usage.cache_read_input_tokens or 0) > 0, f"cache read was not reported: {message.usage!r}" + assert (message.usage.cache_read_input_tokens or 0) > 0, f"cache read was not reported: {message.usage!r}" - -@pytest.mark.parametrize("backend", BEDROCK_BACKENDS) -@pytest.mark.covers( - "llm.messages.bedrock_converse.basic.stream.works", - "llm.messages.bedrock_invoke.basic.stream.works", -) -@meta( - Subject( - domain=Domain.LLM_TRANSLATION, - route=Route.MESSAGES, - providers=(Provider.BEDROCK,), - models=(BEDROCK_MODEL_ID,), - mode=Mode.STREAM, + @pytest.mark.parametrize("backend", BEDROCK_BACKENDS) + @pytest.mark.covers( + "llm.messages.bedrock_converse.basic.stream.works", + "llm.messages.bedrock_invoke.basic.stream.works", ) -) -def test_bedrock_messages_streaming_prompt_caching_reads_cache( - proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, backend: str -) -> None: - model, key = _register_deployment(proxy, resources, _deployment_params(backend), "e2e-prompt-cache-stream") - client = sdk.anthropic(key) - messages: Final = _cached_messages(unique_marker()) - client.messages.create(model=model, max_tokens=100, messages=messages, extra_body=NO_PROXY_CACHE) - - with client.messages.stream( - model=model, - max_tokens=100, - messages=messages, - extra_body=NO_PROXY_CACHE, - ) as stream: - message: Final = stream.get_final_message() - - assert (message.usage.cache_read_input_tokens or 0) > 0, f"streaming cache read was not reported: {message.usage!r}" - assert _message_text(message).strip(), f"streaming response contained no text: {message.content!r}" - - -@pytest.mark.parametrize("backend", TOOL_SEARCH_BACKENDS) -@pytest.mark.covers( - "llm.messages.anthropic.tool_search.nonstream.works", - "llm.messages.bedrock_converse.tool_search.nonstream.works", - "llm.messages.bedrock_invoke.tool_search.nonstream.works", -) -@meta( - Subject( - domain=Domain.LLM_TRANSLATION, - route=Route.MESSAGES, - providers=(Provider.ANTHROPIC, Provider.BEDROCK), - models=(ANTHROPIC_MODEL, BEDROCK_MODEL_ID), - mode=Mode.NONSTREAM, + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.MESSAGES, + providers=(Provider.BEDROCK,), + models=(BEDROCK_MODEL_ID,), + mode=Mode.STREAM, + ) ) -) -def test_anthropic_tool_search_returns_a_message( - proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, backend: str -) -> None: - model, key = _register_deployment(proxy, resources, _deployment_params(backend), "e2e-tool-search") - message = sdk.anthropic(key).messages.create( - model=model, - max_tokens=1024, - messages=[{"role": "user", "content": f"Find the weather tool for San Francisco. {unique_marker()}"}], - tools=_tool_search_tools(), - extra_headers={"anthropic-beta": "tool-search-tool-2025-10-19,advanced-tool-use-2025-11-20"}, - extra_body=NO_PROXY_CACHE, + def test_bedrock_messages_streaming_prompt_caching_reads_cache( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, backend: str + ) -> None: + model, key = _register_deployment(proxy, resources, _deployment_params(backend), "e2e-prompt-cache-stream") + client = sdk.anthropic(key) + messages: Final = _cached_messages(unique_marker()) + client.messages.create(model=model, max_tokens=100, messages=messages, extra_body=NO_PROXY_CACHE) + + with client.messages.stream( + model=model, + max_tokens=100, + messages=messages, + extra_body=NO_PROXY_CACHE, + ) as stream: + message: Final = stream.get_final_message() + + assert (message.usage.cache_read_input_tokens or 0) > 0, ( + f"streaming cache read was not reported: {message.usage!r}" + ) + assert _message_text(message).strip(), f"streaming response contained no text: {message.content!r}" + + +class TestAnthropicToolSearch: + @pytest.mark.parametrize("backend", TOOL_SEARCH_BACKENDS) + @pytest.mark.covers( + "llm.messages.anthropic.tool_search.nonstream.works", + "llm.messages.bedrock_converse.tool_search.nonstream.works", + "llm.messages.bedrock_invoke.tool_search.nonstream.works", ) - - assert message.role == "assistant", f"unexpected role: {message.role!r}" - assert message.content, "tool search response contained no content" - assert message.usage.input_tokens > 0, f"tool search response had no input usage: {message.usage!r}" - - -@pytest.mark.parametrize("backend", TOOL_SEARCH_BACKENDS) -@pytest.mark.covers( - "llm.messages.anthropic.tool_search.nonstream.works", - "llm.messages.bedrock_converse.tool_search.nonstream.works", - "llm.messages.bedrock_invoke.tool_search.nonstream.works", -) -@meta( - Subject( - domain=Domain.LLM_TRANSLATION, - route=Route.MESSAGES, - providers=(Provider.ANTHROPIC, Provider.BEDROCK), - models=(ANTHROPIC_MODEL, BEDROCK_MODEL_ID), - mode=Mode.STREAM, + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.MESSAGES, + providers=(Provider.ANTHROPIC, Provider.BEDROCK), + models=(ANTHROPIC_MODEL, BEDROCK_MODEL_ID), + mode=Mode.NONSTREAM, + ) ) -) -def test_anthropic_tool_search_streaming_returns_a_message( - proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, backend: str -) -> None: - model, key = _register_deployment(proxy, resources, _deployment_params(backend), "e2e-tool-search-stream") - client = sdk.anthropic(key) + def test_anthropic_tool_search_returns_a_message( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, backend: str + ) -> None: + model, key = _register_deployment(proxy, resources, _deployment_params(backend), "e2e-tool-search") + message = sdk.anthropic(key).messages.create( + model=model, + max_tokens=1024, + messages=[{"role": "user", "content": f"Find the weather tool for San Francisco. {unique_marker()}"}], + tools=_tool_search_tools(), + extra_headers={"anthropic-beta": "tool-search-tool-2025-10-19,advanced-tool-use-2025-11-20"}, + extra_body=NO_PROXY_CACHE, + ) - with client.messages.stream( - model=model, - max_tokens=1024, - messages=[{"role": "user", "content": f"Find the weather tool for Tokyo. {unique_marker()}"}], - tools=_tool_search_tools(), - extra_headers={"anthropic-beta": "tool-search-tool-2025-10-19,advanced-tool-use-2025-11-20"}, - extra_body=NO_PROXY_CACHE, - ) as stream: - message: Final = stream.get_final_message() + assert message.role == "assistant", f"unexpected role: {message.role!r}" + assert message.content, "tool search response contained no content" + assert message.usage.input_tokens > 0, f"tool search response had no input usage: {message.usage!r}" - assert message.role == "assistant", f"unexpected role: {message.role!r}" - assert message.content, "streaming tool search response contained no content" - assert message.usage.input_tokens > 0, f"streaming tool search response had no input usage: {message.usage!r}" - - -@pytest.mark.covers("llm.messages.bedrock_invoke.basic.nonstream.works") -@meta( - Subject( - domain=Domain.LLM_TRANSLATION, - route=Route.MESSAGES, - providers=(Provider.BEDROCK,), - models=(BEDROCK_MODEL_ID,), - mode=Mode.NONSTREAM, + @pytest.mark.parametrize("backend", TOOL_SEARCH_BACKENDS) + @pytest.mark.covers( + "llm.messages.anthropic.tool_search.nonstream.works", + "llm.messages.bedrock_converse.tool_search.nonstream.works", + "llm.messages.bedrock_invoke.tool_search.nonstream.works", ) -) -def test_bedrock_invoke_messages_accepts_forwarded_headers( - proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients -) -> None: - params: Final = _deployment_params(f"bedrock/invoke/{BEDROCK_MODEL_ID}") - model, key = _register_deployment(proxy, resources, params, "e2e-bedrock-forwarded-headers") - message = sdk.anthropic(key).messages.create( - model=model, - max_tokens=32, - messages=[{"role": "user", "content": f"Reply with one word. {unique_marker()}"}], - extra_headers={ - "x-forwarded-for": "10.11.232.194", - "x-forwarded-port": "443", - "x-forwarded-proto": "https", - "x-app": "cli", - }, - extra_body=NO_PROXY_CACHE, + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.MESSAGES, + providers=(Provider.ANTHROPIC, Provider.BEDROCK), + models=(ANTHROPIC_MODEL, BEDROCK_MODEL_ID), + mode=Mode.STREAM, + ) ) + def test_anthropic_tool_search_streaming_returns_a_message( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, backend: str + ) -> None: + model, key = _register_deployment(proxy, resources, _deployment_params(backend), "e2e-tool-search-stream") + client = sdk.anthropic(key) - assert message.role == "assistant", f"unexpected role: {message.role!r}" - assert _message_text(message).strip(), f"forwarded-header response contained no text: {message.content!r}" + with client.messages.stream( + model=model, + max_tokens=1024, + messages=[{"role": "user", "content": f"Find the weather tool for Tokyo. {unique_marker()}"}], + tools=_tool_search_tools(), + extra_headers={"anthropic-beta": "tool-search-tool-2025-10-19,advanced-tool-use-2025-11-20"}, + extra_body=NO_PROXY_CACHE, + ) as stream: + message: Final = stream.get_final_message() + + assert message.role == "assistant", f"unexpected role: {message.role!r}" + assert message.content, "streaming tool search response contained no content" + assert message.usage.input_tokens > 0, f"streaming tool search response had no input usage: {message.usage!r}" + + +class TestBedrockInvokeForwardedHeaders: + @pytest.mark.covers("llm.messages.bedrock_invoke.basic.nonstream.works") + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.MESSAGES, + providers=(Provider.BEDROCK,), + models=(BEDROCK_MODEL_ID,), + mode=Mode.NONSTREAM, + ) + ) + def test_bedrock_invoke_messages_accepts_forwarded_headers( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + params: Final = _deployment_params(f"bedrock/invoke/{BEDROCK_MODEL_ID}") + model, key = _register_deployment(proxy, resources, params, "e2e-bedrock-forwarded-headers") + message = sdk.anthropic(key).messages.create( + model=model, + max_tokens=32, + messages=[{"role": "user", "content": f"Reply with one word. {unique_marker()}"}], + extra_headers={ + "x-forwarded-for": "10.11.232.194", + "x-forwarded-port": "443", + "x-forwarded-proto": "https", + "x-app": "cli", + }, + extra_body=NO_PROXY_CACHE, + ) + + assert message.role == "assistant", f"unexpected role: {message.role!r}" + assert _message_text(message).strip(), f"forwarded-header response contained no text: {message.content!r}" diff --git a/tests/e2e/secret_manager/test_aws_secret_manager_e2e.py b/tests/e2e/secret_manager/test_aws_secret_manager_e2e.py index 0c6ca0cacaf..5bd67abd2f6 100644 --- a/tests/e2e/secret_manager/test_aws_secret_manager_e2e.py +++ b/tests/e2e/secret_manager/test_aws_secret_manager_e2e.py @@ -24,11 +24,11 @@ _STRING_MAP: Final = TypeAdapter(dict[str, str]) _SECRET_STRING: Final = TypeAdapter(str) -@meta(Subject(domain=Domain.DEPLOY_OPS)) -def test_aws_secret_manager() -> None: - AWSSecretsManagerV2.load_aws_secret_manager(use_aws_secret_manager=True) - secret_value: Final = _STRING_MAP.validate_json(_SECRET_STRING.validate_python(get_secret("litellm_master_key"))) - assert ( - hashlib.sha256(secret_value["litellm_master_key"].encode()).hexdigest() - == _AWS_FIXTURE_MASTER_KEY_SHA256 - ) +class TestAwsSecretManager: + @meta(Subject(domain=Domain.DEPLOY_OPS)) + def test_aws_secret_manager(self) -> None: + AWSSecretsManagerV2.load_aws_secret_manager(use_aws_secret_manager=True) + secret_value: Final = _STRING_MAP.validate_json( + _SECRET_STRING.validate_python(get_secret("litellm_master_key")) + ) + assert hashlib.sha256(secret_value["litellm_master_key"].encode()).hexdigest() == _AWS_FIXTURE_MASTER_KEY_SHA256 diff --git a/tests/integration/sdk/test_ocr_matrix_sdk.py b/tests/integration/sdk/test_ocr_matrix_sdk.py index e1082adebc8..d0e9329b30c 100644 --- a/tests/integration/sdk/test_ocr_matrix_sdk.py +++ b/tests/integration/sdk/test_ocr_matrix_sdk.py @@ -25,6 +25,7 @@ pytestmark = pytest.mark.asyncio Document = Mapping[str, object] AuthMode = Literal["explicit", "env"] CallStyle = Literal["sync", "async"] +PriceField = Literal["ocr_cost_per_page", "ocr_cost_per_credit"] PAGE_TEXT: Final = "scripted ocr page" PDF_BYTES: Final = b"%PDF-1.4\n%ocr matrix document\n%%EOF\n" @@ -37,6 +38,11 @@ DOCUMENTS: Final = { } DI_ANALYZE_PATH: Final = "/documentintelligence/documentModels/prebuilt-layout" OCR_CALL_TYPES: Final = frozenset({"ocr", "aocr"}) +UPLOADED_FILE_ID: Final = "reducto://scripted-upload" +BILLED_PAGES: Final = 2 +BILLED_CREDITS: Final = 3 +RATE: Final = 0.0125 +PAGE_NUMBERS: Final = tuple(range(1, BILLED_PAGES + 1)) _JSON: Final = TypeAdapter(dict[str, JsonValue]) _STRING_KEYED: Final = TypeAdapter(dict[str, object]) _FLOAT: Final = TypeAdapter(float) @@ -61,36 +67,30 @@ def _respond(request: Request) -> Reply: content, content_type = DOCUMENTS[path] return Reply(body=content, content_type=content_type) if path == "/upload": - return _json_reply({"file_id": "reducto://scripted-upload"}) + return _json_reply({"file_id": UPLOADED_FILE_ID}) if path == "/parse": + chunks: Final = [ + {"content": PAGE_TEXT, "blocks": [{"content": PAGE_TEXT, "bbox": {"page": page}}]} for page in PAGE_NUMBERS + ] return _json_reply( { "job_id": "scripted-job", - "usage": {"num_pages": 1, "credits": 1}, - "result": {"chunks": [{"content": PAGE_TEXT, "blocks": [{"content": PAGE_TEXT, "bbox": {"page": 1}}]}]}, + "usage": {"num_pages": BILLED_PAGES, "credits": BILLED_CREDITS}, + "result": {"chunks": chunks}, } ) if path == f"{DI_ANALYZE_PATH}:analyze": operation: Final = f"http://{request.headers['host']}{DI_ANALYZE_PATH}/analyzeResults/scripted?api-version=1" return Reply(status=202, body=b"", headers={"operation-location": operation}) if path.startswith(f"{DI_ANALYZE_PATH}/analyzeResults/"): - return _json_reply( - { - "status": "succeeded", - "analyzeResult": { - "content": PAGE_TEXT, - "pages": [{"pageNumber": 1, "lines": [{"content": PAGE_TEXT}]}], - }, - } - ) + pages: Final = [{"pageNumber": page, "lines": [{"content": PAGE_TEXT}]} for page in PAGE_NUMBERS] + return _json_reply({"status": "succeeded", "analyzeResult": {"content": PAGE_TEXT, "pages": pages}}) if path == "/v2/parse": - return _json_reply( - {"pages": [{"index": 0, "markdown": {"content": PAGE_TEXT}}], "meta": {"billed_units": {"pages": 1}}} - ) + cohere_pages: Final = [{"index": page - 1, "markdown": {"content": PAGE_TEXT}} for page in PAGE_NUMBERS] + return _json_reply({"pages": cohere_pages, "meta": {"billed_units": {"pages": BILLED_PAGES}}}) model: Final = _JSON.validate_json(request.body)["model"] - return _json_reply( - {"pages": [{"index": 0, "markdown": PAGE_TEXT}], "model": model, "usage_info": {"pages_processed": 1}} - ) + mistral_pages: Final = [{"index": page - 1, "markdown": PAGE_TEXT} for page in PAGE_NUMBERS] + return _json_reply({"pages": mistral_pages, "model": model, "usage_info": {"pages_processed": BILLED_PAGES}}) @dataclass(frozen=True, slots=True) @@ -202,10 +202,16 @@ class Provider: auth_header: str base_env: str | None = None downloads_urls: bool = False + parse_input_field: str | None = None + price_field: PriceField = "ocr_cost_per_page" def credential_header(self, key: str) -> str: return key if self.auth_header == "ocp-apim-subscription-key" else f"Bearer {key}" + @property + def billed_units(self) -> int: + return BILLED_CREDITS if self.price_field == "ocr_cost_per_credit" else BILLED_PAGES + MISTRAL: Final = Provider("mistral", "mistral/mistral-ocr-latest", "MISTRAL_API_KEY", "authorization") AZURE_AI_MISTRAL: Final = Provider( @@ -224,8 +230,22 @@ AZURE_DOC_INTELLIGENCE: Final = Provider( base_env="AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT", ) COHERE: Final = Provider("cohere", "cohere/parse-v5.0", "COHERE_API_KEY", "authorization") -REDUCTO_V3: Final = Provider("reducto_v3", "reducto/parse-v3", "REDUCTO_API_KEY", "authorization") -REDUCTO_LEGACY: Final = Provider("reducto_legacy", "reducto/parse-legacy", "REDUCTO_API_KEY", "authorization") +REDUCTO_V3: Final = Provider( + "reducto_v3", + "reducto/parse-v3", + "REDUCTO_API_KEY", + "authorization", + parse_input_field="input", + price_field="ocr_cost_per_credit", +) +REDUCTO_LEGACY: Final = Provider( + "reducto_legacy", + "reducto/parse-legacy", + "REDUCTO_API_KEY", + "authorization", + parse_input_field="document_url", + price_field="ocr_cost_per_credit", +) @dataclass(frozen=True, slots=True) @@ -256,10 +276,11 @@ class Case: return {} async def run(self, document: Document, credentials: Mapping[str, str]) -> OCRResponse: + kwargs: Final = {**credentials, self.provider.price_field: RATE} response: Final = ( - await _OCR_CLIENT.aocr(model=self.provider.model, document=document, **credentials) + await _OCR_CLIENT.aocr(model=self.provider.model, document=document, **kwargs) if self.call == "async" - else _OCR_CLIENT.ocr(model=self.provider.model, document=document, **credentials) + else _OCR_CLIENT.ocr(model=self.provider.model, document=document, **kwargs) ) assert isinstance(response, OCRResponse) return response @@ -284,6 +305,10 @@ CASES: Final = ( ) +def _path(request: Request) -> str: + return request.target.split("?", 1)[0] + + def _decoded(value: str) -> bytes | str: if not value.startswith("data:"): return value @@ -316,46 +341,50 @@ def _sent_document(request: Request) -> bytes | str: def _assert_sent(case: Case, requests: tuple[Request, ...], document: Document) -> None: - provider_calls: Final = tuple(request for request in requests if request.target.split("?", 1)[0] not in DOCUMENTS) + provider_calls: Final = tuple(request for request in requests if _path(request) not in DOCUMENTS) assert provider_calls, requests assert {request.headers.get(case.provider.auth_header) for request in provider_calls} == { case.provider.credential_header(case.key) } - downloads: Final = tuple(request for request in requests if request.target.split("?", 1)[0] in DOCUMENTS) + downloads: Final = tuple(_path(request) for request in requests if _path(request) in DOCUMENTS) url: Final = document.get("document_url") or document.get("image_url") passes_url_through: Final = case.document.url_path is not None and not case.provider.downloads_urls assert _sent_document(provider_calls[0]) == (url if passes_url_through else case.document.content) - assert [request.target for request in downloads] == ( - [case.document.url_path] if case.document.url_path is not None and case.provider.downloads_urls else [] + assert downloads == ( + (case.document.url_path,) if case.document.url_path is not None and case.provider.downloads_urls else () ) + if case.provider.parse_input_field is not None: + parse_bodies: Final = tuple( + _JSON.validate_json(request.body) for request in provider_calls if _path(request) == "/parse" + ) + assert parse_bodies == ({case.provider.parse_input_field: UPLOADED_FILE_ID},) def _assert_ocr_response(response: OCRResponse, model: str) -> None: assert response.object == "ocr" assert response.model == model.split("/", 1)[1] - assert [page.index for page in response.pages] == list(range(len(response.pages))) - assert PAGE_TEXT in " ".join(page.markdown for page in response.pages) + assert [page.index for page in response.pages] == list(range(BILLED_PAGES)) + assert [page.markdown for page in response.pages] == [PAGE_TEXT] * BILLED_PAGES assert response.usage_info is not None - assert response.usage_info.pages_processed == len(response.pages) + assert response.usage_info.pages_processed == BILLED_PAGES -def _assert_logged(logged: LoggedCall, response: OCRResponse, model: str, call: CallStyle) -> None: +def _assert_logged(logged: LoggedCall, response: OCRResponse, case: Case) -> None: assert logged.response is response assert logged.payload["status"] == "success" - assert logged.payload["call_type"] == ("aocr" if call == "async" else "ocr") - assert logged.payload["custom_llm_provider"] == model.split("/", 1)[0] + assert logged.payload["call_type"] == ("aocr" if case.call == "async" else "ocr") + assert logged.payload["custom_llm_provider"] == case.provider.model.split("/", 1)[0] assert logged.payload["model"] == response.model - assert _FLOAT.validate_python(logged.payload["response_cost"]) > 0 + assert _FLOAT.validate_python(logged.payload["response_cost"]) == pytest.approx(case.provider.billed_units * RATE) @pytest.mark.parametrize("case", CASES, ids=[case.id for case in CASES]) -async def test_ocr_sends_the_bound_credential_and_document_then_logs_a_priced_call( +async def test_ocr_sends_the_bound_credential_and_document_then_logs_the_billed_cost( case: Case, monkeypatch: pytest.MonkeyPatch, logger: RecordingLogger, tmp_path: Path, ) -> None: - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") with wire_server(_respond) as wire: monkeypatch.setattr(litellm, "user_url_allowed_hosts", [wire.url.removeprefix("http://")]) credentials: Final = case.bind_credentials(monkeypatch, wire.url) @@ -366,4 +395,4 @@ async def test_ocr_sends_the_bound_credential_and_document_then_logs_a_priced_ca _assert_ocr_response(response, case.provider.model) _assert_sent(case, requests, document) - _assert_logged(logged, response, case.provider.model, case.call) + _assert_logged(logged, response, case)