scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-17 04:34:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 172

test: pin OCR parse bodies and billed cost, group new e2e cases into classes

Browse source
This commit is contained in:
Yuneng Jiang 2026-10-10 11:54:11 -07:00
parent 067f77eb0d
commit 6c10a93fde
No known key found for this signature in database
4 changed files with 300 additions and 263 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

90
tests/e2e/llm_translation/test_messages_beta_headers_e2e.py
View file

@ -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

356
tests/e2e/llm_translation/test_messages_prompt_caching_and_tool_search_e2e.py
View file

@ -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}"

16
tests/e2e/secret_manager/test_aws_secret_manager_e2e.py
View file

@ -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

101
tests/integration/sdk/test_ocr_matrix_sdk.py
View file

@ -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)